From edadfd0daa414ebf8afa5a5f495217cded80c809 Mon Sep 17 00:00:00 2001
From: Marco Pennekamp <marco.pennekamp@jetbrains.com>
Date: Mon, 3 Jul 2023 17:12:16 +0200
Subject: [PATCH] [LL FIR] KT-58257 Fix session project leak in `invokeLater`
 lambda

- The lambda passed to `invokeLater` in `LLFirSession.invalidate` might
  survive beyond project disposal (especially in tests), so me must not
  capture a hard reference to the session that can leak its `project`.
---
 .../level/api/fir/sessions/LLFirSession.kt    | 26 ++++++++++++++-----
 1 file changed, 19 insertions(+), 7 deletions(-)

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSession.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSession.kt
index 4f661dd3a39..dfe2a92a7a6 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSession.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSession.kt
@@ -17,6 +17,7 @@ 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 java.lang.ref.WeakReference
 import java.util.concurrent.atomic.AtomicBoolean
 import java.util.concurrent.atomic.AtomicLong
 
@@ -58,8 +59,12 @@ abstract class LLFirSession(
                 "Outside a write action, a session may only be invalidated directly when FIR guards are turned off."
             }
 
+            // The lambda passed to `invokeLater` might survive beyond project disposal (especially in tests), so me must not capture a hard
+            // reference to `this` that can leak `project`. A weak reference is appropriate because we do not need to invalidate the session
+            // when it has already been collected.
+            val weakSession = WeakReference(this)
             application.invokeLater(
-                { application.runWriteAction { invalidateInWriteAction() } },
+                { weakSession.invalidateIfAlive() },
 
                 // `ModalityState.any()` can be used because session invalidation does not modify PSI, VFS, or the project model.
                 ModalityState.any(),
@@ -67,12 +72,6 @@ abstract class LLFirSession(
         }
     }
 
-    private fun invalidateInWriteAction() {
-        if (!isValid) return
-
-        LLFirSessionInvalidationService.getInstance(project).invalidate(ktModule)
-    }
-
     /**
      * Creates a [ModificationTracker] which tracks the validity of this session via [isValid].
      */
@@ -93,6 +92,19 @@ abstract class LLFirSession(
     }
 }
 
+private fun LLFirSession.invalidateInWriteAction() {
+    if (!isValid) return
+
+    LLFirSessionInvalidationService.getInstance(project).invalidate(ktModule)
+}
+
+private fun WeakReference<LLFirSession>.invalidateIfAlive() = ApplicationManager.getApplication().runWriteAction {
+    val session = get() ?: return@runWriteAction
+    if (session.project.isDisposed) return@runWriteAction
+
+    session.invalidateInWriteAction()
+}
+
 abstract class LLFirModuleSession(
     ktModule: KtModule,
     builtinTypes: BuiltinTypes,