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()