[LL FIR] KT-58257 Implement active session invalidation

- This commit replaces session modification trackers with event-based session invalidation. The removal of modification trackers should improve overall performance, because session invalidation events happen less frequently than sessions are accessed. Getting rid of modification trackers also allows sessions to refer to other sessions lazily, which is essential when cyclic dependencies occur. - The new `LLFirSession` validity tracker has constant complexity and will not cause the same kind of performance issues as dependency modification trackers. It is a bridge to support modification trackers in certain parts of the code (e.g. for `CachedValue`s), while being backed by event-based invalidation. - `LLFirSessionInvalidationService` is the bridge between modification events and `LLFirSessionCache`. It finds out which modules should be invalidated and instructs the session cache to remove the associated sessions. - Session invalidation must always happen in a write action to preserve consistency between sessions. Otherwise, while a session A is already removed from the cache, it might still be referenced by a dependent session B which is in the process of being invalidated. Such a session must never be returned from `getSession`. ^KT-57515 fixed
2023-05-24 14:09:51 +02:00
parent 1f556589a2
commit e76692ef8d
20 changed files with 225 additions and 212 deletions
@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.analysis.api.fir

 import com.intellij.openapi.project.Project
-import com.intellij.openapi.roots.ProjectRootModificationTracker
 import com.intellij.openapi.util.LowMemoryWatcher
 import com.intellij.psi.util.CachedValue
 import com.intellij.psi.util.CachedValueProvider
@@ -21,7 +20,6 @@ import org.jetbrains.kotlin.analysis.api.session.KtAnalysisSessionProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getFirResolveSession
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
-import org.jetbrains.kotlin.analysis.providers.createProjectWideOutOfBlockModificationTracker
 import org.jetbrains.kotlin.psi.KtElement
 import java.util.concurrent.ConcurrentMap
 import kotlin.reflect.KClass
@@ -49,9 +47,7 @@ class KtFirAnalysisSessionProvider(project: Project) : KtAnalysisSessionProvider

                CachedValueProvider.Result(
                    KtFirAnalysisSession.createAnalysisSessionByFirResolveSession(firResolveSession, validityToken),
-                    firResolveSession.useSiteFirSession.modificationTracker,
-                    ProjectRootModificationTracker.getInstance(project),
-                    project.createProjectWideOutOfBlockModificationTracker()
+                    firResolveSession.useSiteFirSession.createValidityTracker(),
                )
            }
        }.value
@@ -1,12 +0,0 @@
-/*
- * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
- * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
- */
-
-package org.jetbrains.kotlin.analysis.api.fir.utils
-
-import com.intellij.openapi.util.ModificationTracker
-import org.jetbrains.kotlin.analysis.utils.trackers.CompositeModificationTracker
-
-public fun List<ModificationTracker>.createCompositeModificationTracker(): ModificationTracker =
-    CompositeModificationTracker.createFlattened(this)
@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.JvmFirD
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.LLFirBuiltinsSessionFactory
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionCache
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionConfigurator
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionInvalidationService
 import org.jetbrains.kotlin.analysis.low.level.api.fir.stubBased.deserialization.JvmStubBasedDeserializedSymbolProviderFactory
 import org.jetbrains.kotlin.asJava.KotlinAsJavaSupport
 import org.jetbrains.kotlin.asJava.finder.JavaElementFinder
@@ -56,6 +57,9 @@ object FirStandaloneServiceRegistrar : AnalysisApiStandaloneServiceRegistrar {
            registerService(JvmFirDeserializedSymbolProviderFactory::class.java, JvmStubBasedDeserializedSymbolProviderFactory::class.java)
            registerService(KotlinReferenceProviderContributor::class.java, KotlinFirReferenceContributor::class.java)
            registerService(ReadWriteAccessChecker::class.java, ReadWriteAccessCheckerFirImpl::class.java)
+
+            registerService(LLFirSessionInvalidationService::class.java)
+            LLFirSessionInvalidationService.getInstance(project).subscribeToModificationEvents()
        }
    }

@@ -114,12 +114,13 @@ public class StandaloneAnalysisAPISessionBuilder(
    ) {
        val project = kotlinCoreProjectEnvironment.project
        project.apply {
+            registerService(KotlinMessageBusProvider::class.java, KotlinProjectMessageBusProvider::class.java)
+
            FirStandaloneServiceRegistrar.registerProjectServices(project)
            FirStandaloneServiceRegistrar.registerProjectExtensionPoints(project)
            FirStandaloneServiceRegistrar.registerProjectModelServices(project, kotlinCoreProjectEnvironment.parentDisposable)

            registerService(KotlinModificationTrackerFactory::class.java, KotlinStaticModificationTrackerFactory::class.java)
-            registerService(KotlinMessageBusProvider::class.java, KotlinProjectMessageBusProvider::class.java)
            registerService(KotlinGlobalModificationService::class.java, KotlinStaticGlobalModificationService::class.java)

            registerService(KtModuleScopeProvider::class.java, KtModuleScopeProviderImpl())
@@ -1,43 +0,0 @@
-/*
- * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
- * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
- */
-
-package org.jetbrains.kotlin.analysis.utils.trackers
-
-import com.intellij.openapi.util.ModificationTracker
-import java.util.*
-
-public class CompositeModificationTracker private constructor(private val trackers: List<ModificationTracker>) : ModificationTracker {
-    override fun getModificationCount(): Long = trackers.sumOf { it.modificationCount }
-
-    public companion object {
-        public fun create(trackers: List<ModificationTracker>): ModificationTracker {
-            return when (trackers.size) {
-                0 -> ModificationTracker.NEVER_CHANGED
-                1 -> trackers.single()
-                else -> CompositeModificationTracker(trackers)
-            }
-        }
-
-        public fun createFlattened(trackers: List<ModificationTracker>): ModificationTracker {
-            val set = Collections.newSetFromMap(IdentityHashMap<ModificationTracker, Boolean>())
-            val flattened = ArrayList<ModificationTracker>()
-
-            fun flatten(tracker: ModificationTracker) {
-                when (tracker) {
-                    is CompositeModificationTracker -> tracker.trackers.forEach(::flatten)
-                    ModificationTracker.NEVER_CHANGED -> {}
-                    else -> {
-                        if (set.add(tracker)) {
-                            flattened += tracker
-                        }
-                    }
-                }
-            }
-
-            trackers.forEach(::flatten)
-            return create(flattened)
-        }
-    }
-}
@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure

 import com.intellij.openapi.project.Project
-import com.intellij.openapi.util.ModificationTracker
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.search.DelegatingGlobalSearchScope
 import com.intellij.psi.util.CachedValue
@@ -70,7 +69,7 @@ class LLFirBuiltinsSessionFactory(private val project: Project) {
        builtinsAndCloneableSessions.getOrPut(platform) {
            CachedValuesManager.getManager(project).createCachedValue {
                val session = createBuiltinsAndCloneableSession(platform)
-                CachedValueProvider.Result(session, session.modificationTracker)
+                CachedValueProvider.Result(session, session.createValidityTracker())
            }
        }.value

@@ -83,7 +82,7 @@ class LLFirBuiltinsSessionFactory(private val project: Project) {
    private fun createBuiltinsAndCloneableSession(platform: TargetPlatform): LLFirBuiltinsAndCloneableSession {
        val builtinsModule = getBuiltinsModule(platform)

-        val session = LLFirBuiltinsAndCloneableSession(builtinsModule, ModificationTracker.NEVER_CHANGED, builtInTypes)
+        val session = LLFirBuiltinsAndCloneableSession(builtinsModule, builtInTypes)
        val moduleData = LLFirModuleData(builtinsModule).apply { bindSession(session) }

        return session.apply {
@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir.sessions

 import com.intellij.openapi.project.Project
-import com.intellij.openapi.util.ModificationTracker
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.search.ProjectScope
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirGlobalResolveComponents
@@ -22,7 +21,6 @@ import org.jetbrains.kotlin.analysis.providers.createDeclarationProvider
 import org.jetbrains.kotlin.analysis.providers.impl.declarationProviders.FileBasedKotlinDeclarationProvider
 import org.jetbrains.kotlin.analysis.providers.impl.util.mergeInto
 import org.jetbrains.kotlin.analysis.utils.errors.withKtModuleEntry
-import org.jetbrains.kotlin.analysis.utils.trackers.CompositeModificationTracker
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.config.LanguageVersionSettingsImpl
@@ -104,9 +102,8 @@ internal abstract class LLFirAbstractSessionFactory(protected val project: Proje
        val components = LLFirModuleResolveComponents(module, globalResolveComponents, scopeProvider)

        val dependencies = collectSourceModuleDependencies(module)
-        val dependencyTracker = createSourceModuleDependencyTracker(module, dependencies)

-        val session = LLFirScriptSession(module, dependencyTracker, components, builtinsSession.builtinTypes)
+        val session = LLFirScriptSession(module, components, builtinsSession.builtinTypes)
        components.session = session

        val moduleData = createModuleData(session)
@@ -180,14 +177,13 @@ internal abstract class LLFirAbstractSessionFactory(protected val project: Proje
        }
    }

-    fun createNotUnderContentRootResolvableSession(module: KtNotUnderContentRootModule): LLFirNonUnderContentRootResolvableModuleSession {
+    fun createNotUnderContentRootResolvableSession(module: KtNotUnderContentRootModule): LLFirNotUnderContentRootResolvableModuleSession {
        val builtinsSession = LLFirBuiltinsSessionFactory.getInstance(project).getBuiltinsSession(JvmPlatforms.unspecifiedJvmPlatform)
        val languageVersionSettings = LanguageVersionSettingsImpl.DEFAULT
        val scopeProvider = FirKotlinScopeProvider(::wrapScopeWithJvmMapped)
        val components = LLFirModuleResolveComponents(module, globalResolveComponents, scopeProvider)

-        val dependencyTracker = builtinsSession.modificationTracker
-        val session = LLFirNonUnderContentRootResolvableModuleSession(module, dependencyTracker, components, builtinsSession.builtinTypes)
+        val session = LLFirNotUnderContentRootResolvableModuleSession(module, components, builtinsSession.builtinTypes)
        components.session = session

        val moduleData = createModuleData(session)
@@ -259,8 +255,7 @@ internal abstract class LLFirAbstractSessionFactory(protected val project: Proje
        val components = LLFirModuleResolveComponents(module, globalResolveComponents, scopeProvider)

        val dependencies = collectSourceModuleDependencies(module)
-        val dependencyTracker = createSourceModuleDependencyTracker(module, dependencies)
-        val session = LLFirSourcesSession(module, dependencyTracker, components, builtinsSession.builtinTypes)
+        val session = LLFirSourcesSession(module, components, builtinsSession.builtinTypes)
        components.session = session

        val moduleData = createModuleData(session)
@@ -346,9 +341,7 @@ internal abstract class LLFirAbstractSessionFactory(protected val project: Proje
            LLFirSessionCache.getInstance(project).getSession(it)
        }

-        val dependencyTracker = createLibraryDependencyTracker(moduleAnchorSession, builtinsSession)
-        val session =
-            LLFirLibraryOrLibrarySourceResolvableModuleSession(module, dependencyTracker, components, builtinsSession.builtinTypes)
+        val session = LLFirLibraryOrLibrarySourceResolvableModuleSession(module, components, builtinsSession.builtinTypes)
        components.session = session

        val moduleData = createModuleData(session)
@@ -423,8 +416,7 @@ internal abstract class LLFirAbstractSessionFactory(protected val project: Proje
        val platform = module.platform
        val builtinsSession = LLFirBuiltinsSessionFactory.getInstance(project).getBuiltinsSession(platform)

-        val dependencyTracker = ModificationTracker.NEVER_CHANGED
-        val session = LLFirLibrarySession(module, dependencyTracker, builtinsSession.builtinTypes)
+        val session = LLFirLibrarySession(module, builtinsSession.builtinTypes)

        val moduleData = createModuleData(session)

@@ -507,27 +499,6 @@ internal abstract class LLFirAbstractSessionFactory(protected val project: Proje
        return dependencyModules.mapNotNull(::getOrCreateSessionForDependency)
    }

-    private fun createSourceModuleDependencyTracker(module: KtModule, exposedDependencies: List<LLFirSession>): ModificationTracker {
-        val llFirSessionCache = LLFirSessionCache.getInstance(project)
-        val friendDependencies = module.directFriendDependencies
-        val trackers = ArrayList<ModificationTracker>(exposedDependencies.size + friendDependencies.size)
-
-        exposedDependencies.forEach { trackers += it.modificationTracker }
-        friendDependencies.forEach { trackers += llFirSessionCache.getSession(it).modificationTracker }
-
-        return CompositeModificationTracker.createFlattened(trackers)
-    }
-
-    private fun createLibraryDependencyTracker(moduleAnchorSession: LLFirSession?, builtinsSession: LLFirBuiltinsAndCloneableSession): ModificationTracker {
-        return moduleAnchorSession?.let {
-            CompositeModificationTracker.createFlattened(buildList {
-                add(moduleAnchorSession.modificationTracker)
-                add(builtinsSession.modificationTracker)
-            })
-        }
-            ?: builtinsSession.modificationTracker
-    }
-
    private fun createModuleData(session: LLFirSession): LLFirModuleData {
        return LLFirModuleData(session.ktModule).apply { bindSession(session) }
    }
@@ -5,13 +5,11 @@

 package org.jetbrains.kotlin.analysis.low.level.api.fir.sessions

-import com.intellij.openapi.util.ModificationTracker
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.fir.BuiltinTypes
 import org.jetbrains.kotlin.fir.PrivateSessionConstructor

 class LLFirBuiltinsAndCloneableSession @PrivateSessionConstructor constructor(
    ktModule: KtModule,
-    dependencyTracker: ModificationTracker,
    builtinTypes: BuiltinTypes,
-) : LLFirLibraryLikeSession(ktModule, dependencyTracker, builtinTypes)
+) : LLFirLibraryLikeSession(ktModule, builtinTypes)
@@ -5,7 +5,6 @@

 package org.jetbrains.kotlin.analysis.low.level.api.fir.sessions

-import com.intellij.openapi.util.ModificationTracker
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.LLFirScopeSessionProvider
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.fir.BuiltinTypes
@@ -13,13 +12,11 @@ import org.jetbrains.kotlin.fir.resolve.ScopeSession

 abstract class LLFirLibraryLikeSession(
    ktModule: KtModule,
-    dependencyTracker: ModificationTracker,
    builtinTypes: BuiltinTypes,
-) : LLFirSession(ktModule, dependencyTracker, builtinTypes, Kind.Library) {
+) : LLFirSession(ktModule, builtinTypes, Kind.Library) {
    private val scopeSessionProvider = LLFirScopeSessionProvider.create(project, invalidationTrackers = emptyList())

    override fun getScopeSession(): ScopeSession {
        return scopeSessionProvider.getScopeSession()
    }
 }
-
@@ -5,7 +5,6 @@

 package org.jetbrains.kotlin.analysis.low.level.api.fir.sessions

-import com.intellij.openapi.util.ModificationTracker
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveComponents
 import org.jetbrains.kotlin.analysis.project.structure.KtLibraryModule
 import org.jetbrains.kotlin.analysis.project.structure.KtLibrarySourceModule
@@ -14,10 +13,9 @@ import org.jetbrains.kotlin.fir.BuiltinTypes

 internal class LLFirLibraryOrLibrarySourceResolvableModuleSession(
    ktModule: KtModule,
-    dependencyTracker: ModificationTracker,
    override val moduleComponents: LLFirModuleResolveComponents,
    builtinTypes: BuiltinTypes,
-) : LLFirResolvableModuleSession(ktModule, dependencyTracker, builtinTypes) {
+) : LLFirResolvableModuleSession(ktModule, builtinTypes) {
    init {
        checkIsValidKtModule(ktModule)
    }
@@ -5,7 +5,6 @@

 package org.jetbrains.kotlin.analysis.low.level.api.fir.sessions

-import com.intellij.openapi.util.ModificationTracker
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.fir.BuiltinTypes
 import org.jetbrains.kotlin.fir.PrivateSessionConstructor
@@ -15,6 +14,5 @@ import org.jetbrains.kotlin.fir.PrivateSessionConstructor
 */
 internal class LLFirLibrarySession @PrivateSessionConstructor constructor(
    ktModule: KtModule,
-    dependencyTracker: ModificationTracker,
    builtinTypes: BuiltinTypes,
-) : LLFirLibraryLikeSession(ktModule, dependencyTracker, builtinTypes)
+) : LLFirLibraryLikeSession(ktModule, builtinTypes)
@@ -5,15 +5,13 @@

 package org.jetbrains.kotlin.analysis.low.level.api.fir.sessions

-import com.intellij.openapi.util.ModificationTracker
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveComponents
 import org.jetbrains.kotlin.analysis.project.structure.KtNotUnderContentRootModule
 import org.jetbrains.kotlin.fir.BuiltinTypes
 import org.jetbrains.kotlin.fir.PrivateSessionConstructor

-internal class LLFirNonUnderContentRootResolvableModuleSession @PrivateSessionConstructor constructor(
+internal class LLFirNotUnderContentRootResolvableModuleSession @PrivateSessionConstructor constructor(
    ktModule: KtNotUnderContentRootModule,
-    dependencyTracker: ModificationTracker,
    override val moduleComponents: LLFirModuleResolveComponents,
    builtinTypes: BuiltinTypes,
-) : LLFirResolvableModuleSession(ktModule, dependencyTracker, builtinTypes)
+) : LLFirResolvableModuleSession(ktModule, builtinTypes)
@@ -5,7 +5,6 @@

 package org.jetbrains.kotlin.analysis.low.level.api.fir.sessions

-import com.intellij.openapi.util.ModificationTracker
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveComponents
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.fir.BuiltinTypes
@@ -15,9 +14,8 @@ import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol

 abstract class LLFirResolvableModuleSession(
    ktModule: KtModule,
-    dependencyTracker: ModificationTracker,
    builtinTypes: BuiltinTypes
-) : LLFirModuleSession(ktModule, dependencyTracker, builtinTypes, Kind.Source) {
+) : LLFirModuleSession(ktModule, builtinTypes, Kind.Source) {
    internal abstract val moduleComponents: LLFirModuleResolveComponents

    final override fun getScopeSession(): ScopeSession {
@@ -5,7 +5,6 @@

 package org.jetbrains.kotlin.analysis.low.level.api.fir.sessions

-import com.intellij.openapi.util.ModificationTracker
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveComponents
 import org.jetbrains.kotlin.analysis.project.structure.KtScriptModule
 import org.jetbrains.kotlin.fir.BuiltinTypes
@@ -13,7 +12,6 @@ import org.jetbrains.kotlin.fir.PrivateSessionConstructor

 internal class LLFirScriptSession @PrivateSessionConstructor constructor(
    ktModule: KtScriptModule,
-    dependencyTracker: ModificationTracker,
    override val moduleComponents: LLFirModuleResolveComponents,
    builtinTypes: BuiltinTypes
-) : LLFirResolvableModuleSession(ktModule, dependencyTracker, builtinTypes)
+) : LLFirResolvableModuleSession(ktModule, builtinTypes)
@@ -5,103 +5,97 @@

 package org.jetbrains.kotlin.analysis.low.level.api.fir.sessions

+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.application.ModalityState
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.ModificationTracker
-import com.intellij.psi.PsiFile
-import com.intellij.psi.SmartPointerManager
-import org.jetbrains.kotlin.analysis.project.structure.*
-import org.jetbrains.kotlin.analysis.providers.KotlinModificationTrackerFactory
-import org.jetbrains.kotlin.analysis.providers.KtModuleStateTracker
-import org.jetbrains.kotlin.analysis.utils.trackers.CompositeModificationTracker
+import org.jetbrains.kotlin.analysis.low.level.api.fir.transformers.USE_STATE_KEEPER
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.fir.BuiltinTypes
 import org.jetbrains.kotlin.fir.FirElementWithResolveState
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.PrivateSessionConstructor
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
-import org.jetbrains.kotlin.utils.addIfNotNull
 import java.util.concurrent.atomic.AtomicBoolean
+import java.util.concurrent.atomic.AtomicLong

@OptIn(PrivateSessionConstructor::class)
 abstract class LLFirSession(
    val ktModule: KtModule,
-    dependencyTracker: ModificationTracker,
    override val builtinTypes: BuiltinTypes,
    kind: Kind
 ) : FirSession(sessionProvider = null, kind) {
    abstract fun getScopeSession(): ScopeSession

-    private val initialModificationCount: Long
-    private val isExplicitlyInvalidated = AtomicBoolean(false)
-
-    val modificationTracker: ModificationTracker
-
    val project: Project
        get() = ktModule.project

-    init {
-        val trackerFactory = KotlinModificationTrackerFactory.getService(ktModule.project)
-        val validityTracker = trackerFactory.createModuleStateTracker(ktModule)
+    private val _isValid = AtomicBoolean(true)

-        val outOfBlockTracker = when (ktModule) {
-            is KtSourceModule -> trackerFactory.createModuleWithoutDependenciesOutOfBlockModificationTracker(ktModule)
-            is KtNotUnderContentRootModule -> ktModule.file?.let(::FileModificationTracker)
-            is KtScriptModule -> FileModificationTracker(ktModule.file)
-            is KtScriptDependencyModule -> ktModule.file?.let(::FileModificationTracker)
-            else -> null
+    /**
+     * Whether the [LLFirSession] is valid. The session should not be used if it is invalid.
+     *
+     * [isValid] should be set to `false` at the same time as the session is removed from [LLFirSessionCache]. Hence, [isValid] should be
+     * managed by [LLFirSessionCache].
+     */
+    var isValid: Boolean
+        get() = _isValid.get()
+        internal set(value) {
+            check(!value) { "An invalid LL FIR session cannot become valid again." }
+            _isValid.set(value)
        }

-        modificationTracker = CompositeModificationTracker.createFlattened(
-            buildList {
-                add(ExplicitInvalidationTracker(ktModule, isExplicitlyInvalidated))
-                add(ModuleStateModificationTracker(ktModule, validityTracker))
-                addIfNotNull(outOfBlockTracker)
-                add(dependencyTracker)
-            }
-        )
-
-        initialModificationCount = modificationTracker.modificationCount
-    }
-
-    private class ModuleStateModificationTracker(val module: KtModule, val tracker: KtModuleStateTracker) : ModificationTracker {
-        override fun getModificationCount(): Long = tracker.rootModificationCount
-        override fun toString(): String = "Module state tracker for module '${module.moduleDescription}'"
-    }
-
-    private class ExplicitInvalidationTracker(val module: KtModule, val isExplicitlyInvalidated: AtomicBoolean) : ModificationTracker {
-        override fun getModificationCount(): Long = if (isExplicitlyInvalidated.get()) 1 else 0
-        override fun toString(): String = "Explicit invalidation tracker for module '${module.moduleDescription}'"
-    }
-
-    private class FileModificationTracker(file: PsiFile) : ModificationTracker {
-        private val pointer = SmartPointerManager.getInstance(file.project).createSmartPsiElementPointer(file)
-
-        override fun getModificationCount(): Long {
-            val file = pointer.element ?: return Long.MAX_VALUE
-            return file.modificationStamp
-        }
-
-        override fun toString(): String {
-            val file = pointer.element ?: return "File tracker for a collected file"
-            val virtualFile = file.virtualFile ?: return "File tracker for a non-physical file '${file.name}'"
-            return "File tracker for path '${virtualFile.path}'"
-        }
-    }
-
    fun invalidate() {
-        isExplicitlyInvalidated.set(true)
+        val application = ApplicationManager.getApplication()
+        if (application.isWriteAccessAllowed) {
+            invalidateInWriteAction()
+        } else {
+            // We have to invalidate the session on the EDT per the contract of `LLFirSessionInvalidationService`. The timing here is not
+            // 100% waterproof, but `LLFirSession.invalidate` is only a workaround for when FIR guards consistency protection (see KT-56503)
+            // is turned off. The check restricts usage of `invalidate` to this scenario.
+            check(!USE_STATE_KEEPER) {
+                "Outside a write action, a session may only be invalidated directly when FIR guards are turned off."
+            }
+
+            application.invokeLater(
+                { application.runWriteAction { invalidateInWriteAction() } },
+
+                // `ModalityState.any()` can be used because session invalidation does not modify PSI, VFS, or the project model.
+                ModalityState.any(),
+            )
+        }
    }

-    val isValid: Boolean
-        get() = modificationTracker.modificationCount == initialModificationCount
+    private fun invalidateInWriteAction() {
+        LLFirSessionInvalidationService.getInstance(project).invalidate(ktModule)
+    }
+
+    /**
+     * Creates a [ModificationTracker] which tracks the validity of this session via [isValid].
+     */
+    fun createValidityTracker(): ModificationTracker = ValidityModificationTracker()
+
+    private inner class ValidityModificationTracker : ModificationTracker {
+        private var count = AtomicLong()
+
+        override fun getModificationCount(): Long {
+            if (isValid) return 0
+
+            // When the session is invalid, we cannot simply return a static modification count of 1. For example, consider situations where
+            // a cached value was created with an already invalid session (so it remembers the modification count of 1). Then, if we return
+            // a static modification count of 1, the modification count never changes and the cached value misses that the session has been
+            // invalidated. Hence, `count` is incremented on each modification count access.
+            return count.incrementAndGet()
+        }
+    }
 }

 abstract class LLFirModuleSession(
    ktModule: KtModule,
-    dependencyTracker: ModificationTracker,
    builtinTypes: BuiltinTypes,
    kind: Kind
-) : LLFirSession(ktModule, dependencyTracker, builtinTypes, kind)
+) : LLFirSession(ktModule, builtinTypes, kind)

 val FirElementWithResolveState.llFirSession: LLFirSession
    get() = moduleData.session as LLFirSession
@@ -5,10 +5,8 @@

 package org.jetbrains.kotlin.analysis.low.level.api.fir.sessions

+import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.project.Project
-import com.intellij.psi.util.CachedValue
-import com.intellij.psi.util.CachedValueProvider
-import com.intellij.psi.util.CachedValuesManager
 import com.intellij.util.containers.CollectionFactory
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirInternals
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkCanceled
@@ -25,6 +23,8 @@ import org.jetbrains.kotlin.platform.konan.NativePlatform
 import org.jetbrains.kotlin.resolve.jvm.platform.JvmPlatformAnalyzerServices
 import java.util.concurrent.ConcurrentMap

+private typealias SessionStorage = ConcurrentMap<KtModule, LLFirSession>
+
@LLFirInternals
 class LLFirSessionCache(private val project: Project) {
    companion object {
@@ -33,8 +33,8 @@ class LLFirSessionCache(private val project: Project) {
        }
    }

-    private val sourceCache: ConcurrentMap<KtModule, CachedValue<LLFirSession>> = CollectionFactory.createConcurrentSoftValueMap()
-    private val binaryCache: ConcurrentMap<KtModule, CachedValue<LLFirSession>> = CollectionFactory.createConcurrentSoftValueMap()
+    private val sourceCache: SessionStorage = CollectionFactory.createConcurrentSoftValueMap()
+    private val binaryCache: SessionStorage = CollectionFactory.createConcurrentSoftValueMap()

    /**
     * Returns the existing session if found, or creates a new session and caches it.
@@ -60,17 +60,54 @@ class LLFirSessionCache(private val project: Project) {

    private fun <T : KtModule> getCachedSession(
        module: T,
-        storage: ConcurrentMap<KtModule, CachedValue<LLFirSession>>,
+        storage: SessionStorage,
        factory: (T) -> LLFirSession
    ): LLFirSession {
        checkCanceled()

-        return storage.computeIfAbsent(module) {
-            CachedValuesManager.getManager(project).createCachedValue {
-                val session = factory(module)
-                CachedValueProvider.Result(session, session.modificationTracker)
-            }
-        }.value
+        return storage.computeIfAbsent(module) { factory(module) }.also { session ->
+            require(session.isValid) { "A session acquired via `getSession` should always be valid. Module: $module" }
+        }
+    }
+
+    /**
+     * Removes the session(s) associated with [module] after it has been invalidated.
+     *
+     * [removeSession] must be called in a write action.
+     */
+    fun removeSession(module: KtModule) {
+        ApplicationManager.getApplication().assertWriteAccessAllowed()
+
+        removeSessionFrom(module, sourceCache)
+        if (module is KtBinaryModule) {
+            removeSessionFrom(module, binaryCache)
+        }
+    }
+
+    private fun removeSessionFrom(module: KtModule, storage: SessionStorage) {
+        val session = storage.remove(module)
+        if (session != null) {
+            session.isValid = false
+        }
+    }
+
+    /**
+     * Removes all sessions after global invalidation. If [includeBinarySessions] is `false`, only source sessions will be removed.
+     *
+     * [removeAllSessions] must be called in a write action.
+     */
+    fun removeAllSessions(includeBinarySessions: Boolean) {
+        ApplicationManager.getApplication().assertWriteAccessAllowed()
+
+        removeAllSessionsFrom(sourceCache)
+        if (includeBinarySessions) removeAllSessionsFrom(binaryCache)
+    }
+
+    private fun removeAllSessionsFrom(storage: SessionStorage) {
+        // Because `removeAllSessionsFrom` is executed in a write action, the order of setting `isValid` and clearing `storage` is not
+        // important.
+        storage.values.forEach { it.isValid = false }
+        storage.clear()
    }

    private fun createSession(module: KtModule): LLFirSession {
@@ -0,0 +1,92 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.analysis.low.level.api.fir.sessions
+
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.analysis.project.structure.KotlinModuleDependentsProvider
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
+import org.jetbrains.kotlin.analysis.providers.analysisMessageBus
+import org.jetbrains.kotlin.analysis.providers.topics.KotlinGlobalModuleStateModificationListener
+import org.jetbrains.kotlin.analysis.providers.topics.KotlinGlobalOutOfBlockModificationListener
+import org.jetbrains.kotlin.analysis.providers.topics.KotlinGlobalSourceModuleStateModificationListener
+import org.jetbrains.kotlin.analysis.providers.topics.KotlinGlobalSourceOutOfBlockModificationListener
+import org.jetbrains.kotlin.analysis.providers.topics.KotlinTopics
+import org.jetbrains.kotlin.analysis.providers.topics.KotlinModuleOutOfBlockModificationListener
+import org.jetbrains.kotlin.analysis.providers.topics.KotlinModuleStateModificationListener
+
+/**
+ * [LLFirSessionInvalidationService] listens to [modification events][KotlinTopics] and invalidates [LLFirSession]s which depend on the
+ * modified [KtModule]. Its invalidation functions should always be invoked in a **write action** because invalidation affects multiple
+ * sessions in [LLFirSessionCache] and the cache has to be kept consistent.
+ */
+class LLFirSessionInvalidationService(private val project: Project) : Disposable {
+    /**
+     * Subscribes to all [modification events][KotlinTopics] via the [analysisMessageBus].
+     *
+     * [subscribeToModificationEvents] must be invoked during setup to allow [LLFirSessionInvalidationService] to listen to events.
+     * Subscribing in `init` is not an option because services are created on demand and there is no guarantee that this service is going to
+     * be requested.
+     */
+    fun subscribeToModificationEvents() {
+        val busConnection = project.analysisMessageBus.connect(this)
+
+        // All modification events the invalidation service subscribes to are guaranteed to be published in a write action. This ensures
+        // that invalidation functions are only called in a write action, per the contract of `LLFirSessionInvalidationService`.
+        busConnection.subscribe(
+            KotlinTopics.MODULE_STATE_MODIFICATION,
+            KotlinModuleStateModificationListener { module, _ -> invalidate(module) },
+        )
+        busConnection.subscribe(
+            KotlinTopics.MODULE_OUT_OF_BLOCK_MODIFICATION,
+            KotlinModuleOutOfBlockModificationListener { module -> invalidate(module) },
+        )
+        busConnection.subscribe(
+            KotlinTopics.GLOBAL_MODULE_STATE_MODIFICATION,
+            KotlinGlobalModuleStateModificationListener { invalidateAll(includeBinaryModules = true) }
+        )
+        busConnection.subscribe(
+            KotlinTopics.GLOBAL_OUT_OF_BLOCK_MODIFICATION,
+            KotlinGlobalOutOfBlockModificationListener { invalidateAll(includeBinaryModules = true) }
+        )
+        busConnection.subscribe(
+            KotlinTopics.GLOBAL_SOURCE_MODULE_STATE_MODIFICATION,
+            KotlinGlobalSourceModuleStateModificationListener { invalidateAll(includeBinaryModules = false) },
+        )
+        busConnection.subscribe(
+            KotlinTopics.GLOBAL_SOURCE_OUT_OF_BLOCK_MODIFICATION,
+            KotlinGlobalSourceOutOfBlockModificationListener { invalidateAll(includeBinaryModules = false) },
+        )
+    }
+
+    /**
+     * Invalidates the session(s) associated with [module].
+     *
+     * Per the contract of [LLFirSessionInvalidationService], [invalidate] may only be called from a write action.
+     */
+    fun invalidate(module: KtModule) {
+        ApplicationManager.getApplication().assertWriteAccessAllowed()
+
+        val sessionCache = LLFirSessionCache.getInstance(project)
+        sessionCache.removeSession(module)
+        KotlinModuleDependentsProvider.getInstance(project).getTransitiveDependents(module).forEach(sessionCache::removeSession)
+    }
+
+    private fun invalidateAll(includeBinaryModules: Boolean) {
+        ApplicationManager.getApplication().assertWriteAccessAllowed()
+
+        LLFirSessionCache.getInstance(project).removeAllSessions(includeBinaryModules)
+    }
+
+    override fun dispose() {
+    }
+
+    companion object {
+        fun getInstance(project: Project): LLFirSessionInvalidationService =
+            project.getService(LLFirSessionInvalidationService::class.java)
+    }
+}
@@ -5,7 +5,6 @@

 package org.jetbrains.kotlin.analysis.low.level.api.fir.sessions

-import com.intellij.openapi.util.ModificationTracker
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveComponents
 import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 import org.jetbrains.kotlin.fir.BuiltinTypes
@@ -13,7 +12,6 @@ import org.jetbrains.kotlin.fir.PrivateSessionConstructor

 internal class LLFirSourcesSession @PrivateSessionConstructor constructor(
    ktModule: KtSourceModule,
-    dependencyTracker: ModificationTracker,
    override val moduleComponents: LLFirModuleResolveComponents,
    builtinTypes: BuiltinTypes,
-) : LLFirResolvableModuleSession(ktModule, dependencyTracker, builtinTypes)
+) : LLFirResolvableModuleSession(ktModule, builtinTypes)
@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir.state

 import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.ModificationTracker
 import com.intellij.psi.util.CachedValueProvider
 import com.intellij.psi.util.CachedValuesManager
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveComponents
@@ -51,7 +52,7 @@ internal abstract class LLFirResolvableResolveSession(

    private val useSiteFirSessionCached = CachedValuesManager.getManager(project).createCachedValue {
        val session = useSiteSessionFactory(useSiteKtModule)
-        CachedValueProvider.Result.create(session, session.modificationTracker)
+        CachedValueProvider.Result.create(session, session.createValidityTracker())
    }

    final override val useSiteFirSession: LLFirSession
@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir.util

 import com.intellij.openapi.progress.ProgressManager
-import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirResolvableModuleSession
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.diagnostics.FirDiagnosticHolder
@@ -19,15 +18,6 @@ import org.jetbrains.kotlin.psi.psiUtil.isObjectLiteral
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.locks.Lock

-internal fun <R : Any> withInvalidationOnException(session: LLFirResolvableModuleSession, action: () -> R): R {
-    try {
-        return action()
-    } catch (e: Throwable) {
-        session.invalidate()
-        throw e
-    }
-}
-
 internal inline fun <T> Lock.lockWithPCECheck(lockingIntervalMs: Long, action: () -> T): T {
    while (true) {
        checkCanceled()