[KAPT] Cache annotation processors' classloaders

Adds experimental feature to cache annotation processors' classloaders
It should increase speed of KAPT for consecutive runs
Works for gradle only (it stores cl in gradle daemon)
#KT-28901 Fixed

gradle/verification-metadata.xml

@@ -1510,6 +1510,12 @@
             <sha256 value="8d175561619289a527573cfbdcf872eb5c1216cf9ab692a0ef3dba6d793f63b4" origin="Generated by Gradle"/>
          </artifact>
       </component>
+      <component group="com.google.errorprone" name="error_prone_annotations" version="2.3.4">
+         <artifact name="error_prone_annotations-2.3.4.jar">
+            <md5 value="67beeee58df00366100061c7da82f4c2"/>
+            <sha256 value="baf7d6ea97ce606c53e11b6854ba5f2ce7ef5c24dddf0afa18d1260bd25b002c"/>
+         </artifact>
+      </component>
       <component group="com.google.errorprone" name="error_prone_parent" version="2.0.18">
          <artifact name="error_prone_parent-2.0.18.pom">
             <md5 value="0456a9fe3a1372c1e924798101ee09ec" origin="Generated by Gradle"/>
@@ -1656,6 +1662,12 @@
             <sha256 value="74af0e8dd935894ce9bcb8e2842dd0001581b84cb6a7751c3ef9314011399ba4" origin="Generated by Gradle"/>
          </artifact>
       </component>
+      <component group="com.google.guava" name="guava" version="29.0-jre">
+         <artifact name="guava-29.0-jre.jar">
+            <md5 value="c87962fe0522b105531f37fe351b3b0c"/>
+            <sha256 value="b22c5fb66d61e7b9522531d04b2f915b5158e80aa0b40ee7282c8bfb07b0da25"/>
+         </artifact>
+      </component>
       <component group="com.google.guava" name="guava-parent" version="13.0.1">
          <artifact name="guava-parent-13.0.1.pom">
             <md5 value="c78a1aecd4642f3f6b3c00b09398b0c8" origin="Generated by Gradle"/>
@@ -8564,6 +8576,12 @@
             <sha256 value="deb10353dd2b1bee0e548b74014d80a1179a8faaf8d0d886c01b4d6ca9300068" origin="Generated by Gradle"/>
          </artifact>
       </component>
+      <component group="org.checkerframework" name="checker-qual" version="2.11.1">
+         <artifact name="checker-qual-2.11.1.jar">
+            <md5 value="936884d7b04ee8f064df137bcdfe470d"/>
+            <sha256 value="015224a4b1dc6de6da053273d4da7d39cfea20e63038169fc45ac0d1dc9c5938"/>
+         </artifact>
+      </component>
       <component group="org.checkerframework" name="checker-qual" version="2.5.2">
          <artifact name="checker-qual-2.5.2.jar">
             <md5 value="04acc78b24bbd365423da357da003cf0" origin="Generated by Gradle"/>

libraries/tools/kotlin-gradle-plugin-integration-tests/src/test/kotlin/org/jetbrains/kotlin/gradle/BaseGradleIT.kt

@@ -258,7 +258,8 @@ abstract class BaseGradleIT {
         val verbose: Boolean,
         val useWorkers: Boolean,
         val incrementalKapt: Boolean = false,
-        val includeCompileClasspath: Boolean = true
+        val includeCompileClasspath: Boolean = true,
+        val classLoadersCacheSize: Int? = null
     )
 
     open inner class Project(
@@ -913,6 +914,9 @@ Finished executing task ':$taskName'|
                 add("-Pkapt.use.worker.api=${kaptOptions.useWorkers}")
                 add("-Pkapt.incremental.apt=${kaptOptions.incrementalKapt}")
                 add("-Pkapt.include.compile.classpath=${kaptOptions.includeCompileClasspath}")
+                kaptOptions.classLoadersCacheSize?.also { cacheSize ->
+                    add("-Pkapt.classloaders.cache.size=$cacheSize")
+                }
             }
 
             options.parallelTasksInProject?.let {

libraries/tools/kotlin-gradle-plugin-integration-tests/src/test/kotlin/org/jetbrains/kotlin/gradle/Kapt3IT.kt

@@ -21,11 +21,14 @@ import org.jetbrains.kotlin.gradle.tasks.USING_JVM_INCREMENTAL_COMPILATION_MESSA
 import org.jetbrains.kotlin.gradle.util.*
 import org.junit.Assert
 import org.junit.Assume
+import org.junit.Ignore
 import org.junit.Test
 import java.io.File
+import java.util.regex.Pattern
 import java.util.zip.ZipFile
 import kotlin.test.assertNotNull
 import kotlin.test.assertNull
+import kotlin.test.assertTrue
 
 abstract class Kapt3BaseIT : BaseGradleIT() {
     companion object {
@@ -43,7 +46,7 @@ abstract class Kapt3BaseIT : BaseGradleIT() {
     }
 }
 
-class Kapt3WorkersIT : Kapt3IT() {
+open class Kapt3WorkersIT : Kapt3IT() {
     override fun kaptOptions(): KaptOptions =
         super.kaptOptions().copy(useWorkers = true)
 
@@ -53,7 +56,9 @@ class Kapt3WorkersIT : Kapt3IT() {
             Project("javacIsLoadedOnce", directoryPrefix = "kapt2")
         project.build("build") {
             assertSuccessful()
-            assertSubstringCount("Loaded com.sun.tools.javac.util.Context from", 1)
+
+            val loadsCount = Pattern.quote("Loaded com.sun.tools.javac.util.Context from").toRegex().findAll(output).count()
+            assertTrue(loadsCount <= 1, "javac is loaded more than once")
         }
     }
 
@@ -95,9 +100,65 @@ class Kapt3WorkersIT : Kapt3IT() {
     }
 }
 
+class Kapt3ClassLoadersCacheIT : Kapt3WorkersIT() {
+    override fun kaptOptions(): KaptOptions =
+        super.kaptOptions().copy(classLoadersCacheSize = 10, includeCompileClasspath = false)
+
+    @Ignore
+    override fun testDisableDiscoveryInCompileClasspath() {
+        //classloaders cache is incompatible with AP discovery in classpath
+    }
+
+    @Test
+    override fun testAnnotationProcessorAsFqName() {
+        val project = Project("annotationProcessorAsFqName", directoryPrefix = "kapt2").also { it.setupWorkingDir() }
+
+        //classloaders caching is not compatible with includeCompileClasspath
+        project.projectDir.getFileByName("build.gradle").modify {
+            it.addBeforeSubstring(
+                "kapt \"org.jetbrains.kotlin:annotation-processor-example:\$kotlin_version\"\n",
+                "implementation \"org.jetbrains.kotlin:annotation-processor-example"
+            )
+        }
+
+        project.build("build") {
+            assertSuccessful()
+            assertKaptSuccessful()
+            assertTasksExecuted(":compileKotlin", ":compileJava")
+            assertFileExists("build/generated/source/kapt/main/example/TestClassGenerated.java")
+            assertFileExists(kotlinClassesDir() + "example/TestClass.class")
+            assertFileExists(javaClassesDir() + "example/TestClassGenerated.class")
+        }
+    }
+
+    @Test
+    fun testAnnotationProcessorClassIsLoadedOnce() {
+        val project = Project("javacIsLoadedOnce", directoryPrefix = "kapt2")
+
+        fun CompiledProject.classLoadingCount() =
+            Pattern.quote("Loaded example.ExampleAnnotationProcessor from").toRegex().findAll(output).count()
+
+        project.build("build") {
+            assertSuccessful()
+            assertTasksExecuted(":module1:kaptKotlin", ":module2:kaptKotlin")
+            assertTrue(classLoadingCount() <= 1, "AP class is loaded more than once")
+        }
+
+        project.projectDir.getFilesByNames("Module1Class.kt", "Module2Class.kt").forEach {
+            it.appendText("\n fun touch() = null")
+        }
+
+        project.build("build") {
+            assertSuccessful()
+            assertTasksExecuted(":module1:kaptKotlin", ":module2:kaptKotlin")
+            assertTrue(classLoadingCount() == 0, "AP class shouldn't be loaded on the second build")
+        }
+    }
+}
+
 open class Kapt3IT : Kapt3BaseIT() {
     @Test
-    fun testAnnotationProcessorAsFqName() {
+    open fun testAnnotationProcessorAsFqName() {
         val project = Project("annotationProcessorAsFqName", directoryPrefix = "kapt2")
 
         project.build("build") {
@@ -487,7 +548,7 @@ open class Kapt3IT : Kapt3BaseIT() {
     }
 
     @Test
-    fun testDisableDiscoveryInCompileClasspath() = with(Project("kaptAvoidance", directoryPrefix = "kapt2")) {
+    open fun testDisableDiscoveryInCompileClasspath() = with(Project("kaptAvoidance", directoryPrefix = "kapt2")) {
         setupWorkingDir()
         val buildGradle = projectDir.resolve("app/build.gradle")
         buildGradle.modify {
@@ -510,6 +571,12 @@ open class Kapt3IT : Kapt3BaseIT() {
 
     @Test
     fun testKaptAvoidance() = with(Project("kaptAvoidance", directoryPrefix = "kapt2")) {
+        setupWorkingDir()
+
+        projectDir.resolve("app/build.gradle").modify {
+            "$it\n\nkapt.includeCompileClasspath = true"
+        }
+
         build("assemble") {
             assertSuccessful()
             assertTasksExecuted(
@@ -538,7 +605,10 @@ open class Kapt3IT : Kapt3BaseIT() {
 
         // enable discovery
         projectDir.resolve("app/build.gradle").modify {
-            "$it\n\nkapt.includeCompileClasspath = false"
+            it.replace(
+                "kapt.includeCompileClasspath = true",
+                "kapt.includeCompileClasspath = false"
+            )
         }
         build("assemble") {
             assertSuccessful()

libraries/tools/kotlin-gradle-plugin-integration-tests/src/test/resources/testProject/kapt2/mpp-kapt-presence/build.gradle

@@ -3,9 +3,12 @@ buildscript {
       'kotlin': [
           'stdlib': [
               'common': "org.jetbrains.kotlin:kotlin-stdlib-common:$kotlin_version",
-              'jdk': "org.jetbrains.kotlin:kotlin-stdlib:$kotlin_version",
+              'jdk': "org.jetbrains.kotlin:kotlin-stdlib:$kotlin_version"
-          ],
+          ]
       ],
+      'other' : [
+          'ap_example': "org.jetbrains.kotlin:annotation-processor-example:$kotlin_version"
+      ]
   ]
 
   repositories {

libraries/tools/kotlin-gradle-plugin-integration-tests/src/test/resources/testProject/kapt2/mpp-kapt-presence/dac/jdk/build.gradle

@@ -7,4 +7,5 @@ targetCompatibility = JavaVersion.VERSION_1_8
 dependencies {
   expectedBy project(':dac')
   implementation deps.kotlin.stdlib.jdk
+  kapt deps.other.ap_example
 }

libraries/tools/kotlin-gradle-plugin/build.gradle.kts

@@ -58,6 +58,7 @@ dependencies {
     embedded(project(":kotlin-gradle-build-metrics"))
 
     implementation("com.google.code.gson:gson:${rootProject.extra["versions.jar.gson"]}")
+    implementation("com.google.guava:guava:${rootProject.extra["versions.jar.guava"]}")
     implementation("de.undercouch:gradle-download-task:4.1.1")
     implementation("com.github.gundy:semver4j:0.16.4:nodeps") {
         exclude(group = "*")
@@ -81,7 +82,7 @@ dependencies {
     runtimeOnly(project(":kotlin-reflect"))
 
     jarContents(compileOnly(intellijDep()) {
-        includeJars("asm-all", "gson", "serviceMessages", rootProject = rootProject)
+        includeJars("asm-all", "gson", "guava", "serviceMessages", rootProject = rootProject)
     })
 
     // com.android.tools.build:gradle has ~50 unneeded transitive dependencies

libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/internal/kapt/Kapt3KotlinGradleSubplugin.kt

@@ -15,6 +15,7 @@ import org.gradle.api.GradleException
 import org.gradle.api.Project
 import org.gradle.api.artifacts.Configuration
 import org.gradle.api.artifacts.Dependency
+import org.gradle.api.artifacts.ExternalDependency
 import org.gradle.api.attributes.Attribute
 import org.gradle.api.attributes.Usage
 import org.gradle.api.file.FileCollection
@@ -78,6 +79,9 @@ class Kapt3GradleSubplugin @Inject internal constructor(private val registry: To
 
         private val KAPT_KOTLIN_GENERATED = "kapt.kotlin.generated"
 
+        private val CLASSLOADERS_CACHE_SIZE = "kapt.classloaders.cache.size"
+        private val CLASSLOADERS_CACHE_DISABLE_FOR_PROCESSORS = "kapt.classloaders.cache.disableForProcessors"
+
         val MAIN_KAPT_CONFIGURATION_NAME = "kapt"
 
         const val KAPT_ARTIFACT_NAME = "kotlin-annotation-processing-gradle"
@@ -118,6 +122,28 @@ class Kapt3GradleSubplugin @Inject internal constructor(private val registry: To
             return getBooleanOptionValue(BooleanOption.KAPT_KEEP_KDOC_COMMENTS_IN_STUBS)
         }
 
+        fun Project.classLoadersCacheSize(): Int = findPropertySafe(CLASSLOADERS_CACHE_SIZE)?.toString()?.toInt() ?: 0
+
+        fun Project.disableClassloaderCacheForProcessors(): Set<String> {
+            val value = findPropertySafe(CLASSLOADERS_CACHE_DISABLE_FOR_PROCESSORS)?.toString() ?: ""
+            return value
+                .split(",")
+                .map { it.trim() }
+                .filter { it.isNotEmpty() }
+                .toSet()
+        }
+
+        /**
+         * In case [Project.findProperty] can throw exception, this version catch it and return null
+         */
+        private fun Project.findPropertySafe(propertyName: String): Any? =
+            try {
+                findProperty(propertyName)
+            } catch (ex: Exception) {
+                logger.warn("Error getting property $propertyName", ex)
+                null
+            }
+
         fun findMainKaptConfiguration(project: Project) = project.findKaptConfiguration(SourceSet.MAIN_SOURCE_SET_NAME)
 
         fun createAptConfigurationIfNeeded(project: Project, sourceSetName: String): Configuration {
@@ -511,7 +537,9 @@ class Kapt3GradleSubplugin @Inject internal constructor(private val registry: To
                     it.isVisible = false
                     it.isCanBeConsumed = false
                 }
+
             kaptTask.kaptClasspath.from(kaptClasspathConfiguration)
+            kaptTask.kaptExternalClasspath.from(*kaptClasspathConfiguration.files { it is ExternalDependency }.toTypedArray())
             kaptTask.kaptClasspathConfigurationNames.set(kaptClasspathConfigurations.map { it.name })
 
             KaptWithAndroid.androidVariantData(this)?.annotationProcessorOptionProviders?.let {
@@ -569,6 +597,15 @@ class Kapt3GradleSubplugin @Inject internal constructor(private val registry: To
                 it.mapDiagnosticLocations = kaptExtension.mapDiagnosticLocations
                 it.annotationProcessorFqNames = kaptExtension.processors.split(',').filter { it.isNotEmpty() }
                 it.javacOptions = dslJavacOptions.get()
+                if (includeCompileClasspath && project.classLoadersCacheSize() > 0) {
+                    project.logger.warn(
+                        "ClassLoaders cache can't be enabled together with AP discovery in compilation classpath."
+                                + "\nSet 'kapt.includeCompileClasspath = false' to disable discovery"
+                    )
+                } else {
+                    it.classLoadersCacheSize = project.classLoadersCacheSize()
+                }
+                it.disableClassloaderCacheForProcessors = project.disableClassloaderCacheForProcessors()
             }
 
             val subpluginOptions = getKaptApOptions()

libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/internal/kapt/KaptTask.kt

@@ -53,6 +53,12 @@ abstract class KaptTask @Inject constructor(
     @get:InputFiles
     abstract val kaptClasspath: ConfigurableFileCollection
 
+    //part of kaptClasspath consisting from external artifacts only
+    //basically kaptClasspath = kaptExternalClasspath + artifacts built locally
+    @get:Classpath
+    @get:InputFiles
+    abstract val kaptExternalClasspath: ConfigurableFileCollection
+
     @get:Classpath
     @get:InputFiles
     val compilerClasspath: List<File> by project.provider {

libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/internal/kapt/KaptWithoutKotlincTask.kt

@@ -17,6 +17,8 @@ import org.gradle.workers.WorkAction
 import org.gradle.workers.WorkParameters
 import org.gradle.workers.WorkerExecutor
 import org.jetbrains.kotlin.gradle.internal.Kapt3GradleSubplugin.Companion.KAPT_WORKER_DEPENDENCIES_CONFIGURATION_NAME
+import org.jetbrains.kotlin.gradle.internal.kapt.classloaders.ClassLoadersCache
+import org.jetbrains.kotlin.gradle.internal.kapt.classloaders.rootOrSelf
 import org.jetbrains.kotlin.gradle.internal.kapt.incremental.KaptIncrementalChanges
 import org.jetbrains.kotlin.gradle.plugin.KotlinAndroidPluginWrapper
 import org.jetbrains.kotlin.gradle.tasks.CompilerPluginOptions
@@ -24,6 +26,7 @@ import org.jetbrains.kotlin.gradle.tasks.findKotlinStdlibClasspath
 import org.jetbrains.kotlin.gradle.tasks.findToolsJar
 import org.jetbrains.kotlin.gradle.utils.isGradleVersionAtLeast
 import org.jetbrains.kotlin.utils.PathUtil
+import org.slf4j.LoggerFactory
 import java.io.File
 import java.io.Serializable
 import java.net.URL
@@ -44,6 +47,12 @@ abstract class KaptWithoutKotlincTask @Inject constructor(
     @get:Input
     var isVerbose: Boolean = false
 
+    @get:Input
+    var classLoadersCacheSize: Int = 0
+
+    @get:Input
+    var disableClassloaderCacheForProcessors: Set<String> = emptySet()
+
     @get:Input
     var mapDiagnosticLocations: Boolean = false
 
@@ -124,12 +133,15 @@ abstract class KaptWithoutKotlincTask @Inject constructor(
             stubsDir.asFile.get(),
 
             kaptClasspath.files.toList(),
+            kaptExternalClasspath.files.toList(),
             annotationProcessorFqNames,
 
             getAnnotationProcessorOptions(),
             javacOptions,
 
-            kaptFlagsForWorker
+            kaptFlagsForWorker,
+
+            disableClassloaderCacheForProcessors
         )
 
         // Skip annotation processing if no annotation processors were provided.
@@ -184,6 +196,7 @@ abstract class KaptWithoutKotlincTask @Inject constructor(
             it.workerOptions.set(optionsForWorker)
             it.toolsJarURLSpec.set(toolsJarURLSpec)
             it.kaptClasspath.setFrom(kaptClasspath)
+            it.classloadersCacheSize.set(classLoadersCacheSize)
         }
     }
 
@@ -198,6 +211,7 @@ abstract class KaptWithoutKotlincTask @Inject constructor(
         val workerOptions: Property<KaptOptionsForWorker>
         val toolsJarURLSpec: Property<String>
         val kaptClasspath: ConfigurableFileCollection
+        val classloadersCacheSize: Property<Int>
     }
 
     internal abstract class KaptExecutionWorkAction : WorkAction<KaptWorkParameters> {
@@ -205,7 +219,8 @@ abstract class KaptWithoutKotlincTask @Inject constructor(
             KaptExecution(
                 parameters.workerOptions.get(),
                 parameters.toolsJarURLSpec.get(),
-                parameters.kaptClasspath.toList()
+                parameters.kaptClasspath.toList(),
+                parameters.classloadersCacheSize.get()
             ).run()
         }
     }
@@ -215,29 +230,40 @@ abstract class KaptWithoutKotlincTask @Inject constructor(
 private class KaptExecution @Inject constructor(
     val optionsForWorker: KaptOptionsForWorker,
     val toolsJarURLSpec: String,
-    val kaptClasspath: List<File>
+    val kaptClasspath: List<File>,
+    val classloadersCacheSize: Int
 ) : Runnable {
     private companion object {
         private const val JAVAC_CONTEXT_CLASS = "com.sun.tools.javac.util.Context"
 
         private fun kaptClass(classLoader: ClassLoader) = Class.forName("org.jetbrains.kotlin.kapt3.base.Kapt", true, classLoader)
-        private var cachedClassLoaderWithToolsJar: ClassLoader? = null
+
+        private var classLoadersCache: ClassLoadersCache? = null
+
         private var cachedKaptClassLoader: ClassLoader? = null
     }
 
+    private val logger = LoggerFactory.getLogger(KaptExecution::class.java)
+
     override fun run(): Unit = with(optionsForWorker) {
         val kaptClasspathUrls = kaptClasspath.map { it.toURI().toURL() }.toTypedArray()
         val rootClassLoader = findRootClassLoader()
 
-        val classLoaderWithToolsJar = cachedClassLoaderWithToolsJar ?: if (!toolsJarURLSpec.isEmpty() && !javacIsAlreadyHere()) {
+        val kaptClassLoader = cachedKaptClassLoader ?: run {
-            URLClassLoader(arrayOf(URL(toolsJarURLSpec)), rootClassLoader)
+            val classLoaderWithToolsJar = if (toolsJarURLSpec.isNotEmpty() && !javacIsAlreadyHere()) {
-        } else {
+                URLClassLoader(arrayOf(URL(toolsJarURLSpec)), rootClassLoader)
-            rootClassLoader
+            } else {
+                rootClassLoader
+            }
+            val result = URLClassLoader(kaptClasspathUrls, classLoaderWithToolsJar)
+            cachedKaptClassLoader = result
+            result
         }
-        cachedClassLoaderWithToolsJar = classLoaderWithToolsJar
 
-        val kaptClassLoader = cachedKaptClassLoader ?: URLClassLoader(kaptClasspathUrls, classLoaderWithToolsJar)
+        if (classLoadersCache == null && classloadersCacheSize > 0) {
-        cachedKaptClassLoader = kaptClassLoader
+            logger.info("Initializing KAPT classloaders cache with size = $classloadersCacheSize")
+            classLoadersCache = ClassLoadersCache(classloadersCacheSize, kaptClassLoader)
+        }
 
         val kaptMethod = kaptClass(kaptClassLoader).declaredMethods.single { it.name == "kapt" }
         kaptMethod.invoke(null, createKaptOptions(kaptClassLoader))
@@ -260,6 +286,15 @@ private class KaptExecution @Inject constructor(
         val detectMemoryLeaksMode = Class.forName("org.jetbrains.kotlin.base.kapt3.DetectMemoryLeaksMode", true, classLoader)
             .enumConstants.single { (it as Enum<*>).name == "NONE" }
 
+        //in case cache was enabled and then disabled
+        //or disabled for some modules
+        val processingClassLoader =
+            if (classloadersCacheSize > 0) {
+                classLoadersCache!!.getForSplitPaths(processingClasspath - processingExternalClasspath, processingExternalClasspath)
+            } else {
+                null
+            }
+
         Class.forName("org.jetbrains.kotlin.base.kapt3.KaptOptions", true, classLoader).constructors.single().newInstance(
             projectBaseDir,
             compileClasspath,
@@ -283,17 +318,14 @@ private class KaptExecution @Inject constructor(
 
             flags,
             mode,
-            detectMemoryLeaksMode
+            detectMemoryLeaksMode,
+
+            processingClassLoader,
+            disableClassloaderCacheForProcessors
         )
     }
 
-    private fun findRootClassLoader(): ClassLoader {
+    private fun findRootClassLoader(): ClassLoader = KaptExecution::class.java.classLoader.rootOrSelf()
-        tailrec fun parentOrSelf(classLoader: ClassLoader): ClassLoader {
-            val parent = classLoader.parent ?: return classLoader
-            return parentOrSelf(parent)
-        }
-        return parentOrSelf(KaptExecution::class.java.classLoader)
-    }
 }
 
 internal data class KaptOptionsForWorker(
@@ -311,10 +343,13 @@ internal data class KaptOptionsForWorker(
     val stubsOutputDir: File,
 
     val processingClasspath: List<File>,
+    val processingExternalClasspath: List<File>,
     val processors: List<String>,
 
     val processingOptions: Map<String, String>,
     val javacOptions: Map<String, String>,
 
-    val flags: Set<String>
+    val flags: Set<String>,
+
+    val disableClassloaderCacheForProcessors: Set<String>
 ) : Serializable

libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/internal/kapt/classloaders/ClassLoadersCache.kt

@@ -0,0 +1,86 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.gradle.internal.kapt.classloaders
+
+import com.google.common.cache.CacheBuilder
+import org.slf4j.LoggerFactory
+import java.io.File
+import java.net.URL
+import java.net.URLClassLoader
+import java.time.Duration
+import java.util.concurrent.ConcurrentMap
+
+/**
+ * LRU cache for [ClassLoader]s by class path.
+ */
+class ClassLoadersCache(
+    size: Int,
+    private val parentClassLoader: ClassLoader = ClassLoader.getSystemClassLoader(),
+    ttl: Duration = Duration.ofHours(1)
+) : AutoCloseable {
+
+    private val logger = LoggerFactory.getLogger(ClassLoadersCache::class.java)
+
+    private val cache: ConcurrentMap<CacheKey, URLClassLoader> =
+        CacheBuilder
+            .newBuilder()
+            .maximumSize(size.toLong())
+            .expireAfterAccess(ttl)
+            .removalListener<CacheKey, URLClassLoader> { (key, cl) ->
+                logger.info("Removing classloader from cache: ${key.entries.map { it.path }}")
+                cl.close()
+            }
+            .build<CacheKey, URLClassLoader>()
+            .asMap()
+
+    fun getForClassPath(files: List<File>): ClassLoader = getForClassPath(files, parentClassLoader)
+
+    private fun getForClassPath(files: List<File>, parent: ClassLoader): ClassLoader {
+        val key = makeKey(files)
+        return cache.getOrPut(key) {
+            makeClassLoader(key, parent)
+        }
+    }
+
+    /**
+     * Gets or creates [ClassLoader] from [bottom] + [top] files.
+     * When creating new [ClassLoader] it tries to get [top] from cache first and then create new ClassLoader from [bottom] files,
+     * providing [top] [ClassLoader] as parent.
+     * Useful when you have internal and external artifacts and internal ones can be references from other internal artefacts only.
+     * So you can safely cache [ClassLoader] from external artifacts and use it for internal ones.
+     */
+    fun getForSplitPaths(bottom: List<File>, top: List<File>): ClassLoader {
+        return if (bottom.isEmpty() || top.isEmpty()) {
+            getForClassPath(bottom + top)
+        } else {
+            val key = makeKey(bottom + top)
+            cache.getOrPut(key) {
+                val parent = getForClassPath(top)
+                makeClassLoader(makeKey(bottom), parent)
+            }
+        }
+    }
+
+    override fun close() {
+        cache.clear()
+    }
+
+    private fun makeClassLoader(key: CacheKey, parent: ClassLoader): URLClassLoader {
+        val cp = key.entries.map { it.path }
+        logger.info("Creating new classloader for classpath: $cp")
+        return URLClassLoader(cp.toTypedArray(), parent)
+    }
+
+    private fun makeKey(files: List<File>): CacheKey {
+        //probably should walk dirs content for actual last modified
+        val entries = files.map { f -> ClasspathEntry(f.toURI().toURL(), f.lastModified()) }
+        return CacheKey(entries)
+    }
+
+    private data class ClasspathEntry(val path: URL, val modificationTimestamp: Long)
+
+    private data class CacheKey(val entries: List<ClasspathEntry>)
+}

libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/internal/kapt/classloaders/classLoadersUtils.kt

@@ -0,0 +1,16 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.gradle.internal.kapt.classloaders
+
+
+internal fun ClassLoader.rootOrSelf(): ClassLoader {
+    tailrec fun parentOrSelf(classLoader: ClassLoader): ClassLoader {
+        val parent = classLoader.parent ?: return classLoader
+        return parentOrSelf(parent)
+    }
+
+    return parentOrSelf(this)
+}

libraries/tools/kotlin-gradle-plugin/src/main/kotlin/org/jetbrains/kotlin/gradle/tasks/jarSearchingUtil.kt

@@ -119,7 +119,7 @@ internal fun findToolsJar(): File? {
     return javacUtilContextClass?.let(::findJarByClass)
 }
 
-private fun findJarByClass(klass: Class<*>): File? {
+internal fun findJarByClass(klass: Class<*>): File? {
     val classFileName = klass.name.substringAfterLast(".") + ".class"
     val resource = klass.getResource(classFileName) ?: return null
     val uri = resource.toString()

libraries/tools/kotlin-gradle-plugin/src/test/kotlin/org/jetbrains/kotlin/gradle/internal/kapt/classloaders/ClassLoadersCacheTest.kt

@@ -0,0 +1,80 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.gradle.internal.kapt.classloaders
+
+import com.google.gson.Gson
+import org.jetbrains.kotlin.gradle.tasks.findJarByClass
+import org.junit.Test
+import kotlin.test.assertNotSame
+import kotlin.test.assertSame
+
+class ClassLoadersCacheTest {
+
+    private val rootClassLoader = this.javaClass.classLoader.rootOrSelf()
+
+    private val someClass = Test::class.java
+    private val someJar = findJarByClass(someClass)!!
+
+    private val otherClass = Gson::class.java
+    private val otherJar = findJarByClass(otherClass)!!
+
+    @Test
+    fun testNewClassLoader() {
+        val cache = ClassLoadersCache(10, rootClassLoader)
+        val cl = cache.getForClassPath(listOf(someJar))
+        val loaded = cl.loadClass(someClass.name)
+        assertNotSame(someClass, loaded, "Class should be from different ClassLoader")
+    }
+
+    @Test
+    fun testCacheClassLoader() {
+        val cache = ClassLoadersCache(10, rootClassLoader)
+        val cp = listOf(someJar)
+
+        val cl1 = cache.getForClassPath(cp)
+        val loaded1 = cl1.loadClass(someClass.name)
+
+        val cl2 = cache.getForClassPath(cp)
+        val loaded2 = cl2.loadClass(someClass.name)
+
+        assertSame(loaded2, loaded1, "Should return the same ClassLoader for same class path")
+    }
+
+    @Test
+    fun testDifferentClassPath() {
+        val cache = ClassLoadersCache(10, rootClassLoader)
+
+        val cl1 = cache.getForClassPath(listOf(someJar))
+        val loaded1 = cl1.loadClass(someClass.name)
+
+        val cl2 = cache.getForClassPath(listOf(someJar, otherJar))
+        val loaded2 = cl2.loadClass(someClass.name)
+
+        assertNotSame(loaded2, loaded1, "Should create different ClassLoaders for different class paths")
+    }
+
+    @Test
+    fun testSplitClassPath() {
+        val cache = ClassLoadersCache(10, rootClassLoader)
+        val topCp = listOf(someJar)
+        val bottomCp1 = listOf(otherJar)
+        val bottomCp2 = listOf(otherJar, findJarByClass(JvmField::class.java)!!)
+
+        val cl1 = cache.getForSplitPaths(bottomCp1, topCp)
+        val cl2 = cache.getForSplitPaths(bottomCp2, topCp)
+
+        assertSame(
+            cl1.loadClass(someClass.name),
+            cl2.loadClass(someClass.name),
+            "Top classpath should be cached separately. ClassLoader shouldn't change if top classpath stays the same"
+        )
+        assertNotSame(
+            cl1.loadClass(otherClass.name),
+            cl2.loadClass(otherClass.name),
+            "Bottom ClassLoader should be recreated as class path changed"
+        )
+    }
+}

plugins/kapt3/kapt3-base/src/org/jetbrains/kotlin/kapt3/base/KaptOptions.kt

@@ -35,7 +35,14 @@ class KaptOptions(
     val flags: KaptFlags,
 
     val mode: AptMode,
-    val detectMemoryLeaks: DetectMemoryLeaksMode
+    val detectMemoryLeaks: DetectMemoryLeaksMode,
+
+    //these two config can be replaced with single function-like interface (ProcessorName -> ClassLoader),
+    // but it is hard to pass function between different classloaders
+    //if defined use it to run processors instead of creating new one
+    val processingClassLoader: ClassLoader?,
+    //construct new classloader for these processors instead of using one defined in processingClassLoader
+    val separateClassloaderForProcessors: Set<String>
 ) : KaptFlags {
     override fun get(flag: KaptFlag) = flags[flag]
 
@@ -76,7 +83,9 @@ class KaptOptions(
                 changedFiles, compiledSources, incrementalCache, classpathChanges,
                 sourcesOutputDir, classesOutputDir, stubsOutputDir, incrementalDataOutputDir,
                 processingClasspath, processors, processingOptions, javacOptions, KaptFlags.fromSet(flags),
-                mode, detectMemoryLeaks
+                mode, detectMemoryLeaks,
+                processingClassLoader = null,
+                separateClassloaderForProcessors = emptySet()
             )
         }
     }

plugins/kapt3/kapt3-base/src/org/jetbrains/kotlin/kapt3/base/ProcessorLoader.kt

@@ -32,6 +32,7 @@ open class ProcessorLoader(private val options: KaptOptions, private val logger:
                 addAll(options.compileClasspath)
             }
         }
+
         val classLoader = URLClassLoader(classpath.map { it.toURI().toURL() }.toTypedArray(), parentClassLoader)
         this.annotationProcessingClassLoader = classLoader
 
@@ -74,7 +75,7 @@ open class ProcessorLoader(private val options: KaptOptions, private val logger:
         }
     }
 
-    open fun doLoadProcessors(classpath: LinkedHashSet<File>, classLoader: URLClassLoader): List<Processor> {
+    open fun doLoadProcessors(classpath: LinkedHashSet<File>, classLoader: ClassLoader): List<Processor> {
         val processorNames = mutableSetOf<String>()
 
         fun processSingleInput(input: InputStream) {
@@ -119,8 +120,17 @@ open class ProcessorLoader(private val options: KaptOptions, private val logger:
     }
 
     private fun tryLoadProcessor(fqName: String, classLoader: ClassLoader): Processor? {
+        val providedClassloader = options.processingClassLoader?.takeIf { !options.separateClassloaderForProcessors.contains(fqName) }
+        val classLoaderToUse = if (providedClassloader != null) {
+            logger.info { "Use provided ClassLoader for processor '$fqName'" }
+            providedClassloader
+        } else {
+            logger.info { "Use own ClassLoader for processor '$fqName'" }
+            classLoader
+        }
+
         val annotationProcessorClass = try {
-            Class.forName(fqName, true, classLoader)
+            Class.forName(fqName, true, classLoaderToUse)
         } catch (e: Throwable) {
             logger.warn("Can't find annotation processor class $fqName: ${e.message}")
             return null

plugins/kapt3/kapt3-compiler/src/org/jetbrains/kotlin/kapt3/Kapt3Extension.kt

@@ -81,9 +81,11 @@ class ClasspathBasedKapt3Extension(
 
     override fun loadProcessors(): LoadedProcessors {
         val efficientProcessorLoader = object : ProcessorLoader(options, logger) {
-            override fun doLoadProcessors(classpath: LinkedHashSet<File>, classLoader: URLClassLoader): List<Processor> {
+            override fun doLoadProcessors(classpath: LinkedHashSet<File>, classLoader: ClassLoader): List<Processor> =
-                return ServiceLoaderLite.loadImplementations(Processor::class.java, classLoader)
+                when (classLoader) {
-            }
+                    is URLClassLoader -> ServiceLoaderLite.loadImplementations(Processor::class.java, classLoader)
+                    else -> super.doLoadProcessors(classpath, classLoader)
+                }
         }
 
         this.processorLoader = efficientProcessorLoader