Add Dispatch Tracking

Author: steve-the-edwards

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/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()