Switch executor instrumentation to use context storage
ADD
dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/State.java
+package datadog.trace.bootstrap.instrumentation.java.concurrent;
+
+import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.context.TraceScope;
+import java.util.concurrent.atomic.AtomicReference;
+import lombok.extern.slf4j.Slf4j;
+
+@Slf4j
+public class State {
+
+ public static ContextStore.Factory<State> FACTORY =
+ new ContextStore.Factory<State>() {
+ @Override
+ public State create() {
+ return new State();
+ }
+ };
+
+ private final AtomicReference<TraceScope.Continuation> continuationRef =
+ new AtomicReference<>(null);
+
+ private State() {}
+
+ public boolean setContinuation(final TraceScope.Continuation continuation) {
+ final boolean result = continuationRef.compareAndSet(null, continuation);
+ if (!result) {
+ log.debug(
+ "Failed to set continuation because another continuation is already set {}: new: {}, old: {}",
+ this,
+ continuation,
+ continuationRef.get());
+ }
+ return result;
+ }
+
+ public void closeContinuation() {
+ final TraceScope.Continuation continuation = continuationRef.getAndSet(null);
+ if (continuation != null) {
+ continuation.close();
+ }
+ }
+
+ public TraceScope.Continuation getAndResetContinuation() {
+ return continuationRef.getAndSet(null);
+ }
+}

MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java
@@ -55,6 +55,9 @@ public class AgentInstaller {
.with(POOL_STRATEGY)
.with(new LoggingListener())
.with(LOCATION_STRATEGY)
+ // FIXME: we cannot enable it yet dur to BB/JVM bug, see
+ // https://github.com/raphw/byte-buddy/issues/558
+ // .with(AgentBuilder.LambdaInstrumentationStrategy.ENABLED)
.ignore(any(), skipClassLoader())
.or(nameStartsWith("datadog.trace."))
.or(nameStartsWith("datadog.opentracing."))

MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelperInjector.java
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelperInjector.java
@@ -40,7 +40,7 @@ public class HelperInjector implements Transformer {
}
public HelperInjector(final Map<String, byte[]> helperMap) {
- this.helperClassNames = helperMap.keySet();
+ helperClassNames = helperMap.keySet();
this.helperMap = new LinkedHashMap<>(helperClassNames.size());
for (final String helperName : helperClassNames) {
final TypeDescription typeDesc =

MODIFY
dd-java-agent/instrumentation/akka-http-10.0/src/lagomTest/groovy/LagomTest.groovy
dd-java-agent/instrumentation/akka-http-10.0/src/lagomTest/groovy/LagomTest.groovy
@@ -16,9 +16,6 @@ import static com.lightbend.lagom.javadsl.testkit.ServiceTest.defaultSetup
import static com.lightbend.lagom.javadsl.testkit.ServiceTest.startServer
class LagomTest extends AgentTestRunner {
- static {
- System.setProperty("dd.integration.akka-http-server.enabled", "true")
- }
@Shared
private TestServer server

MODIFY
dd-java-agent/instrumentation/akka-http-10.0/src/lagomTest/java/EchoServiceImpl.java
dd-java-agent/instrumentation/akka-http-10.0/src/lagomTest/java/EchoServiceImpl.java
@@ -3,6 +3,7 @@ import akka.stream.javadsl.Source;
import com.lightbend.lagom.javadsl.api.ServiceCall;
import datadog.trace.api.Trace;
import java.util.List;
+import java.util.concurrent.Callable;
import java.util.concurrent.CompletableFuture;
public class EchoServiceImpl implements EchoService {
@@ -10,7 +11,17 @@ public class EchoServiceImpl implements EchoService {
@Override
public ServiceCall<Source<String, NotUsed>, Source<String, NotUsed>> echo() {
final CompletableFuture<Source<String, NotUsed>> fut = new CompletableFuture<>();
- ServiceTestModule.executor.submit(() -> fut.complete(Source.from(tracedMethod())));
+ ServiceTestModule.executor.submit(
+ // FIXME: we cannot use lambda here due to BB/JVM bug and RunnableInstrumentation being present.
+ // See https://github.com/raphw/byte-buddy/issues/558. Once that is fixed we can make this lambda again.
+ // Technically this doesn't return object but we do not really care here
+ new Callable<Object>() {
+
+ @Override
+ public Object call() throws Exception {
+ return fut.complete(Source.from(tracedMethod()));
+ }
+ });
return req -> fut;
}
@@ -23,4 +34,5 @@ public class EchoServiceImpl implements EchoService {
public List<String> tracedMethod() {
return java.util.Arrays.asList("msg1", "msg2", "msg3");
}
+
}

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
@@ -9,13 +9,12 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
import com.google.auto.service.AutoService;
import datadog.trace.agent.tooling.Instrumenter;
-import datadog.trace.bootstrap.CallDepthThreadLocalMap;
-import datadog.trace.bootstrap.WeakMap;
+import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.bootstrap.InstrumentationContext;
+import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
import datadog.trace.context.TraceScope;
import io.opentracing.Scope;
import io.opentracing.util.GlobalTracer;
-import java.lang.reflect.Field;
-import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
@@ -23,7 +22,6 @@ import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.concurrent.Callable;
-import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executor;
import java.util.concurrent.Future;
import lombok.extern.slf4j.Slf4j;
@@ -64,7 +62,6 @@ public final class ExecutorInstrumentation extends Instrumenter.Default {
"scala.concurrent.impl.ExecutionContextImpl$$anon$3",
"akka.dispatch.MessageDispatcher",
"akka.dispatch.Dispatcher",
- "akka.dispatch.Dispatcher$LazyExecutorServiceDelegate",
"akka.actor.ActorSystemImpl$$anon$1",
"akka.dispatch.ForkJoinExecutorConfigurator$AkkaForkJoinPool",
"akka.dispatch.forkjoin.ForkJoinPool",
@@ -128,37 +125,46 @@ public final class ExecutorInstrumentation extends Instrumenter.Default {
public String[] helperClassNames() {
return new String[] {
ExecutorInstrumentation.class.getName() + "$ConcurrentUtils",
- ExecutorInstrumentation.class.getName() + "$DatadogWrapper",
- ExecutorInstrumentation.class.getName() + "$CallableWrapper",
- ExecutorInstrumentation.class.getName() + "$RunnableWrapper"
};
}
+ @Override
+ public Map<String, String> contextStore() {
+ final Map<String, String> map = new HashMap<>();
+ map.put("java.lang.Runnable", State.class.getName());
+ map.put("java.util.concurrent.Callable", State.class.getName());
+ map.put("java.util.concurrent.Future", State.class.getName());
+ return Collections.unmodifiableMap(map);
+ }
+
@Override
public Map<ElementMatcher, String> transformers() {
final Map<ElementMatcher, String> transformers = new HashMap<>();
transformers.put(
- named("execute").and(takesArgument(0, Runnable.class)), WrapRunnableAdvice.class.getName());
+ named("execute").and(takesArgument(0, Runnable.class)),
+ SetRunnableStateAdvice.class.getName());
transformers.put(
- named("submit").and(takesArgument(0, Runnable.class)), WrapRunnableAdvice.class.getName());
+ named("submit").and(takesArgument(0, Runnable.class)),
+ SetRunnableStateAdvice.class.getName());
transformers.put(
- named("submit").and(takesArgument(0, Callable.class)), WrapCallableAdvice.class.getName());
+ named("submit").and(takesArgument(0, Callable.class)),
+ SetCallableStateAdvice.class.getName());
transformers.put(
nameMatches("invoke(Any|All)$").and(takesArgument(0, Callable.class)),
- WrapCallableCollectionAdvice.class.getName());
+ SetCallableStateForCallableCollectionAdvice.class.getName());
return transformers;
}
- public static class WrapRunnableAdvice {
+ public static class SetRunnableStateAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
- public static DatadogWrapper enterJobSubmit(
- @Advice.This final Executor executor,
- @Advice.Argument(value = 0, readOnly = false) Runnable task) {
+ public static State enterJobSubmit(
+ @Advice.This final Executor executor, @Advice.Argument(value = 0) final Runnable task) {
final Scope scope = GlobalTracer.get().scopeManager().active();
- if (DatadogWrapper.shouldWrapTask(task, executor)) {
- task = new RunnableWrapper(task, (TraceScope) scope);
- return (RunnableWrapper) task;
+ if (ConcurrentUtils.shouldAttachStateToTask(task, executor)) {
+ final ContextStore<Runnable, State> contextStore =
+ InstrumentationContext.get(Runnable.class, State.class);
+ return ConcurrentUtils.setupState(contextStore, task, (TraceScope) scope);
}
return null;
}
@@ -166,23 +172,22 @@ public final class ExecutorInstrumentation extends Instrumenter.Default {
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void exitJobSubmit(
@Advice.This final Executor executor,
- @Advice.Enter final DatadogWrapper wrapper,
+ @Advice.Enter final State state,
@Advice.Thrown final Throwable throwable) {
- DatadogWrapper.cleanUpOnMethodExit(executor, wrapper, throwable);
+ ConcurrentUtils.cleanUpOnMethodExit(executor, state, throwable);
}
}
- public static class WrapCallableAdvice {
+ public static class SetCallableStateAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
- public static DatadogWrapper enterJobSubmit(
- @Advice.This final Executor executor,
- @Advice.Argument(value = 0, readOnly = false) Callable<?> task) {
-
+ public static State enterJobSubmit(
+ @Advice.This final Executor executor, @Advice.Argument(value = 0) final Callable task) {
final Scope scope = GlobalTracer.get().scopeManager().active();
- if (DatadogWrapper.shouldWrapTask(task, executor)) {
- task = new CallableWrapper<>(task, (TraceScope) scope);
- return (CallableWrapper) task;
+ if (ConcurrentUtils.shouldAttachStateToTask(task, executor)) {
+ final ContextStore<Callable, State> contextStore =
+ InstrumentationContext.get(Callable.class, State.class);
+ return ConcurrentUtils.setupState(contextStore, task, (TraceScope) scope);
}
return null;
}
@@ -190,47 +195,65 @@ public final class ExecutorInstrumentation extends Instrumenter.Default {
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void exitJobSubmit(
@Advice.This final Executor executor,
- @Advice.Enter final DatadogWrapper wrapper,
- @Advice.Thrown final Throwable throwable) {
- DatadogWrapper.cleanUpOnMethodExit(executor, wrapper, throwable);
+ @Advice.Enter final State state,
+ @Advice.Thrown final Throwable throwable,
+ @Advice.Return final Future future) {
+ if (state != null && future != null) {
+ final ContextStore<Future, State> contextStore =
+ InstrumentationContext.get(Future.class, State.class);
+ contextStore.put(future, state);
+ }
+ ConcurrentUtils.cleanUpOnMethodExit(executor, state, throwable);
}
}
- public static class WrapCallableCollectionAdvice {
+ public static class SetCallableStateForCallableCollectionAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
- public static Collection<?> wrapJob(
+ public static void wrapJob(
@Advice.This final Executor executor,
- @Advice.Argument(value = 0, readOnly = false) Collection<? extends Callable<?>> tasks) {
+ @Advice.Argument(value = 0) final Collection<? extends Callable<?>> tasks) {
final Scope scope = GlobalTracer.get().scopeManager().active();
if (scope instanceof TraceScope
&& ((TraceScope) scope).isAsyncPropagating()
- && tasks != null
- && DatadogWrapper.isTopLevelCall(executor)) {
- final Collection<Callable<?>> wrappedTasks = new ArrayList<>(tasks.size());
+ && tasks != null) {
for (final Callable<?> task : tasks) {
- if (task != null && !(task instanceof CallableWrapper)) {
- wrappedTasks.add(new CallableWrapper<>(task, (TraceScope) scope));
+ if (task != null) {
+ final ContextStore<Callable, State> contextStore =
+ InstrumentationContext.get(Callable.class, State.class);
+ ConcurrentUtils.setupState(contextStore, task, (TraceScope) scope);
}
}
- tasks = wrappedTasks;
- return tasks;
}
- return null;
}
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void checkCancel(
@Advice.This final Executor executor,
- @Advice.Enter final Collection<?> wrappedJobs,
+ @Advice.Argument(value = 0) final Collection<? extends Callable<?>> tasks,
@Advice.Thrown final Throwable throwable) {
- if (null != wrappedJobs) {
- DatadogWrapper.resetNestedCalls(executor);
-
+ /*
+ Note1: invokeAny doesn't return any futures so all we need to do for it
+ is to make sure we close all scopes in case of an exception.
+ Note2: invokeAll does return futures - but according to its documentation
+ it actually only returns after all futures have been completed - i.e. it blocks.
+ This means we do not need to setup any hooks on these futures, we just need to clean
+ up any continuations in case of an error.
+ (according to ExecutorService docs and AbstractExecutorService code)
+ */
if (null != throwable) {
- for (final Object wrapper : wrappedJobs) {
- if (wrapper instanceof DatadogWrapper) {
- ((DatadogWrapper) wrapper).cancel();
+ for (final Callable<?> task : tasks) {
+ if (task != null) {
+ final ContextStore<Callable, State> contextStore =
+ InstrumentationContext.get(Callable.class, State.class);
+ final State state = contextStore.get(task);
+ if (state != null) {
+ // Note: this may potentially close somebody else's continuation if we didn't set it up
+ // in setupState because it was already present before us. This should be safe but may lead
+ // to non-attributed async work in some very rare cases.
+ // Alternative is to not close continuation here if we did not set it up in setupState but
+ // this may potentially lead to memory leaks if callers do not properly handle exceptions.
+ state.closeContinuation();
}
}
}
@@ -238,169 +261,64 @@ public final class ExecutorInstrumentation extends Instrumenter.Default {
}
}
- /** Marker interface for tasks which are wrapped to propagate the trace context. */
+ /** Utils for concurrent instrumentations. */
@Slf4j
- public abstract static class DatadogWrapper {
- protected final TraceScope.Continuation continuation;
-
- public DatadogWrapper(final TraceScope scope) {
- continuation = scope.capture();
- log.debug("created continuation {} from scope {}", continuation, scope);
- }
-
- public void cancel() {
- if (null != continuation) {
- continuation.close();
- log.debug("canceled continuation {}", continuation);
- }
- }
+ public static class ConcurrentUtils {
/**
- * Check if given call to executor is nested. We would like to ignore nested calls to execute to
- * avoid wrapping tasks twice. Note: this condition may lead to problems with executors that
- * 'fork' several tasks, but we do not have such executors at the moment. Note: this condition
- * is mutating and needs to be checked right before task is actually wrapped.
+ * Checks if given task should get state attached.
+ *
+ * <p>Warning: this also increments nested call counter!
*
- * @return true iff call is not nested
- */
- 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(final Executor executor) {
- CallDepthThreadLocalMap.reset(executor.getClass());
- }
-
- /**
* @param task task object
+ * @param executor executor this task was scheduled on
* @return true iff given task object should be wrapped
*/
- public static boolean shouldWrapTask(final Object task, final Executor executor) {
+ public static boolean shouldAttachStateToTask(final Object task, final Executor executor) {
final Scope scope = GlobalTracer.get().scopeManager().active();
return (scope instanceof TraceScope
&& ((TraceScope) scope).isAsyncPropagating()
- && task != null
- && !(task instanceof DatadogWrapper)
- && isTopLevelCall(executor)
- && !ConcurrentUtils.isDisabled(executor));
+ && task != null);
}
/**
- * Clean up after job submission method has exited
+ * Create task state given current scope
*
- * @param executor the current executor
- * @param wrapper task wrapper
- * @param throwable throwable that may have been thrown
+ * @param contextStore context storage
+ * @param task task instance
+ * @param scope current scope
+ * @param <T> task class type
+ * @return new state
*/
- public static void cleanUpOnMethodExit(
- final Executor executor, final DatadogWrapper wrapper, final Throwable throwable) {
- if (null != wrapper) {
- resetNestedCalls(executor);
- if (null != throwable) {
- wrapper.cancel();
- }
- }
- }
- }
-
- @Slf4j
- public static class RunnableWrapper extends DatadogWrapper implements Runnable {
- private final Runnable delegatee;
-
- public RunnableWrapper(final Runnable toWrap, final TraceScope scope) {
- super(scope);
- delegatee = toWrap;
- }
-
- @Override
- public void run() {
- final TraceScope context = continuation.activate();
- context.setAsyncPropagation(true);
- try {
- delegatee.run();
- } finally {
- context.close();
- }
- }
- }
-
- @Slf4j
- public static class CallableWrapper<T> extends DatadogWrapper implements Callable<T> {
- private final Callable<T> delegatee;
-
- public CallableWrapper(final Callable<T> toWrap, final TraceScope scope) {
- super(scope);
- delegatee = toWrap;
- }
-
- @Override
- public T call() throws Exception {
- final TraceScope context = continuation.activate();
- context.setAsyncPropagation(true);
- try {
- return delegatee.call();
- } finally {
- context.close();
- }
- }
- }
-
- /** Utils for pulling DatadogWrapper out of Future instances. */
- @Slf4j
- public static class ConcurrentUtils {
- private static final WeakMap<Executor, Boolean> disabledExecutors =
- WeakMap.Provider.newWeakMap();
- private static final Map<Class<?>, Field> fieldCache = new ConcurrentHashMap<>();
- private static final String[] wrapperFields = {"runnable", "callable"};
-
- public static void disableExecutor(final Executor executor) {
- log.debug("Disabling Executor tracing for instance {}", executor);
- disabledExecutors.put(executor, true);
- }
-
- public static boolean isDisabled(final Executor executor) {
- return disabledExecutors.containsKey(executor);
- }
-
- public static DatadogWrapper getDatadogWrapper(final Future<?> f) {
- final Field field;
- if (fieldCache.containsKey(f.getClass())) {
- field = fieldCache.get(f.getClass());
+ public static <T> State setupState(
+ final ContextStore<T, State> contextStore, final T task, final TraceScope scope) {
+ final State state = contextStore.putIfAbsent(task, State.FACTORY);
+ final TraceScope.Continuation continuation = scope.capture();
+ if (state.setContinuation(continuation)) {
+ log.debug("created continuation {} from scope {}, state: {}", continuation, scope, state);
} else {
- field = getWrapperField(f.getClass());
- fieldCache.put(f.getClass(), field);
- }
-
- if (field != null) {
- try {
- field.setAccessible(true);
- final Object o = field.get(f);
- if (o instanceof DatadogWrapper) {
- return (DatadogWrapper) o;
- }
- } catch (final IllegalArgumentException | IllegalAccessException e) {
- } finally {
- field.setAccessible(false);
+ continuation.close();
}
- }
- return null;
+ return state;
}
- private static Field getWrapperField(Class<?> clazz) {
- Field field = null;
- while (field == null && clazz != null) {
- for (int i = 0; i < wrapperFields.length; ++i) {
- try {
- field = clazz.getDeclaredField(wrapperFields[i]);
- break;
- } catch (final Exception e) {
- }
- }
- clazz = clazz.getSuperclass();
+ /**
+ * Clean up after job submission method has exited
+ *
+ * @param executor the current executor
+ * @param state task instrumentation state
+ * @param throwable throwable that may have been thrown
+ */
+ public static void cleanUpOnMethodExit(
+ final Executor executor, final State state, final Throwable throwable) {
+ if (null != state && null != throwable) {
+ // Note: this may potentially close somebody else's continuation if we didn't set it up
+ // in setupState because it was already present before us. This should be safe but may lead
+ // to non-attributed async work in some very rare cases.
+ // Alternative is to not close continuation here if we did not set it up in setupState but
+ // this may potentially lead to memory leaks if callers do not properly handle exceptions.
+ state.closeContinuation();
}
- return field;
}
}
}

MODIFY
dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/FutureInstrumentation.java
dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/FutureInstrumentation.java
@@ -8,8 +8,9 @@ import static net.bytebuddy.matcher.ElementMatchers.returns;
import com.google.auto.service.AutoService;
import datadog.trace.agent.tooling.Instrumenter;
-import datadog.trace.instrumentation.java.concurrent.ExecutorInstrumentation.ConcurrentUtils;
-import datadog.trace.instrumentation.java.concurrent.ExecutorInstrumentation.DatadogWrapper;
+import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.bootstrap.InstrumentationContext;
+import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
@@ -86,12 +87,14 @@ public final class FutureInstrumentation extends Instrumenter.Default {
@Override
public String[] helperClassNames() {
- return new String[] {
- ExecutorInstrumentation.class.getName() + "$ConcurrentUtils",
- ExecutorInstrumentation.class.getName() + "$DatadogWrapper",
- ExecutorInstrumentation.class.getName() + "$CallableWrapper",
- ExecutorInstrumentation.class.getName() + "$RunnableWrapper"
- };
+ return new String[] {ExecutorInstrumentation.class.getName() + "$ConcurrentUtils"};
+ }
+
+ @Override
+ public Map<String, String> contextStore() {
+ final Map<String, String> map = new HashMap<>();
+ map.put("java.util.concurrent.Future", State.class.getName());
+ return Collections.unmodifiableMap(map);
}
@Override
@@ -103,16 +106,16 @@ public final class FutureInstrumentation extends Instrumenter.Default {
}
public static class CanceledFutureAdvice {
- @Advice.OnMethodEnter(suppress = Throwable.class)
- public static DatadogWrapper findWrapper(@Advice.This final Future<?> future) {
- return ConcurrentUtils.getDatadogWrapper(future);
- }
-
@Advice.OnMethodExit(suppress = Throwable.class)
- public static void abortTracing(
- @Advice.Enter final DatadogWrapper wrapper, @Advice.Return final boolean canceled) {
- if (canceled && null != wrapper) {
- wrapper.cancel();
+ public static void exit(
+ @Advice.This final Future<?> future, @Advice.Return final boolean canceled) {
+ if (canceled) {
+ final ContextStore<Future, State> contextStore =
+ InstrumentationContext.get(Future.class, State.class);
+ final State state = contextStore.get(future);
+ if (state != null) {
+ state.closeContinuation();
+ }
}
}
}

ADD
dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/RunnableInstrumentation.java
+package datadog.trace.instrumentation.java.concurrent;
+
+import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
+import static net.bytebuddy.matcher.ElementMatchers.isInterface;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.not;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.bootstrap.InstrumentationContext;
+import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
+import datadog.trace.context.TraceScope;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.Callable;
+import lombok.extern.slf4j.Slf4j;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@Slf4j
+@AutoService(Instrumenter.class)
+public final class RunnableInstrumentation extends Instrumenter.Default {
+
+ public RunnableInstrumentation() {
+ super(ExecutorInstrumentation.EXEC_NAME);
+ }
+
+ @Override
+ public ElementMatcher<TypeDescription> typeMatcher() {
+ return not(isInterface())
+ .and(safeHasSuperType(named(Runnable.class.getName()).or(named(Callable.class.getName()))));
+ }
+
+ @Override
+ public String[] helperClassNames() {
+ return new String[] {
+ RunnableInstrumentation.class.getName() + "$RunnableUtils",
+ };
+ }
+
+ @Override
+ public Map<String, String> contextStore() {
+ final Map<String, String> map = new HashMap<>();
+ map.put("java.lang.Runnable", State.class.getName());
+ map.put("java.util.concurrent.Callable", State.class.getName());
+ return Collections.unmodifiableMap(map);
+ }
+
+ @Override
+ public Map<ElementMatcher, String> transformers() {
+ final Map<ElementMatcher, String> transformers = new HashMap<>();
+ transformers.put(named("run").and(takesArguments(0)), RunnableAdvice.class.getName());
+ transformers.put(named("call").and(takesArguments(0)), CallableAdvice.class.getName());
+ return transformers;
+ }
+
+ public static class RunnableAdvice {
+
+ @Advice.OnMethodEnter(suppress = Throwable.class)
+ public static TraceScope enter(@Advice.This final Runnable thiz) {
+ final ContextStore<Runnable, State> contextStore =
+ InstrumentationContext.get(Runnable.class, State.class);
+ return RunnableUtils.startTaskScope(contextStore, thiz);
+ }
+
+ @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+ public static void exit(@Advice.Enter final TraceScope scope) {
+ RunnableUtils.endTaskScope(scope);
+ }
+ }
+
+ public static class CallableAdvice {
+
+ @Advice.OnMethodEnter(suppress = Throwable.class)
+ public static TraceScope enter(@Advice.This final Callable thiz) {
+ final ContextStore<Callable, State> contextStore =
+ InstrumentationContext.get(Callable.class, State.class);
+ return RunnableUtils.startTaskScope(contextStore, thiz);
+ }
+
+ @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+ public static void exit(@Advice.Enter final TraceScope scope) {
+ RunnableUtils.endTaskScope(scope);
+ }
+ }
+
+ /** Helper utils for Runnable/Callable instrumentation */
+ @Slf4j
+ public static class RunnableUtils {
+
+ /**
+ * Start scope for a given task
+ *
+ * @param contextStore context storage for task's state
+ * @param task task to start scope for
+ * @param <T> task's type
+ * @return scope if scope was started, or null
+ */
+ public static <T> TraceScope startTaskScope(
+ final ContextStore<T, State> contextStore, final T task) {
+ final State state = contextStore.get(task);
+ if (state != null) {
+ final TraceScope.Continuation continuation = state.getAndResetContinuation();
+ if (continuation != null) {
+ final TraceScope scope = continuation.activate();
+ scope.setAsyncPropagation(true);
+ return scope;
+ }
+ }
+ return null;
+ }
+
+ public static void endTaskScope(final TraceScope scope) {
+ if (scope != null) {
+ scope.close();
+ }
+ }
+ }
+}

DELETE
dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/ThreadPoolExecutorInstrumentation.java
-package datadog.trace.instrumentation.java.concurrent;
-
-import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
-import static net.bytebuddy.matcher.ElementMatchers.named;
-import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
-import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
-
-import com.google.auto.service.AutoService;
-import datadog.trace.agent.tooling.Instrumenter;
-import java.util.Collections;
-import java.util.Map;
-import java.util.concurrent.BlockingQueue;
-import java.util.concurrent.ThreadPoolExecutor;
-import lombok.extern.slf4j.Slf4j;
-import net.bytebuddy.asm.Advice;
-import net.bytebuddy.description.type.TypeDescription;
-import net.bytebuddy.matcher.ElementMatcher;
-
-@Slf4j
-@AutoService(Instrumenter.class)
-public class ThreadPoolExecutorInstrumentation extends Instrumenter.Default {
-
- public ThreadPoolExecutorInstrumentation() {
- super(ExecutorInstrumentation.EXEC_NAME);
- }
-
- @Override
- public ElementMatcher<? super TypeDescription> typeMatcher() {
- return named("java.util.concurrent.ThreadPoolExecutor");
- }
-
- @Override
- public String[] helperClassNames() {
- return new String[] {
- ExecutorInstrumentation.class.getName() + "$ConcurrentUtils",
- ThreadPoolExecutorInstrumentation.class.getName() + "$GenericRunnable",
- };
- }
-
- @Override
- public Map<? extends ElementMatcher, String> transformers() {
- return Collections.singletonMap(
- isConstructor()
- .and(takesArgument(4, named("java.util.concurrent.BlockingQueue")))
- .and(takesArguments(7)),
- ThreadPoolExecutorAdvice.class.getName());
- }
-
- public static class ThreadPoolExecutorAdvice {
- @Advice.OnMethodExit(suppress = Throwable.class)
- public static void disableIfQueueWrongType(
- @Advice.This final ThreadPoolExecutor executor,
- @Advice.Argument(4) final BlockingQueue queue) {
-
- if (queue.size() == 0) {
- try {
- queue.offer(new GenericRunnable());
- queue.clear(); // Remove the Runnable we just added.
- } catch (final ClassCastException | IllegalArgumentException e) {
- // These errors indicate the queue is fundamentally incompatible with wrapped runnables.
- // We must disable the executor instance to avoid passing wrapped runnables later.
- ExecutorInstrumentation.ConcurrentUtils.disableExecutor(executor);
- } catch (final Exception e) {
- // Other errors might indicate the queue is not fully initialized.
- // We might want to disable for those too, but for now just ignore.
- }
- }
- }
- }
-
- public static class GenericRunnable implements Runnable {
-
- @Override
- public void run() {}
- }
-}

MODIFY
dd-java-agent/instrumentation/java-concurrent/src/slickTest/groovy/SlickTest.groovy
dd-java-agent/instrumentation/java-concurrent/src/slickTest/groovy/SlickTest.groovy
@@ -6,7 +6,6 @@ import io.opentracing.tag.Tags
class SlickTest extends AgentTestRunner {
// Can't be @Shared, otherwise the work queue is initialized before the instrumentation is applied
- // @Shared
def database = new SlickUtils()
def "Basic statement generates spans"() {

MODIFY
dd-java-agent/instrumentation/java-concurrent/src/test/groovy/ExecutorInstrumentationTest.groovy
dd-java-agent/instrumentation/java-concurrent/src/test/groovy/ExecutorInstrumentationTest.groovy
@@ -73,8 +73,8 @@ class ExecutorInstrumentationTest extends AgentTestRunner {
def "#poolImpl reports after canceled jobs"() {
setup:
def pool = poolImpl
- final AsyncChild child = new AsyncChild(true, true)
- List<Future> jobFutures = new ArrayList<Future>()
+ List<AsyncChild> children = new ArrayList<>()
+ List<Future> jobFutures = new ArrayList<>()
new Runnable() {
@Override
@@ -83,6 +83,14 @@ class ExecutorInstrumentationTest extends AgentTestRunner {
((ContinuableScope) GlobalTracer.get().scopeManager().active()).setAsyncPropagation(true)
try {
for (int i = 0; i < 20; ++i) {
+ // Our current instrumentation instrumentation does not behave very well
+ // if we try to reuse Callable/Runnable. Namely we would be getting 'orphaned'
+ // child traces sometimes since state can contain only one continuation - and
+ // we do not really have a good way for attributing work to correct parent span
+ // if we reuse Callable/Runnable.
+ // Solution for now is to never reuse a Callable/Runnable.
+ final AsyncChild child = new AsyncChild(true, true)
+ children.add(child)
Future f = pool.submit((Callable) child)
jobFutures.add(f)
}
@@ -92,8 +100,10 @@ class ExecutorInstrumentationTest extends AgentTestRunner {
for (Future f : jobFutures) {
f.cancel(false)
}
+ for (AsyncChild child : children) {
child.unblock()
}
+ }
}.run()
TEST_WRITER.waitForTraces(1)

MODIFY
dd-java-agent/instrumentation/java-concurrent/src/test/java/AsyncChild.java
dd-java-agent/instrumentation/java-concurrent/src/test/java/AsyncChild.java
@@ -12,7 +12,7 @@ public class AsyncChild implements Runnable, Callable {
this(true, false);
}
- public AsyncChild(boolean doTraceableWork, boolean blockThread) {
+ public AsyncChild(final boolean doTraceableWork, final boolean blockThread) {
this.doTraceableWork = doTraceableWork;
this.blockThread = new AtomicBoolean(blockThread);
}

Switch executor instrumentation to use context storage