Add WorkflowTracer Through some Runtime Internals

workflow-core/api/workflow-core.api

@@ -35,6 +35,7 @@ public abstract interface class com/squareup/workflow1/BaseRenderContext {
 	public abstract fun eventHandler (Ljava/lang/String;Lkotlin/jvm/functions/Function8;)Lkotlin/jvm/functions/Function7;
 	public abstract fun eventHandler (Ljava/lang/String;Lkotlin/jvm/functions/Function9;)Lkotlin/jvm/functions/Function8;
 	public abstract fun getActionSink ()Lcom/squareup/workflow1/Sink;
+	public abstract fun getWorkflowTracer ()Lcom/squareup/workflow1/WorkflowTracer;
 	public abstract fun renderChild (Lcom/squareup/workflow1/Workflow;Ljava/lang/Object;Ljava/lang/String;Lkotlin/jvm/functions/Function1;)Ljava/lang/Object;
 	public abstract fun runningSideEffect (Ljava/lang/String;Lkotlin/jvm/functions/Function2;)V
 }
@@ -167,6 +168,7 @@ public final class com/squareup/workflow1/StatefulWorkflow$RenderContext : com/s
 	public fun eventHandler (Ljava/lang/String;Lkotlin/jvm/functions/Function8;)Lkotlin/jvm/functions/Function7;
 	public fun eventHandler (Ljava/lang/String;Lkotlin/jvm/functions/Function9;)Lkotlin/jvm/functions/Function8;
 	public fun getActionSink ()Lcom/squareup/workflow1/Sink;
+	public fun getWorkflowTracer ()Lcom/squareup/workflow1/WorkflowTracer;
 	public fun renderChild (Lcom/squareup/workflow1/Workflow;Ljava/lang/Object;Ljava/lang/String;Lkotlin/jvm/functions/Function1;)Ljava/lang/Object;
 	public fun runningSideEffect (Ljava/lang/String;Lkotlin/jvm/functions/Function2;)V
 }
@@ -192,6 +194,7 @@ public final class com/squareup/workflow1/StatelessWorkflow$RenderContext : com/
 	public fun eventHandler (Ljava/lang/String;Lkotlin/jvm/functions/Function8;)Lkotlin/jvm/functions/Function7;
 	public fun eventHandler (Ljava/lang/String;Lkotlin/jvm/functions/Function9;)Lkotlin/jvm/functions/Function8;
 	public fun getActionSink ()Lcom/squareup/workflow1/Sink;
+	public fun getWorkflowTracer ()Lcom/squareup/workflow1/WorkflowTracer;
 	public fun renderChild (Lcom/squareup/workflow1/Workflow;Ljava/lang/Object;Ljava/lang/String;Lkotlin/jvm/functions/Function1;)Ljava/lang/Object;
 	public fun runningSideEffect (Ljava/lang/String;Lkotlin/jvm/functions/Function2;)V
 }
@@ -307,6 +310,16 @@ public final class com/squareup/workflow1/WorkflowOutput {
 	public fun toString ()Ljava/lang/String;
 }
 
+public abstract interface class com/squareup/workflow1/WorkflowTracer {
+	public abstract fun beginSection (Ljava/lang/String;)V
+	public abstract fun endSection ()V
+}
+
+public final class com/squareup/workflow1/WorkflowTracerKt {
+	public static final fun trace (Lcom/squareup/workflow1/WorkflowTracer;Ljava/lang/String;Lkotlin/jvm/functions/Function0;)Ljava/lang/Object;
+	public static final fun trace (Lcom/squareup/workflow1/WorkflowTracer;Lkotlin/jvm/functions/Function0;Lkotlin/jvm/functions/Function0;)Ljava/lang/Object;
+}
+
 public final class com/squareup/workflow1/Workflows {
 	public static final fun RenderContext (Lcom/squareup/workflow1/BaseRenderContext;Lcom/squareup/workflow1/StatefulWorkflow;)Lcom/squareup/workflow1/StatefulWorkflow$RenderContext;
 	public static final fun RenderContext (Lcom/squareup/workflow1/BaseRenderContext;Lcom/squareup/workflow1/StatelessWorkflow;)Lcom/squareup/workflow1/StatelessWorkflow$RenderContext;

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

@@ -55,6 +55,8 @@ public interface BaseRenderContext<out PropsT, StateT, in OutputT> {
    */
   public val actionSink: Sink<WorkflowAction<PropsT, StateT, OutputT>>
 
+  public val workflowTracer: WorkflowTracer?
+
   /**
    * Ensures [child] is running as a child of this workflow, and returns the result of its
    * `render` method.
@@ -437,6 +439,8 @@ internal fun <T, PropsT, StateT, OutputT>
     key: String = "",
     handler: (T) -> WorkflowAction<PropsT, StateT, OutputT>
   ) {
-  val workerWorkflow = WorkerWorkflow<T>(workerType, key)
+  val workerWorkflow = workflowTracer.trace("CreateWorkerWorkflow") {
+    WorkerWorkflow<T>(workerType, key)
+  }
   renderChild(workerWorkflow, props = worker, key = key, handler = handler)
 }

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

@@ -30,11 +30,16 @@ import kotlin.reflect.KType
  */
 internal class WorkerWorkflow<OutputT>(
   val workerType: KType,
-  private val key: String
+  private val key: String,
+  workflowTracer: WorkflowTracer? = null
 ) : StatefulWorkflow<Worker<OutputT>, Int, OutputT, Unit>(),
   ImpostorWorkflow {
 
-  override val realIdentifier: WorkflowIdentifier = unsnapshottableIdentifier(workerType)
+  override val realIdentifier: WorkflowIdentifier =
+    workflowTracer.trace("ComputeRealIdentifier" ) {
+      unsnapshottableIdentifier(workerType)
+    }
+
   override fun describeRealIdentifier(): String = workerType.toString()
 
   override fun initialState(

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

@@ -0,0 +1,36 @@
+package com.squareup.workflow1
+
+/**
+ * This is a very simple tracing interface that can be passed into a workflow runtime in order
+ * to inject span tracing throughout the workflow core and runtime internals.
+ */
+public interface WorkflowTracer {
+  public fun beginSection(label: String): Unit
+  public fun endSection(): Unit
+}
+
+/**
+ * Convenience function to wrap [block] with a trace span as defined by [WorkflowTracer].
+ * Only calls [label] if there is an active [WorkflowTracer] use this for any label other than
+ * a constant.
+ */
+public inline fun <T> WorkflowTracer?.trace(label: () -> String, block: () -> T): T {
+  val optimizedLabel = if (this !== null) {
+    label()
+  } else {
+    ""
+  }
+  return trace(optimizedLabel, block)
+}
+
+/**
+ * Convenience function to wrap [block] with a trace span as defined by [WorkflowTracer].
+ */
+public inline fun <T> WorkflowTracer?.trace(label: String, block: () -> T): T {
+  this?.beginSection(label)
+  try {
+    return block()
+  } finally {
+    this?.endSection()
+  }
+}

workflow-runtime/api/workflow-runtime.api

@@ -10,8 +10,8 @@ public final class com/squareup/workflow1/NoopWorkflowInterceptor : com/squareup
 }
 
 public final class com/squareup/workflow1/RenderWorkflowKt {
-	public static final fun renderWorkflowIn (Lcom/squareup/workflow1/Workflow;Lkotlinx/coroutines/CoroutineScope;Lkotlinx/coroutines/flow/StateFlow;Lcom/squareup/workflow1/TreeSnapshot;Ljava/util/List;Ljava/util/Set;Lkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/flow/StateFlow;
-	public static synthetic fun renderWorkflowIn$default (Lcom/squareup/workflow1/Workflow;Lkotlinx/coroutines/CoroutineScope;Lkotlinx/coroutines/flow/StateFlow;Lcom/squareup/workflow1/TreeSnapshot;Ljava/util/List;Ljava/util/Set;Lkotlin/jvm/functions/Function2;ILjava/lang/Object;)Lkotlinx/coroutines/flow/StateFlow;
+	public static final fun renderWorkflowIn (Lcom/squareup/workflow1/Workflow;Lkotlinx/coroutines/CoroutineScope;Lkotlinx/coroutines/flow/StateFlow;Lcom/squareup/workflow1/TreeSnapshot;Ljava/util/List;Ljava/util/Set;Lcom/squareup/workflow1/WorkflowTracer;Lkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/flow/StateFlow;
+	public static synthetic fun renderWorkflowIn$default (Lcom/squareup/workflow1/Workflow;Lkotlinx/coroutines/CoroutineScope;Lkotlinx/coroutines/flow/StateFlow;Lcom/squareup/workflow1/TreeSnapshot;Ljava/util/List;Ljava/util/Set;Lcom/squareup/workflow1/WorkflowTracer;Lkotlin/jvm/functions/Function2;ILjava/lang/Object;)Lkotlinx/coroutines/flow/StateFlow;
 }
 
 public final class com/squareup/workflow1/RenderingAndSnapshot {
@@ -104,6 +104,7 @@ public abstract interface class com/squareup/workflow1/WorkflowInterceptor$Workf
 	public abstract fun getRenderKey ()Ljava/lang/String;
 	public abstract fun getRuntimeConfig ()Ljava/util/Set;
 	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/RenderWorkflow.kt

@@ -110,12 +110,21 @@ public fun <PropsT, OutputT, RenderingT> renderWorkflowIn(
   initialSnapshot: TreeSnapshot? = null,
   interceptors: List<WorkflowInterceptor> = emptyList(),
   runtimeConfig: RuntimeConfig = RuntimeConfigOptions.DEFAULT_CONFIG,
+  workflowTracer: WorkflowTracer? = null,
   onOutput: suspend (OutputT) -> Unit
 ): StateFlow<RenderingAndSnapshot<RenderingT>> {
   val chainedInterceptor = interceptors.chained()
 
   val runner =
-    WorkflowRunner(scope, workflow, props, initialSnapshot, chainedInterceptor, runtimeConfig)
+    WorkflowRunner(
+      scope,
+      workflow,
+      props,
+      initialSnapshot,
+      chainedInterceptor,
+      runtimeConfig,
+      workflowTracer
+    )
 
   // Rendering is synchronous, so we can run the first render pass before launching the runtime
   // coroutine to calculate the initial rendering.

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

@@ -160,6 +160,9 @@ public interface WorkflowInterceptor {
 
     /** The [RuntimeConfig] of the runtime this session is executing in. */
     public val runtimeConfig: RuntimeConfig
+
+    /** The optional [WorkflowTracer] of the runtime this session is executing in. */
+    public val workflowTracer: WorkflowTracer?
   }
 
   /**
@@ -314,6 +317,7 @@ private class InterceptedRenderContext<P, S, O>(
   private val interceptor: RenderContextInterceptor<P, S, O>
 ) : BaseRenderContext<P, S, O>, Sink<WorkflowAction<P, S, O>> {
   override val actionSink: Sink<WorkflowAction<P, S, O>> get() = this
+  override val workflowTracer: WorkflowTracer? = baseRenderContext.workflowTracer
 
   override fun send(value: WorkflowAction<P, S, O>) {
     interceptor.onActionSent(value) { interceptedAction ->

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

@@ -4,13 +4,15 @@ import com.squareup.workflow1.BaseRenderContext
 import com.squareup.workflow1.Sink
 import com.squareup.workflow1.Workflow
 import com.squareup.workflow1.WorkflowAction
+import com.squareup.workflow1.WorkflowTracer
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.channels.SendChannel
 
 internal class RealRenderContext<out PropsT, StateT, OutputT>(
   private val renderer: Renderer<PropsT, StateT, OutputT>,
   private val sideEffectRunner: SideEffectRunner,
-  private val eventActionsChannel: SendChannel<WorkflowAction<PropsT, StateT, OutputT>>
+  private val eventActionsChannel: SendChannel<WorkflowAction<PropsT, StateT, OutputT>>,
+  override val workflowTracer: WorkflowTracer?
 ) : BaseRenderContext<PropsT, StateT, OutputT>, Sink<WorkflowAction<PropsT, StateT, OutputT>> {
 
   interface Renderer<PropsT, StateT, OutputT> {

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

@@ -9,7 +9,9 @@ import com.squareup.workflow1.Workflow
 import com.squareup.workflow1.WorkflowAction
 import com.squareup.workflow1.WorkflowInterceptor
 import com.squareup.workflow1.WorkflowInterceptor.WorkflowSession
+import com.squareup.workflow1.WorkflowTracer
 import com.squareup.workflow1.identifier
+import com.squareup.workflow1.trace
 import kotlinx.coroutines.selects.SelectBuilder
 import kotlin.coroutines.CoroutineContext
 
@@ -91,6 +93,7 @@ internal class SubtreeManager<PropsT, StateT, OutputT>(
     childResult: ActionApplied<*>
   ) -> ActionProcessingResult,
   private val runtimeConfig: RuntimeConfig,
+  private val workflowTracer: WorkflowTracer?,
   private val workflowSession: WorkflowSession? = null,
   private val interceptor: WorkflowInterceptor = NoopWorkflowInterceptor,
   private val idCounter: IdCounter? = null
@@ -121,17 +124,22 @@ internal class SubtreeManager<PropsT, StateT, OutputT>(
     handler: (ChildOutputT) -> WorkflowAction<PropsT, StateT, OutputT>
   ): ChildRenderingT {
     // Prevent duplicate workflows with the same key.
-    children.forEachStaging {
-      require(!(it.matches(child, key))) {
-        "Expected keys to be unique for ${child.identifier}: key=\"$key\""
+    workflowTracer.trace("CheckingUniqueMatches") {
+      children.forEachStaging {
+        require(!(it.matches(child, key))) {
+          "Expected keys to be unique for ${child.identifier}: key=\"$key\""
+        }
       }
     }
 
     // Start tracking this case so we can be ready to render it.
-    val stagedChild = children.retainOrCreate(
-      predicate = { it.matches(child, key) },
-      create = { createChildNode(child, props, key, handler) }
-    )
+    val stagedChild =
+      workflowTracer.trace("RetainingChildren") {
+        children.retainOrCreate(
+          predicate = { it.matches(child, key) },
+          create = { createChildNode(child, props, key, handler) }
+        )
+      }
     stagedChild.setHandler(handler)
     return stagedChild.render(child.asStatefulWorkflow(), props)
   }
@@ -188,6 +196,7 @@ internal class SubtreeManager<PropsT, StateT, OutputT>(
       snapshot = childTreeSnapshots,
       baseContext = contextForChildren,
       runtimeConfig = runtimeConfig,
+      workflowTracer = workflowTracer,
       emitAppliedActionToParent = ::acceptChildActionResult,
       parent = workflowSession,
       interceptor = interceptor,

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

@@ -14,9 +14,11 @@ import com.squareup.workflow1.WorkflowExperimentalApi
 import com.squareup.workflow1.WorkflowIdentifier
 import com.squareup.workflow1.WorkflowInterceptor
 import com.squareup.workflow1.WorkflowInterceptor.WorkflowSession
+import com.squareup.workflow1.WorkflowTracer
 import com.squareup.workflow1.applyTo
 import com.squareup.workflow1.intercept
 import com.squareup.workflow1.internal.RealRenderContext.SideEffectRunner
+import com.squareup.workflow1.trace
 import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CoroutineName
 import kotlinx.coroutines.CoroutineScope
@@ -51,6 +53,7 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
   baseContext: CoroutineContext,
   // Providing default value so we don't need to specify in test.
   override val runtimeConfig: RuntimeConfig = RuntimeConfigOptions.DEFAULT_CONFIG,
+  override val workflowTracer: WorkflowTracer? = null,
   private val emitAppliedActionToParent: (ActionApplied<OutputT>) -> ActionProcessingResult =
     { it },
   override val parent: WorkflowSession? = null,
@@ -74,6 +77,7 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
     contextForChildren = coroutineContext,
     emitActionToParent = ::applyAction,
     runtimeConfig = runtimeConfig,
+    workflowTracer = workflowTracer,
     workflowSession = this,
     interceptor = interceptor,
     idCounter = idCounter
@@ -87,7 +91,8 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
   private val baseRenderContext = RealRenderContext(
     renderer = subtreeManager,
     sideEffectRunner = this,
-    eventActionsChannel = eventActionsChannel
+    eventActionsChannel = eventActionsChannel,
+    workflowTracer = workflowTracer,
   )
   private val context = RenderContext(baseRenderContext, workflow)
 
@@ -212,12 +217,14 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
       .render(props, state, context)
     baseRenderContext.freeze()
 
-    // Tear down workflows and workers that are obsolete.
-    subtreeManager.commitRenderedChildren()
-    // Side effect jobs are launched lazily, since they can send actions to the sink, and can only
-    // be started after context is frozen.
-    sideEffects.forEachStaging { it.job.start() }
-    sideEffects.commitStaging { it.job.cancel() }
+    workflowTracer.trace("UpdateRuntimeTree") {
+      // Tear down workflows and workers that are obsolete.
+      subtreeManager.commitRenderedChildren()
+      // Side effect jobs are launched lazily, since they can send actions to the sink, and can only
+      // be started after context is frozen.
+      sideEffects.forEachStaging { it.job.start() }
+      sideEffects.commitStaging { it.job.cancel() }
+    }
 
     return rendering
   }
@@ -261,8 +268,10 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
     key: String,
     sideEffect: suspend CoroutineScope.() -> Unit
   ): SideEffectNode {
-    val scope = this + CoroutineName("sideEffect[$key] for $id")
-    val job = scope.launch(start = LAZY, block = sideEffect)
-    return SideEffectNode(key, job)
+    return workflowTracer.trace("CreateSideEffectNode") {
+      val scope = this + CoroutineName("sideEffect[$key] for $id")
+      val job = scope.launch(start = LAZY, block = sideEffect)
+      SideEffectNode(key, job)
+    }
   }
 }

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

@@ -10,6 +10,7 @@ import com.squareup.workflow1.TreeSnapshot
 import com.squareup.workflow1.Workflow
 import com.squareup.workflow1.WorkflowExperimentalRuntime
 import com.squareup.workflow1.WorkflowInterceptor
+import com.squareup.workflow1.WorkflowTracer
 import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.DelicateCoroutinesApi
@@ -28,7 +29,8 @@ internal class WorkflowRunner<PropsT, OutputT, RenderingT>(
   props: StateFlow<PropsT>,
   snapshot: TreeSnapshot?,
   private val interceptor: WorkflowInterceptor,
-  private val runtimeConfig: RuntimeConfig
+  private val runtimeConfig: RuntimeConfig,
+  private val workflowTracer: WorkflowTracer?
 ) {
   private val workflow = protoWorkflow.asStatefulWorkflow()
   private val idCounter = IdCounter()
@@ -55,6 +57,7 @@ internal class WorkflowRunner<PropsT, OutputT, RenderingT>(
     snapshot = snapshot,
     baseContext = scope.coroutineContext,
     runtimeConfig = runtimeConfig,
+    workflowTracer = workflowTracer,
     interceptor = interceptor,
     idCounter = idCounter
   )