From c2de8712f2ecbd442f98aefa674866773a374414 Mon Sep 17 00:00:00 2001
From: Victor Petukhov <victor.petukhov@jetbrains.com>
Date: Tue, 25 May 2021 20:55:13 +0300
Subject: [PATCH] Don't store strong references on thread local cache and
 clean-up request caches properly

^KT-46744 Fixed
---
 .../load/kotlin/KotlinBinaryClassCache.kt     | 22 ++++++++++---------
 1 file changed, 12 insertions(+), 10 deletions(-)

diff --git a/compiler/frontend.java/src/org/jetbrains/kotlin/load/kotlin/KotlinBinaryClassCache.kt b/compiler/frontend.java/src/org/jetbrains/kotlin/load/kotlin/KotlinBinaryClassCache.kt
index 49ff93c8e88..c43b2826ad0 100644
--- a/compiler/frontend.java/src/org/jetbrains/kotlin/load/kotlin/KotlinBinaryClassCache.kt
+++ b/compiler/frontend.java/src/org/jetbrains/kotlin/load/kotlin/KotlinBinaryClassCache.kt
@@ -23,9 +23,12 @@ import com.intellij.openapi.components.ServiceManager
 import com.intellij.openapi.util.Computable
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.PsiJavaModule
+import java.lang.ref.WeakReference
 import java.util.concurrent.CopyOnWriteArrayList
 
 class KotlinBinaryClassCache : Disposable {
+    private val requestCaches = CopyOnWriteArrayList<WeakReference<RequestCache>>()
+
     private class RequestCache {
         var virtualFile: VirtualFile? = null
         var modificationStamp: Long = 0
@@ -44,22 +47,21 @@ class KotlinBinaryClassCache : Disposable {
     }
 
     private val cache = object : ThreadLocal<RequestCache>() {
-        private val requestCaches = CopyOnWriteArrayList<RequestCache>()
-
         override fun initialValue(): RequestCache {
-            return RequestCache().also { requestCaches.add(it) }
-        }
-
-        override fun remove() {
-            for (cache in requestCaches) {
-                cache.result = null
-                cache.virtualFile = null
+            return RequestCache().also {
+                requestCaches.add(WeakReference(it))
             }
-            super.remove()
         }
     }
 
     override fun dispose() {
+        for (cache in requestCaches) {
+            cache.get()?.run {
+                result = null
+                virtualFile = null
+            }
+        }
+        requestCaches.clear()
         // This is only relevant for tests. We create a new instance of Application for each test, and so a new instance of this service is
         // also created for each test. However all tests share the same event dispatch thread, which would collect all instances of this
         // thread-local if they're not removed properly. Each instance would transitively retain VFS resulting in OutOfMemoryError