From 6bf8193a536d6570dc113d9ee5f4a308ffb69fb4 Mon Sep 17 00:00:00 2001 From: Jiaxiang Chen Date: Thu, 26 Oct 2023 15:28:15 -0700 Subject: [PATCH] isolate static members into thread local --- .../devtools/ksp/KSObjectCacheManager.kt | 21 ++++++++++++++----- .../google/devtools/ksp/gradle/KspAATask.kt | 16 ++++++++++---- .../devtools/ksp/impl/KSPCoreEnvironment.kt | 8 ++++++- .../ksp/impl/KotlinSymbolProcessing.kt | 4 ++-- .../devtools/ksp/impl/ResolverAAImpl.kt | 19 ++++++++++++----- 5 files changed, 51 insertions(+), 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 index 08a8bd7fbc..7de1337911 100644 --- 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,7 +19,13 @@ package com.google.devtools.ksp class KSObjectCacheManager { companion object { - val caches = arrayListOf>() + private val caches_prop = object : ThreadLocal>>() { + override fun initialValue(): ArrayList> { + return ArrayList() + } + } + val caches + get() = caches_prop.get() fun register(cache: KSObjectCache<*, *>) = caches.add(cache) fun clear() = caches.forEach { it.clear() } @@ -27,11 +33,16 @@ class KSObjectCacheManager { } abstract class KSObjectCache { - val cache = mutableMapOf() + private val cache_prop = ThreadLocal>() - init { - KSObjectCacheManager.register(this) - } + val cache: MutableMap + 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 index 8525b3bfd7..90d712a733 100644 --- 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() + abstract class KspAAWorkerAction : WorkAction { 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 index 1a4387dedd..ccff832b12 100644 --- 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 = 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 index 87bb2a58f0..98bd705662 100644 --- 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 index b3b497bea5..e609a5a4c6 100644 --- 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, KSType> - lateinit var functionAsMemberOfCache: MutableMap, KSFunction> + val instance_prop: ThreadLocal = ThreadLocal() + private val ktModule_prop: ThreadLocal = 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, KSType> + lateinit var functionAsMemberOfCache: MutableMap, KSFunction> val javaFileManager = project.getService(JavaFileManager::class.java) as KotlinCliJavaFileManagerImpl val classBinaryCache = ClsKotlinBinaryClassCache()