Add Dispatch Tracking

.github/workflows/publish-snapshot.yml

@@ -10,7 +10,7 @@ jobs:
   publish-snapshot:
     runs-on: macos-latest
     if: github.repository == 'square/workflow-kotlin'
-    timeout-minutes: 45
+    timeout-minutes: 90
 
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4

workflow-runtime/api/workflow-runtime.api

@@ -123,6 +123,7 @@ public abstract interface class com/squareup/workflow1/WorkflowInterceptor$Workf
 	public abstract fun getParent ()Lcom/squareup/workflow1/WorkflowInterceptor$WorkflowSession;
 	public abstract fun getRenderKey ()Ljava/lang/String;
 	public abstract fun getRuntimeConfig ()Ljava/util/Set;
+	public abstract fun getRuntimeContext ()Lkotlin/coroutines/CoroutineContext;
 	public abstract fun getSessionId ()J
 	public abstract fun getWorkflowTracer ()Lcom/squareup/workflow1/WorkflowTracer;
 	public abstract fun isRootWorkflow ()Z

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

@@ -223,6 +223,9 @@ public interface WorkflowInterceptor {
     /** The [RuntimeConfig] of the runtime this session is executing in. */
     public val runtimeConfig: RuntimeConfig
 
+    /** The [CoroutineContext] of the runtime this session is executing in. */
+    public val runtimeContext: CoroutineContext
+
     /** The optional [WorkflowTracer] of the runtime this session is executing in. */
     public val workflowTracer: WorkflowTracer?
   }

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

@@ -71,6 +71,9 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
    */
   override val coroutineContext = baseContext + Job(baseContext[Job]) + CoroutineName(id.toString())
 
+  override val runtimeContext: CoroutineContext
+    get() = coroutineContext
+
   // WorkflowInstance properties
   override val identifier: WorkflowIdentifier get() = id.identifier
   override val renderKey: String get() = id.name

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

@@ -3,6 +3,7 @@ package com.squareup.workflow1
 import com.squareup.workflow1.WorkflowInterceptor.WorkflowSession
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.cancel
+import kotlin.coroutines.CoroutineContext
 import kotlin.coroutines.EmptyCoroutineContext
 import kotlin.reflect.KType
 import kotlin.reflect.typeOf
@@ -90,6 +91,7 @@ internal class SimpleLoggingWorkflowInterceptorTest {
     override val parent: WorkflowSession? get() = null
     override val runtimeConfig: RuntimeConfig = RuntimeConfigOptions.DEFAULT_CONFIG
     override val workflowTracer: WorkflowTracer? = null
+    override val runtimeContext: CoroutineContext = EmptyCoroutineContext
   }
 
   private object FakeRenderContext : BaseRenderContext<Unit, Unit, Nothing> {

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

@@ -183,6 +183,7 @@ internal class WorkflowInterceptorTest {
       override val parent: WorkflowSession? = null
       override val runtimeConfig: RuntimeConfig = RuntimeConfigOptions.DEFAULT_CONFIG
       override val workflowTracer: WorkflowTracer? = null
+      override val runtimeContext: CoroutineContext = EmptyCoroutineContext
     }
 
   private object TestWorkflow : StatefulWorkflow<String, String, String, String>() {

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

@@ -24,6 +24,7 @@ import kotlinx.coroutines.Job
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.test.advanceUntilIdle
 import kotlinx.coroutines.test.runTest
+import kotlin.coroutines.CoroutineContext
 import kotlin.coroutines.EmptyCoroutineContext
 import kotlin.reflect.KType
 import kotlin.test.Test
@@ -359,5 +360,6 @@ internal class ChainedWorkflowInterceptorTest {
     override val parent: WorkflowSession? = null
     override val runtimeConfig: RuntimeConfig = RuntimeConfigOptions.DEFAULT_CONFIG
     override val workflowTracer: WorkflowTracer? = null
+    override val runtimeContext: CoroutineContext = EmptyCoroutineContext
   }
 }

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

@@ -42,6 +42,7 @@ import kotlinx.coroutines.suspendCancellableCoroutine
 import kotlinx.coroutines.test.runTest
 import kotlinx.coroutines.withTimeout
 import kotlin.coroutines.CoroutineContext
+import kotlin.coroutines.EmptyCoroutineContext
 import kotlin.reflect.typeOf
 import kotlin.test.AfterTest
 import kotlin.test.Test
@@ -1418,5 +1419,6 @@ internal class WorkflowNodeTest {
     override val parent: WorkflowSession? = null
     override val runtimeConfig: RuntimeConfig = RuntimeConfigOptions.DEFAULT_CONFIG
     override val workflowTracer: WorkflowTracer? = null
+    override val runtimeContext: CoroutineContext = EmptyCoroutineContext
   }
 }

workflow-tracing-papa/src/test/java/com/squareup/workflow1/tracing/papa/WorkflowPapaTracerTest.kt

@@ -16,6 +16,8 @@ import com.squareup.workflow1.tracing.RuntimeTraceContext
 import com.squareup.workflow1.tracing.RuntimeUpdateLogLine
 import com.squareup.workflow1.tracing.WorkflowSessionInfo
 import kotlinx.coroutines.test.TestScope
+import kotlin.coroutines.CoroutineContext
+import kotlin.coroutines.EmptyCoroutineContext
 import kotlin.test.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertNotNull
@@ -246,7 +248,8 @@ internal class WorkflowPapaTracerTest {
     private val workflow: TestWorkflow,
     override val sessionId: Long,
     override val renderKey: String,
-    override val parent: WorkflowSession?
+    override val parent: WorkflowSession?,
+    override val runtimeContext: CoroutineContext = EmptyCoroutineContext
   ) : WorkflowSession {
     override val identifier = workflow.identifier
     override val runtimeConfig = TestRuntimeConfig()

workflow-tracing/api/workflow-tracing.api

@@ -25,8 +25,10 @@ public final class com/squareup/workflow1/tracing/ChainedWorkflowRuntimeTracerKt
 }
 
 public final class com/squareup/workflow1/tracing/ConfigSnapshot {
-	public fun <init> (Ljava/util/Set;)V
+	public fun <init> (Ljava/util/Set;Lkotlinx/coroutines/CoroutineDispatcher;)V
+	public synthetic fun <init> (Ljava/util/Set;Lkotlinx/coroutines/CoroutineDispatcher;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public final fun getConfigAsString ()Ljava/lang/String;
+	public final fun getRuntimeDispatch ()Lkotlinx/coroutines/CoroutineDispatcher;
 	public final fun getShortConfigAsString ()Ljava/lang/String;
 }
 

workflow-tracing/src/main/java/com/squareup/workflow1/tracing/ConfigSnapshot.kt

@@ -8,13 +8,18 @@ import com.squareup.workflow1.RuntimeConfigOptions.RENDER_ONLY_WHEN_STATE_CHANGE
 import com.squareup.workflow1.RuntimeConfigOptions.STABLE_EVENT_HANDLERS
 import com.squareup.workflow1.RuntimeConfigOptions.WORK_STEALING_DISPATCHER
 import com.squareup.workflow1.WorkflowExperimentalRuntime
+import kotlinx.coroutines.CoroutineDispatcher
 
 /**
  * Snapshot of the current [RuntimeConfig]
  */
 @OptIn(WorkflowExperimentalRuntime::class)
-public class ConfigSnapshot(config: RuntimeConfig) {
-  public val configAsString: String = config.toString()
+public class ConfigSnapshot(
+  config: RuntimeConfig,
+  public val runtimeDispatch: CoroutineDispatcher? = null
+) {
+
+  public val configAsString: String = "$config, $runtimeDispatch"
 
   public val shortConfigAsString: String by lazy {
     buildString {
@@ -40,6 +45,7 @@ public class ConfigSnapshot(config: RuntimeConfig) {
       if (config.isEmpty()) {
         append("Base, ")
       }
+      append("Dispatch: ${runtimeDispatch?.toString()?.take(6)}")
     }
   }
 }

workflow-tracing/src/main/java/com/squareup/workflow1/tracing/WorkflowRuntimeMonitor.kt

@@ -29,6 +29,7 @@ import com.squareup.workflow1.tracing.RenderCause.RootPropsChanged
 import com.squareup.workflow1.tracing.RenderCause.WaitingForOutput
 import com.squareup.workflow1.tracing.WorkflowRuntimeMonitor.ActionType.CascadeAction
 import com.squareup.workflow1.tracing.WorkflowRuntimeMonitor.ActionType.QueuedAction
+import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Job
 import kotlin.time.Duration.Companion.nanoseconds
@@ -105,7 +106,7 @@ public class WorkflowRuntimeMonitor(
 
     if (session.isRootWorkflow) {
       // Cache the config snapshot for this whole runtime.
-      configSnapshot = ConfigSnapshot(session.runtimeConfig)
+      configSnapshot = ConfigSnapshot(session.runtimeConfig, session.runtimeContext[CoroutineDispatcher])
       check(renderIncomingCauses.isEmpty()) {
         "Workflow runtime for $runtimeName already has incoming render on creation triggered by " +
           "${renderIncomingCauses.lastOrNull()}"

workflow-tracing/src/test/java/com/squareup/workflow1/tracing/WorkflowRuntimeMonitorTest.kt

@@ -21,7 +21,11 @@ import com.squareup.workflow1.identifier
 import com.squareup.workflow1.tracing.RenderCause.RootCreation
 import com.squareup.workflow1.tracing.RenderCause.RootPropsChanged
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlin.coroutines.CoroutineContext
+import kotlin.coroutines.EmptyCoroutineContext
 import kotlin.reflect.KType
 import kotlin.test.Test
 import kotlin.test.assertContains
@@ -74,6 +78,23 @@ internal class WorkflowRuntimeMonitorTest {
     assertTrue(monitor.renderIncomingCauses.first() is RootCreation)
   }
 
+  @OptIn(ExperimentalCoroutinesApi::class)
+  @Test
+  fun `onSessionStarted captures runtime context`() {
+    val monitor = WorkflowRuntimeMonitor(
+      runtimeName = runtimeName,
+      workflowRuntimeTracers = listOf(fakeRuntimeTracer)
+    )
+    val testWorkflow = TestWorkflow()
+    val testCoroutineDispatcher = UnconfinedTestDispatcher()
+    val rootSession = testWorkflow.createRootSession(testCoroutineDispatcher)
+    val testScope = TestScope(testCoroutineDispatcher)
+
+    monitor.onSessionStarted(testScope, rootSession)
+
+    assertEquals(testCoroutineDispatcher, monitor.configSnapshot.runtimeDispatch)
+  }
+
   @Test
   fun `onSessionStarted handles child workflow session`() {
     val monitor = WorkflowRuntimeMonitor(
@@ -647,26 +668,33 @@ internal class WorkflowRuntimeMonitorTest {
 
     override fun snapshotState(state: String): Snapshot = Snapshot.of(state)
 
-    fun createRootSession(): WorkflowSession = TestWorkflowSession(
-      workflow = this,
-      sessionId = 1L,
-      renderKey = "root",
-      parent = null
-    )
-
-    fun createChildSession(parent: WorkflowSession): WorkflowSession = TestWorkflowSession(
+    fun createRootSession(context: CoroutineContext = EmptyCoroutineContext): WorkflowSession =
+      TestWorkflowSession(
+        workflow = this,
+        sessionId = 1L,
+        renderKey = "root",
+        parent = null,
+        runtimeContext = context
+      )
+
+    fun createChildSession(
+      parent: WorkflowSession,
+      context: CoroutineContext = EmptyCoroutineContext
+    ): WorkflowSession = TestWorkflowSession(
       workflow = this,
       sessionId = 2L,
       renderKey = "child",
-      parent = parent
+      parent = parent,
+      runtimeContext = context
     )
   }
 
   private class TestWorkflowSession(
     private val workflow: TestWorkflow,
     override val sessionId: Long,
     override val renderKey: String,
-    override val parent: WorkflowSession?
+    override val parent: WorkflowSession?,
+    override val runtimeContext: CoroutineContext = EmptyCoroutineContext
   ) : WorkflowSession {
     override val identifier = workflow.identifier
     override val runtimeConfig = TestRuntimeConfig()

workflow-ui/core-android/src/main/java/com/squareup/workflow1/ui/WorkflowLayout.kt

@@ -2,6 +2,7 @@ package com.squareup.workflow1.ui
 
 import android.content.Context
 import android.os.Build.VERSION
+import android.os.Looper
 import android.os.Parcel
 import android.os.Parcelable
 import android.os.Parcelable.Creator
@@ -16,10 +17,10 @@ import androidx.lifecycle.coroutineScope
 import androidx.lifecycle.repeatOnLifecycle
 import com.squareup.workflow1.ui.androidx.OnBackPressedDispatcherOwnerKey
 import com.squareup.workflow1.ui.androidx.WorkflowAndroidXSupport.onBackPressedDispatcherOwnerOrNull
-import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
 import kotlin.coroutines.CoroutineContext
 import kotlin.coroutines.EmptyCoroutineContext
 
@@ -89,8 +90,8 @@ public class WorkflowLayout(
    * @param [repeatOnLifecycle] the lifecycle state in which renderings should be actively
    * updated. Defaults to STARTED, which is appropriate for Activity and Fragment.
    * @param [collectionContext] additional [CoroutineContext] we want for the coroutine that is
-   * launched to collect the renderings. This should not override the [CoroutineDispatcher][kotlinx.coroutines.CoroutineDispatcher]
-   * but may include some other instrumentation elements.
+   * launched to collect the renderings, can include a different dispatcher - but it should be
+   * a main thread dispatcher!
    *
    * @return the [Job] started to collect [renderings], to give callers the option to
    * [cancel][Job.cancel] collection -- e.g., before calling [take] again with a new
@@ -104,16 +105,15 @@ public class WorkflowLayout(
     repeatOnLifecycle: State = STARTED,
     collectionContext: CoroutineContext = EmptyCoroutineContext
   ): Job {
-    // We remove the dispatcher as we want to use what is provided by the lifecycle.coroutineScope.
-    val contextWithoutDispatcher = collectionContext.minusKey(CoroutineDispatcher.Key)
-    val lifecycleDispatcher = lifecycle.coroutineScope.coroutineContext[CoroutineDispatcher.Key]
     // Just like https://medium.com/androiddevelopers/a-safer-way-to-collect-flows-from-android-uis-23080b1f8bda
-    return lifecycle.coroutineScope.launch(contextWithoutDispatcher) {
+    return lifecycle.coroutineScope.launch {
       lifecycle.repeatOnLifecycle(repeatOnLifecycle) {
-        require(coroutineContext[CoroutineDispatcher.Key] == lifecycleDispatcher) {
-          "Collection dispatch should happen on the lifecycle's dispatcher."
+        withContext(collectionContext) {
+          require(Looper.myLooper() == Looper.getMainLooper()) {
+            "Collection dispatch should happen on the main thread!"
+          }
+          renderings.collect { show(it) }
         }
-        renderings.collect { show(it) }
       }
     }
   }

workflow-ui/core-android/src/test/java/com/squareup/workflow1/ui/WorkflowLayoutTest.kt

@@ -17,14 +17,12 @@ import com.squareup.workflow1.ui.androidx.OnBackPressedDispatcherOwnerKey
 import com.squareup.workflow1.ui.navigation.WrappedScreen
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.MutableSharedFlow
-import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.robolectric.RobolectricTestRunner
 import org.robolectric.annotation.Config
-import kotlin.coroutines.CoroutineContext
 
 @RunWith(RobolectricTestRunner::class)
 // SDK 28 required for the four-arg constructor we use in our custom view classes.
@@ -92,23 +90,6 @@ internal class WorkflowLayoutTest {
     unoriginal.show(BScreen(), env)
   }
 
-  @Test fun usesLifecycleDispatcher() {
-    val lifecycleDispatcher = UnconfinedTestDispatcher()
-    val collectionContext: CoroutineContext = UnconfinedTestDispatcher()
-    val testLifecycle = TestLifecycleOwner(
-      Lifecycle.State.RESUMED,
-      lifecycleDispatcher
-    )
-
-    workflowLayout.take(
-      lifecycle = testLifecycle.lifecycle,
-      renderings = flowOf(WrappedScreen(), WrappedScreen()),
-      collectionContext = collectionContext
-    )
-
-    // No crash then we safely removed the dispatcher.
-  }
-
   @Test fun takes() {
     val lifecycleDispatcher = UnconfinedTestDispatcher()
     val testLifecycle = TestLifecycleOwner(