feat: implement SuiteEnd listener for sink flushing

Author: daniel-mohedano

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/domain/TestSuiteImpl.java

@@ -6,6 +6,7 @@
 
 import datadog.trace.api.Config;
 import datadog.trace.api.civisibility.DDTestSuite;
+import datadog.trace.api.civisibility.InstrumentationTestBridge;
 import datadog.trace.api.civisibility.config.LibraryCapability;
 import datadog.trace.api.civisibility.coverage.CoverageStore;
 import datadog.trace.api.civisibility.execution.TestStatus;
@@ -220,6 +221,8 @@ public void end(@Nullable Long endTime) {
       AgentTracer.closeActive();
     }
 
+    InstrumentationTestBridge.fireBeforeSuiteEnd();
+
     onSpanFinish.accept(span);
 
     if (endTime != null) {

dd-java-agent/agent-ci-visibility/src/testFixtures/groovy/datadog/trace/civisibility/CiVisibilitySmokeTest.groovy

@@ -52,4 +52,38 @@ abstract class CiVisibilitySmokeTest extends Specification {
     // an even more basic smoke check for distributions: assert that we received some
     assert !receivedTelemetryDistributions.isEmpty()
   }
+
+  protected static verifySnapshotLogs(List<Map<String, Object>> receivedLogs, int expectedSnapshots) {
+    def logsPerSnapshot = 4 // 3 probe statuses + 1 snapshot log are expected per snapshot
+
+    assert receivedLogs.size() == logsPerSnapshot * expectedSnapshots
+
+    def probeStatusLogs = receivedLogs.findAll { it.containsKey("message") }
+    def snapshotLogs = receivedLogs.findAll { !it.containsKey("message") }
+
+    verifyProbeStatuses(probeStatusLogs, expectedSnapshots)
+    verifySnapshots(snapshotLogs)
+  }
+
+  private static verifyProbeStatuses(List<Map<String, Object>> logs, int expectedCount) {
+    assert logs.findAll { log -> ((String) log.message).startsWith("Received probe") }.size() == expectedCount
+    assert logs.findAll { log -> ((String) log.message).startsWith("Installed probe") }.size() == expectedCount
+    assert logs.findAll { log -> ((String) log.message).endsWith("is emitting.") }.size() == expectedCount
+  }
+
+  private static verifySnapshots(List<Map<String, Object>> logs) {
+    def requiredLogFields = ["logger.name", "logger.method", "dd.spanid", "dd.traceid"]
+    def requiredSnapshotFields = ["captures", "exceptionId", "probe", "stack"]
+
+    logs.each { log ->
+      assert log.product == "test_optimization"
+      requiredLogFields.each { field -> log.containsKey(field) }
+
+      Map<String, Object> debuggerMap = log.debugger as Map<String, Object>
+      Map<String, Object> snapshotContent = debuggerMap.snapshot as Map<String, Object>
+
+      assert snapshotContent != null
+      requiredSnapshotFields.each { field -> snapshotContent.containsKey(field) }
+    }
+  }
 }

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/sink/DebuggerSink.java

@@ -4,6 +4,8 @@
 import com.datadog.debugger.uploader.BatchUploader;
 import com.datadog.debugger.util.DebuggerMetrics;
 import datadog.trace.api.Config;
+import datadog.trace.api.civisibility.InstrumentationTestBridge;
+import datadog.trace.api.civisibility.domain.TestContext;
 import datadog.trace.bootstrap.debugger.DebuggerContext.SkipCause;
 import datadog.trace.bootstrap.debugger.ProbeId;
 import datadog.trace.util.AgentTaskScheduler;
@@ -68,6 +70,11 @@ public DebuggerSink(
     this.snapshotSink = snapshotSink;
     this.symbolSink = symbolSink;
     this.uploadFlushInterval = config.getDynamicInstrumentationUploadFlushInterval();
+
+    if (config.isCiVisibilityFailedTestReplayActive()) {
+      // register test listener to flush snapshots on suite end
+      InstrumentationTestBridge.registerListener(new DebuggerTestListener(this));
+    }
   }
 
   public void start() {
@@ -235,4 +242,23 @@ public void skipSnapshot(String probeId, SkipCause cause) {
   long getCurrentLowRateFlushInterval() {
     return currentLowRateFlushInterval;
   }
+
+  static class DebuggerTestListener implements InstrumentationTestBridge.TestListener {
+    private final DebuggerSink sink;
+
+    DebuggerTestListener(DebuggerSink sink) {
+      this.sink = sink;
+    }
+
+    @Override
+    public void beforeTestEnd(TestContext ignored) {
+      // noop
+    }
+
+    @Override
+    public void beforeSuiteEnd() {
+      LOGGER.debug("CiVisibility BeforeSuiteEnd fired, flushing sink");
+      sink.lowRateFlush(sink);
+    }
+  }
 }

dd-java-agent/instrumentation/selenium/src/main/java/datadog/trace/instrumentation/selenium/SeleniumTestListener.java

@@ -25,4 +25,9 @@ public void beforeTestEnd(TestContext testContext) {
 
     SeleniumUtils.beforePageClose(driver);
   }
+
+  @Override
+  public void beforeSuiteEnd() {
+    // noop
+  }
 }

dd-smoke-tests/junit-console/src/test/groovy/datadog/smoketest/JUnitConsoleSmokeTest.groovy

@@ -61,8 +61,7 @@ class JUnitConsoleSmokeTest extends CiVisibilitySmokeTest {
 
     def additionalDynamicTags = ["content.meta.['_dd.debug.error.6.snapshot_id']", "content.meta.['_dd.debug.error.exception_id']"]
     verifyEventsAndCoverages(projectName, "junit-console", "headless", mockBackend.waitForEvents(5), mockBackend.waitForCoverages(0), additionalDynamicTags)
-    //TODO: add verification of the logs payload
-    //mockBackend.waitForLogs(8)
+    verifySnapshotLogs(mockBackend.waitForLogs(4), 1)
 
     where:
     projectName = "test_junit_console_failed_test_replay"

dd-smoke-tests/maven/src/test/groovy/datadog/smoketest/MavenSmokeTest.groovy

@@ -250,8 +250,7 @@ class MavenSmokeTest extends CiVisibilitySmokeTest {
 
     def additionalDynamicTags = ["content.meta.['_dd.debug.error.3.snapshot_id']", "content.meta.['_dd.debug.error.exception_id']"]
     verifyEventsAndCoverages(projectName, "maven", mavenVersion, mockBackend.waitForEvents(5), mockBackend.waitForCoverages(0), additionalDynamicTags)
-    //TODO: add verification of the logs payload
-    //mockBackend.waitForLogs(8)
+    verifySnapshotLogs(mockBackend.waitForLogs(4), 1)
 
     where:
     projectName                            | mavenVersion

dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java

@@ -189,7 +189,7 @@ public final class ConfigDefaults {
   static final boolean DEFAULT_DYNAMIC_INSTRUMENTATION_ENABLED = false;
   static final int DEFAULT_DYNAMIC_INSTRUMENTATION_UPLOAD_TIMEOUT = 30; // seconds
   static final int DEFAULT_DYNAMIC_INSTRUMENTATION_UPLOAD_FLUSH_INTERVAL = 0; // ms, 0 = dynamic
-  static final int DEFAULT_FAILED_TEST_REPLAY_UPLOAD_FLUSH_INTERVAL = 200; // ms, 0 = dynamic
+  static final int DEFAULT_FAILED_TEST_REPLAY_UPLOAD_FLUSH_INTERVAL = 1000; // ms, 0 = dynamic
   static final boolean DEFAULT_DYNAMIC_INSTRUMENTATION_CLASSFILE_DUMP_ENABLED = false;
   static final int DEFAULT_DYNAMIC_INSTRUMENTATION_POLL_INTERVAL = 1; // seconds
   static final int DEFAULT_DYNAMIC_INSTRUMENTATION_DIAGNOSTICS_INTERVAL = 60 * 60; // seconds

internal-api/src/main/java/datadog/trace/api/civisibility/InstrumentationTestBridge.java

@@ -38,11 +38,19 @@ public static void fireBeforeTestEnd(TestContext testContext) {
     }
   }
 
+  public static void fireBeforeSuiteEnd() {
+    for (TestListener testListener : TEST_LISTENERS) {
+      testListener.beforeSuiteEnd();
+    }
+  }
+
   public static void registerListener(TestListener listener) {
     TEST_LISTENERS.addIfAbsent(listener);
   }
 
   public interface TestListener {
     void beforeTestEnd(TestContext testContext);
+
+    void beforeSuiteEnd();
   }
 }