Handle ForkJoinTask in concurrent instrumentation
ADD
dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/AdviceUtils.java
+package datadog.trace.instrumentation.java.concurrent;
+
+import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
+import datadog.trace.context.TraceScope;
+
+/** Helper utils for Runnable/Callable instrumentation */
+public class AdviceUtils {
+
+ /**
+ * 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();
+ }
+ }
+}

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
@@ -27,6 +27,7 @@ import java.util.HashSet;
import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.Executor;
+import java.util.concurrent.ForkJoinTask;
import java.util.concurrent.Future;
import lombok.extern.slf4j.Slf4j;
import net.bytebuddy.asm.Advice;
@@ -140,6 +141,7 @@ public final class ExecutorInstrumentation extends Instrumenter.Default {
final Map<String, String> map = new HashMap<>();
map.put(Runnable.class.getName(), State.class.getName());
map.put(Callable.class.getName(), State.class.getName());
+ map.put(ForkJoinTask.class.getName(), State.class.getName());
map.put(Future.class.getName(), State.class.getName());
return Collections.unmodifiableMap(map);
}
@@ -150,15 +152,24 @@ public final class ExecutorInstrumentation extends Instrumenter.Default {
transformers.put(
named("execute").and(takesArgument(0, Runnable.class)),
SetExecuteRunnableStateAdvice.class.getName());
+ transformers.put(
+ named("execute").and(takesArgument(0, ForkJoinTask.class)),
+ SetExecuteForkJoinStateAdvice.class.getName());
transformers.put(
named("submit").and(takesArgument(0, Runnable.class)),
SetSubmitRunnableStateAdvice.class.getName());
transformers.put(
named("submit").and(takesArgument(0, Callable.class)),
SetCallableStateAdvice.class.getName());
+ transformers.put(
+ named("submit").and(takesArgument(0, ForkJoinTask.class)),
+ SetExecuteForkJoinStateAdvice.class.getName());
transformers.put(
nameMatches("invoke(Any|All)$").and(takesArgument(0, Callable.class)),
SetCallableStateForCallableCollectionAdvice.class.getName());
+ transformers.put(
+ nameMatches("invoke").and(takesArgument(0, ForkJoinTask.class)),
+ SetExecuteForkJoinStateAdvice.class.getName());
transformers.put( // kotlinx.coroutines.scheduling.CoroutineScheduler
named("dispatch")
.and(takesArgument(0, Runnable.class))
@@ -192,6 +203,30 @@ public final class ExecutorInstrumentation extends Instrumenter.Default {
}
}
+ public static class SetExecuteForkJoinStateAdvice {
+
+ @Advice.OnMethodEnter(suppress = Throwable.class)
+ public static State enterJobSubmit(
+ @Advice.This final Executor executor,
+ @Advice.Argument(value = 0, readOnly = false) final ForkJoinTask task) {
+ final Scope scope = GlobalTracer.get().scopeManager().active();
+ if (ConcurrentUtils.shouldAttachStateToTask(task, executor)) {
+ final ContextStore<ForkJoinTask, State> contextStore =
+ InstrumentationContext.get(ForkJoinTask.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) {
+ ConcurrentUtils.cleanUpOnMethodExit(executor, state, throwable);
+ }
+ }
+
public static class SetSubmitRunnableStateAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)

ADD
dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/ForkJoinTaskInstrumentation.java
+package datadog.trace.instrumentation.java.concurrent;
+
+import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
+import static net.bytebuddy.matcher.ElementMatchers.isAbstract;
+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 java.util.concurrent.ForkJoinPool;
+import java.util.concurrent.ForkJoinTask;
+import lombok.extern.slf4j.Slf4j;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+/**
+ * Instrument {@link ForkJoinTask}.
+ *
+ * <p>Note: There are quite a few separate implementations of {@code ForkJoinTask}/{@code
+ * ForkJoinPool}: JVM, akka, scala, netty to name a few. For now we only deal with JVM one because
+ * there are known cases when JVM {@code ForkJoinTask} is supplied as {@code Runnable} into {@code
+ * ForkJoinPool}.
+ */
+@Slf4j
+@AutoService(Instrumenter.class)
+public final class ForkJoinTaskInstrumentation extends Instrumenter.Default {
+
+ public ForkJoinTaskInstrumentation() {
+ super(ExecutorInstrumentation.EXEC_NAME);
+ }
+
+ @Override
+ public ElementMatcher<TypeDescription> typeMatcher() {
+ return not(isInterface()).and(safeHasSuperType(named(ForkJoinTask.class.getName())));
+ }
+
+ @Override
+ public String[] helperClassNames() {
+ return new String[] {
+ AdviceUtils.class.getName(),
+ };
+ }
+
+ @Override
+ public Map<String, String> contextStore() {
+ final Map<String, String> map = new HashMap<>();
+ map.put(Runnable.class.getName(), State.class.getName());
+ map.put(Callable.class.getName(), State.class.getName());
+ map.put(ForkJoinTask.class.getName(), State.class.getName());
+ return Collections.unmodifiableMap(map);
+ }
+
+ @Override
+ public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+ final Map<ElementMatcher<? super MethodDescription>, String> transformers = new HashMap<>();
+ transformers.put(
+ named("exec").and(takesArguments(0)).and(not(isAbstract())),
+ ForkJoinTaskAdvice.class.getName());
+ return transformers;
+ }
+
+ public static class ForkJoinTaskAdvice {
+
+ /**
+ * When {@link ForkJoinTask} object is submitted to {@link ForkJoinPool} as {@link Runnable} or
+ * {@link Callable} it will not get wrapped, instead it will be casted to {@code ForkJoinTask}
+ * directly. This means state is still stored in {@code Runnable} or {@code Callable} and we
+ * need to use that state.
+ */
+ @Advice.OnMethodEnter(suppress = Throwable.class)
+ public static TraceScope enter(@Advice.This final ForkJoinTask thiz) {
+ final ContextStore<ForkJoinTask, State> contextStore =
+ InstrumentationContext.get(ForkJoinTask.class, State.class);
+ TraceScope scope = AdviceUtils.startTaskScope(contextStore, thiz);
+ if (thiz instanceof Runnable) {
+ final ContextStore<Runnable, State> runnableContextStore =
+ InstrumentationContext.get(Runnable.class, State.class);
+ final TraceScope newScope =
+ AdviceUtils.startTaskScope(runnableContextStore, (Runnable) thiz);
+ if (null != newScope) {
+ if (null != scope) {
+ newScope.close();
+ } else {
+ scope = newScope;
+ }
+ }
+ }
+ if (thiz instanceof Callable) {
+ final ContextStore<Callable, State> callableContextStore =
+ InstrumentationContext.get(Callable.class, State.class);
+ final TraceScope newScope =
+ AdviceUtils.startTaskScope(callableContextStore, (Callable) thiz);
+ if (null != newScope) {
+ if (null != scope) {
+ newScope.close();
+ } else {
+ scope = newScope;
+ }
+ }
+ }
+ return scope;
+ }
+
+ @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+ public static void exit(@Advice.Enter final TraceScope scope) {
+ AdviceUtils.endTaskScope(scope);
+ }
+ }
+}

MODIFY
dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/RunnableCallableInstrumentation.java
dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/RunnableCallableInstrumentation.java
@@ -2,6 +2,7 @@ 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.isPublic;
import static net.bytebuddy.matcher.ElementMatchers.named;
import static net.bytebuddy.matcher.ElementMatchers.not;
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
@@ -22,6 +23,7 @@ import net.bytebuddy.description.method.MethodDescription;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;
+/** Instrument {@link Runnable} and {@Callable} */
@Slf4j
@AutoService(Instrumenter.class)
public final class RunnableCallableInstrumentation extends Instrumenter.Default {
@@ -39,7 +41,7 @@ public final class RunnableCallableInstrumentation extends Instrumenter.Default
@Override
public String[] helperClassNames() {
return new String[] {
- RunnableCallableInstrumentation.class.getName() + "$RunnableUtils",
+ AdviceUtils.class.getName(),
};
}
@@ -54,8 +56,10 @@ public final class RunnableCallableInstrumentation extends Instrumenter.Default
@Override
public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
final Map<ElementMatcher<? super MethodDescription>, String> transformers = new HashMap<>();
- transformers.put(named("run").and(takesArguments(0)), RunnableAdvice.class.getName());
- transformers.put(named("call").and(takesArguments(0)), CallableAdvice.class.getName());
+ transformers.put(
+ named("run").and(takesArguments(0)).and(isPublic()), RunnableAdvice.class.getName());
+ transformers.put(
+ named("call").and(takesArguments(0)).and(isPublic()), CallableAdvice.class.getName());
return transformers;
}
@@ -65,12 +69,12 @@ public final class RunnableCallableInstrumentation extends Instrumenter.Default
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);
+ return AdviceUtils.startTaskScope(contextStore, thiz);
}
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void exit(@Advice.Enter final TraceScope scope) {
- RunnableUtils.endTaskScope(scope);
+ AdviceUtils.endTaskScope(scope);
}
}
@@ -80,45 +84,12 @@ public final class RunnableCallableInstrumentation extends Instrumenter.Default
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);
+ return AdviceUtils.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();
- }
+ AdviceUtils.endTaskScope(scope);
}
}
}

MODIFY
dd-java-agent/instrumentation/java-concurrent/src/test/groovy/ExecutorInstrumentationTest.groovy
dd-java-agent/instrumentation/java-concurrent/src/test/groovy/ExecutorInstrumentationTest.groovy
@@ -12,6 +12,7 @@ import java.util.concurrent.Callable
import java.util.concurrent.Executor
import java.util.concurrent.ExecutorService
import java.util.concurrent.ForkJoinPool
+import java.util.concurrent.ForkJoinTask
import java.util.concurrent.Future
import java.util.concurrent.RejectedExecutionException
import java.util.concurrent.ThreadPoolExecutor
@@ -19,16 +20,25 @@ import java.util.concurrent.TimeUnit
class ExecutorInstrumentationTest extends AgentTestRunner {
@Shared
- Method executeMethod
+ Method executeRunnableMethod
+ @Shared
+ Method executeForkJoinTaskMethod
@Shared
Method submitRunnableMethod
@Shared
Method submitCallableMethod
+ @Shared
+ Method submitForkJoinTaskMethod
+ @Shared
+ Method invokeForkJoinTaskMethod
def setupSpec() {
- executeMethod = Executor.getMethod("execute", Runnable)
+ executeRunnableMethod = Executor.getMethod("execute", Runnable)
+ executeForkJoinTaskMethod = ForkJoinPool.getMethod("execute", ForkJoinTask)
submitRunnableMethod = ExecutorService.getMethod("submit", Runnable)
submitCallableMethod = ExecutorService.getMethod("submit", Callable)
+ submitForkJoinTaskMethod = ForkJoinPool.getMethod("submit", ForkJoinTask)
+ invokeForkJoinTaskMethod = ForkJoinPool.getMethod("invoke", ForkJoinTask)
}
// more useful name breaks java9 javac
@@ -66,12 +76,15 @@ class ExecutorInstrumentationTest extends AgentTestRunner {
// Unfortunately, there's no simple way to test the cross product of methods/pools.
where:
poolImpl | method
+ new ForkJoinPool() | executeRunnableMethod
+ new ForkJoinPool() | executeForkJoinTaskMethod
new ForkJoinPool() | submitRunnableMethod
new ForkJoinPool() | submitCallableMethod
- new ForkJoinPool() | executeMethod
+ new ForkJoinPool() | submitForkJoinTaskMethod
+ new ForkJoinPool() | invokeForkJoinTaskMethod
+ new ThreadPoolExecutor(1, 1, 1000, TimeUnit.NANOSECONDS, new ArrayBlockingQueue<Runnable>(1)) | executeRunnableMethod
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
}
// more useful name breaks java9 javac

MODIFY
dd-java-agent/instrumentation/java-concurrent/src/test/java/AsyncChild.java
dd-java-agent/instrumentation/java-concurrent/src/test/java/AsyncChild.java
import datadog.trace.api.Trace;
import java.util.concurrent.Callable;
+import java.util.concurrent.ForkJoinTask;
import java.util.concurrent.atomic.AtomicBoolean;
-public class AsyncChild implements Runnable, Callable {
+public class AsyncChild extends ForkJoinTask implements Runnable, Callable {
private final AtomicBoolean blockThread;
private final boolean doTraceableWork;
@@ -10,6 +11,20 @@ public class AsyncChild implements Runnable, Callable {
this(true, false);
}
+ @Override
+ public Object getRawResult() {
+ return null;
+ }
+
+ @Override
+ protected void setRawResult(final Object value) {}
+
+ @Override
+ protected boolean exec() {
+ runImpl();
+ return true;
+ }
+
public AsyncChild(final boolean doTraceableWork, final boolean blockThread) {
this.doTraceableWork = doTraceableWork;
this.blockThread = new AtomicBoolean(blockThread);

Handle ForkJoinTask in concurrent instrumentation