Move ksp task cache to a gradle build service

Author: hfmehmed

gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspAAService.kt

@@ -0,0 +1,28 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.devtools.ksp.gradle
+
+import org.gradle.api.services.BuildService
+import org.gradle.api.services.BuildServiceParameters
+import java.net.URLClassLoader
+import java.util.concurrent.ConcurrentHashMap
+
+abstract class KspAAService : BuildService<BuildServiceParameters.None>, AutoCloseable {
+    val isolatedClassLoaderCache = ConcurrentHashMap<String, URLClassLoader>()
+
+    override fun close() = isolatedClassLoaderCache.clear()
+}

gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspAATask.kt

@@ -39,6 +39,7 @@ import org.gradle.api.provider.MapProperty
 import org.gradle.api.provider.Property
 import org.gradle.api.provider.Provider
 import org.gradle.api.provider.SetProperty
+import org.gradle.api.services.ServiceReference
 import org.gradle.api.tasks.CacheableTask
 import org.gradle.api.tasks.Classpath
 import org.gradle.api.tasks.IgnoreEmptyDirectories
@@ -95,6 +96,9 @@ abstract class KspAATask @Inject constructor(
     @get:Nested
     abstract val commandLineArgumentProviders: ListProperty<CommandLineArgumentProvider>
 
+    @get:ServiceReference("KspAAService")
+    abstract val kspAAService: Property<KspAAService>
+
     @TaskAction
     fun execute(inputChanges: InputChanges) {
         // FIXME: Create a class loader with clean classpath instead of shadowing existing ones. It'll require either:
@@ -155,6 +159,7 @@ abstract class KspAATask @Inject constructor(
             it.removedSources = removedSources
             it.isInputChangeIncremental = inputChanges.isIncremental
             it.changedClasses = changedClasses
+            it.kspAAService.set(kspAAService)
         }
     }
 
@@ -167,6 +172,11 @@ abstract class KspAATask @Inject constructor(
             kspExtension: KspExtension,
         ): TaskProvider<KspAATask> {
             val project = kotlinCompilation.target.project
+            val kspAAServiceProvider = project.gradle.sharedServices.registerIfAbsent(
+                "KspAAService",
+                KspAAService::class.java
+            ) {}
+
             val target = kotlinCompilation.target.name
             val sourceSetName = kotlinCompilation.defaultSourceSet.name
             val kspTaskName = kotlinCompileProvider.name.replaceFirst("compile", "ksp")
@@ -190,6 +200,7 @@ abstract class KspAATask @Inject constructor(
                 kspAATask.onlyIf {
                     !incomingProcessors.isEmpty
                 }
+                kspAATask.kspAAService.set(kspAAServiceProvider)
                 kspAATask.kspClasspath.from(kspAADepCfg.incoming.artifactView { }.files)
                 kspAATask.kspConfig.let { cfg ->
                     cfg.processorClasspath.from(incomingProcessors)
@@ -539,25 +550,23 @@ interface KspAAWorkParameter : WorkParameters {
     var removedSources: List<File>
     var changedClasses: List<String>
     var isInputChangeIncremental: Boolean
+    var kspAAService: Property<KspAAService>
 }
 
-var isolatedClassLoaderCache = mutableMapOf<String, URLClassLoader>()
 val doNotGC = mutableSetOf<Any>()
 
 abstract class KspAAWorkerAction : WorkAction<KspAAWorkParameter> {
     override fun execute() {
         val gradleCfg = parameters.config
         val kspClasspath = parameters.kspClasspath
+        val isolatedClassLoaderCache = parameters.kspAAService.get().isolatedClassLoaderCache
         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.computeIfAbsent(key) {
+            URLClassLoader(
+                kspClasspath.files.map { it.toURI().toURL() }.toTypedArray(),
+                ClassLoader.getPlatformClassLoader()
+            )
         }
-        val isolatedClassLoader = isolatedClassLoaderCache[key]!!
 
         // Clean stale files for now.
         // TODO: support incremental processing.

gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspSubplugin.kt

@@ -195,7 +195,7 @@ class KspGradleSubplugin @Inject internal constructor(private val registry: Tool
             kotlinCompilation,
             kotlinCompileProvider,
             processorClasspath,
-            kspExtension
+            kspExtension,
         )
 
         val generatedSources = arrayOf(