[kairos] propagate transaction errors gracefully

Flag: EXEMPT unused
Test: atest kairos-tests
Change-Id: I0d48b5abf0de27a0008b3a69e49692dcca8a2fba

3 files changed, 42 insertions, 24 deletions

diff --git a/packages/SystemUI/utils/kairos/src/com/android/systemui/kairos/FrpNetwork.kt b/packages/SystemUI/utils/kairos/src/com/android/systemui/kairos/FrpNetwork.kt
index 97252b4a199a..cec76886c06d 100644
--- a/packages/SystemUI/utils/kairos/src/com/android/systemui/kairos/FrpNetwork.kt
+++ b/packages/SystemUI/utils/kairos/src/com/android/systemui/kairos/FrpNetwork.kt
@@ -24,7 +24,6 @@ import com.android.systemui.kairos.internal.util.childScope
 import kotlin.coroutines.CoroutineContext
 import kotlin.coroutines.EmptyCoroutineContext
 import kotlin.coroutines.coroutineContext
-import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.CoroutineName
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Job
@@ -134,19 +133,8 @@ internal class LocalFrpNetwork(
     private val scope: CoroutineScope,
     private val endSignal: TFlow<Any>,
 ) : FrpNetwork {
-    override suspend fun <R> transact(block: suspend FrpTransactionScope.() -> R): R {
-        val result = CompletableDeferred<R>(coroutineContext[Job])
-        @Suppress("DeferredResultUnused")
-        network.transaction("FrpNetwork.transact") {
-            val buildScope =
-                BuildScopeImpl(
-                    stateScope = StateScopeImpl(evalScope = this, endSignal = endSignal),
-                    coroutineScope = scope,
-                )
-            buildScope.runInBuildScope { effect { result.complete(block()) } }
-        }
-        return result.await()
-    }
+    override suspend fun <R> transact(block: suspend FrpTransactionScope.() -> R): R =
+        network.transaction("FrpNetwork.transact") { runInTransactionScope { block() } }.await()
 
     override suspend fun activateSpec(spec: FrpSpec<*>) {
         val job =
diff --git a/packages/SystemUI/utils/kairos/src/com/android/systemui/kairos/internal/Network.kt b/packages/SystemUI/utils/kairos/src/com/android/systemui/kairos/internal/Network.kt
index cc36fda0c3ae..b2b3ca3001ad 100644
--- a/packages/SystemUI/utils/kairos/src/com/android/systemui/kairos/internal/Network.kt
+++ b/packages/SystemUI/utils/kairos/src/com/android/systemui/kairos/internal/Network.kt
@@ -101,17 +101,27 @@ internal class Network(val coroutineScope: CoroutineScope) : NetworkScope {
                 actions.add(func)
             }
             transactionMutex.withLock {
-                // Run all actions
-                evalScope {
-                    for (action in actions) {
-                        launch { action.started(evalScope = this@evalScope) }
+                try {
+                    // Run all actions
+                    evalScope {
+                        for (action in actions) {
+                            launch { action.started(evalScope = this@evalScope) }
+                        }
+                    }
+                    // Step through the network
+                    doTransaction()
+                } catch (e: Exception) {
+                    // Signal failure
+                    while (actions.isNotEmpty()) {
+                        actions.removeLast().fail(e)
+                    }
+                    // re-throw, cancelling this coroutine
+                    throw e
+                } finally {
+                    // Signal completion
+                    while (actions.isNotEmpty()) {
+                        actions.removeLast().completed()
                     }
-                }
-                // Step through the network
-                doTransaction()
-                // Signal completion
-                while (actions.isNotEmpty()) {
-                    actions.removeLast().completed()
                 }
             }
         }
@@ -234,6 +244,11 @@ internal class ScheduledAction<T>(
         result = just(onStartTransaction(evalScope))
     }
 
+    fun fail(ex: Exception) {
+        result = none
+        onResult?.completeExceptionally(ex)
+    }
+
     fun completed() {
         if (onResult != null) {
             when (val result = result) {
diff --git a/packages/SystemUI/utils/kairos/test/com/android/systemui/kairos/KairosTests.kt b/packages/SystemUI/utils/kairos/test/com/android/systemui/kairos/KairosTests.kt
index 688adae8fcae..f3303f697fc9 100644
--- a/packages/SystemUI/utils/kairos/test/com/android/systemui/kairos/KairosTests.kt
+++ b/packages/SystemUI/utils/kairos/test/com/android/systemui/kairos/KairosTests.kt
@@ -53,6 +53,7 @@ import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Assert.assertNull
 import org.junit.Assert.assertTrue
+import org.junit.Assert.fail
 import org.junit.Test
 
 class KairosTests {
@@ -1057,6 +1058,20 @@ class KairosTests {
     }
 
     @Test
+    fun propagateError() {
+        try {
+            runFrpTest { network ->
+                runCurrent()
+                try {
+                    network.transact<Unit> { error("message") }
+                    fail("caller did not throw exception")
+                } catch (_: IllegalStateException) {}
+            }
+            fail("scheduler did not throw exception")
+        } catch (_: IllegalStateException) {}
+    }
+
+    @Test
     fun fanOutLateSubscribe() = runFrpTest { network ->
         val e = network.mutableTFlow<Map<String, Int>>()
         val barFlow =