Prevent double reporting of Scalatest events when using SBT with test forking (#8682)

Author: nikita-tkachenko-datadog

.circleci/config.continue.yml.j2

@@ -36,7 +36,7 @@ instrumentation_modules: &instrumentation_modules "dd-java-agent/instrumentation
 debugger_modules: &debugger_modules "dd-java-agent/agent-debugger|dd-java-agent/agent-bootstrap|dd-java-agent/agent-builder|internal-api|communication|dd-trace-core"
 profiling_modules: &profiling_modules "dd-java-agent/agent-profiling"
 
-default_system_tests_commit: &default_system_tests_commit 2799fa982318da14c9d3e5f722abdc670d2802c3
+default_system_tests_commit: &default_system_tests_commit 761b9e7a82ffb136c4653a4d1623d120d67b005b
 
 parameters:
   nightly:
@@ -932,7 +932,7 @@ jobs:
           command: |
             cd system-tests
             DD_SITE=datadoghq.com DD_API_KEY=$SYSTEM_TESTS_E2E_DD_API_KEY DD_APPLICATION_KEY=$SYSTEM_TESTS_E2E_DD_APP_KEY ./run.sh INTEGRATIONS
-        
+
       - run:
           name: Run IDM Crossed Tracing Libraries propagation tests for messaging
           environment:
@@ -975,7 +975,7 @@ jobs:
           no_output_timeout: 5m
           command: |
             cd system-tests
-            export DD_API_KEY=$SYSTEM_TESTS_E2E_DD_API_KEY 
+            export DD_API_KEY=$SYSTEM_TESTS_E2E_DD_API_KEY
             ./run.sh DEBUGGER_SCENARIOS
 
       - run:

dd-java-agent/instrumentation/scalatest/src/main/java/datadog/trace/instrumentation/scalatest/DatadogReporter.java

@@ -86,14 +86,17 @@ private static void stop(Event event) {
 
   private static void onSuiteStart(SuiteStarting event) {
     int runStamp = event.ordinal().runStamp();
-    RunContext context = RunContext.getOrCreate(runStamp);
-    TestEventsHandler<TestSuiteDescriptor, TestDescriptor> eventHandler = context.getEventHandler();
+    RunContext context = RunContext.get(runStamp);
+    if (context == null) {
+      return;
+    }
 
     String testSuiteName = event.suiteId();
     Class<?> testClass = ScalatestUtils.getClass(event.suiteClassName());
     Collection<String> categories = Collections.emptyList();
     boolean parallelized = true;
 
+    TestEventsHandler<TestSuiteDescriptor, TestDescriptor> eventHandler = context.getEventHandler();
     eventHandler.onTestSuiteStart(
         new TestSuiteDescriptor(testSuiteName, testClass),
         testSuiteName,
@@ -108,30 +111,40 @@ private static void onSuiteStart(SuiteStarting event) {
 
   private static void onSuiteFinish(SuiteCompleted event) {
     int runStamp = event.ordinal().runStamp();
-    RunContext context = RunContext.getOrCreate(runStamp);
-    TestEventsHandler<TestSuiteDescriptor, TestDescriptor> eventHandler = context.getEventHandler();
+    RunContext context = RunContext.get(runStamp);
+    if (context == null) {
+      return;
+    }
 
     String testSuiteName = event.suiteId();
     Class<?> testClass = ScalatestUtils.getClass(event.suiteClassName());
+
+    TestEventsHandler<TestSuiteDescriptor, TestDescriptor> eventHandler = context.getEventHandler();
     eventHandler.onTestSuiteFinish(new TestSuiteDescriptor(testSuiteName, testClass), null);
   }
 
   private static void onSuiteAbort(SuiteAborted event) {
     int runStamp = event.ordinal().runStamp();
-    RunContext context = RunContext.getOrCreate(runStamp);
-    TestEventsHandler<TestSuiteDescriptor, TestDescriptor> eventHandler = context.getEventHandler();
+    RunContext context = RunContext.get(runStamp);
+    if (context == null) {
+      return;
+    }
 
     String testSuiteName = event.suiteId();
     Class<?> testClass = ScalatestUtils.getClass(event.suiteClassName());
     Throwable throwable = event.throwable().getOrElse(null);
+
+    TestEventsHandler<TestSuiteDescriptor, TestDescriptor> eventHandler = context.getEventHandler();
     eventHandler.onTestSuiteFailure(new TestSuiteDescriptor(testSuiteName, testClass), throwable);
     eventHandler.onTestSuiteFinish(new TestSuiteDescriptor(testSuiteName, testClass), null);
   }
 
   private static void onTestStart(TestStarting event) {
     int runStamp = event.ordinal().runStamp();
-    RunContext context = RunContext.getOrCreate(runStamp);
-    TestEventsHandler<TestSuiteDescriptor, TestDescriptor> eventHandler = context.getEventHandler();
+    RunContext context = RunContext.get(runStamp);
+    if (context == null) {
+      return;
+    }
 
     String testSuiteName = event.suiteId();
     String testName = event.testName();
@@ -146,6 +159,7 @@ private static void onTestStart(TestStarting event) {
     }
     Class<?> testClass = ScalatestUtils.getClass(event.suiteClassName());
 
+    TestEventsHandler<TestSuiteDescriptor, TestDescriptor> eventHandler = context.getEventHandler();
     eventHandler.onTestStart(
         new TestSuiteDescriptor(testSuiteName, testClass),
         new TestDescriptor(testSuiteName, testClass, testName, testParameters, testQualifier),
@@ -161,8 +175,10 @@ private static void onTestStart(TestStarting event) {
 
   private static void onTestSuccess(TestSucceeded event) {
     int runStamp = event.ordinal().runStamp();
-    RunContext context = RunContext.getOrCreate(runStamp);
-    TestEventsHandler<TestSuiteDescriptor, TestDescriptor> eventHandler = context.getEventHandler();
+    RunContext context = RunContext.get(runStamp);
+    if (context == null) {
+      return;
+    }
 
     String testSuiteName = event.suiteId();
     Class<?> testClass = ScalatestUtils.getClass(event.suiteClassName());
@@ -175,13 +191,16 @@ private static void onTestSuccess(TestSucceeded event) {
     TestIdentifier testIdentifier = new TestIdentifier(testSuiteName, testName, null);
     TestExecutionHistory executionHistory = context.popExecutionHistory(testIdentifier);
 
+    TestEventsHandler<TestSuiteDescriptor, TestDescriptor> eventHandler = context.getEventHandler();
     eventHandler.onTestFinish(testDescriptor, null, executionHistory);
   }
 
   private static void onTestFailure(TestFailed event) {
     int runStamp = event.ordinal().runStamp();
-    RunContext context = RunContext.getOrCreate(runStamp);
-    TestEventsHandler<TestSuiteDescriptor, TestDescriptor> eventHandler = context.getEventHandler();
+    RunContext context = RunContext.get(runStamp);
+    if (context == null) {
+      return;
+    }
 
     String testSuiteName = event.suiteId();
     Class<?> testClass = ScalatestUtils.getClass(event.suiteClassName());
@@ -195,14 +214,17 @@ private static void onTestFailure(TestFailed event) {
     TestIdentifier testIdentifier = new TestIdentifier(testSuiteName, testName, null);
     TestExecutionHistory executionHistory = context.popExecutionHistory(testIdentifier);
 
+    TestEventsHandler<TestSuiteDescriptor, TestDescriptor> eventHandler = context.getEventHandler();
     eventHandler.onTestFailure(testDescriptor, throwable);
     eventHandler.onTestFinish(testDescriptor, null, executionHistory);
   }
 
   private static void onTestIgnore(TestIgnored event) {
     int runStamp = event.ordinal().runStamp();
-    RunContext context = RunContext.getOrCreate(runStamp);
-    TestEventsHandler<TestSuiteDescriptor, TestDescriptor> eventHandler = context.getEventHandler();
+    RunContext context = RunContext.get(runStamp);
+    if (context == null) {
+      return;
+    }
 
     String testSuiteName = event.suiteId();
     String testName = event.testName();
@@ -214,6 +236,7 @@ private static void onTestIgnore(TestIgnored event) {
     TestIdentifier skippableTest = new TestIdentifier(testSuiteName, testName, null);
     SkipReason reason = context.getSkipReason(skippableTest);
 
+    TestEventsHandler<TestSuiteDescriptor, TestDescriptor> eventHandler = context.getEventHandler();
     eventHandler.onTestIgnore(
         new TestSuiteDescriptor(testSuiteName, testClass),
         new TestDescriptor(testSuiteName, testClass, testName, testParameters, testQualifier),
@@ -228,8 +251,10 @@ private static void onTestIgnore(TestIgnored event) {
 
   private static void onTestCancel(TestCanceled event) {
     int runStamp = event.ordinal().runStamp();
-    RunContext context = RunContext.getOrCreate(runStamp);
-    TestEventsHandler<TestSuiteDescriptor, TestDescriptor> eventHandler = context.getEventHandler();
+    RunContext context = RunContext.get(runStamp);
+    if (context == null) {
+      return;
+    }
 
     String testSuiteName = event.suiteId();
     String testName = event.testName();
@@ -241,6 +266,7 @@ private static void onTestCancel(TestCanceled event) {
 
     TestDescriptor testDescriptor =
         new TestDescriptor(testSuiteName, testClass, testName, testParameters, testQualifier);
+    TestEventsHandler<TestSuiteDescriptor, TestDescriptor> eventHandler = context.getEventHandler();
     if (throwable instanceof SuppressedTestFailedException) {
       eventHandler.onTestFailure(testDescriptor, throwable.getCause());
     } else {
@@ -255,8 +281,10 @@ private static void onTestCancel(TestCanceled event) {
 
   private static void onTestPending(TestPending event) {
     int runStamp = event.ordinal().runStamp();
-    RunContext context = RunContext.getOrCreate(runStamp);
-    TestEventsHandler<TestSuiteDescriptor, TestDescriptor> eventHandler = context.getEventHandler();
+    RunContext context = RunContext.get(runStamp);
+    if (context == null) {
+      return;
+    }
 
     String testSuiteName = event.suiteId();
     String testName = event.testName();
@@ -267,6 +295,7 @@ private static void onTestPending(TestPending event) {
 
     TestDescriptor testDescriptor =
         new TestDescriptor(testSuiteName, testClass, testName, testParameters, testQualifier);
+    TestEventsHandler<TestSuiteDescriptor, TestDescriptor> eventHandler = context.getEventHandler();
     eventHandler.onTestSkip(testDescriptor, reason);
 
     TestIdentifier testIdentifier = new TestIdentifier(testSuiteName, testName, null);

dd-java-agent/instrumentation/scalatest/src/main/java/datadog/trace/instrumentation/scalatest/RunContext.java

@@ -31,9 +31,15 @@ public static RunContext getOrCreate(int runStamp) {
     return CONTEXTS.computeIfAbsent(runStamp, RunContext::new);
   }
 
+  public static RunContext get(int runStamp) {
+    return CONTEXTS.get(runStamp);
+  }
+
   public static void destroy(int runStamp) {
     RunContext context = CONTEXTS.remove(runStamp);
-    context.destroy();
+    if (context != null) {
+      context.destroy();
+    }
   }
 
   private final int runStamp;

dd-java-agent/instrumentation/scalatest/src/main/java/datadog/trace/instrumentation/scalatest/ScalatestForkInstrumentation.java

@@ -0,0 +1,80 @@
+package datadog.trace.instrumentation.scalatest;
+
+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 datadog.trace.api.Config;
+import datadog.trace.api.config.CiVisibilityConfig;
+import datadog.trace.bootstrap.CallDepthThreadLocalMap;
+import java.util.Set;
+import net.bytebuddy.asm.Advice;
+import org.scalatest.Reporter;
+
+@AutoService(InstrumenterModule.class)
+public class ScalatestForkInstrumentation extends InstrumenterModule.CiVisibility
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+  public ScalatestForkInstrumentation() {
+    super("ci-visibility", "scalatest");
+  }
+
+  @Override
+  public boolean isApplicable(Set<TargetSystem> enabledSystems) {
+    return super.isApplicable(enabledSystems)
+        && Config.get().isCiVisibilityScalatestForkMonitorEnabled();
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "org.scalatest.tools.Framework";
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {};
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        named("runner").and(takesArgument(1, String[].class)),
+        ScalatestForkInstrumentation.class.getName() + "$RunnerAdvice");
+  }
+
+  /**
+   * This instrumentation is needed to prevent double-reporting of Scalatest events when using SBT
+   * and Test/fork. Current version of the code cannot detect such cases automatically, so it has to
+   * be activated explicitly by setting {@link
+   * CiVisibilityConfig#CIVISIBILITY_SCALATEST_FORK_MONITOR_ENABLED}.
+   *
+   * <p>When using SBT with test forking, Scalatest reporter is activated both in the parent and in
+   * the child process. We only want to instrument the child process (since it is actually executing
+   * the tests) and not the parent one. The way to distinguish between the two is by examining the
+   * "remoteArgs" passed to the runner method: the args are non-empty in the child process, and
+   * empty in the parent one.
+   *
+   * <p>The instrumentation works by increasing the Reporter.class counter in the call depth thread
+   * local map. The counter is decreased when the runner method exits.
+   *
+   * <p>Since the tracing instrumentation is checking the counter value to avoid nested calls,
+   * increasing it here results in effectively disabling the tracing.
+   */
+  public static class RunnerAdvice {
+    @Advice.OnMethodEnter
+    public static void beforeRunnerStart(@Advice.Argument(value = 1) String[] remoteArgs) {
+      if (remoteArgs.length == 0) {
+        CallDepthThreadLocalMap.incrementCallDepth(Reporter.class);
+      }
+    }
+
+    @Advice.OnMethodExit
+    public static void afterRunnerStart(@Advice.Argument(value = 1) String[] remoteArgs) {
+      if (remoteArgs.length == 0) {
+        CallDepthThreadLocalMap.decrementCallDepth(Reporter.class);
+      }
+    }
+  }
+}

dd-trace-api/src/main/java/datadog/trace/api/config/CiVisibilityConfig.java

@@ -79,6 +79,8 @@ public final class CiVisibilityConfig {
   public static final String CIVISIBILITY_REMOTE_ENV_VARS_PROVIDER_KEY =
       "civisibility.remote.env.vars.provider.key";
   public static final String CIVISIBILITY_TEST_ORDER = "civisibility.test.order";
+  public static final String CIVISIBILITY_SCALATEST_FORK_MONITOR_ENABLED =
+      "civisibility.scalatest.fork.monitor.enabled";
   public static final String TEST_MANAGEMENT_ENABLED = "test.management.enabled";
   public static final String TEST_MANAGEMENT_ATTEMPT_TO_FIX_RETRIES =
       "test.management.attempt.to.fix.retries";

internal-api/src/main/java/datadog/trace/api/Config.java

@@ -383,6 +383,7 @@ public static String getHostName() {
   private final String ciVisibilityTestOrder;
   private final boolean ciVisibilityTestManagementEnabled;
   private final Integer ciVisibilityTestManagementAttemptToFixRetries;
+  private final boolean ciVisibilityScalatestForkMonitorEnabled;
 
   private final boolean remoteConfigEnabled;
   private final boolean remoteConfigIntegrityCheckEnabled;
@@ -1607,6 +1608,8 @@ PROFILING_DATADOG_PROFILER_ENABLED, isDatadogProfilerSafeInCurrentEnvironment())
     ciVisibilityTestManagementEnabled = configProvider.getBoolean(TEST_MANAGEMENT_ENABLED, true);
     ciVisibilityTestManagementAttemptToFixRetries =
         configProvider.getInteger(TEST_MANAGEMENT_ATTEMPT_TO_FIX_RETRIES);
+    ciVisibilityScalatestForkMonitorEnabled =
+        configProvider.getBoolean(CIVISIBILITY_SCALATEST_FORK_MONITOR_ENABLED, false);
 
     remoteConfigEnabled =
         configProvider.getBoolean(
@@ -3114,6 +3117,10 @@ public boolean isCiVisibilityExecutionPoliciesEnabled() {
         || ciVisibilityTestManagementEnabled;
   }
 
+  public boolean isCiVisibilityScalatestForkMonitorEnabled() {
+    return ciVisibilityScalatestForkMonitorEnabled;
+  }
+
   public int getCiVisibilityFlakyRetryCount() {
     return ciVisibilityFlakyRetryCount;
   }