WIP: Intercept to freeze for sideEffects

Author: steve-the-edwards

workflow-runtime-android/src/androidTest/java/com/squareup/workflow1/AndroidRenderWorkflowInTest.kt

@@ -14,15 +14,14 @@ import kotlinx.coroutines.plus
 import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import kotlinx.coroutines.test.runTest
-import org.junit.Ignore
 import org.junit.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertTrue
 
 @OptIn(WorkflowExperimentalRuntime::class, ExperimentalCoroutinesApi::class)
 class AndroidRenderWorkflowInTest {
 
-  @Ignore("#1311: Does not yet work with immediate dispatcher.")
+  // @Ignore("#1311: Does not yet work with immediate dispatcher.")
   @Test
   fun with_conflate_we_conflate_stacked_actions_into_one_rendering() =
     runTest(UnconfinedTestDispatcher()) {

workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/RealRenderContext.kt

@@ -98,7 +98,7 @@ internal class RealRenderContext<out PropsT, StateT, OutputT>(
    * Freezes this context so that any further calls to this context will throw.
    */
   fun freeze() {
-    checkNotFrozen("freeze") { "freeze" }
+    // checkNotFrozen("freeze") { "freeze" }
     frozen = true
   }
 

workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/WorkflowNode.kt

@@ -26,7 +26,7 @@ import com.squareup.workflow1.trace
 import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CoroutineName
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.CoroutineStart.LAZY
+import kotlinx.coroutines.CoroutineStart.DEFAULT
 import kotlinx.coroutines.DelicateCoroutinesApi
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.Job
@@ -36,7 +36,10 @@ import kotlinx.coroutines.channels.Channel.Factory.UNLIMITED
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.plus
 import kotlinx.coroutines.selects.SelectBuilder
+import kotlin.coroutines.Continuation
+import kotlin.coroutines.ContinuationInterceptor
 import kotlin.coroutines.CoroutineContext
+import kotlin.coroutines.CoroutineContext.Key
 import kotlin.reflect.KType
 
 /**
@@ -279,9 +282,6 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
       workflowTracer.trace("UpdateRuntimeTree") {
         // Tear down workflows and workers that are obsolete.
         subtreeManager.commitRenderedChildren()
-        // Side effect jobs are launched lazily, since they can send actions to the sink, and can only
-        // be started after context is frozen.
-        sideEffects.forEachStaging { it.job.start() }
         sideEffects.commitStaging { it.job.cancel() }
         remembered.commitStaging { /* Nothing to clean up. */ }
       }
@@ -351,13 +351,36 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
     }
   }
 
+  inner class FrozenContextContinuationInterceptor : ContinuationInterceptor {
+    override val key: Key<*>
+      get() = ContinuationInterceptor.Key
+
+    override fun <T> interceptContinuation(continuation: Continuation<T>): Continuation<T> =
+      object : Continuation<T> {
+        override val context: CoroutineContext
+          get() = continuation.context
+
+        override fun resumeWith(result: Result<T>) {
+          // Freeze the render context each time we are resuming the side effect continuation.
+          baseRenderContext.freeze()
+          continuation.resumeWith(result)
+        }
+      }
+
+    override fun releaseInterceptedContinuation(continuation: Continuation<*>) {
+      baseRenderContext.unfreeze()
+    }
+  }
+
   private fun createSideEffectNode(
     key: String,
     sideEffect: suspend CoroutineScope.() -> Unit
   ): SideEffectNode {
     return workflowTracer.trace("CreateSideEffectNode") {
-      val scope = this + CoroutineName("sideEffect[$key] for $id")
-      val job = scope.launch(start = LAZY, block = sideEffect)
+      val scope = this +
+        CoroutineName("sideEffect[$key] for $id") +
+        FrozenContextContinuationInterceptor()
+      val job = scope.launch(start = DEFAULT, block = sideEffect)
       SideEffectNode(key, job)
     }
   }