Propagate context into promise callbacks

Author: patbeqo

temporal-sdk/src/main/java/io/temporal/internal/sync/CompletablePromiseImpl.java

@@ -1,6 +1,7 @@
 package io.temporal.internal.sync;
 
 import io.temporal.failure.TemporalFailure;
+import io.temporal.internal.context.ContextThreadLocal;
 import io.temporal.workflow.CancellationScope;
 import io.temporal.workflow.CompletablePromise;
 import io.temporal.workflow.Functions;
@@ -9,6 +10,7 @@
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 
@@ -246,7 +248,19 @@ private <U> Promise<U> then(Functions.Proc1<CompletablePromise<U>> proc) {
       proc.apply(resultPromise);
       unregisterWithRunner();
     } else {
-      handlers.add(() -> proc.apply(resultPromise));
+      // Handlers run in a callback thread created by the runner, which inherits its propagated
+      // contexts from the workflow header rather than from the thread registering the handler.
+      Map<String, Object> contexts = ContextThreadLocal.getCurrentContextForPropagation();
+      handlers.add(
+          () -> {
+            Map<String, Object> previous = ContextThreadLocal.getCurrentContextForPropagation();
+            ContextThreadLocal.propagateContextToCurrentThread(contexts);
+            try {
+              proc.apply(resultPromise);
+            } finally {
+              ContextThreadLocal.propagateContextToCurrentThread(previous);
+            }
+          });
     }
     return resultPromise;
   }

temporal-sdk/src/test/java/io/temporal/workflow/ContextPropagationTest.java

@@ -104,6 +104,20 @@ public void testThreadContextPropagation() {
     assertEquals("asynctesting123", result);
   }
 
+  @Test
+  public void testPromiseCallbackContextPropagation() {
+    Worker worker = testEnvironment.newWorker(TASK_QUEUE);
+    worker.registerWorkflowImplementationTypes(ContextPropagationCallbackWorkflowImpl.class);
+    testEnvironment.start();
+    MDC.put("test", "testing123");
+    WorkflowClient client = testEnvironment.getWorkflowClient();
+    WorkflowOptions options = WorkflowOptions.newBuilder().setTaskQueue(TASK_QUEUE).build();
+    TestWorkflows.TestWorkflow1 workflow =
+        client.newWorkflowStub(TestWorkflows.TestWorkflow1.class, options);
+    String result = workflow.execute("input1");
+    assertEquals("asynccallback123", result);
+  }
+
   @Test
   public void testActivityContextPropagation() {
     Worker worker = testEnvironment.newWorker(TASK_QUEUE);
@@ -339,6 +353,25 @@ private String async() {
     }
   }
 
+  public static class ContextPropagationCallbackWorkflowImpl
+      implements TestWorkflows.TestWorkflow1 {
+
+    @Override
+    public String execute(String input) {
+      Promise<String> asyncPromise = Async.function(this::async);
+      // The callback must see the context of the thread that registered it, not the context the
+      // thread completing the promise happens to carry.
+      MDC.put("test", "callback123");
+      Promise<String> callback = asyncPromise.thenApply(result -> result + MDC.get("test"));
+      MDC.put("test", "testing123");
+      return callback.get();
+    }
+
+    private String async() {
+      return "async";
+    }
+  }
+
   public static class ContextActivityImpl implements TestActivities.TestActivity1 {
     @Override
     public String execute(String input) {