Allow executor tracing to flow across distinct executors
Otherwise any async work is lost after the first executor jump.
MODIFY
dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/ExecutorInstrumentation.java
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<? 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 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();
}

Allow executor tracing to flow across distinct executors

Otherwise any async work is lost after the first executor jump.