diff --git a/dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/ExecutorInstrumentation.java b/dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/ExecutorInstrumentation.java index 8c1add7530..36db882e3a 100644 --- a/dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/ExecutorInstrumentation.java +++ b/dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/ExecutorInstrumentation.java @@ -25,7 +25,6 @@ import java.util.Map; 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 final class ExecutorInstrumentation extends Instrumenter.Default { @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,8 +196,10 @@ public final class ExecutorInstrumentation extends Instrumenter.Default { @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); } } @@ -204,12 +207,13 @@ public final class ExecutorInstrumentation extends Instrumenter.Default { @Advice.OnMethodEnter(suppress = Throwable.class) public static Collection wrapJob( + @Advice.This final Executor executor, @Advice.Argument(value = 0, readOnly = false) Collection> tasks) { final Scope scope = GlobalTracer.get().scopeManager().active(); if (scope instanceof TraceScope && ((TraceScope) scope).isAsyncPropagating() && tasks != null - && DatadogWrapper.isTopLevelCall()) { + && DatadogWrapper.isTopLevelCall(executor)) { final Collection> wrappedTasks = new ArrayList<>(tasks.size()); for (final Callable task : tasks) { if (task != null && !(task instanceof CallableWrapper)) { @@ -224,9 +228,11 @@ public final class ExecutorInstrumentation extends Instrumenter.Default { @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 final class ExecutorInstrumentation extends Instrumenter.Default { * * @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 final class ExecutorInstrumentation extends Instrumenter.Default { && ((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(); }