[IC] Let CompilationTransaction to reset in-memory changes in IC caches

#KT-56052 In Progress

build-common/src/org/jetbrains/kotlin/incremental/CompilationTransaction.kt

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.build.report.debug
 import org.jetbrains.kotlin.build.report.metrics.BuildTime
 import org.jetbrains.kotlin.build.report.metrics.measure
 import org.jetbrains.kotlin.compilerRunner.OutputItemsCollector
+import org.jetbrains.kotlin.incremental.storage.InMemoryStorageWrapper
 import org.jetbrains.kotlin.konan.file.use
 import java.io.Closeable
 import java.io.File
@@ -45,6 +46,8 @@ interface CompilationTransaction : Closeable {
      * A throwable that was thrown during the transaction execution
      */
     var executionThrowable: Throwable?
+
+    fun registerInMemoryStorageWrapper(inMemoryStorageWrapper: InMemoryStorageWrapper<*, *>)
 }
 
 fun CompilationTransaction.write(file: Path, writeAction: () -> Unit) {
@@ -81,10 +84,54 @@ inline fun <R> CompilationTransaction.runWithin(
     }.recover { exceptionTransformer(it) }.getOrThrow() // some exceptions may be transformed into results
 }
 
+abstract class BaseCompilationTransaction : CompilationTransaction {
+    private val inMemoryStorageWrappers = hashSetOf<InMemoryStorageWrapper<*, *>>()
+
+    protected var isSuccessful: Boolean = false
+
+    override fun markAsSuccessful() {
+        isSuccessful = true
+    }
+
+    override fun registerInMemoryStorageWrapper(inMemoryStorageWrapper: InMemoryStorageWrapper<*, *>) {
+        inMemoryStorageWrappers.add(inMemoryStorageWrapper)
+    }
+
+    override var cachesManager: Closeable? = null
+        set(value) {
+            check(field == null) {
+                "cachesManager is already set"
+            }
+            field = value
+        }
+
+    override var executionThrowable: Throwable? = null
+        set(value) {
+            check(field == null) {
+                "executionThrowable is already set"
+            }
+            field = value
+        }
+
+    protected fun closeCachesManager() = runCatching {
+        if (!isSuccessful) {
+            for (wrapper in inMemoryStorageWrappers) {
+                wrapper.resetInMemoryChanges()
+            }
+        }
+        cachesManager?.close()
+    }.exceptionOrNull()?.run {
+        isSuccessful = false
+        val exception = CachesManagerCloseException(this)
+        executionThrowable?.addSuppressed(exception)
+        executionThrowable ?: exception
+    }
+}
+
 /**
  * A non-recoverable implementation of compilation transaction. Changes reverting on failure should be performed externally if needed.
  */
-class NonRecoverableCompilationTransaction : CompilationTransaction {
+class NonRecoverableCompilationTransaction : CompilationTransaction, BaseCompilationTransaction() {
     override fun registerAddedOrChangedFile(outputFile: Path) {
         // do nothing
     }
@@ -95,23 +142,9 @@ class NonRecoverableCompilationTransaction : CompilationTransaction {
         }
     }
 
-    override fun markAsSuccessful() {
-        // do nothing
-    }
-
-    override var cachesManager: Closeable? = null
-
-    override var executionThrowable: Throwable? = null
-
     override fun close() {
-        try {
+        closeCachesManager()?.let {
-            cachesManager?.close()
+            throw it
-        } catch (t: Throwable) {
-            if (executionThrowable != null) {
-                executionThrowable?.addSuppressed(CachesManagerCloseException(t))
-            } else {
-                throw CachesManagerCloseException(t)
-            }
         }
     }
 }
@@ -125,10 +158,9 @@ class NonRecoverableCompilationTransaction : CompilationTransaction {
 class RecoverableCompilationTransaction(
     private val reporter: BuildReporter,
     private val stashDir: Path,
-) : CompilationTransaction {
+) : CompilationTransaction, BaseCompilationTransaction() {
     private val fileRelocationRegistry = hashMapOf<Path, Path?>()
     private var filesCounter = 0
-    private var successful = false
 
     /**
      * Moves the original [outputFile] before change to the [stashDir].
@@ -206,26 +238,11 @@ class RecoverableCompilationTransaction(
         }
     }
 
-    override fun markAsSuccessful() {
-        successful = true
-    }
-
-    override var cachesManager: Closeable? = null
-
-    override var executionThrowable: Throwable? = null
-
     override fun close() {
-        val mainException = runCatching {
+        val mainException = closeCachesManager()
-            cachesManager?.close()
-        }.exceptionOrNull()?.run {
-            successful = false
-            val exception = CachesManagerCloseException(this)
-            executionThrowable?.addSuppressed(exception)
-            executionThrowable ?: exception
-        }
 
         try {
-            if (successful) {
+            if (isSuccessful) {
                 cleanupStash()
             } else {
                 revertChanges()

build-common/src/org/jetbrains/kotlin/incremental/storage/BasicMap.kt

@@ -32,7 +32,9 @@ abstract class BasicMap<K : Comparable<K>, V>(
 ) {
     protected val storage: LazyStorage<K, V> = CachingLazyStorage(storageFile, keyDescriptor, valueExternalizer).let {
         if (icContext.keepIncrementalCompilationCachesInMemory) {
-            InMemoryStorageWrapper(it)
+            InMemoryStorageWrapper(it).also { wrapper ->
+                icContext.transaction.registerInMemoryStorageWrapper(wrapper)
+            }
         } else {
             it
         }

build-common/src/org/jetbrains/kotlin/incremental/storage/InMemoryStorageWrapper.kt

@@ -13,7 +13,7 @@ import kotlin.collections.LinkedHashMap
  * Flushes all the changes to the [origin] on [flush] invocation.
  * [resetInMemoryChanges] should be called to reset in-memory changes of this wrapper.
  */
-class InMemoryStorageWrapper<K, V>(val origin: LazyStorage<K, V>) : LazyStorage<K, V> {
+class InMemoryStorageWrapper<K, V>(private val origin: LazyStorage<K, V>) : LazyStorage<K, V> {
     private val inMemoryStorage = LinkedHashMap<K, ValueWrapper<V>>()
     private val removedKeys = hashSetOf<K>()
     private var isCleanRequested = false