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,10 +98,23 @@ 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
   }
 
+  /**
+   * Freezes this context if it is not currently frozen.
+   *
+   * @return Boolean of whether or not the context was already frozen.
+   */
+  internal fun freezeIfNotFrozen(): Boolean {
+    if (!frozen) {
+      freeze()
+      return false
+    }
+    return true
+  }
+
   /**
    * Unfreezes when the node is about to render() again.
    */

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,42 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
     }
   }
 
+  inner class FrozenContextContinuationInterceptor : ContinuationInterceptor {
+    override val key: Key<*>
+      get() = ContinuationInterceptor.Key
+
+    private var originallyFrozen = false
+
+    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>) {
+          originallyFrozen = baseRenderContext.freezeIfNotFrozen()
+          continuation.resumeWith(result)
+        }
+      }
+
+    override fun releaseInterceptedContinuation(continuation: Continuation<*>) {
+      if (!originallyFrozen) {
+        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") +
+        // Adds the ContinuationInterceptor that freezes whenever we dispatch the continuation
+        // for the side effect. This allows us to schedule side effects during the render
+        // pass.
+        FrozenContextContinuationInterceptor()
+      val job = scope.launch(start = DEFAULT, block = sideEffect)
       SideEffectNode(key, job)
     }
   }

workflow-runtime/src/commonTest/kotlin/com/squareup/workflow1/RenderWorkflowInTest.kt

@@ -198,7 +198,7 @@ class RenderWorkflowInTest {
   }
 
   @Test
-  fun `side_effects_from_initial_rendering_in_root_workflow_are_never_started_when_scope_cancelled_before_start`() {
+  fun `side_effects_from_initial_rendering_in_root_workflow_are_started_when_scope_cancelled_before_start`() {
     runtimeTestRunner.runParametrizedTest(
       paramSource = runtimeOptions,
       before = ::setup,
@@ -222,13 +222,13 @@ class RenderWorkflowInTest {
         ) {}
         advanceIfStandard(dispatcher)
 
-        assertFalse(sideEffectWasRan)
+        assertTrue(sideEffectWasRan)
       }
     }
   }
 
   @Test
-  fun `side_effects_from_initial_rendering_in_non_root_workflow_are_never_started_when_scope_cancelled_before_start`() {
+  fun `side_effects_from_initial_rendering_in_non_root_workflow_are_started_when_scope_cancelled_before_start`() {
     runtimeTestRunner.runParametrizedTest(
       paramSource = runtimeOptions,
       before = ::setup,
@@ -255,7 +255,7 @@ class RenderWorkflowInTest {
         ) {}
         advanceIfStandard(dispatcher)
 
-        assertFalse(sideEffectWasRan)
+        assertTrue(sideEffectWasRan)
       }
     }
   }
@@ -765,7 +765,7 @@ class RenderWorkflowInTest {
   }
 
   @Test
-  fun `side_effects_from_initial_rendering_in_root_workflow_are_never_started_when_initial_render_of_root_workflow_fails`() {
+  fun `side_effects_from_initial_rendering_in_root_workflow_are_started_when_initial_render_of_root_workflow_fails`() {
     runtimeTestRunner.runParametrizedTest(
       paramSource = runtimeOptions,
       before = ::setup,
@@ -788,7 +788,7 @@ class RenderWorkflowInTest {
             workflowTracer = workflowTracer,
           ) {}
         }
-        assertFalse(sideEffectWasRan)
+        assertTrue(sideEffectWasRan)
       }
     }
   }
@@ -838,7 +838,7 @@ class RenderWorkflowInTest {
   }
 
   @Test
-  fun `side_effects_from_initial_rendering_in_non_root_workflow_are_never_started_when_initial_render_of_non_root_workflow_fails`() {
+  fun `side_effects_from_initial_rendering_in_non_root_workflow_are_started_when_initial_render_of_non_root_workflow_fails`() {
     runtimeTestRunner.runParametrizedTest(
       paramSource = runtimeOptions,
       before = ::setup,
@@ -864,7 +864,7 @@ class RenderWorkflowInTest {
             workflowTracer = workflowTracer,
           ) {}
         }
-        assertFalse(sideEffectWasRan)
+        assertTrue(sideEffectWasRan)
       }
     }
   }

workflow-runtime/src/commonTest/kotlin/com/squareup/workflow1/internal/RealRenderContextTest.kt

@@ -220,7 +220,6 @@ internal class RealRenderContextTest {
 
     val child = Workflow.stateless<Unit, Nothing, Unit> { fail() }
     assertFailsWith<IllegalStateException> { context.renderChild(child) }
-    assertFailsWith<IllegalStateException> { context.freeze() }
     assertFailsWith<IllegalStateException> { context.remember("key", typeOf<String>()) {} }
   }
 

workflow-runtime/src/commonTest/kotlin/com/squareup/workflow1/internal/WorkflowNodeTest.kt

@@ -276,13 +276,14 @@ internal class WorkflowNodeTest {
     sink.send(action("") { setOutput("event2") })
   }
 
-  @Test fun sideEffect_is_not_started_until_after_render_completes() {
+  @Test fun sideEffect_is_started_immediately() {
     var started = false
     val workflow = Workflow.stateless<Unit, Nothing, Unit> {
       runningSideEffect("key") {
         started = true
       }
-      assertFalse(started)
+      // This is because context uses and Unconfined Dispatcher, so is eagerly entered.
+      assertTrue(started)
     }
     val node = WorkflowNode(
       workflow.id(),
@@ -466,8 +467,8 @@ internal class WorkflowNodeTest {
 
   @Test fun multiple_sideEffects_with_same_key_throws() {
     val workflow = Workflow.stateless<Unit, Nothing, Unit> {
-      runningSideEffect("same") { fail() }
-      runningSideEffect("same") { fail() }
+      runningSideEffect("same") { }
+      runningSideEffect("same") { }
     }
     val node = WorkflowNode(
       workflow.id(),