Merge pull request #443 from DataDog/tyler/executor-instrumentation

Allow executor tracing to flow across distinct executors
DataDog · Sep 23, 2018 · 76876e7 · 76876e7
2 parents bf24e5d + 2174f21
commit 76876e7
Showing 1 changed file with 23 additions and 15 deletions.
diff --git a/.../src/main/java/datadog/trace/instrumentation/java/concurrent/ExecutorInstrumentation.java b/.../src/main/java/datadog/trace/instrumentation/java/concurrent/ExecutorInstrumentation.java
@@ -25,7 +25,6 @@
 import java.util.concurrent.Callable;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executor;
-import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Future;
 import lombok.extern.slf4j.Slf4j;
 import net.bytebuddy.asm.Advice;
@@ -173,8 +172,10 @@ public static DatadogWrapper enterJobSubmit(
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void exitJobSubmit(
-        @Advice.Enter final DatadogWrapper wrapper, @Advice.Thrown final Throwable throwable) {
-      DatadogWrapper.cleanUpOnMethodExit(wrapper, throwable);
+        @Advice.This final Executor executor,
+        @Advice.Enter final DatadogWrapper wrapper,
+        @Advice.Thrown final Throwable throwable) {
+      DatadogWrapper.cleanUpOnMethodExit(executor, wrapper, throwable);
     }
   }
 
@@ -195,21 +196,24 @@ public static DatadogWrapper enterJobSubmit(
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void exitJobSubmit(
-        @Advice.Enter final DatadogWrapper wrapper, @Advice.Thrown final Throwable throwable) {
-      DatadogWrapper.cleanUpOnMethodExit(wrapper, throwable);
+        @Advice.This final Executor executor,
+        @Advice.Enter final DatadogWrapper wrapper,
+        @Advice.Thrown final Throwable throwable) {
+      DatadogWrapper.cleanUpOnMethodExit(executor, wrapper, throwable);
     }
   }
 
   public static class WrapCallableCollectionAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static Collection<?> wrapJob(
+        @Advice.This final Executor executor,
         @Advice.Argument(value = 0, readOnly = false) Collection<? extends Callable<?>> tasks) {
       final Scope scope = GlobalTracer.get().scopeManager().active();
       if (scope instanceof TraceScope
           && ((TraceScope) scope).isAsyncPropagating()
           && tasks != null
-          && DatadogWrapper.isTopLevelCall()) {
+          && DatadogWrapper.isTopLevelCall(executor)) {
         final Collection<Callable<?>> wrappedTasks = new ArrayList<>(tasks.size());
         for (final Callable<?> task : tasks) {
           if (task != null && !(task instanceof CallableWrapper)) {
@@ -224,9 +228,11 @@ public static Collection<?> wrapJob(
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void checkCancel(
-        @Advice.Enter final Collection<?> wrappedJobs, @Advice.Thrown final Throwable throwable) {
+        @Advice.This final Executor executor,
+        @Advice.Enter final Collection<?> wrappedJobs,
+        @Advice.Thrown final Throwable throwable) {
       if (null != wrappedJobs) {
-        DatadogWrapper.resetNestedCalls();
+        DatadogWrapper.resetNestedCalls(executor);
 
         if (null != throwable) {
           for (final Object wrapper : wrappedJobs) {
@@ -264,13 +270,14 @@ public void cancel() {
      *
      * @return true iff call is not nested
      */
-    public static boolean isTopLevelCall() {
-      return CallDepthThreadLocalMap.incrementCallDepth(ExecutorService.class) <= 0;
+    public static boolean isTopLevelCall(final Executor executor) {
+      final int i = CallDepthThreadLocalMap.incrementCallDepth(executor.getClass());
+      return i <= 0;
     }
 
     /** Reset nested calls to executor. */
-    public static void resetNestedCalls() {
-      CallDepthThreadLocalMap.reset(ExecutorService.class);
+    public static void resetNestedCalls(final Executor executor) {
+      CallDepthThreadLocalMap.reset(executor.getClass());
     }
 
     /**
@@ -283,20 +290,21 @@ public static boolean shouldWrapTask(final Object task, final Executor executor)
           && ((TraceScope) scope).isAsyncPropagating()
           && task != null
           && !(task instanceof DatadogWrapper)
-          && isTopLevelCall()
+          && isTopLevelCall(executor)
           && !ConcurrentUtils.isDisabled(executor));
     }
 
     /**
      * Clean up after job submission method has exited
      *
+     * @param executor the current executor
      * @param wrapper task wrapper
      * @param throwable throwable that may have been thrown
      */
     public static void cleanUpOnMethodExit(
-        final DatadogWrapper wrapper, final Throwable throwable) {
+        final Executor executor, final DatadogWrapper wrapper, final Throwable throwable) {
       if (null != wrapper) {
-        resetNestedCalls();
+        resetNestedCalls(executor);
         if (null != throwable) {
           wrapper.cancel();
         }