fix: proper handling of bazel test runner suite events

Author: daniel-mohedano

dd-java-agent/instrumentation/junit/junit-4/junit-4.10/src/main/java/datadog/trace/instrumentation/junit4/JUnit4TracingListener.java

@@ -8,11 +8,8 @@
 import datadog.trace.bootstrap.ContextStore;
 import java.lang.reflect.Method;
 import java.util.List;
-import java.util.Set;
-import java.util.concurrent.ConcurrentHashMap;
 import org.junit.Ignore;
 import org.junit.runner.Description;
-import org.junit.runner.Result;
 import org.junit.runner.notification.Failure;
 
 public class JUnit4TracingListener extends TracingListener {
@@ -22,26 +19,6 @@ public class JUnit4TracingListener extends TracingListener {
 
   private final ContextStore<Description, TestExecutionTracker> executionTrackers;
 
-  /**
-   * Suites for which {@code onTestSuiteStart} has been fired (from either the normal
-   * ParentRunner-based flow or via lazy-registration in {@link #testStarted}). Used to keep
-   * lifecycle events idempotent and to know which auto-started suite still needs closing.
-   */
-  private final Set<TestSuiteDescriptor> startedSuites = ConcurrentHashMap.newKeySet();
-
-  /**
-   * Last suite lazy-started from {@link #testStarted} because no {@link #testSuiteStarted} event
-   * was observed for it first. This has been seen under {@code
-   * com.google.testing.junit.runner.BazelTestRunner}, where the suite-start advice in {@code
-   * JUnit4SuiteEventsInstrumentation} does not fire for reasons still to be pinpointed (likely a
-   * classloader or runner-wrapping quirk specific to the Bazel test launcher). Closed when the next
-   * test belongs to a different suite, or when the whole test run finishes.
-   *
-   * <p>TODO: investigate the exact cause under {@code BazelTestRunner} and add a dedicated
-   * instrumentation that emits proper suite-lifecycle events instead of relying on this fallback.
-   */
-  private volatile TestSuiteDescriptor autoStartedSuite;
-
   public JUnit4TracingListener(ContextStore<Description, TestExecutionTracker> executionTrackers) {
     this.executionTrackers = executionTrackers;
   }
@@ -55,9 +32,6 @@ public void testSuiteStarted(final Description description) {
     }
 
     TestSuiteDescriptor suiteDescriptor = JUnit4Utils.toSuiteDescriptor(description);
-    if (!startedSuites.add(suiteDescriptor)) {
-      return; // already started (idempotent vs. lazy-registration or duplicate events)
-    }
     Class<?> testClass = description.getTestClass();
     String testSuiteName = JUnit4Utils.getSuiteName(testClass, description);
     List<String> categories = JUnit4Utils.getCategories(testClass, null);
@@ -84,9 +58,6 @@ public void testSuiteFinished(final Description description) {
     }
 
     TestSuiteDescriptor suiteDescriptor = JUnit4Utils.toSuiteDescriptor(description);
-    if (!startedSuites.remove(suiteDescriptor)) {
-      return; // never started
-    }
     TestEventsHandlerHolder.HANDLERS
         .get(TestFrameworkInstrumentation.JUNIT4)
         .onTestSuiteFinish(suiteDescriptor, null);
@@ -102,8 +73,6 @@ public void testStarted(final Description description) {
     TestDescriptor testDescriptor = JUnit4Utils.toTestDescriptor(description);
     TestSourceData testSourceData = JUnit4Utils.toTestSourceData(description);
 
-    lazyStartSuiteIfNeeded(suiteDescriptor, description, testSourceData);
-
     String testName = JUnit4Utils.getTestName(description, testSourceData.getTestMethod());
     String testParameters = JUnit4Utils.getParameters(description);
     List<String> categories =
@@ -124,50 +93,6 @@ public void testStarted(final Description description) {
             executionTrackers.get(description));
   }
 
-  @Override
-  public void testRunFinished(Result result) {
-    closeAutoStartedSuite();
-  }
-
-  private void lazyStartSuiteIfNeeded(
-      TestSuiteDescriptor newSuite, Description description, TestSourceData testSourceData) {
-    if (startedSuites.contains(newSuite)) {
-      return;
-    }
-    closeAutoStartedSuite();
-
-    Class<?> testClass = testSourceData.getTestClass();
-    String testSuiteName = JUnit4Utils.getSuiteName(testClass, description);
-    List<String> categories = JUnit4Utils.getCategories(testClass, null);
-    TestEventsHandlerHolder.HANDLERS
-        .get(TestFrameworkInstrumentation.JUNIT4)
-        .onTestSuiteStart(
-            newSuite,
-            testSuiteName,
-            FRAMEWORK_NAME,
-            FRAMEWORK_VERSION,
-            testClass,
-            categories,
-            false,
-            TestFrameworkInstrumentation.JUNIT4,
-            null);
-    startedSuites.add(newSuite);
-    autoStartedSuite = newSuite;
-  }
-
-  private void closeAutoStartedSuite() {
-    TestSuiteDescriptor suite = autoStartedSuite;
-    if (suite == null) {
-      return;
-    }
-    autoStartedSuite = null;
-    if (startedSuites.remove(suite)) {
-      TestEventsHandlerHolder.HANDLERS
-          .get(TestFrameworkInstrumentation.JUNIT4)
-          .onTestSuiteFinish(suite, null);
-    }
-  }
-
   @Override
   public void testFinished(final Description description) {
     if (JUnit4Utils.isJUnitPlatformRunnerTest(description)) {

dd-java-agent/instrumentation/junit/junit-4/junit-4.10/src/main/java/datadog/trace/instrumentation/junit4/JUnit4Utils.java

@@ -117,10 +117,11 @@ public static List<RunListener> runListenersFromRunNotifier(final RunNotifier ru
   }
 
   /**
-   * Walks through {@link RunNotifier} wrappers (e.g. Bazel's {@code RunNotifierWrapper}) so the
-   * effective {@code listeners} field is read, not the wrapper's own (forwarded) one.
+   * Walks through {@link RunNotifier} wrappers (e.g. Bazel's {@code RunNotifierWrapper}) and
+   * returns the inner notifier whose {@code listeners} field actually receives {@code addListener}
+   * calls. Returns the input untouched when it is not a known wrapper.
    */
-  private static RunNotifier unwrapRunNotifier(RunNotifier notifier) {
+  public static RunNotifier unwrapRunNotifier(RunNotifier notifier) {
     RunNotifier current = notifier;
     for (int i = 0; i < 8 && current != null; i++) {
       if (!isBazelRunNotifierWrapper(current.getClass())) {

dd-java-agent/instrumentation/junit/junit-4/junit-4.13/src/main/java/datadog/trace/instrumentation/junit4/BazelRunNotifierWrapperInstrumentation.java

@@ -0,0 +1,92 @@
+package datadog.trace.instrumentation.junit4;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+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 net.bytebuddy.asm.Advice;
+import org.junit.runner.Description;
+import org.junit.runner.notification.RunNotifier;
+
+/**
+ * Restores suite lifecycle events when JUnit 4.13+ tests run under Bazel's {@code
+ * com.google.testing.junit.runner.BazelTestRunner}.
+ *
+ * <p>Bazel's {@code com.google.testing.junit.junit4.runner.RunNotifierWrapper} (last touched in
+ * 2015) explicitly delegates {@code addListener}, {@code fireTestStarted}, etc. to the inner
+ * notifier, but does not override {@link RunNotifier#fireTestSuiteStarted(Description)} and {@link
+ * RunNotifier#fireTestSuiteFinished(Description)} — those were added in JUnit 4.13. {@link
+ * org.junit.runners.ParentRunner#run} therefore fires the suite-lifecycle events on the wrapper's
+ * own (always empty) listener list, and our tracing listener — installed on the inner notifier via
+ * the wrapper's delegating {@code addListener} — never sees them.
+ *
+ * <p>This advice intercepts {@link RunNotifier#fireTestSuiteStarted(Description)} and {@link
+ * RunNotifier#fireTestSuiteFinished(Description)} on the wrapper instance and re-fires the event on
+ * the inner notifier, where the listener actually lives. Calls on a non-wrapper notifier are a
+ * no-op.
+ */
+@AutoService(InstrumenterModule.class)
+public class BazelRunNotifierWrapperInstrumentation extends InstrumenterModule.CiVisibility
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+  public BazelRunNotifierWrapperInstrumentation() {
+    super("ci-visibility", "junit-4");
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "org.junit.runner.notification.RunNotifier";
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".JUnit4Utils",
+      packageName + ".TracingListener",
+      packageName + ".SkippedByDatadog",
+    };
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        named("fireTestSuiteStarted").and(takesArgument(0, named("org.junit.runner.Description"))),
+        BazelRunNotifierWrapperInstrumentation.class.getName() + "$FireSuiteStartedAdvice");
+    transformer.applyAdvice(
+        named("fireTestSuiteFinished").and(takesArgument(0, named("org.junit.runner.Description"))),
+        BazelRunNotifierWrapperInstrumentation.class.getName() + "$FireSuiteFinishedAdvice");
+  }
+
+  public static class FireSuiteStartedAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void fireOnInnerNotifier(
+        @Advice.This final RunNotifier self, @Advice.Argument(0) final Description description) {
+      RunNotifier inner = JUnit4Utils.unwrapRunNotifier(self);
+      if (inner != null && inner != self) {
+        inner.fireTestSuiteStarted(description);
+      }
+    }
+
+    // JUnit 4.13 muzzle marker: fireTestSuiteStarted exists from 4.13.
+    public static void muzzleCheck(final RunNotifier notifier) {
+      notifier.fireTestSuiteStarted(null);
+    }
+  }
+
+  public static class FireSuiteFinishedAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void fireOnInnerNotifier(
+        @Advice.This final RunNotifier self, @Advice.Argument(0) final Description description) {
+      RunNotifier inner = JUnit4Utils.unwrapRunNotifier(self);
+      if (inner != null && inner != self) {
+        inner.fireTestSuiteFinished(description);
+      }
+    }
+
+    public static void muzzleCheck(final RunNotifier notifier) {
+      notifier.fireTestSuiteFinished(null);
+    }
+  }
+}