Use lifecycle methods on ContextPropagator so as not to leak OpenTracing implementation

Author: sdwr98

build.gradle

@@ -58,10 +58,13 @@ dependencies {
     compile group: 'com.google.guava', name: 'guava', version: '28.1-jre'
     compile group: 'com.cronutils', name: 'cron-utils', version: '9.0.0'
     compile group: 'io.micrometer', name: 'micrometer-core', version: '1.1.2'
+    compile group: 'io.opentracing', name: 'opentracing-api', version: '0.33.0'
+    compile group: 'io.opentracing', name: 'opentracing-util', version: '0.33.0'
 
     testCompile group: 'junit', name: 'junit', version: '4.12'
     testCompile group: 'com.googlecode.junit-toolbox', name: 'junit-toolbox', version: '2.4'
     testCompile group: 'ch.qos.logback', name: 'logback-classic', version: '1.2.3'
+    testCompile group: 'io.opentracing', name: 'opentracing-mock', version: '0.33.0'
 }
 
 license {

src/main/java/com/uber/cadence/activity/LocalActivityOptions.java

@@ -101,15 +101,19 @@ public LocalActivityOptions validateAndBuildWithDefaults() {
       if (retryOptions != null) {
         ro = new RetryOptions.Builder(retryOptions).validateBuildWithDefaults();
       }
-      return new LocalActivityOptions(roundUpToSeconds(scheduleToCloseTimeout), ro, contextPropagators);
+      return new LocalActivityOptions(
+          roundUpToSeconds(scheduleToCloseTimeout), ro, contextPropagators);
     }
   }
 
   private final Duration scheduleToCloseTimeout;
   private final RetryOptions retryOptions;
   private final List<ContextPropagator> contextPropagators;
 
-  private LocalActivityOptions(Duration scheduleToCloseTimeout, RetryOptions retryOptions, List<ContextPropagator> contextPropagators) {
+  private LocalActivityOptions(
+      Duration scheduleToCloseTimeout,
+      RetryOptions retryOptions,
+      List<ContextPropagator> contextPropagators) {
     this.scheduleToCloseTimeout = scheduleToCloseTimeout;
     this.retryOptions = retryOptions;
     this.contextPropagators = contextPropagators;

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

@@ -23,6 +23,9 @@
  * Context Propagators are used to propagate information from workflow to activity, workflow to
  * child workflow, and workflow to child thread (using {@link com.uber.cadence.workflow.Async}).
  *
+ * <p>It is important to note that all threads share one ContextPropagator instance, so your
+ * implementation <b>must</b> be thread-safe and store any state in ThreadLocal variables.
+ *
  * <p>A sample <code>ContextPropagator</code> that copies all {@link org.slf4j.MDC} entries starting
  * with a given prefix along the code path looks like this:
  *
@@ -136,4 +139,31 @@ public interface ContextPropagator {
 
   /** Sets the current context */
   void setCurrentContext(Object context);
+
+  /**
+   * This is a lifecycle method, called after the context has been propagated to the
+   * workflow/activity thread but the workflow/activity has not yet started.
+   */
+  default void setUp() {
+    // No-op
+  }
+
+  /**
+   * This is a lifecycle method, called after the workflow/activity has completed. If the method
+   * finished without exception, {@code successful} will be true. Otherwise, it will be false and
+   * {@link #onError(Throwable)} will have already been called.
+   */
+  default void finish(boolean successful) {
+    // No-op
+  }
+
+  /**
+   * This is a lifecycle method, called when the workflow/activity finishes by throwing an unhandled
+   * exception. {@link #finish(boolean)} is called after this method.
+   *
+   * @param t The unhandled exception that caused the workflow/activity to terminate
+   */
+  default void onError(Throwable t) {
+    // No-op
+  }
 }

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

@@ -17,21 +17,28 @@
 
 package com.uber.cadence.context;
 
+import com.uber.cadence.internal.logging.LoggerTag;
+import io.opentracing.Scope;
 import io.opentracing.Span;
 import io.opentracing.SpanContext;
 import io.opentracing.Tracer;
+import io.opentracing.log.Fields;
 import io.opentracing.propagation.Format;
 import io.opentracing.propagation.TextMap;
+import io.opentracing.tag.Tags;
 import io.opentracing.util.GlobalTracer;
 import java.nio.charset.Charset;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.Map;
+import org.slf4j.MDC;
 
 /** Support for OpenTracing spans */
 public class OpenTracingContextPropagator implements ContextPropagator {
 
   private static ThreadLocal<SpanContext> currentOpenTracingSpanContext = new ThreadLocal<>();
+  private static ThreadLocal<Span> currentOpenTracingSpan = new ThreadLocal<>();
+  private static ThreadLocal<Scope> currentOpenTracingScope = new ThreadLocal<>();
 
   public static void setCurrentOpenTracingSpanContext(SpanContext ctx) {
     if (ctx != null) {
@@ -91,6 +98,46 @@ public void setCurrentContext(Object context) {
     }
   }
 
+  @Override
+  public void setUp() {
+    Tracer openTracingTracer = GlobalTracer.get();
+    Tracer.SpanBuilder builder =
+        openTracingTracer
+            .buildSpan("cadence.workflow")
+            .withTag("resource.name", MDC.get(LoggerTag.WORKFLOW_TYPE));
+
+    if (getCurrentOpenTracingSpanContext() != null) {
+      builder.asChildOf(getCurrentOpenTracingSpanContext());
+    }
+
+    Span span = builder.start();
+    openTracingTracer.activateSpan(span);
+    currentOpenTracingSpan.set(span);
+    Scope scope = openTracingTracer.activateSpan(span);
+    currentOpenTracingScope.set(scope);
+  }
+
+  @Override
+  public void onError(Throwable t) {
+    Span span = currentOpenTracingSpan.get();
+    Tags.ERROR.set(span, true);
+    Map<String, Object> errorData = new HashMap<>();
+    errorData.put(Fields.EVENT, "error");
+    if (t != null) {
+      errorData.put(Fields.ERROR_OBJECT, t);
+      errorData.put(Fields.MESSAGE, t.getMessage());
+    }
+    span.log(errorData);
+  }
+
+  @Override
+  public void finish(boolean successful) {
+    Scope currentScope = currentOpenTracingScope.get();
+    Span currentSpan = currentOpenTracingSpan.get();
+    currentScope.close();
+    currentSpan.finish();
+  }
+
   private class HashMapTextMap implements TextMap {
 
     private final HashMap<String, String> backingMap = new HashMap<>();

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

@@ -24,10 +24,14 @@
 import java.util.List;
 import java.util.Map;
 import java.util.function.Supplier;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /** This class holds the current set of context propagators */
 public class ContextThreadLocal {
 
+  private static final Logger log = LoggerFactory.getLogger(ContextThreadLocal.class);
+
   private static WorkflowThreadLocal<List<ContextPropagator>> contextPropagators =
       WorkflowThreadLocal.withInitial(
           new Supplier<List<ContextPropagator>>() {
@@ -57,6 +61,11 @@ public static Map<String, Object> getCurrentContextForPropagation() {
     return contextData;
   }
 
+  /**
+   * Injects the context data into the thread for each configured context propagator
+   *
+   * @param contextData The context data received from the server
+   */
   public static void propagateContextToCurrentThread(Map<String, Object> contextData) {
     if (contextData == null || contextData.isEmpty()) {
       return;
@@ -67,4 +76,49 @@ public static void propagateContextToCurrentThread(Map<String, Object> contextDa
       }
     }
   }
+
+  /** Calls {@link ContextPropagator#setUp()} for each propagator */
+  public static void setUpContextPropagators() {
+    for (ContextPropagator propagator : contextPropagators.get()) {
+      try {
+        propagator.setUp();
+      } catch (Throwable t) {
+        // Don't let an error in one propagator block the others
+        log.error("Error calling setUp() on a contextpropagator", t);
+      }
+    }
+  }
+
+  /**
+   * Calls {@link ContextPropagator#onError(Throwable)} for each propagator
+   *
+   * @param t The Throwable that caused the workflow/activity to finish
+   */
+  public static void onErrorContextPropagators(Throwable t) {
+    for (ContextPropagator propagator : contextPropagators.get()) {
+      try {
+        propagator.onError(t);
+      } catch (Throwable t1) {
+        // Don't let an error in one propagator block the others
+        log.error("Error calling onError() on a contextpropagator", t1);
+      }
+    }
+  }
+
+  /**
+   * Calls {@link ContextPropagator#finish(boolean))} for each propagator
+   *
+   * @param successful True if the workflow/activity completed without unhandled exception, false
+   *     otherwise
+   */
+  public static void finishContextPropagators(boolean successful) {
+    for (ContextPropagator propagator : contextPropagators.get()) {
+      try {
+        propagator.finish(successful);
+      } catch (Throwable t) {
+        // Don't let an error in one propagator block the others
+        log.error("Error calling finish() on a contextpropagator", t);
+      }
+    }
+  }
 }

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

@@ -323,10 +323,11 @@ private ExecuteActivityParameters constructExecuteActivityParameters(
 
   private ExecuteLocalActivityParameters constructExecuteLocalActivityParameters(
       String name, LocalActivityOptions options, byte[] input, long elapsed, int attempt) {
-    ExecuteLocalActivityParameters parameters = new ExecuteLocalActivityParameters()
-        .withActivityType(new ActivityType().setName(name))
-        .withInput(input)
-        .withScheduleToCloseTimeoutSeconds(options.getScheduleToCloseTimeout().getSeconds());
+    ExecuteLocalActivityParameters parameters =
+        new ExecuteLocalActivityParameters()
+            .withActivityType(new ActivityType().setName(name))
+            .withInput(input)
+            .withScheduleToCloseTimeoutSeconds(options.getScheduleToCloseTimeout().getSeconds());
 
     RetryOptions retryOptions = options.getRetryOptions();
     if (retryOptions != null) {
@@ -337,8 +338,8 @@ private ExecuteLocalActivityParameters constructExecuteLocalActivityParameters(
     parameters.setWorkflowDomain(this.context.getDomain());
     parameters.setWorkflowExecution(this.context.getWorkflowExecution());
 
-    List<ContextPropagator> propagators = Optional.ofNullable(options.getContextPropagators())
-        .orElse(contextPropagators);
+    List<ContextPropagator> propagators =
+        Optional.ofNullable(options.getContextPropagators()).orElse(contextPropagators);
     parameters.setContext(extractContextsAndConvertToBytes(propagators));
 
     return parameters;

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

@@ -19,19 +19,12 @@
 
 import com.google.common.util.concurrent.RateLimiter;
 import com.uber.cadence.context.ContextPropagator;
-import com.uber.cadence.context.OpenTracingContextPropagator;
 import com.uber.cadence.internal.context.ContextThreadLocal;
 import com.uber.cadence.internal.logging.LoggerTag;
 import com.uber.cadence.internal.metrics.MetricsType;
 import com.uber.cadence.internal.replay.DeciderCache;
 import com.uber.cadence.internal.replay.DecisionContext;
 import com.uber.cadence.workflow.Promise;
-import io.opentracing.Scope;
-import io.opentracing.Span;
-import io.opentracing.Tracer;
-import io.opentracing.log.Fields;
-import io.opentracing.tag.Tags;
-import io.opentracing.util.GlobalTracer;
 import java.io.PrintWriter;
 import java.io.StringWriter;
 import java.util.HashMap;
@@ -98,27 +91,16 @@ public void run() {
       // Repopulate the context(s)
       ContextThreadLocal.setContextPropagators(this.contextPropagators);
       ContextThreadLocal.propagateContextToCurrentThread(this.propagatedContexts);
+      ContextThreadLocal.setUpContextPropagators();
 
-      // Set up an opentracing span
-      Tracer openTracingTracer = GlobalTracer.get();
-      Tracer.SpanBuilder builder =
-          openTracingTracer
-              .buildSpan("cadence.workflow")
-              .withTag("resource.name", decisionContext.getWorkflowType().getName());
-
-      if (OpenTracingContextPropagator.getCurrentOpenTracingSpanContext() != null) {
-        builder.asChildOf(OpenTracingContextPropagator.getCurrentOpenTracingSpanContext());
-      }
-      Span span = builder.start();
-
-      try (Scope scope = openTracingTracer.activateSpan(span)) {
+      try {
         // initialYield blocks thread until the first runUntilBlocked is called.
         // Otherwise r starts executing without control of the sync.
         threadContext.initialYield();
         cancellationScope.run();
       } catch (DestroyWorkflowThreadError e) {
         if (!threadContext.isDestroyRequested()) {
-          setSpanError(span, e);
+          ContextThreadLocal.onErrorContextPropagators(e);
           threadContext.setUnhandledException(e);
         }
       } catch (Error e) {
@@ -131,11 +113,11 @@ public void run() {
           log.error(
               String.format("Workflow thread \"%s\" run failed with Error:\n%s", name, stackTrace));
         }
-        setSpanError(span, e);
+        ContextThreadLocal.onErrorContextPropagators(e);
         threadContext.setUnhandledException(e);
       } catch (CancellationException e) {
         if (!isCancelRequested()) {
-          setSpanError(span, e);
+          ContextThreadLocal.onErrorContextPropagators(e);
           threadContext.setUnhandledException(e);
         }
         if (log.isDebugEnabled()) {
@@ -152,27 +134,16 @@ public void run() {
                   "Workflow thread \"%s\" run failed with unhandled exception:\n%s",
                   name, stackTrace));
         }
-        setSpanError(span, e);
+        ContextThreadLocal.onErrorContextPropagators(e);
         threadContext.setUnhandledException(e);
       } finally {
+        ContextThreadLocal.finishContextPropagators(threadContext.getUnhandledException() == null);
         DeterministicRunnerImpl.setCurrentThreadInternal(null);
         threadContext.setStatus(Status.DONE);
         thread.setName(originalName);
         thread = null;
         MDC.clear();
-        span.finish();
-      }
-    }
-
-    private void setSpanError(Span span, Throwable ex) {
-      Tags.ERROR.set(span, true);
-      Map<String, Object> errorData = new HashMap<>();
-      errorData.put(Fields.EVENT, "error");
-      if (ex != null) {
-        errorData.put(Fields.ERROR_OBJECT, ex);
-        errorData.put(Fields.MESSAGE, ex.getMessage());
       }
-      span.log(errorData);
     }
 
     public String getName() {

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

@@ -271,7 +271,9 @@ private void propagateContext(ExecuteLocalActivityParameters params) {
   }
 
   private void restoreContext(Map<String, byte[]> context) {
-    options.getContextPropagators()
-        .forEach(propagator -> propagator.setCurrentContext(propagator.deserializeContext(context)));
+    options
+        .getContextPropagators()
+        .forEach(
+            propagator -> propagator.setCurrentContext(propagator.deserializeContext(context)));
   }
 }

src/test/java/com/uber/cadence/activity/LocalActivityContextPropagationTest.java

@@ -59,7 +59,8 @@ public class LocalActivityContextPropagationTest {
 
   private final WrapperContext wrapperContext = new WrapperContext(EXPECTED_CONTEXT_NAME);
 
-  //let's add safe TestWorkflowEnvironment closing and make configurable propagation enabling/disabling
+  // let's add safe TestWorkflowEnvironment closing and make configurable propagation
+  // enabling/disabling
   private class TestEnvAutoCloseable implements AutoCloseable {
 
     private TestWorkflowEnvironment testEnv;