isolate static members into thread local

google · Nov 1, 2023 · 6bf8193 · 6bf8193
1 parent 85c5005
commit 6bf8193
Show file tree

Hide file tree

Showing 5 changed files with 51 additions and 17 deletions.
diff --git a/common-util/src/main/kotlin/com/google/devtools/ksp/KSObjectCacheManager.kt b/common-util/src/main/kotlin/com/google/devtools/ksp/KSObjectCacheManager.kt
@@ -19,19 +19,30 @@ package com.google.devtools.ksp
 
 class KSObjectCacheManager {
     companion object {
-        val caches = arrayListOf<KSObjectCache<*, *>>()
+        private val caches_prop = object : ThreadLocal<ArrayList<KSObjectCache<*, *>>>() {
+            override fun initialValue(): ArrayList<KSObjectCache<*, *>> {
+                return ArrayList()
+            }
+        }
+        val caches
+            get() = caches_prop.get()
 
         fun register(cache: KSObjectCache<*, *>) = caches.add(cache)
         fun clear() = caches.forEach { it.clear() }
     }
 }
 
 abstract class KSObjectCache<K, V> {
-    val cache = mutableMapOf<K, V>()
+    private val cache_prop = ThreadLocal<MutableMap<K, V>>()
 
-    init {
-        KSObjectCacheManager.register(this)
-    }
+    val cache: MutableMap<K, V>
+        get() {
+            if (cache_prop.get() == null) {
+                KSObjectCacheManager.register(this)
+                cache_prop.set(mutableMapOf())
+            }
+            return cache_prop.get()
+        }
 
     open fun clear() = cache.clear()
 }
diff --git a/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspAATask.kt b/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspAATask.kt
@@ -267,14 +267,22 @@ interface KspAAWorkParameter : WorkParameters {
     var kspClassPath: ConfigurableFileCollection
 }
 
+var isolatedClassLoaderCache = mutableMapOf<String, URLClassLoader>()
+
 abstract class KspAAWorkerAction : WorkAction<KspAAWorkParameter> {
     override fun execute() {
         val gradleCfg = parameters.config
         val kspClassPath = parameters.kspClassPath
-        val isolatedClassLoader = URLClassLoader(
-            kspClassPath.files.map { it.toURI().toURL() }.toTypedArray(),
-            ClassLoader.getPlatformClassLoader()
-        )
+        val key = kspClassPath.files.map { it.toURI().toURL() }.joinToString { it.path }
+        synchronized(isolatedClassLoaderCache) {
+            if (isolatedClassLoaderCache[key] == null) {
+                isolatedClassLoaderCache[key] = URLClassLoader(
+                    kspClassPath.files.map { it.toURI().toURL() }.toTypedArray(),
+                    ClassLoader.getPlatformClassLoader()
+                )
+            }
+        }
+        val isolatedClassLoader = isolatedClassLoaderCache[key]!!
 
         // Clean stale files for now.
         // TODO: support incremental processing.

diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KSPCoreEnvironment.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KSPCoreEnvironment.kt
@@ -23,7 +23,13 @@ import com.intellij.psi.PsiDocumentManager
 class KSPCoreEnvironment(internal val project: MockProject) {
     companion object {
         // TODO: get rid of singleton.
-        lateinit var instance: KSPCoreEnvironment
+        val instance_prop: ThreadLocal<KSPCoreEnvironment> = ThreadLocal()
+
+        var instance: KSPCoreEnvironment
+            get() = instance_prop.get()
+            set(value) {
+                instance_prop.set(value)
+            }
     }
     init {
         instance = this

diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KotlinSymbolProcessing.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KotlinSymbolProcessing.kt
@@ -478,8 +478,8 @@ class KotlinSymbolProcessing(
                 project
             )
             ResolverAAImpl.instance = resolver
-            ResolverAAImpl.functionAsMemberOfCache = mutableMapOf()
-            ResolverAAImpl.propertyAsMemberOfCache = mutableMapOf()
+            ResolverAAImpl.instance.functionAsMemberOfCache = mutableMapOf()
+            ResolverAAImpl.instance.propertyAsMemberOfCache = mutableMapOf()
 
             processors.forEach {
                 deferredSymbols[it] =

diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/ResolverAAImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/ResolverAAImpl.kt
@@ -65,12 +65,21 @@ class ResolverAAImpl(
     val project: Project
 ) : Resolver {
     companion object {
-        lateinit var instance: ResolverAAImpl
-        lateinit var ktModule: KtModule
-        lateinit var propertyAsMemberOfCache: MutableMap<Pair<KSPropertyDeclaration, KSType>, KSType>
-        lateinit var functionAsMemberOfCache: MutableMap<Pair<KSFunctionDeclaration, KSType>, KSFunction>
+        val instance_prop: ThreadLocal<ResolverAAImpl> = ThreadLocal()
+        private val ktModule_prop: ThreadLocal<KtModule> = ThreadLocal()
+        var instance
+            get() = instance_prop.get()
+            set(value) {
+                instance_prop.set(value)
+            }
+        var ktModule: KtModule
+            get() = ktModule_prop.get()
+            set(value) {
+                ktModule_prop.set(value)
+            }
     }
-
+    lateinit var propertyAsMemberOfCache: MutableMap<Pair<KSPropertyDeclaration, KSType>, KSType>
+    lateinit var functionAsMemberOfCache: MutableMap<Pair<KSFunctionDeclaration, KSType>, KSFunction>
     val javaFileManager = project.getService(JavaFileManager::class.java) as KotlinCliJavaFileManagerImpl
     val classBinaryCache = ClsKotlinBinaryClassCache()