diff --git a/dd-java-agent/instrumentation/weaver-0.9/src/main/java/datadog/trace/instrumentation/weaver/WeaverInstrumentation.java b/dd-java-agent/instrumentation/weaver-0.9/src/main/java/datadog/trace/instrumentation/weaver/WeaverInstrumentation.java index fe8e60e4fd2..be033fb009a 100644 --- a/dd-java-agent/instrumentation/weaver-0.9/src/main/java/datadog/trace/instrumentation/weaver/WeaverInstrumentation.java +++ b/dd-java-agent/instrumentation/weaver-0.9/src/main/java/datadog/trace/instrumentation/weaver/WeaverInstrumentation.java @@ -1,12 +1,12 @@ package datadog.trace.instrumentation.weaver; +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; 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 de.thetaphi.forbiddenapis.SuppressForbidden; -import java.lang.reflect.Field; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.LinkedBlockingQueue; import net.bytebuddy.asm.Advice; @@ -38,34 +38,34 @@ public String[] helperClassNames() { @Override public void methodAdvice(MethodTransformer transformer) { + // disneystreaming/weaver-test (0.8.4+) uses a ConcurrentLinkedQueue transformer.applyAdvice( - isConstructor(), WeaverInstrumentation.class.getName() + "$SbtTaskCreationAdvice"); + isConstructor().and(takesArgument(5, named("java.util.concurrent.ConcurrentLinkedQueue"))), + WeaverInstrumentation.class.getName() + "$ConcurrentLinkedQueueAdvice"); + // typelevel/weaver-test (0.9+) uses a LinkedBlockingQueue + transformer.applyAdvice( + isConstructor().and(takesArgument(5, named("java.util.concurrent.LinkedBlockingQueue"))), + WeaverInstrumentation.class.getName() + "$LinkedBlockingQueueAdvice"); + } + + public static class ConcurrentLinkedQueueAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void wrapQueue( + @Advice.Argument(0) TaskDef taskDef, + @Advice.Argument(value = 5, readOnly = false) ConcurrentLinkedQueue queue) { + if (!(queue instanceof TaskDefAwareConcurrentLinkedQueueProxy)) { + queue = new TaskDefAwareConcurrentLinkedQueueProxy<>(taskDef, queue); + } + } } - public static class SbtTaskCreationAdvice { - // TODO: JEP 500 - avoid mutating final fields - @SuppressForbidden - @Advice.OnMethodExit(suppress = Throwable.class) - public static void onTaskCreation( - @Advice.This Object sbtTask, @Advice.FieldValue("taskDef") TaskDef taskDef) { - try { - Field queueField = sbtTask.getClass().getDeclaredField("queue"); - queueField.setAccessible(true); - Object queue = queueField.get(sbtTask); - if (queue instanceof ConcurrentLinkedQueue) { - // disney's implementation (0.8.4+) uses a ConcurrentLinkedQueue for the field - queueField.set( - sbtTask, - new TaskDefAwareConcurrentLinkedQueueProxy( - taskDef, (ConcurrentLinkedQueue) queue)); - } else if (queue instanceof LinkedBlockingQueue) { - // typelevel's implementation (0.9+) uses a LinkedBlockingQueue for the field - queueField.set( - sbtTask, - new TaskDefAwareLinkedBlockingQueueProxy( - taskDef, (LinkedBlockingQueue) queue)); - } - } catch (Exception ignored) { + public static class LinkedBlockingQueueAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void wrapQueue( + @Advice.Argument(0) TaskDef taskDef, + @Advice.Argument(value = 5, readOnly = false) LinkedBlockingQueue queue) { + if (!(queue instanceof TaskDefAwareLinkedBlockingQueueProxy)) { + queue = new TaskDefAwareLinkedBlockingQueueProxy<>(taskDef, queue); } } }