feat(ContextPropagator): Replace imperative set/unset context propagation with lexically-scoped runWithContext

src/main/java/com/uber/cadence/context/ContextPropagator.java

@@ -139,4 +139,37 @@ public interface ContextPropagator {
 
   /** Unsets the current context. This is called when the context is no longer needed */
   void unsetCurrentContext();
+
+  /** A context task that may throw a checked exception. */
+  @FunctionalInterface
+  interface ContextRunnable {
+    void run() throws Exception;
+  }
+
+  /**
+   * Executes {@code task} with {@code context} installed as the current context.
+   *
+   * <p>The default implementation preserves the legacy imperative context lifecycle. Propagators
+   * that use lexical context, such as Java Scoped Values, should override this method and execute
+   * {@code task} inside their lexical binding rather than implementing {@link #setCurrentContext}
+   * and {@link #unsetCurrentContext}.
+   *
+   * <p><b>Contract:</b> implementations must call {@code task.run()} exactly once and let any
+   * exception it throws propagate to the caller unchanged. Only wrap the call to {@code task.run()}
+   * in {@code try}/{@code finally} for cleanup (as the default implementation does) -- never in a
+   * {@code try}/{@code catch} that suppresses or replaces the exception, and never skip or retry
+   * the call. Cadence relies on exceptions thrown by the wrapped task (including workflow
+   * cancellation and thread-destruction signals) reaching the caller in order to function
+   * correctly; a propagator that violates this contract causes {@link
+   * com.uber.cadence.internal.context.ContextThreadLocal} to throw an internal error rather than
+   * silently continue as if {@code task} had succeeded.
+   */
+  default void runWithContext(Object context, ContextRunnable task) throws Exception {
+    setCurrentContext(context);
+    try {
+      task.run();
+    } finally {
+      unsetCurrentContext();
+    }
+  }
 }

src/main/java/com/uber/cadence/internal/context/ContextPropagatorContractViolationError.java

@@ -0,0 +1,65 @@
+/*
+ *  Modifications Copyright (c) 2017-2020 Uber Technologies Inc.
+ *  Portions of the Software are attributed to Copyright (c) 2020 Temporal Technologies Inc.
+ *  Copyright 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"). You may not
+ *  use this file except in compliance with the License. A copy of the License is
+ *  located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ *  or in the "license" file accompanying this file. This file is distributed on
+ *  an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ *  express or implied. See the License for the specific language governing
+ *  permissions and limitations under the License.
+ */
+
+package com.uber.cadence.internal.context;
+
+import com.uber.cadence.context.ContextPropagator;
+import java.util.List;
+import java.util.stream.Collectors;
+
+/**
+ * Thrown when a {@link ContextPropagator#runWithContext} implementation violates its contract:
+ * either by catching and suppressing an exception thrown by the wrapped task instead of letting it
+ * propagate, or by invoking the task a number of times other than exactly once (e.g. retrying it
+ * after a failure, or never invoking it at all). This indicates a bug in the propagator
+ * implementation, not a workflow or activity failure, so it is modeled as an {@link Error} rather
+ * than a checked/unchecked exception.
+ */
+public final class ContextPropagatorContractViolationError extends Error {
+
+  private ContextPropagatorContractViolationError(String message, Throwable cause) {
+    super(message, cause);
+  }
+
+  static ContextPropagatorContractViolationError swallowedException(
+      List<ContextPropagator> appliedPropagators, Throwable swallowed) {
+    return new ContextPropagatorContractViolationError(
+        "A ContextPropagator swallowed an exception thrown by the task it wraps instead of "
+            + "propagating it. ContextPropagator#runWithContext must not catch and suppress "
+            + "exceptions from the task it is given -- only wrap the call in try/finally, never "
+            + "try/catch. One of these configured propagators must be fixed: "
+            + propagatorNames(appliedPropagators),
+        swallowed);
+  }
+
+  static ContextPropagatorContractViolationError unexpectedInvocationCount(
+      List<ContextPropagator> appliedPropagators, int invocationCount) {
+    return new ContextPropagatorContractViolationError(
+        "A ContextPropagator invoked the task it wraps "
+            + invocationCount
+            + " time(s) instead of exactly once. ContextPropagator#runWithContext must call "
+            + "task.run() exactly once -- it must not skip the call, and it must not retry the "
+            + "task after catching an exception from it. One of these configured propagators "
+            + "must be fixed: "
+            + propagatorNames(appliedPropagators),
+        null);
+  }
+
+  private static String propagatorNames(List<ContextPropagator> propagators) {
+    return propagators.stream().map(p -> p.getClass().getName()).collect(Collectors.joining(", "));
+  }
+}

src/main/java/com/uber/cadence/internal/context/ContextThreadLocal.java

@@ -18,11 +18,14 @@
 package com.uber.cadence.internal.context;
 
 import com.uber.cadence.context.ContextPropagator;
+import com.uber.cadence.context.ContextPropagator.ContextRunnable;
 import com.uber.cadence.workflow.WorkflowThreadLocal;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
 
 /** This class holds the current set of context propagators */
@@ -57,20 +60,68 @@ public static Map<String, Object> getCurrentContextForPropagation() {
     return contextData;
   }
 
-  public static void propagateContextToCurrentThread(Map<String, Object> contextData) {
-    if (contextData == null || contextData.isEmpty()) {
+  public static void runWithContext(Map<String, Object> contextData, ContextRunnable task)
+      throws Exception {
+    runWithContext(contextPropagators.get(), contextData, task);
+  }
+
+  /**
+   * Executes {@code task} inside every applicable propagator context.
+   *
+   * <p>Propagators are composed in configuration order, so the first propagator is the outermost
+   * context and cleanup occurs in reverse order. Legacy propagators retain their existing set/unset
+   * behavior through {@link ContextPropagator#runWithContext(Object, ContextRunnable)}.
+   *
+   * <p>Every {@link ContextPropagator#runWithContext(Object, ContextRunnable)} implementation in
+   * the chain is required to call {@code task.run()} exactly once and propagate any exception it
+   * throws, rather than skipping the call, retrying it, or catching and suppressing the exception.
+   * This method verifies that contract: if {@code task} is invoked a number of times other than
+   * exactly one, or if it throws but no exception escapes the composed propagator chain, a {@link
+   * ContextPropagatorContractViolationError} is thrown instead of silently continuing as if {@code
+   * task} had succeeded normally.
+   */
+  public static void runWithContext(
+      List<ContextPropagator> propagators, Map<String, Object> contextData, ContextRunnable task)
+      throws Exception {
+    if (propagators == null
+        || propagators.isEmpty()
+        || contextData == null
+        || contextData.isEmpty()) {
+      task.run();
       return;
     }
-    for (ContextPropagator propagator : contextPropagators.get()) {
+
+    List<ContextPropagator> applied = new ArrayList<>();
+    AtomicInteger invocationCount = new AtomicInteger();
+    AtomicReference<Throwable> thrown = new AtomicReference<>();
+    ContextRunnable invocation =
+        () -> {
+          invocationCount.incrementAndGet();
+          try {
+            task.run();
+          } catch (Throwable t) {
+            thrown.set(t);
+            throw t;
+          }
+        };
+    for (int i = propagators.size() - 1; i >= 0; i--) {
+      ContextPropagator propagator = propagators.get(i);
       if (contextData.containsKey(propagator.getName())) {
-        propagator.setCurrentContext(contextData.get(propagator.getName()));
+        applied.add(0, propagator);
+        Object context = contextData.get(propagator.getName());
+        ContextRunnable next = invocation;
+        invocation = () -> propagator.runWithContext(context, next);
       }
     }
-  }
+    invocation.run();
 
-  public static void unsetCurrentContext() {
-    for (ContextPropagator propagator : contextPropagators.get()) {
-      propagator.unsetCurrentContext();
+    if (invocationCount.get() != 1) {
+      throw ContextPropagatorContractViolationError.unexpectedInvocationCount(
+          applied, invocationCount.get());
+    }
+    Throwable swallowed = thrown.get();
+    if (swallowed != null) {
+      throw ContextPropagatorContractViolationError.swallowedException(applied, swallowed);
     }
   }
 }

src/main/java/com/uber/cadence/internal/sync/WorkflowThreadImpl.java

@@ -88,15 +88,13 @@ public void run() {
       MDC.put(LoggerTag.TASK_LIST, decisionContext.getTaskList());
       MDC.put(LoggerTag.DOMAIN, decisionContext.getDomain());
 
-      // Repopulate the context(s)
       ContextThreadLocal.setContextPropagators(this.contextPropagators);
-      ContextThreadLocal.propagateContextToCurrentThread(this.propagatedContexts);
 
       try {
         // initialYield blocks thread until the first runUntilBlocked is called.
         // Otherwise r starts executing without control of the sync.
         threadContext.initialYield();
-        cancellationScope.run();
+        ContextThreadLocal.runWithContext(this.propagatedContexts, cancellationScope::run);
       } catch (DestroyWorkflowThreadError e) {
         if (!threadContext.isDestroyRequested()) {
           threadContext.setUnhandledException(e);
@@ -132,7 +130,6 @@ public void run() {
         }
         threadContext.setUnhandledException(e);
       } finally {
-        ContextThreadLocal.unsetCurrentContext();
         DeterministicRunnerImpl.setCurrentThreadInternal(null);
         threadContext.setStatus(Status.DONE);
         thread.setName(originalName);

src/main/java/com/uber/cadence/internal/worker/ActivityWorker.java

@@ -20,6 +20,7 @@
 import com.uber.cadence.*;
 import com.uber.cadence.context.ContextPropagator;
 import com.uber.cadence.internal.common.RpcRetryer;
+import com.uber.cadence.internal.context.ContextThreadLocal;
 import com.uber.cadence.internal.logging.LoggerTag;
 import com.uber.cadence.internal.metrics.HistogramBuckets;
 import com.uber.cadence.internal.metrics.MetricsEmit;
@@ -142,7 +143,15 @@ public void handle(ActivityTask task) throws Exception {
       MDC.put(LoggerTag.RUN_ID, response.getWorkflowExecution().getRunId());
       MDC.put(LoggerTag.ATTEMPT, String.valueOf(response.getAttempt()));
 
-      propagateContext(response);
+      ContextThreadLocal.runWithContext(
+          options.getContextPropagators(),
+          deserializeContext(response),
+          () -> handleWithContext(task, response, metricsScope));
+    }
+
+    private void handleWithContext(
+        ActivityTask task, PollForActivityTaskResponse response, Scope metricsScope)
+        throws CadenceError {
       Span span = spanFactory.spanForExecuteActivity(response);
       ActivityTaskHandler.Result handlerResponse = null;
       try (io.opentracing.Scope scope = tracer.activateSpan(span)) {
@@ -186,7 +195,6 @@ public void handle(ActivityTask task) throws Exception {
         MDC.remove(LoggerTag.WORKFLOW_ID);
         MDC.remove(LoggerTag.RUN_ID);
         MDC.remove(LoggerTag.ATTEMPT);
-        unsetCurrentContext();
         // Apply completion handle if task has been completed synchronously or is async and manual
         // completion hasn't been requested.
         if (handlerResponse != null && !handlerResponse.isManualCompletion()) {
@@ -195,33 +203,24 @@ public void handle(ActivityTask task) throws Exception {
       }
     }
 
-    void propagateContext(PollForActivityTaskResponse response) {
+    private Map<String, Object> deserializeContext(PollForActivityTaskResponse response) {
       if (options.getContextPropagators() == null || options.getContextPropagators().isEmpty()) {
-        return;
+        return new HashMap<>();
       }
 
       Header headers = response.getHeader();
       if (headers == null) {
-        return;
+        return new HashMap<>();
       }
 
       Map<String, byte[]> headerData = new HashMap<>();
-      headers
-          .getFields()
-          .forEach(
-              (k, v) -> {
-                headerData.put(k, v);
-              });
-
-      for (ContextPropagator propagator : options.getContextPropagators()) {
-        propagator.setCurrentContext(propagator.deserializeContext(headerData));
-      }
-    }
+      headers.getFields().forEach((k, v) -> headerData.put(k, v));
 
-    void unsetCurrentContext() {
+      Map<String, Object> contextData = new HashMap<>();
       for (ContextPropagator propagator : options.getContextPropagators()) {
-        propagator.unsetCurrentContext();
+        contextData.put(propagator.getName(), propagator.deserializeContext(headerData));
       }
+      return contextData;
     }
 
     @Override

src/main/java/com/uber/cadence/internal/worker/LocalActivityWorker.java

@@ -24,6 +24,7 @@
 import com.uber.cadence.common.RetryOptions;
 import com.uber.cadence.context.ContextPropagator;
 import com.uber.cadence.internal.common.LocalActivityMarkerData;
+import com.uber.cadence.internal.context.ContextThreadLocal;
 import com.uber.cadence.internal.metrics.HistogramBuckets;
 import com.uber.cadence.internal.metrics.MetricsEmit;
 import com.uber.cadence.internal.metrics.MetricsTag;
@@ -36,9 +37,9 @@
 import io.opentracing.Span;
 import io.opentracing.Tracer;
 import java.time.Duration;
+import java.util.Collections;
 import java.util.Map;
 import java.util.Objects;
-import java.util.Optional;
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
 import java.util.function.LongSupplier;
@@ -132,8 +133,13 @@ private TaskHandlerImpl(ActivityTaskHandler handler) {
 
     @Override
     public void handle(Task task) throws Exception {
-      propagateContext(task.params);
+      ContextThreadLocal.runWithContext(
+          options.getContextPropagators(),
+          deserializeContext(task.params),
+          () -> handleWithContext(task));
+    }
 
+    private void handleWithContext(Task task) throws InterruptedException {
       // start and activate span for local activities
       Span span = spanFactory.spanForExecuteLocalActivity(task);
       try (io.opentracing.Scope scope = tracer.activateSpan(span)) {
@@ -171,7 +177,6 @@ public void handle(Task task) throws Exception {
         task.eventConsumer.accept(event);
       } finally {
         span.finish();
-        unsetCurrentContext();
       }
     }
 
@@ -240,24 +245,20 @@ private ActivityTaskHandler.Result handleLocalActivity(Task task) throws Interru
     }
   }
 
-  private void propagateContext(ExecuteLocalActivityParameters params) {
+  private Map<String, Object> deserializeContext(ExecuteLocalActivityParameters params) {
     if (options.getContextPropagators() == null || options.getContextPropagators().isEmpty()) {
-      return;
+      return Collections.emptyMap();
     }
 
-    Optional.ofNullable(params.getContext())
-        .filter(context -> !context.isEmpty())
-        .ifPresent(this::restoreContext);
-  }
-
-  private void unsetCurrentContext() {
-    options.getContextPropagators().forEach(ContextPropagator::unsetCurrentContext);
-  }
+    Map<String, byte[]> context = params.getContext();
+    if (context == null || context.isEmpty()) {
+      return Collections.emptyMap();
+    }
 
-  private void restoreContext(Map<String, byte[]> context) {
-    options
-        .getContextPropagators()
-        .forEach(
-            propagator -> propagator.setCurrentContext(propagator.deserializeContext(context)));
+    Map<String, Object> contextData = new java.util.HashMap<>();
+    for (ContextPropagator propagator : options.getContextPropagators()) {
+      contextData.put(propagator.getName(), propagator.deserializeContext(context));
+    }
+    return contextData;
   }
 }