diff --git a/dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/RunnableCallableInstrumentation.java b/dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/CallableInstrumentation.java similarity index 69% rename from dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/RunnableCallableInstrumentation.java rename to dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/CallableInstrumentation.java index 4795056c84..bced21fffc 100644 --- a/dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/RunnableCallableInstrumentation.java +++ b/dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/CallableInstrumentation.java @@ -24,15 +24,15 @@ import net.bytebuddy.matcher.ElementMatcher; /** Instrument {@link Runnable} and {@Callable} */ @Slf4j @AutoService(Instrumenter.class) -public final class RunnableCallableInstrumentation extends Instrumenter.Default { +public final class CallableInstrumentation extends Instrumenter.Default { - public RunnableCallableInstrumentation() { + public CallableInstrumentation() { super(AbstractExecutorInstrumentation.EXEC_NAME); } @Override public ElementMatcher typeMatcher() { - return implementsInterface(named(Runnable.class.getName()).or(named(Callable.class.getName()))); + return implementsInterface(named(Callable.class.getName())); } @Override @@ -45,7 +45,6 @@ public final class RunnableCallableInstrumentation extends Instrumenter.Default @Override public Map contextStore() { final Map map = new HashMap<>(); - map.put(Runnable.class.getName(), State.class.getName()); map.put(Callable.class.getName(), State.class.getName()); return Collections.unmodifiableMap(map); } @@ -53,30 +52,12 @@ public final class RunnableCallableInstrumentation extends Instrumenter.Default @Override public Map, String> transformers() { final Map, String> transformers = new HashMap<>(); - transformers.put( - named("run").and(takesArguments(0)).and(isPublic()), - RunnableCallableInstrumentation.class.getName() + "$RunnableAdvice"); transformers.put( named("call").and(takesArguments(0)).and(isPublic()), - RunnableCallableInstrumentation.class.getName() + "$CallableAdvice"); + CallableInstrumentation.class.getName() + "$CallableAdvice"); return transformers; } - public static class RunnableAdvice { - - @Advice.OnMethodEnter(suppress = Throwable.class) - public static TraceScope enter(@Advice.This final Runnable thiz) { - final ContextStore contextStore = - InstrumentationContext.get(Runnable.class, State.class); - return AdviceUtils.startTaskScope(contextStore, thiz); - } - - @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) - public static void exit(@Advice.Enter final TraceScope scope) { - AdviceUtils.endTaskScope(scope); - } - } - public static class CallableAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) diff --git a/dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/RunnableInstrumentation.java b/dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/RunnableInstrumentation.java new file mode 100644 index 0000000000..2ae88a96f7 --- /dev/null +++ b/dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/RunnableInstrumentation.java @@ -0,0 +1,74 @@ +package datadog.trace.instrumentation.java.concurrent; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.DDElementMatchers.implementsInterface; +import static net.bytebuddy.matcher.ElementMatchers.isPublic; +import static net.bytebuddy.matcher.ElementMatchers.named; +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 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 Runnable} and {@Callable} */ +@Slf4j +@AutoService(Instrumenter.class) +public final class RunnableInstrumentation extends Instrumenter.Default { + + public RunnableInstrumentation() { + super(AbstractExecutorInstrumentation.EXEC_NAME); + } + + @Override + public ElementMatcher typeMatcher() { + return implementsInterface(named(Runnable.class.getName())); + } + + @Override + public String[] helperClassNames() { + return new String[] { + AdviceUtils.class.getName(), + }; + } + + @Override + public Map contextStore() { + final Map map = new HashMap<>(); + map.put(Runnable.class.getName(), State.class.getName()); + return Collections.unmodifiableMap(map); + } + + @Override + public Map, String> transformers() { + final Map, String> transformers = new HashMap<>(); + transformers.put( + named("run").and(takesArguments(0)).and(isPublic()), + RunnableInstrumentation.class.getName() + "$RunnableAdvice"); + return transformers; + } + + public static class RunnableAdvice { + + @Advice.OnMethodEnter(suppress = Throwable.class) + public static TraceScope enter(@Advice.This final Runnable thiz) { + final ContextStore contextStore = + InstrumentationContext.get(Runnable.class, State.class); + return AdviceUtils.startTaskScope(contextStore, thiz); + } + + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) + public static void exit(@Advice.Enter final TraceScope scope) { + AdviceUtils.endTaskScope(scope); + } + } +}