Use WorkStealingDispatcher in runtime, behind a flag.

Author: zach-klippenstein

workflow-core/api/workflow-core.api

@@ -167,6 +167,7 @@ public final class com/squareup/workflow1/RuntimeConfigOptions : java/lang/Enum
 	public static final field PARTIAL_TREE_RENDERING Lcom/squareup/workflow1/RuntimeConfigOptions;
 	public static final field RENDER_ONLY_WHEN_STATE_CHANGES Lcom/squareup/workflow1/RuntimeConfigOptions;
 	public static final field STABLE_EVENT_HANDLERS Lcom/squareup/workflow1/RuntimeConfigOptions;
+	public static final field WORK_STEALING_DISPATCHER Lcom/squareup/workflow1/RuntimeConfigOptions;
 	public static fun getEntries ()Lkotlin/enums/EnumEntries;
 	public static fun valueOf (Ljava/lang/String;)Lcom/squareup/workflow1/RuntimeConfigOptions;
 	public static fun values ()[Lcom/squareup/workflow1/RuntimeConfigOptions;
@@ -179,6 +180,7 @@ public final class com/squareup/workflow1/RuntimeConfigOptions$Companion {
 }
 
 public final class com/squareup/workflow1/RuntimeConfigOptions$Companion$RuntimeOptions : java/lang/Enum {
+	public static final field ALL Lcom/squareup/workflow1/RuntimeConfigOptions$Companion$RuntimeOptions;
 	public static final field CONFLATE Lcom/squareup/workflow1/RuntimeConfigOptions$Companion$RuntimeOptions;
 	public static final field DEFAULT Lcom/squareup/workflow1/RuntimeConfigOptions$Companion$RuntimeOptions;
 	public static final field RENDER_ONLY Lcom/squareup/workflow1/RuntimeConfigOptions$Companion$RuntimeOptions;

workflow-core/src/commonMain/kotlin/com/squareup/workflow1/RuntimeConfig.kt

@@ -74,6 +74,13 @@ public enum class RuntimeConfigOptions {
   //  */
   // @WorkflowExperimentalRuntime
   // DRAIN_EXCLUSIVE_ACTIONS,
+
+  /**
+   * Wrap the dispatcher passed to the runtime with a special dispatcher that can be advanced
+   * explicitly, to allow any tasks scheduled by the workflow runtime to run before certain phases.
+   */
+  @WorkflowExperimentalRuntime
+  WORK_STEALING_DISPATCHER,
   ;
 
   public companion object {
@@ -136,6 +143,7 @@ public enum class RuntimeConfigOptions {
           STABLE_EVENT_HANDLERS,
         )
       ),
+
       // RENDER_ONLY_CONFLATE_PARTIAL_DEA(
       //   setOf(
       //     RENDER_ONLY_WHEN_STATE_CHANGES,
@@ -161,6 +169,8 @@ public enum class RuntimeConfigOptions {
       //     DRAIN_EXCLUSIVE_ACTIONS,
       //   )
       // ),
+
+      ALL(RuntimeConfigOptions.entries.toSet())
     }
   }
 }

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

@@ -4,6 +4,7 @@ import com.squareup.workflow1.RuntimeConfigOptions.CONFLATE_STALE_RENDERINGS
 import com.squareup.workflow1.RuntimeConfigOptions.RENDER_ONLY_WHEN_STATE_CHANGES
 import com.squareup.workflow1.WorkflowInterceptor.RenderPassSkipped
 import com.squareup.workflow1.WorkflowInterceptor.RenderPassesComplete
+import com.squareup.workflow1.internal.WorkStealingDispatcher
 import com.squareup.workflow1.internal.WorkflowRunner
 import com.squareup.workflow1.internal.chained
 import kotlinx.coroutines.CancellationException
@@ -14,6 +15,7 @@ import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.isActive
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.plus
 
 /**
  * Launches the [workflow] in a new coroutine in [scope] and returns a [StateFlow] of its
@@ -117,6 +119,15 @@ public fun <PropsT, OutputT, RenderingT> renderWorkflowIn(
 ): StateFlow<RenderingAndSnapshot<RenderingT>> {
   val chainedInterceptor = interceptors.chained()
 
+  val dispatcher = if (RuntimeConfigOptions.WORK_STEALING_DISPATCHER in runtimeConfig) {
+    WorkStealingDispatcher.wrapDispatcherFrom(scope.coroutineContext)
+  } else {
+    null
+  }
+
+  @Suppress("NAME_SHADOWING")
+  val scope = dispatcher?.let { scope + dispatcher } ?: scope
+
   val runner = WorkflowRunner(
     scope,
     workflow,
@@ -202,6 +213,13 @@ public fun <PropsT, OutputT, RenderingT> renderWorkflowIn(
           conflationHasChangedState = conflationHasChangedState || actionResult.stateChanged
           // We may have more actions we can process, this rendering could be stale.
           // This will check for any actions that are immediately available and apply them.
+          // We advance the dispatcher first to allow any coroutines that were launched by the last
+          // render pass to start up and potentially enqueue actions.
+          dispatcher?.let {
+            workflowTracer.trace("AdvancingWorkflowDispatcher") {
+              dispatcher.advanceUntilIdle()
+            }
+          }
           actionResult = runner.applyNextAvailableTreeAction()
 
           // If no actions processed, then no new rendering needed. Pass on to UI.

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

@@ -6,22 +6,29 @@ import com.squareup.workflow1.RuntimeConfigOptions.Companion.RuntimeOptions
 import com.squareup.workflow1.RuntimeConfigOptions.Companion.RuntimeOptions.DEFAULT
 import com.squareup.workflow1.RuntimeConfigOptions.PARTIAL_TREE_RENDERING
 import com.squareup.workflow1.RuntimeConfigOptions.RENDER_ONLY_WHEN_STATE_CHANGES
+import com.squareup.workflow1.RuntimeConfigOptions.WORK_STEALING_DISPATCHER
 import com.squareup.workflow1.WorkflowInterceptor.RenderPassSkipped
 import com.squareup.workflow1.WorkflowInterceptor.RenderPassesComplete
 import com.squareup.workflow1.WorkflowInterceptor.RuntimeLoopOutcome
 import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.CoroutineExceptionHandler
+import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.awaitCancellation
 import kotlinx.coroutines.cancel
+import kotlinx.coroutines.cancelAndJoin
 import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.produceIn
 import kotlinx.coroutines.flow.receiveAsFlow
 import kotlinx.coroutines.isActive
+import kotlinx.coroutines.job
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.plus
 import kotlinx.coroutines.suspendCancellableCoroutine
 import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.test.StandardTestDispatcher
@@ -45,7 +52,7 @@ import kotlin.test.assertTrue
 @Burst
 class RenderWorkflowInTest(
   useTracer: Boolean = false,
-  useUnconfined: Boolean = true,
+  private val useUnconfined: Boolean = true,
   private val runtime: RuntimeOptions = DEFAULT
 ) {
 
@@ -1494,7 +1501,9 @@ class RenderWorkflowInTest(
 
   @Test
   fun for_conflate_we_do_not_conflate_stacked_actions_into_one_rendering_if_output() {
-    if (runtimeConfig.contains(CONFLATE_STALE_RENDERINGS)) {
+    if (CONFLATE_STALE_RENDERINGS in runtimeConfig &&
+      WORK_STEALING_DISPATCHER !in runtimeConfig
+    ) {
       runTest(dispatcherUsed) {
         check(runtimeConfig.contains(CONFLATE_STALE_RENDERINGS))
 
@@ -1738,6 +1747,126 @@ class RenderWorkflowInTest(
     }
   }
 
+  /**
+   * When the [CONFLATE_STALE_RENDERINGS] flag is specified, the runtime will repeatedly run all
+   * enqueued WorkflowActions after a render pass, before emitting the rendering to the external
+   * flow. When the [WORK_STEALING_DISPATCHER] flag is specified at the same time, any coroutines
+   * launched (or even resumed) since the render pass will be allowed to run _before_ checking for
+   * actions. This means that any new side effects or workers started by the render pass will be
+   * allowed to run to their first suspension point before the rendering is emitted. And if they
+   * happen to emit more actions as part of that, then those actions will also be processed, etc.
+   * until no more actions are available – only then will the rendering actually be emitted.
+   */
+  @Test
+  fun new_effect_coroutines_dispatched_before_rendering_emitted_when_work_stealing_dispatcher() {
+    // This tests is specifically for standard dispatching behavior. It currently only works when
+    // CSR is enabled, although an additional test for DEA should be added.
+    if (WORK_STEALING_DISPATCHER !in runtimeConfig ||
+      CONFLATE_STALE_RENDERINGS !in runtimeConfig ||
+      useUnconfined
+    ) {
+      return
+    }
+
+    runTest(dispatcherUsed) {
+      val workflow = Workflow.stateful<Int, Nothing, Unit>(initialState = 0) { effectCount ->
+        // Because of the WSD, this effect will be allowed to run after the render pass but before
+        // emitting the rendering OR checking for new actions, in the CSR loop. Since it emits an
+        // action, that action will be processed and trigger a second render pass.
+        runningSideEffect("sender") {
+          actionSink.send(
+            action("0") {
+              expect(2)
+              this.state++
+            }
+          )
+        }
+
+        if (effectCount >= 1) {
+          // This effect will be started by the first action and cancelled only when the runtime
+          // is cancelled.
+          // It will also start in the CSR loop, and trigger a third render pass before emitting the
+          // rendering.
+          runningSideEffect("0") {
+            expect(3)
+            actionSink.send(
+              action("1") {
+                expect(4)
+                this.state++
+              }
+            )
+            awaitCancellation {
+              expect(9)
+            }
+          }
+        }
+
+        if (effectCount >= 2) {
+          // This effect will be started by the second action, and cancelled by its own action in
+          // the same run of the CSR loop again.
+          runningSideEffect("1") {
+            expect(5)
+            actionSink.send(
+              action("-1") {
+                expect(6)
+                this.state--
+              }
+            )
+            awaitCancellation {
+              expect(7)
+            }
+          }
+        }
+      }
+
+      // We collect the renderings flow to a channel to drive the runtime loop by receiving from the
+      // channel. We can't use testScheduler.advanceUntilIdle() et al because we only want the test
+      // scheduler to run tasks until a rendering is available, not indefinitely.
+      val renderings = renderWorkflowIn(
+        workflow = workflow,
+        // Run in this scope so it is advanced by advanceUntilIdle.
+        scope = backgroundScope,
+        props = MutableStateFlow(Unit),
+        runtimeConfig = runtimeConfig,
+        workflowTracer = testTracer,
+        onOutput = {}
+      ).produceIn(backgroundScope + Dispatchers.Unconfined)
+
+      expect(0)
+      // Receiving the first rendering allows the runtime coroutine to start. The first rendering
+      // is returned synchronously.
+      renderings.receive()
+      expect(1)
+      // Receiving the second rendering will allow the runtime to continue until the rendering is
+      // emitted. Since the CSR loop will start all our effects before emitting the next rendering,
+      // only one rendering will be emitted for all those render passes.
+      renderings.receive()
+      expect(8)
+
+      // No more renderings should be produced.
+      testScheduler.advanceUntilIdle()
+      assertTrue(renderings.isEmpty)
+
+      // Cancel the whole workflow runtime, including all effects.
+      backgroundScope.coroutineContext.job.cancelAndJoin()
+      expect(10)
+    }
+  }
+
+  private suspend fun awaitCancellation(onFinally: () -> Unit) {
+    try {
+      awaitCancellation()
+    } finally {
+      onFinally()
+    }
+  }
+
+  private var expectCounter = 0
+  private fun expect(expected: Int) {
+    assertEquals(expected, expectCounter)
+    expectCounter++
+  }
+
   private class ExpectedException : RuntimeException()
 
   companion object {