diff --git a/dd-java-agent/instrumentation/spring-scheduling-3.1/src/main/java/datadog/trace/instrumentation/springscheduling/SpringSchedulingInstrumentation.java b/dd-java-agent/instrumentation/spring-scheduling-3.1/src/main/java/datadog/trace/instrumentation/springscheduling/SpringSchedulingInstrumentation.java index 3786cd909c0..c2b29f6b8d2 100644 --- a/dd-java-agent/instrumentation/spring-scheduling-3.1/src/main/java/datadog/trace/instrumentation/springscheduling/SpringSchedulingInstrumentation.java +++ b/dd-java-agent/instrumentation/spring-scheduling-3.1/src/main/java/datadog/trace/instrumentation/springscheduling/SpringSchedulingInstrumentation.java @@ -1,30 +1,24 @@ package datadog.trace.instrumentation.springscheduling; -import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; -import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.nameStartsWith; -import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static net.bytebuddy.matcher.ElementMatchers.isMethod; +import static net.bytebuddy.matcher.ElementMatchers.isConstructor; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; -import datadog.trace.bootstrap.CallDepthThreadLocalMap; import net.bytebuddy.asm.Advice; -import net.bytebuddy.description.type.TypeDescription; -import net.bytebuddy.matcher.ElementMatcher; -import org.springframework.scheduling.TaskScheduler; @AutoService(InstrumenterModule.class) public final class SpringSchedulingInstrumentation extends InstrumenterModule.Tracing - implements Instrumenter.ForTypeHierarchy { + implements Instrumenter.ForSingleType { public SpringSchedulingInstrumentation() { super("spring-scheduling"); } - public String hierarchyMarkerType() { - return "org.springframework.scheduling.TaskScheduler"; + @Override + public String instrumentedType() { + return "org.springframework.scheduling.config.Task"; } @Override @@ -34,36 +28,18 @@ public String[] helperClassNames() { }; } - @Override - public ElementMatcher hierarchyMatcher() { - return implementsInterface(named(hierarchyMarkerType())); - } - @Override public void methodAdvice(MethodTransformer transformer) { transformer.applyAdvice( - isMethod().and(nameStartsWith("schedule")).and(takesArgument(0, Runnable.class)), + isConstructor().and(takesArgument(0, Runnable.class)), getClass().getName() + "$SpringSchedulingAdvice"); } public static class SpringSchedulingAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static boolean beforeSchedule( + public static void onConstruction( @Advice.Argument(value = 0, readOnly = false) Runnable runnable) { - if (CallDepthThreadLocalMap.incrementCallDepth(TaskScheduler.class) > 0) { - return false; - } runnable = SpringSchedulingRunnableWrapper.wrapIfNeeded(runnable); - return true; - } - - @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) - public static void afterSchedule( - @Advice.Enter boolean reset, - @Advice.Argument(value = 0, readOnly = false) Runnable runnable) { - if (reset) { - CallDepthThreadLocalMap.reset(TaskScheduler.class); - } } } } diff --git a/dd-java-agent/instrumentation/spring-scheduling-3.1/src/main/java/datadog/trace/instrumentation/springscheduling/SpringSchedulingRunnableWrapper.java b/dd-java-agent/instrumentation/spring-scheduling-3.1/src/main/java/datadog/trace/instrumentation/springscheduling/SpringSchedulingRunnableWrapper.java index c2f3e4af818..166494c03cb 100644 --- a/dd-java-agent/instrumentation/spring-scheduling-3.1/src/main/java/datadog/trace/instrumentation/springscheduling/SpringSchedulingRunnableWrapper.java +++ b/dd-java-agent/instrumentation/spring-scheduling-3.1/src/main/java/datadog/trace/instrumentation/springscheduling/SpringSchedulingRunnableWrapper.java @@ -44,4 +44,9 @@ public static Runnable wrapIfNeeded(final Runnable task) { } return new SpringSchedulingRunnableWrapper(task); } + + @Override + public String toString() { + return runnable.toString(); + } }