fix: update weaver instrumentation to handle both queue implementations

Author: daniel-mohedano

dd-java-agent/instrumentation/weaver/build.gradle

@@ -2,24 +2,32 @@ apply from: "$rootDir/gradle/java.gradle"
 apply plugin: 'scala'
 
 muzzle {
+  pass {
+    group = 'com.disneystreaming'
+    module = 'weaver-cats_3'
+    versions = '[0.8.4,)'
+  }
+
   pass {
     group = 'org.typelevel'
     module = 'weaver-cats_3'
-    versions = '[0.9.2,)'
+    versions = '[0.9.0,)'
   }
 }
 
 addTestSuiteForDir('latestDepTest', 'test')
+// weaver 0.8.4 is the earliest supported version before the org migration
+addTestSuiteForDir('weaver084Test', 'test')
 
 dependencies {
-  compileOnly group: 'org.typelevel', name: 'weaver-cats_3', version: '0.9.2'
+  compileOnly group: 'org.typelevel', name: 'weaver-cats_3', version: '0.9.0'
 
   testImplementation testFixtures(project(':dd-java-agent:agent-ci-visibility'))
 
   testImplementation group: 'org.scala-lang', name: 'scala-library', version: '2.12.20'
-  testImplementation group: 'org.typelevel', name: 'weaver-cats_3', version: '0.9.2'
-
-  testImplementation group: 'org.typelevel', name: 'weaver-cats_3', version: '+'
+  testImplementation group: 'org.typelevel', name: 'weaver-cats_3', version: '0.9.0'
+  latestDepTestImplementation group: 'org.typelevel', name: 'weaver-cats_3', version: '+'
+  weaver084TestImplementation group: 'com.disneystreaming', name: 'weaver-cats_3', version: '0.8.4'
 }
 
 compileTestGroovy {
@@ -31,3 +39,8 @@ compileLatestDepTestGroovy {
   dependsOn compileLatestDepTestScala
   classpath += files(sourceSets.latestDepTest.scala.destinationDirectory)
 }
+
+compileWeaver084TestGroovy {
+  dependsOn compileWeaver084TestScala
+  classpath += files(sourceSets.weaver084Test.scala.destinationDirectory)
+}

dd-java-agent/instrumentation/weaver/src/main/java/datadog/trace/instrumentation/weaver/DatadogWeaverReporter.java

@@ -14,6 +14,7 @@
 import sbt.testing.TaskDef;
 import weaver.Result;
 import weaver.TestOutcome;
+import weaver.framework.RunEvent;
 import weaver.framework.SuiteFinished;
 import weaver.framework.SuiteStarted;
 import weaver.framework.TestFinished;
@@ -42,6 +43,19 @@ public static synchronized void stop() {
     }
   }
 
+  public static void processEvent(Object event, TaskDef taskDef) {
+    if (event instanceof RunEvent) {
+      // handle event here, using taskDef reference to get suite details
+      if (event instanceof SuiteStarted) {
+        onSuiteStart((SuiteStarted) event);
+      } else if (event instanceof SuiteFinished) {
+        onSuiteFinish((SuiteFinished) event);
+      } else if (event instanceof TestFinished) {
+        onTestFinished((TestFinished) event, taskDef);
+      }
+    }
+  }
+
   public static void onSuiteStart(SuiteStarted event) {
     String testSuiteName = event.name();
     Class<?> testClass = WeaverUtils.getClass(testSuiteName);

dd-java-agent/instrumentation/weaver/src/main/java/datadog/trace/instrumentation/weaver/TaskDefAwareConcurrentLinkedQueueProxy.java

@@ -0,0 +1,24 @@
+package datadog.trace.instrumentation.weaver;
+
+import java.util.concurrent.ConcurrentLinkedQueue;
+import sbt.testing.TaskDef;
+
+public final class TaskDefAwareConcurrentLinkedQueueProxy<T> extends ConcurrentLinkedQueue<T> {
+
+  private final TaskDef taskDef;
+  private final ConcurrentLinkedQueue<T> delegate;
+
+  public TaskDefAwareConcurrentLinkedQueueProxy(TaskDef taskDef, ConcurrentLinkedQueue<T> delegate) {
+    super();
+    this.taskDef = taskDef;
+    this.delegate = delegate;
+    DatadogWeaverReporter.start();
+  }
+
+  @Override
+  public T poll() {
+    T event = delegate.poll();
+    DatadogWeaverReporter.processEvent(event, taskDef);
+    return event;
+  }
+}

dd-java-agent/instrumentation/weaver/src/main/java/datadog/trace/instrumentation/weaver/TaskDefAwareLinkedBlockingQueueProxy.java

@@ -0,0 +1,34 @@
+package datadog.trace.instrumentation.weaver;
+
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.TimeUnit;
+import sbt.testing.TaskDef;
+
+public final class TaskDefAwareLinkedBlockingQueueProxy<T> extends LinkedBlockingQueue<T> {
+
+  private final TaskDef taskDef;
+  private final LinkedBlockingQueue<T> delegate;
+
+  public TaskDefAwareLinkedBlockingQueueProxy(TaskDef taskDef, LinkedBlockingQueue<T> delegate) {
+    super();
+    this.taskDef = taskDef;
+    this.delegate = delegate;
+    DatadogWeaverReporter.start();
+  }
+
+  @Override
+  public T poll() {
+    T event = delegate.poll();
+    DatadogWeaverReporter.processEvent(event, taskDef);
+    return event;
+  }
+
+  @Override
+  public T poll(long timeout, TimeUnit unit) throws InterruptedException {
+    T event = delegate.poll(timeout, unit);
+    if (event != null) {
+      DatadogWeaverReporter.processEvent(event, taskDef);
+    }
+    return event;
+  }
+}

dd-java-agent/instrumentation/weaver/src/main/java/datadog/trace/instrumentation/weaver/TaskDefAwareQueueProxy.java

@@ -5,7 +5,9 @@
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import java.lang.reflect.Field;
 import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.LinkedBlockingQueue;
 import net.bytebuddy.asm.Advice;
 import sbt.testing.TaskDef;
 import weaver.framework.SuiteEvent;
@@ -28,7 +30,8 @@ public String[] helperClassNames() {
     return new String[] {
       packageName + ".WeaverUtils",
       packageName + ".DatadogWeaverReporter",
-      packageName + ".TaskDefAwareQueueProxy",
+      packageName + ".TaskDefAwareLinkedBlockingQueueProxy",
+      packageName + ".TaskDefAwareConcurrentLinkedQueueProxy",
     };
   }
 
@@ -41,10 +44,26 @@ public void methodAdvice(MethodTransformer transformer) {
   public static class SbtTaskCreationAdvice {
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void onTaskCreation(
-        @Advice.FieldValue(value = "queue", readOnly = false)
-            ConcurrentLinkedQueue<SuiteEvent> queue,
-        @Advice.FieldValue("taskDef") TaskDef taskDef) {
-      queue = new TaskDefAwareQueueProxy<SuiteEvent>(taskDef, queue);
+        @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<SuiteEvent>(
+                  taskDef, (ConcurrentLinkedQueue<SuiteEvent>) queue));
+        } else if (queue instanceof LinkedBlockingQueue) {
+          // typelevel's implementation (0.9+) uses a LinkedBlockingQueue for the field
+          queueField.set(
+              sbtTask,
+              new TaskDefAwareLinkedBlockingQueueProxy<SuiteEvent>(
+                  taskDef, (LinkedBlockingQueue<SuiteEvent>) queue));
+        }
+      } catch (Exception ignored) {
+      }
     }
   }
 }