Use WorkStealingDispatcher in runtime, behind a flag.

Author: zach-klippenstein

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,7 @@ 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,13 @@ 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 +211,9 @@ 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?.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,6 +6,7 @@ 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
@@ -1494,7 +1495,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))