From f353199b7b390aac3d80056d3798d6a2977f18a2 Mon Sep 17 00:00:00 2001
From: Yan Zhulanow <yan.zhulanow@jetbrains.com>
Date: Mon, 27 Feb 2023 23:02:42 +0900
Subject: [PATCH] [LL API] Ensure that the session is valid right after
 entering the lock

---
 .../level/api/fir/file/builder/LLFirLockProvider.kt  | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/builder/LLFirLockProvider.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/builder/LLFirLockProvider.kt
index 5cf35e8fab8..c40fbb30146 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/builder/LLFirLockProvider.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/builder/LLFirLockProvider.kt
@@ -5,6 +5,8 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder
 
+import com.intellij.openapi.progress.ProcessCanceledException
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.llFirSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.lockWithPCECheck
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import java.util.concurrent.locks.ReentrantLock
@@ -18,11 +20,17 @@ internal class LLFirLockProvider {
     private val globalLock = ReentrantLock()
 
     inline fun <R> withLock(
-        @Suppress("UNUSED_PARAMETER") key: FirFile,
+        key: FirFile,
         lockingIntervalMs: Long = DEFAULT_LOCKING_INTERVAL,
         action: () -> R
     ): R {
-        return globalLock.lockWithPCECheck(lockingIntervalMs) { action() }
+        return globalLock.lockWithPCECheck(lockingIntervalMs) {
+            val session = key.llFirSession
+            if (!session.isValid) {
+                throw ProcessCanceledException()
+            }
+            action()
+        }
     }
 }