[LL API] Do not throw PCE on invalid sessions, patch specific cases

This is a temporary solution to fix the problem with sessions that
turned to be invalid at the time of analysis.

Currently, 'FirSession's and associated 'FirFile's are reused
between threads. When an exception, which might be an ordinary one
or 'ProcessCancelledException', occurs during an analysis of a module
file, a 'FirSession' for that module becomes invalid.

As there is currently a single resolution lock in the K2 IDE, other
threads might be waiting for resolution on the same session.
These threads are unaware that the session cannot be safely used
anymore, as it was correct before the lock.

The previous approach with cancelling analysis in all waiting threads
with manual 'ProcessCancelledException' throwing right inside the lock
showed it's incorrect. While some actions might be prepared to sudden
PCEs and can restart themselves, such behavior is not granted
by the platform.

The fix patches a few common places so the diagnostic list and
a file structure tree can be built safely. It is expected that in
some places exceptions caused by an inconsistent FIR tree can revive.

The right, ultimate fix to the problem involves modifying transformers,
so the FIR tree will mutate only when the analysis is complete. Although
it's planned work, it is expected to take quite a time.

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/builder/LLFirLockProvider.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder
 
+import com.intellij.openapi.diagnostic.Logger
 import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.LLFirLazyResolveContractChecker
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkCanceled
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkPhase
@@ -15,7 +16,9 @@ import java.util.concurrent.CountDownLatch
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicReferenceFieldUpdater
 import java.util.concurrent.locks.ReentrantLock
-import com.intellij.openapi.progress.ProcessCanceledException
+import kotlin.contracts.ExperimentalContracts
+import kotlin.contracts.InvocationKind
+import kotlin.contracts.contract
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.llFirSession
 import org.jetbrains.kotlin.fir.declarations.FirFile
 
@@ -36,11 +39,14 @@ internal class LLFirLockProvider(private val checker: LLFirLazyResolveContractCh
         action: () -> R
     ): R = globalLock.lockWithPCECheck(lockingIntervalMs) {
         val session = key.llFirSession
-        if (!session.isValid) {
-            throw ProcessCanceledException()
+        if (!session.isValid && shouldRetryFlag.get()) {
+            val description = session.ktModule.moduleDescription
+            throw InvalidSessionException("Session '$description' is invalid", description)
         }
 
-        action()
+        // Normally, analysis should not be allowed on an invalid session.
+        // However, there isn't an easy way to cancel or redo it in general case, as it must then be supported on use-site.
+        withRetryFlag(false, action)
     }
 
     fun withGlobalPhaseLock(
@@ -177,4 +183,49 @@ private val resolveStateFieldUpdater = AtomicReferenceFieldUpdater.newUpdater(
     "resolveState"
 )
 
-private const val DEFAULT_LOCKING_INTERVAL = 50L
+private const val DEFAULT_LOCKING_INTERVAL = 50L
+
+internal class InvalidSessionException(message: String, val moduleDescription: String) : RuntimeException(message)
+
+/*
+    The flag specifies whether the analysis action should be repeated in case if it was originally started on an invalid session.
+
+    Possible values:
+        - `true` – throw the marker [InvalidSessionException] to trigger the retry.
+        - `false` – process analysis as usual (default).
+ */
+private val shouldRetryFlag: ThreadLocal<Boolean> = ThreadLocal.withInitial { false }
+
+private val LOG = Logger.getInstance(LLFirLockProvider::class.java)
+
+/**
+ * Retry the `action` calculation with a new FIR session if session passed to [LLFirLockProvider.withLock] turns to be invalid.
+ * This is a temporary solution to fix inconsistent analysis state in common cases of idempotent analysis.
+ * The right solution would be to modify the FIR tree after the analysis is done, so the tree will always be in consistent state.
+ */
+internal inline fun <R> retryOnInvalidSession(action: () -> R): R {
+    withRetryFlag(true) {
+        while (true) {
+            try {
+                return action()
+            } catch (e: InvalidSessionException) {
+                LOG.warn("Processing with invalid module '${e.moduleDescription}'")
+            }
+        }
+    }
+}
+
+@OptIn(ExperimentalContracts::class)
+private inline fun <R> withRetryFlag(value: Boolean, action: () -> R): R {
+    contract {
+        callsInPlace(action, InvocationKind.EXACTLY_ONCE)
+    }
+
+    val oldValue = shouldRetryFlag.get()
+    shouldRetryFlag.set(value)
+    try {
+        return action()
+    } finally {
+        shouldRetryFlag.set(oldValue)
+    }
+}

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/LLFirModuleLazyDeclarationResolver.kt

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignationWithFil
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirClassWithAllMembersResolveTarget
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirResolveTarget
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirSingleResolveTarget
+import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.InvalidSessionException
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.llFirModuleData
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.llFirSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.transformers.LLFirLazyResolverRunner
@@ -204,6 +205,10 @@ private fun handleExceptionFromResolve(
     fromPhase: FirResolvePhase,
     toPhase: FirResolvePhase?
 ): Nothing {
+    if (exception is InvalidSessionException) {
+        throw exception
+    }
+
     firDeclarationToResolve.llFirSession.invalidate()
     rethrowExceptionWithDetails(
         buildString {
@@ -231,6 +236,10 @@ private fun handleExceptionFromResolve(
     designation: LLFirResolveTarget,
     toPhase: FirResolvePhase?
 ): Nothing {
+    if (exception is InvalidSessionException) {
+        throw exception
+    }
+
     val llFirSession = designation.firFile.llFirSession
     llFirSession.invalidate()
     val moduleData = llFirSession.llFirModuleData

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/state/LLFirResolvableResolveSession.kt

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveCompone
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirTowerContextProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.getNonLocalContainingOrThisDeclaration
+import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.retryOnInvalidSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirResolvableModuleSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionCache
@@ -76,13 +77,17 @@ internal abstract class LLFirResolvableResolveSession(
     }
 
     override fun getOrBuildFirFor(element: KtElement): FirElement? {
-        val moduleComponents = getModuleComponentsForElement(element)
-        return moduleComponents.elementsBuilder.getOrBuildFirFor(element, this)
+        retryOnInvalidSession {
+            val moduleComponents = getModuleComponentsForElement(element)
+            return moduleComponents.elementsBuilder.getOrBuildFirFor(element, this)
+        }
     }
 
     override fun getOrBuildFirFile(ktFile: KtFile): FirFile {
-        val moduleComponents = getModuleComponentsForElement(ktFile)
-        return moduleComponents.firFileBuilder.buildRawFirFileWithCaching(ktFile)
+        retryOnInvalidSession {
+            val moduleComponents = getModuleComponentsForElement(ktFile)
+            return moduleComponents.firFileBuilder.buildRawFirFileWithCaching(ktFile)
+        }
     }
 
     protected fun getModuleComponentsForElement(element: KtElement): LLFirModuleResolveComponents {
@@ -98,9 +103,11 @@ internal abstract class LLFirResolvableResolveSession(
             return findFirCompiledSymbol(ktDeclaration)
         }
         val module = ktDeclaration.getKtModule()
-        return when (getModuleKind(module)) {
-            ModuleKind.RESOLVABLE_MODULE -> findSourceFirSymbol(ktDeclaration, module).also { resolveFirToPhase(it.fir, phase) }
-            ModuleKind.BINARY_MODULE -> findFirCompiledSymbol(ktDeclaration)
+        retryOnInvalidSession {
+            return when (getModuleKind(module)) {
+                ModuleKind.RESOLVABLE_MODULE -> findSourceFirSymbol(ktDeclaration, module).also { resolveFirToPhase(it.fir, phase) }
+                ModuleKind.BINARY_MODULE -> findFirCompiledSymbol(ktDeclaration)
+            }
         }
     }
 

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/state/LLFirScriptResolveSession.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir.state
 
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.DiagnosticCheckerFilter
+import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.retryOnInvalidSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
 import org.jetbrains.kotlin.analysis.project.structure.*
 import org.jetbrains.kotlin.analysis.utils.errors.unexpectedElementError
@@ -18,13 +19,17 @@ internal class LLFirScriptResolveSession(
     useSiteSessionFactory: (KtModule) -> LLFirSession
 ) : LLFirResolvableResolveSession(useSiteKtModule, useSiteSessionFactory) {
     override fun getDiagnostics(element: KtElement, filter: DiagnosticCheckerFilter): List<KtPsiDiagnostic> {
-        val moduleComponents = getModuleComponentsForElement(element)
-        return moduleComponents.diagnosticsCollector.getDiagnosticsFor(element, filter)
+        retryOnInvalidSession {
+            val moduleComponents = getModuleComponentsForElement(element)
+            return moduleComponents.diagnosticsCollector.getDiagnosticsFor(element, filter)
+        }
     }
 
     override fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<KtPsiDiagnostic> {
-        val moduleComponents = getModuleComponentsForElement(ktFile)
-        return moduleComponents.diagnosticsCollector.collectDiagnosticsForFile(ktFile, filter)
+        retryOnInvalidSession {
+            val moduleComponents = getModuleComponentsForElement(ktFile)
+            return moduleComponents.diagnosticsCollector.collectDiagnosticsForFile(ktFile, filter)
+        }
     }
 
     override fun getModuleKind(module: KtModule): ModuleKind {

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/state/LLFirSourceResolveSession.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir.state
 
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.DiagnosticCheckerFilter
+import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.retryOnInvalidSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
 import org.jetbrains.kotlin.analysis.project.structure.KtBuiltinsModule
 import org.jetbrains.kotlin.analysis.project.structure.KtLibraryModule
@@ -21,13 +22,17 @@ internal class LLFirSourceResolveSession(
     useSiteSessionFactory: (KtModule) -> LLFirSession
 ) : LLFirResolvableResolveSession(useSiteKtModule, useSiteSessionFactory) {
     override fun getDiagnostics(element: KtElement, filter: DiagnosticCheckerFilter): List<KtPsiDiagnostic> {
-        val moduleComponents = getModuleComponentsForElement(element)
-        return moduleComponents.diagnosticsCollector.getDiagnosticsFor(element, filter)
+        retryOnInvalidSession {
+            val moduleComponents = getModuleComponentsForElement(element)
+            return moduleComponents.diagnosticsCollector.getDiagnosticsFor(element, filter)
+        }
     }
 
     override fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<KtPsiDiagnostic> {
-        val moduleComponents = getModuleComponentsForElement(ktFile)
-        return moduleComponents.diagnosticsCollector.collectDiagnosticsForFile(ktFile, filter)
+        retryOnInvalidSession {
+            val moduleComponents = getModuleComponentsForElement(ktFile)
+            return moduleComponents.diagnosticsCollector.collectDiagnosticsForFile(ktFile, filter)
+        }
     }
 
     override fun getModuleKind(module: KtModule): ModuleKind {