[IC] Add unit tests for resetting in-memory caches wrappers in a transaction

#KT-56052 In Progress

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

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

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

@@ -8,12 +8,16 @@ package org.jetbrains.kotlin.incremental.storage
 import java.util.*
 import kotlin.collections.LinkedHashMap
 
+interface InMemoryStorageWrapper<K, V> : LazyStorage<K, V> {
+    fun resetInMemoryChanges()
+}
+
 /**
  * An in-memory wrapper for [origin] that keeps all the write operations in-memory.
  * 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>(private val origin: LazyStorage<K, V>) : LazyStorage<K, V> {
+class DefaultInMemoryStorageWrapper<K, V>(private val origin: LazyStorage<K, V>) : InMemoryStorageWrapper<K, V> {
     private val inMemoryStorage = LinkedHashMap<K, ValueWrapper<V>>()
     private val removedKeys = hashSetOf<K>()
     private var isCleanRequested = false
@@ -21,7 +25,7 @@ class InMemoryStorageWrapper<K, V>(private val origin: LazyStorage<K, V>) : Lazy
     override val keys: Collection<K>
         get() = if (isCleanRequested) inMemoryStorage.keys else (origin.keys - removedKeys) + inMemoryStorage.keys
 
-    fun resetInMemoryChanges() {
+    override fun resetInMemoryChanges() {
         isCleanRequested = false
         inMemoryStorage.clear()
         removedKeys.clear()

build-common/test/org/jetbrains/kotlin/incremental/CompilationTransactionTest.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.incremental
 
 import org.jetbrains.kotlin.build.report.DoNothingBuildReporter
+import org.jetbrains.kotlin.incremental.storage.InMemoryStorageWrapper
 import org.junit.jupiter.api.Assertions.assertEquals
 import org.junit.jupiter.api.Assertions.assertFalse
 import org.junit.jupiter.api.Assertions.assertTrue
@@ -26,6 +27,32 @@ private class CacheMock(private val throwsException: Boolean = false) : Closeabl
     }
 }
 
+private class InMemoryStorageWrapperMock : InMemoryStorageWrapper<Any, Any> {
+    var reset = false
+
+    override fun resetInMemoryChanges() {
+        reset = true
+    }
+
+    override val keys: Collection<Any> = emptyList()
+
+    override fun clean() {}
+
+    override fun flush(memoryCachesOnly: Boolean) {}
+
+    override fun close() {}
+
+    override fun append(key: Any, value: Any) {}
+
+    override fun remove(key: Any) {}
+
+    override fun set(key: Any, value: Any) {}
+
+    override fun get(key: Any) = null
+
+    override fun contains(key: Any) = false
+}
+
 abstract class BaseCompilationTransactionTest {
     @TempDir
     protected lateinit var stashDir: Path
@@ -84,6 +111,38 @@ abstract class BaseCompilationTransactionTest {
             }
         }
     }
+
+    @Test
+    fun testInMemoryWrappersAreResetOnUnsuccessfulTransaction() {
+        val inMemoryStorageWrapperMock = InMemoryStorageWrapperMock()
+        useTransaction {
+            registerInMemoryStorageWrapper(inMemoryStorageWrapperMock)
+        }
+        assertTrue(inMemoryStorageWrapperMock.reset)
+    }
+
+    @Test
+    fun testInMemoryWrappersAreResetOnExecutionException() {
+        val inMemoryStorageWrapperMock = InMemoryStorageWrapperMock()
+        assertThrows<Exception> {
+            useTransaction {
+                registerInMemoryStorageWrapper(inMemoryStorageWrapperMock)
+                markAsSuccessful()
+                throw Exception()
+            }
+        }
+        assertTrue(inMemoryStorageWrapperMock.reset)
+    }
+
+    @Test
+    fun testInMemoryWrappersAreNotResetOnSuccessfulTransaction() {
+        val inMemoryStorageWrapperMock = InMemoryStorageWrapperMock()
+        useTransaction {
+            registerInMemoryStorageWrapper(inMemoryStorageWrapperMock)
+            markAsSuccessful()
+        }
+        assertFalse(inMemoryStorageWrapperMock.reset)
+    }
 }
 
 class NonRecoverableCompilationTransactionTest : BaseCompilationTransactionTest() {