Close ksp caches on exception and error

Previously, LookupCache doesn't need to be explicitly closed. With newer intellij dependencies that comes with Kotlin 2.0.20 it becomes a problem.
google · Oct 2, 2024 · 50e9a96 · 50e9a96
1 parent 6276135
commit 50e9a96
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 23 deletions.
diff --git a/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/KotlinSymbolProcessingExtension.kt b/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/KotlinSymbolProcessingExtension.kt
@@ -310,6 +310,7 @@ abstract class AbstractKotlinSymbolProcessingExtension(
                         processor.process(resolver).filter { it.origin == Origin.KOTLIN || it.origin == Origin.JAVA }
                 }?.let {
                     resolver.tearDown()
+                    incrementalContext.closeFiles()
                     return it
                 }
                 if (logger.hasError()) {
@@ -336,16 +337,19 @@ abstract class AbstractKotlinSymbolProcessingExtension(
                     processor.onError()
                 }?.let {
                     resolver.tearDown()
+                    incrementalContext.closeFiles()
                     return it
                 }
             }
+            incrementalContext.closeFiles()
         } else {
             if (finished) {
                 processors.forEach { processor ->
                     handleException(module, project) {
                         processor.finish()
                     }?.let {
                         resolver.tearDown()
+                        incrementalContext.closeFiles()
                         return it
                     }
                 }
@@ -364,6 +368,8 @@ abstract class AbstractKotlinSymbolProcessingExtension(
                         codeGenerator.outputs,
                         codeGenerator.sourceToOutputs
                     )
+                } else {
+                    incrementalContext.closeFiles()
                 }
             }
         }

diff --git a/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/common/IncrementalContextBase.kt b/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/common/IncrementalContextBase.kt
@@ -211,13 +211,13 @@ abstract class IncrementalContextBase(
     // Beware: no side-effects here; Caches should only be touched in updateCaches.
     fun calcDirtyFiles(ksFiles: List<KSFile>): Collection<KSFile> = closeFilesOnException {
         if (!isIncremental) {
-            return ksFiles
+            return@closeFilesOnException ksFiles
         }
 
         if (rebuild) {
             collectDefinedSymbols(ksFiles)
             logDirtyFiles(ksFiles, ksFiles)
-            return ksFiles
+            return@closeFilesOnException ksFiles
         }
 
         val newSyms = mutableSetOf<LookupSymbolWrapper>()
@@ -285,7 +285,7 @@ abstract class IncrementalContextBase(
             dirtyFilesByNewSyms,
             dirtyFilesBySealed
         )
-        return ksFiles.filter { it.relativeFile in dirtyFiles }
+        return@closeFilesOnException ksFiles.filter { it.relativeFile in dirtyFiles }
     }
 
     // Loop detection isn't needed because of overwritten checks in CodeGeneratorImpl
@@ -403,27 +403,31 @@ abstract class IncrementalContextBase(
         collectDefinedSymbols(newFiles)
     }
 
-    private inline fun <T> closeFilesOnException(f: () -> T): T {
+    fun <T> closeFilesOnException(f: () -> T): T {
         try {
             return f()
         } catch (e: Exception) {
-            symbolsMap.flush()
-            sealedMap.flush()
-            symbolLookupCache.close()
-            classLookupCache.close()
-            sourceToOutputsMap.flush()
+            closeFiles()
             throw e
         }
     }
 
+    fun closeFiles() {
+        symbolsMap.flush()
+        sealedMap.flush()
+        symbolLookupCache.close()
+        classLookupCache.close()
+        sourceToOutputsMap.flush()
+    }
+
     // TODO: add a wildcard for outputs with no source and get rid of the outputs parameter.
     fun updateCachesAndOutputs(
         dirtyFiles: Collection<KSFile>,
         outputs: Set<File>,
         sourceToOutputs: Map<File, Set<File>>,
     ) = closeFilesOnException {
         if (!isIncremental)
-            return
+            return@closeFilesOnException
 
         cachesUpToDateFile.delete()
         assert(!cachesUpToDateFile.exists())

diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/common/IncrementalContextBase.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/common/IncrementalContextBase.kt
@@ -212,13 +212,13 @@ abstract class IncrementalContextBase(
     // Beware: no side-effects here; Caches should only be touched in updateCaches.
     fun calcDirtyFiles(ksFiles: List<KSFile>): Collection<KSFile> = closeFilesOnException {
         if (!isIncremental) {
-            return ksFiles
+            return@closeFilesOnException ksFiles
         }
 
         if (rebuild) {
             collectDefinedSymbols(ksFiles)
             logDirtyFiles(ksFiles, ksFiles)
-            return ksFiles
+            return@closeFilesOnException ksFiles
         }
 
         val newSyms = mutableSetOf<LookupSymbolWrapper>()
@@ -286,7 +286,7 @@ abstract class IncrementalContextBase(
             dirtyFilesByNewSyms,
             dirtyFilesBySealed
         )
-        return ksFiles.filter { it.relativeFile in dirtyFiles }
+        return@closeFilesOnException ksFiles.filter { it.relativeFile in dirtyFiles }
     }
 
     // Loop detection isn't needed because of overwritten checks in CodeGeneratorImpl
@@ -404,27 +404,31 @@ abstract class IncrementalContextBase(
         collectDefinedSymbols(newFiles)
     }
 
-    private inline fun <T> closeFilesOnException(f: () -> T): T {
+    fun <T> closeFilesOnException(f: () -> T): T {
         try {
             return f()
         } catch (e: Exception) {
-            symbolsMap.flush()
-            sealedMap.flush()
-            symbolLookupCache.close()
-            classLookupCache.close()
-            sourceToOutputsMap.flush()
+            closeFiles()
             throw e
         }
     }
 
+    fun closeFiles() {
+        symbolsMap.flush()
+        sealedMap.flush()
+        symbolLookupCache.close()
+        classLookupCache.close()
+        sourceToOutputsMap.flush()
+    }
+
     // TODO: add a wildcard for outputs with no source and get rid of the outputs parameter.
     fun updateCachesAndOutputs(
         dirtyFiles: Collection<KSFile>,
         outputs: Set<File>,
         sourceToOutputs: Map<File, Set<File>>,
     ) = closeFilesOnException {
         if (!isIncremental)
-            return
+            return@closeFilesOnException
 
         cachesUpToDateFile.delete()
         assert(!cachesUpToDateFile.exists())

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
@@ -548,9 +548,11 @@ class KotlinSymbolProcessing(
                 ResolverAAImpl.instance.propertyAsMemberOfCache = mutableMapOf()
 
                 processors.forEach {
-                    deferredSymbols[it] =
-                        it.process(resolver).filter { it.origin == Origin.KOTLIN || it.origin == Origin.JAVA }
-                            .filterIsInstance<Deferrable>().mapNotNull(Deferrable::defer)
+                    incrementalContext.closeFilesOnException {
+                        deferredSymbols[it] =
+                            it.process(resolver).filter { it.origin == Origin.KOTLIN || it.origin == Origin.JAVA }
+                                .filterIsInstance<Deferrable>().mapNotNull(Deferrable::defer)
+                    }
                     if (!deferredSymbols.containsKey(it) || deferredSymbols[it]!!.isEmpty()) {
                         deferredSymbols.remove(it)
                     }
@@ -591,6 +593,8 @@ class KotlinSymbolProcessing(
                     codeGenerator.outputs,
                     codeGenerator.sourceToOutputs
                 )
+            } else {
+                incrementalContext.closeFiles()
             }
 
             dropCaches()