Add code to wrap future returned by submit(Runnable) call
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
@@ -146,10 +146,10 @@ public final class ExecutorInstrumentation extends Instrumenter.Default {
final Map<ElementMatcher, String> transformers = new HashMap<>();
transformers.put(
named("execute").and(takesArgument(0, Runnable.class)),
- SetRunnableStateAdvice.class.getName());
+ SetExecuteRunnableStateAdvice.class.getName());
transformers.put(
named("submit").and(takesArgument(0, Runnable.class)),
- SetRunnableStateAdvice.class.getName());
+ SetSubmitRunnableStateAdvice.class.getName());
transformers.put(
named("submit").and(takesArgument(0, Callable.class)),
SetCallableStateAdvice.class.getName());
@@ -159,7 +159,7 @@ public final class ExecutorInstrumentation extends Instrumenter.Default {
return transformers;
}
- public static class SetRunnableStateAdvice {
+ public static class SetExecuteRunnableStateAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static State enterJobSubmit(
@@ -184,6 +184,37 @@ public final class ExecutorInstrumentation extends Instrumenter.Default {
}
}
+ public static class SetSubmitRunnableStateAdvice {
+
+ @Advice.OnMethodEnter(suppress = Throwable.class)
+ public static State enterJobSubmit(
+ @Advice.This final Executor executor,
+ @Advice.Argument(value = 0, readOnly = false) Runnable task) {
+ final Scope scope = GlobalTracer.get().scopeManager().active();
+ if (ConcurrentUtils.shouldAttachStateToTask(task, executor)) {
+ task = RunnableWrapper.wrapIfNeeded(task);
+ final ContextStore<Runnable, State> contextStore =
+ InstrumentationContext.get(Runnable.class, State.class);
+ return ConcurrentUtils.setupState(contextStore, task, (TraceScope) scope);
+ }
+ return null;
+ }
+
+ @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+ public static void exitJobSubmit(
+ @Advice.This final Executor executor,
+ @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 SetCallableStateAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)

MODIFY
dd-java-agent/instrumentation/java-concurrent/src/test/groovy/ExecutorInstrumentationTest.groovy
dd-java-agent/instrumentation/java-concurrent/src/test/groovy/ExecutorInstrumentationTest.groovy
@@ -5,6 +5,7 @@ import datadog.trace.api.Trace
import io.opentracing.util.GlobalTracer
import spock.lang.Shared
+import java.lang.reflect.InvocationTargetException
import java.lang.reflect.Method
import java.util.concurrent.ArrayBlockingQueue
import java.util.concurrent.Callable
@@ -17,14 +18,17 @@ import java.util.concurrent.ThreadPoolExecutor
import java.util.concurrent.TimeUnit
class ExecutorInstrumentationTest extends AgentTestRunner {
- @Shared
- Method submitMethod
@Shared
Method executeMethod
+ @Shared
+ Method submitRunnableMethod
+ @Shared
+ Method submitCallableMethod
def setupSpec() {
executeMethod = Executor.getMethod("execute", Runnable)
- submitMethod = ExecutorService.getMethod("submit", Callable)
+ submitRunnableMethod = ExecutorService.getMethod("submit", Runnable)
+ submitCallableMethod = ExecutorService.getMethod("submit", Callable)
}
// more useful name breaks java9 javac
@@ -62,9 +66,11 @@ class ExecutorInstrumentationTest extends AgentTestRunner {
// Unfortunately, there's no simple way to test the cross product of methods/pools.
where:
poolImpl | method
- new ForkJoinPool() | submitMethod
+ new ForkJoinPool() | submitRunnableMethod
+ new ForkJoinPool() | submitCallableMethod
new ForkJoinPool() | executeMethod
- new ThreadPoolExecutor(1, 1, 1000, TimeUnit.NANOSECONDS, new ArrayBlockingQueue<Runnable>(1)) | submitMethod
+ new ThreadPoolExecutor(1, 1, 1000, TimeUnit.NANOSECONDS, new ArrayBlockingQueue<Runnable>(1)) | submitRunnableMethod
+ new ThreadPoolExecutor(1, 1, 1000, TimeUnit.NANOSECONDS, new ArrayBlockingQueue<Runnable>(1)) | submitCallableMethod
new ThreadPoolExecutor(1, 1, 1000, TimeUnit.NANOSECONDS, new ArrayBlockingQueue<Runnable>(1)) | executeMethod
}
@@ -73,6 +79,7 @@ class ExecutorInstrumentationTest extends AgentTestRunner {
def "#poolImpl reports after canceled jobs"() {
setup:
def pool = poolImpl
+ def m = method
List<AsyncChild> children = new ArrayList<>()
List<Future> jobFutures = new ArrayList<>()
@@ -91,8 +98,12 @@ class ExecutorInstrumentationTest extends AgentTestRunner {
// 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)
+ try {
+ Future f = m.invoke(pool, new AsyncChild())
jobFutures.add(f)
+ } catch (InvocationTargetException e) {
+ throw e.getCause()
+ }
}
} catch (RejectedExecutionException e) {
}
@@ -109,11 +120,14 @@ class ExecutorInstrumentationTest extends AgentTestRunner {
TEST_WRITER.waitForTraces(1)
expect:
+ // FIXME: we should improve this test to make sure continuations are actually closed
TEST_WRITER.size() == 1
where:
- poolImpl | _
- new ForkJoinPool() | _
- new ThreadPoolExecutor(1, 1, 1000, TimeUnit.NANOSECONDS, new ArrayBlockingQueue<Runnable>(1)) | _
+ poolImpl | method
+ new ForkJoinPool() | submitRunnableMethod
+ new ForkJoinPool() | submitCallableMethod
+ new ThreadPoolExecutor(1, 1, 1000, TimeUnit.NANOSECONDS, new ArrayBlockingQueue<Runnable>(1)) | submitRunnableMethod
+ new ThreadPoolExecutor(1, 1, 1000, TimeUnit.NANOSECONDS, new ArrayBlockingQueue<Runnable>(1)) | submitCallableMethod
}
}

Add code to wrap future returned by submit(Runnable) call