From 70f85dc0f4015859909cd071306fd78c6d7f3abd Mon Sep 17 00:00:00 2001
From: Zach Klippenstein <zachklipp@squareup.com>
Date: Mon, 10 Feb 2020 17:28:44 -0800
Subject: [PATCH] Plumb workerContext to child workflows.

PRs #940 and #943 have a bug where the worker context isn't passed from a parent workflow
to child `WorkflowNode`s. This fixes that.
---
 .../workflow/internal/SubtreeManager.kt       |  7 ++--
 .../workflow/internal/WorkflowNode.kt         |  2 +-
 .../WorkerCompositionIntegrationTest.kt       | 34 ++++++++++++++++++-
 3 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/kotlin/workflow-runtime/src/main/java/com/squareup/workflow/internal/SubtreeManager.kt b/kotlin/workflow-runtime/src/main/java/com/squareup/workflow/internal/SubtreeManager.kt
index f90835350..f068e572d 100644
--- a/kotlin/workflow-runtime/src/main/java/com/squareup/workflow/internal/SubtreeManager.kt
+++ b/kotlin/workflow-runtime/src/main/java/com/squareup/workflow/internal/SubtreeManager.kt
@@ -23,6 +23,7 @@ import com.squareup.workflow.diagnostic.WorkflowDiagnosticListener
 import kotlinx.coroutines.selects.SelectBuilder
 import okio.ByteString
 import kotlin.coroutines.CoroutineContext
+import kotlin.coroutines.EmptyCoroutineContext
 
 /**
  * Responsible for tracking child workflows, starting them and tearing them down when necessary.
@@ -95,7 +96,8 @@ internal class SubtreeManager<StateT, OutputT : Any>(
   private val emitActionToParent: (WorkflowAction<StateT, OutputT>) -> Any?,
   private val parentDiagnosticId: Long,
   private val diagnosticListener: WorkflowDiagnosticListener? = null,
-  private val idCounter: IdCounter? = null
+  private val idCounter: IdCounter? = null,
+  private val workerContext: CoroutineContext = EmptyCoroutineContext
 ) : RealRenderContext.Renderer<StateT, OutputT> {
 
   /**
@@ -197,7 +199,8 @@ internal class SubtreeManager<StateT, OutputT : Any>(
         ::acceptChildOutput,
         parentDiagnosticId,
         diagnosticListener,
-        idCounter
+        idCounter,
+        workerContext = workerContext
     )
     return WorkflowChildNode(child, handler, workflowNode)
         .also { node = it }
diff --git a/kotlin/workflow-runtime/src/main/java/com/squareup/workflow/internal/WorkflowNode.kt b/kotlin/workflow-runtime/src/main/java/com/squareup/workflow/internal/WorkflowNode.kt
index 3c5e81e0d..bf423894a 100644
--- a/kotlin/workflow-runtime/src/main/java/com/squareup/workflow/internal/WorkflowNode.kt
+++ b/kotlin/workflow-runtime/src/main/java/com/squareup/workflow/internal/WorkflowNode.kt
@@ -78,7 +78,7 @@ internal class WorkflowNode<PropsT, StateT, OutputT : Any, RenderingT>(
   internal val diagnosticId = idCounter.createId()
 
   private val subtreeManager = SubtreeManager<StateT, OutputT>(
-      coroutineContext, ::applyAction, diagnosticId, diagnosticListener, idCounter
+      coroutineContext, ::applyAction, diagnosticId, diagnosticListener, idCounter, workerContext
   )
 
   private val workers = ActiveStagingList<WorkerChildNode<*, *, *>>()
diff --git a/kotlin/workflow-testing/src/test/java/com/squareup/workflow/WorkerCompositionIntegrationTest.kt b/kotlin/workflow-testing/src/test/java/com/squareup/workflow/WorkerCompositionIntegrationTest.kt
index b78bcbb57..93e0ded25 100644
--- a/kotlin/workflow-testing/src/test/java/com/squareup/workflow/WorkerCompositionIntegrationTest.kt
+++ b/kotlin/workflow-testing/src/test/java/com/squareup/workflow/WorkerCompositionIntegrationTest.kt
@@ -21,10 +21,14 @@ import com.squareup.workflow.WorkflowAction.Companion.noAction
 import com.squareup.workflow.testing.WorkerSink
 import com.squareup.workflow.testing.testFromStart
 import kotlinx.coroutines.CancellationException
+import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineName
+import kotlinx.coroutines.Dispatchers.Unconfined
 import kotlinx.coroutines.Job
+import kotlinx.coroutines.Runnable
 import kotlinx.coroutines.TimeoutCancellationException
 import kotlinx.coroutines.channels.Channel
+import kotlin.coroutines.ContinuationInterceptor
 import kotlin.coroutines.CoroutineContext
 import kotlin.coroutines.coroutineContext
 import kotlin.test.Test
@@ -32,6 +36,7 @@ import kotlin.test.assertEquals
 import kotlin.test.assertFailsWith
 import kotlin.test.assertFalse
 import kotlin.test.assertNotSame
+import kotlin.test.assertSame
 import kotlin.test.assertTrue
 import kotlin.test.fail
 
@@ -259,9 +264,12 @@ class WorkerCompositionIntegrationTest {
 
   @Test fun `worker context job is ignored`() {
     val worker = Worker.from { coroutineContext }
-    val workflow = Workflow.stateless<Unit, CoroutineContext, Unit> {
+    val leafWorkflow = Workflow.stateless<Unit, CoroutineContext, Unit> {
       runningWorker(worker) { context -> action { setOutput(context) } }
     }
+    val workflow = Workflow.stateless<Unit, CoroutineContext, Unit> {
+      renderChild(leafWorkflow) { action { setOutput(it) } }
+    }
     val job: Job = Job()
 
     workflow.testFromStart(context = job) {
@@ -269,4 +277,28 @@ class WorkerCompositionIntegrationTest {
       assertNotSame(job, actualWorkerContext[Job])
     }
   }
+
+  @Test fun `worker context is used for workers`() {
+    val worker = Worker.from { coroutineContext }
+    val leafWorkflow = Workflow.stateless<Unit, CoroutineContext, Unit> {
+      runningWorker(worker) { context -> action { setOutput(context) } }
+    }
+    val workflow = Workflow.stateless<Unit, CoroutineContext, Unit> {
+      renderChild(leafWorkflow) { action { setOutput(it) } }
+    }
+    val dispatcher: CoroutineDispatcher = object : CoroutineDispatcher() {
+      override fun isDispatchNeeded(context: CoroutineContext): Boolean =
+        Unconfined.isDispatchNeeded(context)
+
+      override fun dispatch(
+        context: CoroutineContext,
+        block: Runnable
+      ) = Unconfined.dispatch(context, block)
+    }
+
+    workflow.testFromStart(context = dispatcher) {
+      val actualWorkerContext = awaitNextOutput()
+      assertSame(dispatcher, actualWorkerContext[ContinuationInterceptor])
+    }
+  }
 }