Add Dispatch Tracking

steve-the-edwards · steve-the-edwards · commit ee8b977b5efe · 2025-10-14T14:14:45.000-04:00
diff --git a/workflow-runtime/api/workflow-runtime.api b/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
diff --git a/workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/WorkflowInterceptor.kt b/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?
   }
diff --git a/workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/WorkflowNode.kt b/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
diff --git a/workflow-runtime/src/commonTest/kotlin/com/squareup/workflow1/SimpleLoggingWorkflowInterceptorTest.kt b/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> {
diff --git a/workflow-runtime/src/commonTest/kotlin/com/squareup/workflow1/WorkflowInterceptorTest.kt b/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>() {
diff --git a/workflow-runtime/src/commonTest/kotlin/com/squareup/workflow1/internal/ChainedWorkflowInterceptorTest.kt b/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
   }
 }
diff --git a/workflow-runtime/src/commonTest/kotlin/com/squareup/workflow1/internal/WorkflowNodeTest.kt b/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
   }
 }
diff --git a/workflow-tracing-papa/src/test/java/com/squareup/workflow1/tracing/papa/WorkflowPapaTracerTest.kt b/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()
diff --git a/workflow-tracing/api/workflow-tracing.api b/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;
 }
 
diff --git a/workflow-tracing/src/main/java/com/squareup/workflow1/tracing/ConfigSnapshot.kt b/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)}")
     }
   }
 }
diff --git a/workflow-tracing/src/main/java/com/squareup/workflow1/tracing/WorkflowRuntimeMonitor.kt b/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()}"
diff --git a/workflow-tracing/src/test/java/com/squareup/workflow1/tracing/WorkflowRuntimeMonitorTest.kt b/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()

Original file line number	Diff line number	Diff line change
`@@ -223,6 +223,9 @@ public interface WorkflowInterceptor {`
`223`	`223`	`/** The [RuntimeConfig] of the runtime this session is executing in. */`
`224`	`224`	`public val runtimeConfig: RuntimeConfig`
`225`	`225`
	`226`	`+ /** The [CoroutineContext] of the runtime this session is executing in. */`
	`227`	`+ public val runtimeContext: CoroutineContext`
	`228`	`+`
`226`	`229`	`/** The optional [WorkflowTracer] of the runtime this session is executing in. */`
`227`	`230`	`public val workflowTracer: WorkflowTracer?`
`228`	`231`	`}`
Original file line number	Diff line number	Diff line change
`@@ -183,6 +183,7 @@ internal class WorkflowInterceptorTest {`
`183`	`183`	`override val parent: WorkflowSession? = null`
`184`	`184`	`override val runtimeConfig: RuntimeConfig = RuntimeConfigOptions.DEFAULT_CONFIG`
`185`	`185`	`override val workflowTracer: WorkflowTracer? = null`
	`186`	`+ override val runtimeContext: CoroutineContext = EmptyCoroutineContext`
`186`	`187`	`}`
`187`	`188`
`188`	`189`	`private object TestWorkflow : StatefulWorkflow<String, String, String, String>() {`
Original file line number	Diff line number	Diff line change
`@@ -24,6 +24,7 @@ import kotlinx.coroutines.Job`
`24`	`24`	`import kotlinx.coroutines.launch`
`25`	`25`	`import kotlinx.coroutines.test.advanceUntilIdle`
`26`	`26`	`import kotlinx.coroutines.test.runTest`
	`27`	`+import kotlin.coroutines.CoroutineContext`
`27`	`28`	`import kotlin.coroutines.EmptyCoroutineContext`
`28`	`29`	`import kotlin.reflect.KType`
`29`	`30`	`import kotlin.test.Test`
`@@ -359,5 +360,6 @@ internal class ChainedWorkflowInterceptorTest {`
`359`	`360`	`override val parent: WorkflowSession? = null`
`360`	`361`	`override val runtimeConfig: RuntimeConfig = RuntimeConfigOptions.DEFAULT_CONFIG`
`361`	`362`	`override val workflowTracer: WorkflowTracer? = null`
	`363`	`+ override val runtimeContext: CoroutineContext = EmptyCoroutineContext`
`362`	`364`	`}`
`363`	`365`	`}`
Original file line number	Diff line number	Diff line change
`@@ -42,6 +42,7 @@ import kotlinx.coroutines.suspendCancellableCoroutine`
`42`	`42`	`import kotlinx.coroutines.test.runTest`
`43`	`43`	`import kotlinx.coroutines.withTimeout`
`44`	`44`	`import kotlin.coroutines.CoroutineContext`
	`45`	`+import kotlin.coroutines.EmptyCoroutineContext`
`45`	`46`	`import kotlin.reflect.typeOf`
`46`	`47`	`import kotlin.test.AfterTest`
`47`	`48`	`import kotlin.test.Test`
`@@ -1418,5 +1419,6 @@ internal class WorkflowNodeTest {`
`1418`	`1419`	`override val parent: WorkflowSession? = null`
`1419`	`1420`	`override val runtimeConfig: RuntimeConfig = RuntimeConfigOptions.DEFAULT_CONFIG`
`1420`	`1421`	`override val workflowTracer: WorkflowTracer? = null`
	`1422`	`+ override val runtimeContext: CoroutineContext = EmptyCoroutineContext`
`1421`	`1423`	`}`
`1422`	`1424`	`}`
Original file line number	Diff line number	Diff line change
`@@ -25,8 +25,10 @@ public final class com/squareup/workflow1/tracing/ChainedWorkflowRuntimeTracerKt`
`25`	`25`	`}`
`26`	`26`
`27`	`27`	`public final class com/squareup/workflow1/tracing/ConfigSnapshot {`
`28`		`- public fun <init> (Ljava/util/Set;)V`
	`28`	`+ public fun <init> (Ljava/util/Set;Lkotlinx/coroutines/CoroutineDispatcher;)V`
	`29`	`+ public synthetic fun <init> (Ljava/util/Set;Lkotlinx/coroutines/CoroutineDispatcher;ILkotlin/jvm/internal/DefaultConstructorMarker;)V`
`29`	`30`	`public final fun getConfigAsString ()Ljava/lang/String;`
	`31`	`+ public final fun getRuntimeDispatch ()Lkotlinx/coroutines/CoroutineDispatcher;`
`30`	`32`	`public final fun getShortConfigAsString ()Ljava/lang/String;`
`31`	`33`	`}`
`32`	`34`