Add Tests for WorkflowTracer

square · Dec 10, 2024 · 4a0b76f · 4a0b76f
1 parent b2e9f1e
commit 4a0b76f
Show file tree

Hide file tree

Showing 7 changed files with 201 additions and 100 deletions.
diff --git a/workflow-core/api/workflow-core.api b/workflow-core/api/workflow-core.api
@@ -317,7 +317,6 @@ public abstract interface class com/squareup/workflow1/WorkflowTracer {
 
 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 {

diff --git a/workflow-core/src/commonMain/kotlin/com/squareup/workflow1/WorkerWorkflow.kt b/workflow-core/src/commonMain/kotlin/com/squareup/workflow1/WorkerWorkflow.kt
@@ -36,7 +36,7 @@ internal class WorkerWorkflow<OutputT>(
   ImpostorWorkflow {
 
   override val realIdentifier: WorkflowIdentifier =
-    workflowTracer.trace("ComputeRealIdentifier" ) {
+    workflowTracer.trace("ComputeRealIdentifier") {
       unsnapshottableIdentifier(workerType)
     }
 

diff --git a/workflow-core/src/commonMain/kotlin/com/squareup/workflow1/WorkflowTracer.kt b/workflow-core/src/commonMain/kotlin/com/squareup/workflow1/WorkflowTracer.kt
@@ -10,27 +10,22 @@ public interface WorkflowTracer {
 }
 
 /**
- * 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.
+ * Convenience function to wrap [block] with a trace span as defined by [WorkflowTracer]. This
+ * wraps very frequently evaluated code and we should only use constants for [label], with no
+ * interpolation.
  */
-public inline fun <T> WorkflowTracer?.trace(label: () -> String, block: () -> T): T {
-  val optimizedLabel = if (this !== null) {
-    label()
+public inline fun <T> WorkflowTracer?.trace(
+  label: String,
+  block: () -> T
+): T {
+  return if (this == null) {
+    block()
   } 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()
+    beginSection(label)
+    try {
+      return block()
+    } finally {
+      endSection()
+    }
   }
 }
diff --git a/workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/RenderWorkflow.kt b/workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/RenderWorkflow.kt
@@ -115,16 +115,15 @@ public fun <PropsT, OutputT, RenderingT> renderWorkflowIn(
 ): StateFlow<RenderingAndSnapshot<RenderingT>> {
   val chainedInterceptor = interceptors.chained()
 
-  val runner =
-    WorkflowRunner(
-      scope,
-      workflow,
-      props,
-      initialSnapshot,
-      chainedInterceptor,
-      runtimeConfig,
-      workflowTracer
-    )
+  val runner = 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.

diff --git a/workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/SubtreeManager.kt b/workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/SubtreeManager.kt
@@ -126,7 +126,7 @@ internal class SubtreeManager<PropsT, StateT, OutputT>(
     // Prevent duplicate workflows with the same key.
     workflowTracer.trace("CheckingUniqueMatches") {
       children.forEachStaging {
-        require(!(it.matches(child, key))) {
+        require(!(it.matches(child, key, workflowTracer))) {
           "Expected keys to be unique for ${child.identifier}: key=\"$key\""
         }
       }
@@ -136,7 +136,7 @@ internal class SubtreeManager<PropsT, StateT, OutputT>(
     val stagedChild =
       workflowTracer.trace("RetainingChildren") {
         children.retainOrCreate(
-          predicate = { it.matches(child, key) },
+          predicate = { it.matches(child, key, workflowTracer) },
           create = { createChildNode(child, props, key, handler) }
         )
       }

diff --git a/workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/WorkflowChildNode.kt b/workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/WorkflowChildNode.kt
@@ -3,7 +3,9 @@ package com.squareup.workflow1.internal
 import com.squareup.workflow1.StatefulWorkflow
 import com.squareup.workflow1.Workflow
 import com.squareup.workflow1.WorkflowAction
+import com.squareup.workflow1.WorkflowTracer
 import com.squareup.workflow1.internal.InlineLinkedList.InlineListNode
+import com.squareup.workflow1.trace
 
 /**
  * Representation of a child workflow that has been rendered by another workflow.
@@ -32,8 +34,9 @@ internal class WorkflowChildNode<
    */
   fun matches(
     otherWorkflow: Workflow<*, *, *>,
-    key: String
-  ): Boolean = id.matches(otherWorkflow, key)
+    key: String,
+    workflowTracer: WorkflowTracer?
+  ): Boolean = workflowTracer.trace("matches") { id.matches(otherWorkflow, key) }
 
   /**
    * Updates the handler function that will be invoked by [acceptChildOutput].