-
Notifications
You must be signed in to change notification settings - Fork 286
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #2334 from DataDog/rgs/tpe-instrumentation
Targeted Instrumentation for ThreadPoolExecutor
- Loading branch information
Showing
8 changed files
with
218 additions
and
45 deletions.
There are no files selected for viewing
30 changes: 3 additions & 27 deletions
30
...main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/ComparableRunnable.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,38 +1,14 @@ | ||
| package datadog.trace.bootstrap.instrumentation.java.concurrent; | ||
|
|
||
| import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope; | ||
|
|
||
| import datadog.trace.context.TraceScope; | ||
|
|
||
| public final class ComparableRunnable<T extends Runnable & Comparable<T>> | ||
| implements Runnable, Comparable<ComparableRunnable<T>> { | ||
|
|
||
| private final T delegate; | ||
| private final TraceScope.Continuation continuation; | ||
| public final class ComparableRunnable<T extends Runnable & Comparable<T>> extends Wrapper<T> | ||
| implements Comparable<ComparableRunnable<T>> { | ||
|
|
||
| public ComparableRunnable(T delegate) { | ||
| this.delegate = delegate; | ||
| TraceScope.Continuation continuation = null; | ||
| TraceScope scope = activeScope(); | ||
| if (null != scope) { | ||
| continuation = scope.capture(); | ||
| } | ||
| this.continuation = continuation; | ||
| super(delegate); | ||
| } | ||
|
|
||
| @Override | ||
| public int compareTo(ComparableRunnable<T> o) { | ||
| return delegate.compareTo(o.delegate); | ||
| } | ||
|
|
||
| @Override | ||
| public void run() { | ||
| try (TraceScope scope = activate()) { | ||
| delegate.run(); | ||
| } | ||
| } | ||
|
|
||
| private TraceScope activate() { | ||
| return null == continuation ? null : continuation.activate(); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
42 changes: 42 additions & 0 deletions
42
...tstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/Wrapper.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,42 @@ | ||
| package datadog.trace.bootstrap.instrumentation.java.concurrent; | ||
|
|
||
| import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope; | ||
|
|
||
| import datadog.trace.context.TraceScope; | ||
|
|
||
| public class Wrapper<T extends Runnable> implements Runnable { | ||
|
|
||
| protected final T delegate; | ||
| private final TraceScope.Continuation continuation; | ||
|
|
||
| public Wrapper(T delegate) { | ||
| this.delegate = delegate; | ||
| TraceScope.Continuation continuation = null; | ||
| TraceScope scope = activeScope(); | ||
| if (null != scope) { | ||
| continuation = scope.capture(); | ||
| } | ||
| this.continuation = continuation; | ||
| } | ||
|
|
||
| @Override | ||
| public void run() { | ||
| try (TraceScope scope = activate()) { | ||
| delegate.run(); | ||
| } | ||
| } | ||
|
|
||
| public void cancel() { | ||
| if (null != continuation) { | ||
| continuation.cancel(); | ||
| } | ||
| } | ||
|
|
||
| public T unwrap() { | ||
| return delegate; | ||
| } | ||
|
|
||
| private TraceScope activate() { | ||
| return null == continuation ? null : continuation.activate(); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
103 changes: 103 additions & 0 deletions
103
...java/datadog/trace/instrumentation/java/concurrent/ThreadPoolExecutorInstrumentation.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,103 @@ | ||
| package datadog.trace.instrumentation.java.concurrent; | ||
|
|
||
| import static datadog.trace.agent.tooling.bytebuddy.matcher.DDElementMatchers.extendsClass; | ||
| import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; | ||
| import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE; | ||
| import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.exclude; | ||
| import static datadog.trace.instrumentation.java.concurrent.AbstractExecutorInstrumentation.EXEC_NAME; | ||
| import static net.bytebuddy.matcher.ElementMatchers.isMethod; | ||
| import static net.bytebuddy.matcher.ElementMatchers.returns; | ||
| import static net.bytebuddy.matcher.ElementMatchers.takesArgument; | ||
|
|
||
| import com.google.auto.service.AutoService; | ||
| import datadog.trace.agent.tooling.Instrumenter; | ||
| import datadog.trace.bootstrap.instrumentation.java.concurrent.ComparableRunnable; | ||
| import datadog.trace.bootstrap.instrumentation.java.concurrent.Wrapper; | ||
| import java.util.Collections; | ||
| import java.util.HashMap; | ||
| import java.util.Map; | ||
| import java.util.concurrent.RunnableFuture; | ||
| import net.bytebuddy.asm.Advice; | ||
| import net.bytebuddy.description.method.MethodDescription; | ||
| import net.bytebuddy.description.type.TypeDescription; | ||
| import net.bytebuddy.matcher.ElementMatcher; | ||
|
|
||
| @AutoService(Instrumenter.class) | ||
| public final class ThreadPoolExecutorInstrumentation extends Instrumenter.Tracing { | ||
|
|
||
| public ThreadPoolExecutorInstrumentation() { | ||
| super(EXEC_NAME); | ||
| } | ||
|
|
||
| @Override | ||
| public ElementMatcher<? super TypeDescription> typeMatcher() { | ||
| return extendsClass(named("java.util.concurrent.ThreadPoolExecutor")); | ||
| } | ||
|
|
||
| @Override | ||
| public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() { | ||
| Map<ElementMatcher<MethodDescription>, String> transformers = new HashMap<>(8); | ||
| transformers.put(named("execute").and(isMethod()), getClass().getName() + "$Execute"); | ||
| transformers.put( | ||
| named("beforeExecute") | ||
| .and(isMethod()) | ||
| .and(takesArgument(1, named(Runnable.class.getName()))), | ||
| getClass().getName() + "$BeforeExecute"); | ||
| transformers.put( | ||
| named("afterExecute") | ||
| .and(isMethod()) | ||
| .and(takesArgument(0, named(Runnable.class.getName()))), | ||
| getClass().getName() + "$AfterExecute"); | ||
| transformers.put( | ||
| named("remove").and(isMethod()).and(returns(Runnable.class)), | ||
| getClass().getName() + "$Remove"); | ||
| return Collections.unmodifiableMap(transformers); | ||
| } | ||
|
|
||
| public static final class Execute { | ||
| @SuppressWarnings({"rawtypes", "unchecked"}) | ||
| @Advice.OnMethodEnter | ||
| public static void wrap(@Advice.Argument(readOnly = false, value = 0) Runnable task) { | ||
| if (task instanceof RunnableFuture || null == task || exclude(RUNNABLE, task)) { | ||
| return; | ||
| } | ||
| if (task instanceof Comparable) { | ||
| task = new ComparableRunnable(task); | ||
| } else { | ||
| task = new Wrapper<>(task); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| public static final class BeforeExecute { | ||
|
|
||
| @Advice.OnMethodEnter | ||
| public static void unwrap(@Advice.Argument(readOnly = false, value = 1) Runnable task) { | ||
| if (task instanceof Wrapper) { | ||
| task = ((Wrapper<?>) task).unwrap(); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| public static final class AfterExecute { | ||
|
|
||
| @Advice.OnMethodEnter | ||
| public static void unwrap(@Advice.Argument(readOnly = false, value = 0) Runnable task) { | ||
| if (task instanceof Wrapper) { | ||
| task = ((Wrapper<?>) task).unwrap(); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| public static final class Remove { | ||
|
|
||
| @Advice.OnMethodExit | ||
| public static void unwrap(@Advice.Return(readOnly = false) Runnable removed) { | ||
| if (removed instanceof Wrapper) { | ||
| Wrapper<?> wrapper = ((Wrapper<?>) removed); | ||
| wrapper.cancel(); | ||
| removed = wrapper.unwrap(); | ||
| } | ||
| } | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
30 changes: 30 additions & 0 deletions
30
dd-java-agent/instrumentation/java-concurrent/src/test/java/TypeAwareThreadPoolExecutor.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,30 @@ | ||
| import datadog.trace.bootstrap.instrumentation.java.concurrent.Wrapper; | ||
| import java.util.concurrent.ArrayBlockingQueue; | ||
| import java.util.concurrent.ThreadPoolExecutor; | ||
| import java.util.concurrent.TimeUnit; | ||
|
|
||
| public class TypeAwareThreadPoolExecutor extends ThreadPoolExecutor { | ||
| public TypeAwareThreadPoolExecutor() { | ||
| super(2, 2, 0, TimeUnit.MICROSECONDS, new ArrayBlockingQueue<Runnable>(100)); | ||
| } | ||
|
|
||
| @Override | ||
| public boolean remove(Runnable task) { | ||
| assertNotWrapper(task); | ||
| return super.remove(task); | ||
| } | ||
|
|
||
| @Override | ||
| protected void beforeExecute(Thread t, Runnable r) { | ||
| assertNotWrapper(r); | ||
| } | ||
|
|
||
| @Override | ||
| protected void afterExecute(Runnable r, Throwable t) { | ||
| assertNotWrapper(r); | ||
| } | ||
|
|
||
| private void assertNotWrapper(Runnable r) { | ||
| assert !(r instanceof Wrapper); | ||
| } | ||
| } |