[LL FIR] fix race in ValueWithPostCompute which happened then the computation finished with a non-recoverable exception

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/caches/ValueWithPostCompute.kt

@@ -5,9 +5,6 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.fir.caches
 
-import com.intellij.openapi.progress.ProcessCanceledException
-import com.intellij.openapi.project.IndexNotReadyException
-import com.intellij.openapi.diagnostic.ControlFlowException
 import org.jetbrains.kotlin.analysis.utils.errors.shouldIjPlatformExceptionBeRethrown
 
 /**
@@ -62,12 +59,20 @@ internal class ValueWithPostCompute<KEY, VALUE, DATA>(
                     return stateSnapshot.value as VALUE
                 } else {
                     synchronized(this) { // wait until other thread which holds the lock now computes the value
-                        if (value == ValueIsNotComputed) {
-                            // if we have a PCE during value computation, then we will enter the critical section with `value == ValueIsNotComputed`
-                            // in this case, we should try to recalculate the value
-                            return computeValueWithoutLock()
-                        } else {
-                            return value as VALUE
+                        when (val newStateSnapshot = value) {
+                            ValueIsNotComputed -> {
+                                // if we have a PCE during value computation, then we will enter the critical section with `value == ValueIsNotComputed`
+                                // in this case, we should try to recalculate the value
+                                return computeValueWithoutLock()
+                            }
+                            is ExceptionWasThrownDuringValueComputation -> {
+                                // if some other thread tried to compute the value but failed with the exception
+                                throw newStateSnapshot.error
+                            }
+                            else -> {
+                                // other thread computed the value for us
+                                return value as VALUE
+                            }
                         }
                     }
                 }
@@ -87,12 +92,24 @@ internal class ValueWithPostCompute<KEY, VALUE, DATA>(
     @Suppress("UNCHECKED_CAST")
     // should be called under a synchronized section
     private fun computeValueWithoutLock(): VALUE {
+        require(Thread.holdsLock(this))
         // if we entered synchronized section that's mean that the value is not yet calculated and was not started to be calculated
         // or the some other thread calculated the value while we were waiting to acquire the lock
 
-        if (value != ValueIsNotComputed) { // some other thread calculated our value
-            return value as VALUE
+        when (val newStateSnapshot = value) {
+            ValueIsNotComputed -> {
+                // will be computed next
+            }
+            is ExceptionWasThrownDuringValueComputation -> {
+                // if some other thread tried to compute the value but failed with the exception
+                throw newStateSnapshot.error
+            }
+            else -> {
+                // other thread computed the value for us
+                return value as VALUE
+            }
         }
+
         val calculatedValue = try {
             val (calculated, data) = recursiveGuarded {
                 _calculate!!(key)

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/caches/ValueWithPostComputeTest.kt

@@ -71,4 +71,64 @@ class ValueWithPostComputeTest {
             }
         }
     }
+
+    /**
+     * Tests the following scenario:
+     *  - thread `t1` access the cache and executes `calculate()` and then `postCompute()` under a lock hold
+     *  - while the lock hold  by `t1`, `t2` tries to also access the value and waits for the lock to be released by `t1`
+     *  - t1: during the post compute, some non-recoverable exception happens inside the `postCompute()` and exception is saved in the cache
+     *  - t1 releases the lock with the `value` set to `ExceptionWasThrownDuringValueComputation`
+     *  - t2 acquires the lock and should correctly see `ExceptionWasThrownDuringValueComputation` and rethrow the exception
+     */
+    @Test
+    fun testSavedExceptionInCacheFromPostCompute() {
+        class ExceptionWhichWillBeSavedInCache : Exception()
+        for (i in 1..100) {
+            val t1CalledCalculate = CountDownLatch(1)
+            val t2AccessedTheCache = CountDownLatch(1)
+
+            val resultRef = AtomicReference<Any?>(null)
+
+            val valueWithPostCompute = ValueWithPostCompute(
+                key = 1,
+                calculate = {
+                    if (Thread.currentThread().name == "t1") {
+                        t1CalledCalculate.countDown()
+                    }
+                    Thread.currentThread().name to Unit
+                },
+                postCompute = { _, _, _ ->
+                    t2AccessedTheCache.await()
+                    if (Thread.currentThread().name == "t1") {
+                        throw ExceptionWhichWillBeSavedInCache()
+                    }
+                }
+            )
+
+            val t1 = thread(name = "t1") {
+                try {
+                    valueWithPostCompute.getValue()
+                } catch (_: Throwable) {
+                }
+            }
+
+            val t2 = thread(name = "t2") {
+                t1CalledCalculate.await()
+                t2AccessedTheCache.countDown()
+
+                try {
+                    resultRef.set(valueWithPostCompute.getValue())
+                } catch (e: Throwable) {
+                    resultRef.set(e)
+                }
+            }
+            t2.join()
+            t1.join()
+            when (val result = resultRef.get()) {
+                is ExceptionWhichWillBeSavedInCache -> {}
+                is Throwable -> throw result
+                else -> Assertions.fail("Unexpected value $result")
+            }
+        }
+    }
 }