Align retry logic for all test framework instrumentations (#8803)

Author: daniel-mohedano

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

@@ -3,6 +3,7 @@
 import static datadog.trace.civisibility.Constants.CI_VISIBILITY_INSTRUMENTATION_NAME;
 
 import datadog.trace.api.Config;
+import datadog.trace.api.civisibility.execution.TestStatus;
 import datadog.trace.api.civisibility.telemetry.CiVisibilityCountMetric;
 import datadog.trace.api.civisibility.telemetry.CiVisibilityMetricCollector;
 import datadog.trace.api.civisibility.telemetry.tag.EventType;

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

@@ -7,6 +7,7 @@
 import datadog.trace.api.DDTraceId;
 import datadog.trace.api.IdGenerationStrategy;
 import datadog.trace.api.civisibility.CIConstants;
+import datadog.trace.api.civisibility.execution.TestStatus;
 import datadog.trace.api.civisibility.telemetry.CiVisibilityCountMetric;
 import datadog.trace.api.civisibility.telemetry.CiVisibilityMetricCollector;
 import datadog.trace.api.civisibility.telemetry.TagValue;

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

@@ -14,6 +14,7 @@
 import datadog.trace.api.civisibility.coverage.CoveragePerTestBridge;
 import datadog.trace.api.civisibility.coverage.CoverageStore;
 import datadog.trace.api.civisibility.domain.TestContext;
+import datadog.trace.api.civisibility.execution.TestStatus;
 import datadog.trace.api.civisibility.telemetry.CiVisibilityCountMetric;
 import datadog.trace.api.civisibility.telemetry.CiVisibilityMetricCollector;
 import datadog.trace.api.civisibility.telemetry.TagValue;
@@ -30,6 +31,7 @@
 import datadog.trace.api.civisibility.telemetry.tag.SkipReason;
 import datadog.trace.api.civisibility.telemetry.tag.TestFrameworkInstrumentation;
 import datadog.trace.api.gateway.RequestContextSlot;
+import datadog.trace.api.time.SystemTimeSource;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
@@ -45,6 +47,7 @@
 import java.lang.reflect.Method;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -64,6 +67,7 @@ public class TestImpl implements DDTest {
   private final Consumer<AgentSpan> onSpanFinish;
   private final TestContext context;
   private final TestIdentifier identifier;
+  private final long startMicros;
 
   public TestImpl(
       AgentSpanContext moduleSpanContext,
@@ -111,7 +115,10 @@ public TestImpl(
             .withRequestContextData(RequestContextSlot.CI_VISIBILITY, context);
 
     if (startTime != null) {
+      startMicros = startTime;
       spanBuilder = spanBuilder.withStartTimestamp(startTime);
+    } else {
+      startMicros = SystemTimeSource.INSTANCE.getCurrentTimeMicros();
     }
 
     span = spanBuilder.start();
@@ -201,10 +208,6 @@ public TestIdentifier getIdentifier() {
     return identifier;
   }
 
-  public boolean hasFailed() {
-    return span.isError();
-  }
-
   @Override
   public void setTag(String key, Object value) {
     span.setTag(key, value);
@@ -231,6 +234,17 @@ public void setSkipReason(String skipReason) {
     }
   }
 
+  public TestStatus getStatus() {
+    return (TestStatus) span.getTag(Tags.TEST_STATUS);
+  }
+
+  public long getDuration(@Nullable Long endMicros) {
+    if (endMicros == null) {
+      endMicros = SystemTimeSource.INSTANCE.getCurrentTimeMicros();
+    }
+    return TimeUnit.MICROSECONDS.toMillis(endMicros - startMicros);
+  }
+
   @Override
   public void end(@Nullable Long endTime) {
     closeOutstandingSpans();

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

@@ -8,6 +8,7 @@
 import datadog.trace.api.civisibility.DDTestSuite;
 import datadog.trace.api.civisibility.config.LibraryCapability;
 import datadog.trace.api.civisibility.coverage.CoverageStore;
+import datadog.trace.api.civisibility.execution.TestStatus;
 import datadog.trace.api.civisibility.telemetry.CiVisibilityCountMetric;
 import datadog.trace.api.civisibility.telemetry.CiVisibilityMetricCollector;
 import datadog.trace.api.civisibility.telemetry.tag.EventType;

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/events/BuildEventsHandlerImpl.java

@@ -5,13 +5,13 @@
 import datadog.trace.api.civisibility.domain.BuildSessionSettings;
 import datadog.trace.api.civisibility.domain.JavaAgent;
 import datadog.trace.api.civisibility.events.BuildEventsHandler;
+import datadog.trace.api.civisibility.execution.TestStatus;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.Tags;
 import datadog.trace.civisibility.config.JvmInfo;
 import datadog.trace.civisibility.config.JvmInfoFactory;
 import datadog.trace.civisibility.domain.BuildSystemModule;
 import datadog.trace.civisibility.domain.BuildSystemSession;
-import datadog.trace.civisibility.domain.TestStatus;
 import java.nio.file.Path;
 import java.util.Collection;
 import java.util.Map;

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/events/NoOpTestEventsHandler.java

@@ -89,7 +89,8 @@ public void onTestIgnore(
       @Nullable String testParameters,
       @Nullable Collection<String> categories,
       @Nonnull TestSourceData testSourceData,
-      @Nullable String reason) {
+      @Nullable String reason,
+      @Nullable TestExecutionHistory testExecutionHistory) {
     // do nothing
   }
 

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/events/TestEventsHandlerImpl.java

@@ -10,6 +10,7 @@
 import datadog.trace.api.civisibility.events.TestEventsHandler;
 import datadog.trace.api.civisibility.execution.TestExecutionHistory;
 import datadog.trace.api.civisibility.execution.TestExecutionPolicy;
+import datadog.trace.api.civisibility.execution.TestStatus;
 import datadog.trace.api.civisibility.telemetry.CiVisibilityCountMetric;
 import datadog.trace.api.civisibility.telemetry.CiVisibilityMetricCollector;
 import datadog.trace.api.civisibility.telemetry.tag.EventType;
@@ -255,11 +256,14 @@ public void onTestFinish(
 
     TestIdentifier thisTest = test.getIdentifier();
     if (testExecutionHistory != null) {
-      if (test.hasFailed() && testExecutionHistory.hasFailedAllRetries()) {
+      TestStatus testStatus = test.getStatus();
+      testExecutionHistory.registerExecution(
+          testStatus != null ? testStatus : TestStatus.skip, test.getDuration(endTime));
+
+      if (testExecutionHistory.hasFailedAllRetries()) {
         test.setTag(Tags.TEST_HAS_FAILED_ALL_RETRIES, true);
-      } else if (!test.hasFailed()
-          && testModule.isAttemptToFix(thisTest)
-          && testExecutionHistory.hasSucceededAllRetries()) {
+      } else if (testExecutionHistory.hasSucceededAllRetries()
+          && testModule.isAttemptToFix(thisTest)) {
         test.setTag(Tags.TEST_TEST_MANAGEMENT_ATTEMPT_TO_FIX_PASSED, true);
       }
     }
@@ -277,7 +281,8 @@ public void onTestIgnore(
       final @Nullable String testParameters,
       final @Nullable Collection<String> categories,
       @Nonnull TestSourceData testSourceData,
-      final @Nullable String reason) {
+      final @Nullable String reason,
+      @Nullable TestExecutionHistory testExecutionHistory) {
     onTestStart(
         suiteDescriptor,
         testDescriptor,
@@ -288,9 +293,9 @@ public void onTestIgnore(
         categories,
         testSourceData,
         null,
-        null);
+        testExecutionHistory);
     onTestSkip(testDescriptor, reason);
-    onTestFinish(testDescriptor, null, null);
+    onTestFinish(testDescriptor, null, testExecutionHistory);
   }
 
   @Override

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/execution/Regular.java

@@ -1,6 +1,7 @@
 package datadog.trace.civisibility.execution;
 
 import datadog.trace.api.civisibility.execution.TestExecutionPolicy;
+import datadog.trace.api.civisibility.execution.TestStatus;
 import datadog.trace.api.civisibility.telemetry.tag.RetryReason;
 import javax.annotation.Nullable;
 
@@ -12,17 +13,20 @@ public class Regular implements TestExecutionPolicy {
   private Regular() {}
 
   @Override
-  public boolean applicable() {
-    return false;
+  public void registerExecution(TestStatus status, long durationMillis) {}
+
+  @Override
+  public boolean wasLastExecution() {
+    return true;
   }
 
   @Override
-  public boolean suppressFailures() {
+  public boolean applicable() {
     return false;
   }
 
   @Override
-  public boolean retry(boolean successful, long durationMillis) {
+  public boolean suppressFailures() {
     return false;
   }
 

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/execution/RetryUntilSuccessful.java

@@ -1,6 +1,7 @@
 package datadog.trace.civisibility.execution;
 
 import datadog.trace.api.civisibility.execution.TestExecutionPolicy;
+import datadog.trace.api.civisibility.execution.TestStatus;
 import datadog.trace.api.civisibility.telemetry.tag.RetryReason;
 import java.util.concurrent.atomic.AtomicInteger;
 import javax.annotation.Nullable;
@@ -13,42 +14,47 @@ public class RetryUntilSuccessful implements TestExecutionPolicy {
   private int executions;
   private boolean successfulExecutionSeen;
 
-  /** Total execution counter that is shared by all retry policies */
-  private final AtomicInteger totalExecutions;
+  /** Total retry counter that is shared by all retry until successful policies (currently ATR) */
+  private final AtomicInteger totalRetryCount;
 
   public RetryUntilSuccessful(
-      int maxExecutions, boolean suppressFailures, AtomicInteger totalExecutions) {
+      int maxExecutions, boolean suppressFailures, AtomicInteger totalRetryCount) {
     this.maxExecutions = maxExecutions;
     this.suppressFailures = suppressFailures;
-    this.totalExecutions = totalExecutions;
+    this.totalRetryCount = totalRetryCount;
     this.executions = 0;
   }
 
   @Override
-  public boolean applicable() {
-    return !currentExecutionIsLast() || suppressFailures;
+  public void registerExecution(TestStatus status, long durationMillis) {
+    ++executions;
+    successfulExecutionSeen |= (status != TestStatus.fail);
+    if (executions > 1) {
+      totalRetryCount.incrementAndGet();
+    }
   }
 
   @Override
-  public boolean suppressFailures() {
-    // do not suppress failures for last execution
-    // (unless flag to suppress all failures is set)
-    return !currentExecutionIsLast() || suppressFailures;
+  public boolean wasLastExecution() {
+    return successfulExecutionSeen || executions == maxExecutions;
   }
 
   private boolean currentExecutionIsLast() {
     return executions == maxExecutions - 1;
   }
 
   @Override
-  public boolean retry(boolean successful, long durationMillis) {
-    successfulExecutionSeen |= successful;
-    if (!successful && ++executions < maxExecutions) {
-      totalExecutions.incrementAndGet();
-      return true;
-    } else {
-      return false;
-    }
+  public boolean applicable() {
+    // executions must always be registered, therefore consider it applicable as long as there are
+    // retries left
+    return !wasLastExecution();
+  }
+
+  @Override
+  public boolean suppressFailures() {
+    // do not suppress failures for last execution
+    // (unless flag to suppress all failures is set)
+    return !currentExecutionIsLast() || suppressFailures;
   }
 
   @Nullable
@@ -63,7 +69,7 @@ private boolean currentExecutionIsRetry() {
 
   @Override
   public boolean hasFailedAllRetries() {
-    return currentExecutionIsLast() && !successfulExecutionSeen;
+    return wasLastExecution() && !successfulExecutionSeen;
   }
 
   @Override

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/execution/RunNTimes.java

@@ -1,6 +1,7 @@
 package datadog.trace.civisibility.execution;
 
 import datadog.trace.api.civisibility.execution.TestExecutionPolicy;
+import datadog.trace.api.civisibility.execution.TestStatus;
 import datadog.trace.api.civisibility.telemetry.tag.RetryReason;
 import datadog.trace.civisibility.config.ExecutionsByDuration;
 import java.util.List;
@@ -13,9 +14,9 @@ public class RunNTimes implements TestExecutionPolicy {
   private final List<ExecutionsByDuration> executionsByDuration;
   private int executions;
   private int maxExecutions;
-  private int totalExecutionsSeen;
   private int successfulExecutionsSeen;
   private final RetryReason retryReason;
+  private TestStatus lastStatus;
 
   public RunNTimes(
       List<ExecutionsByDuration> executionsByDuration,
@@ -25,18 +26,32 @@ public RunNTimes(
     this.executionsByDuration = executionsByDuration;
     this.executions = 0;
     this.maxExecutions = getExecutions(0);
-    this.totalExecutionsSeen = 0;
     this.successfulExecutionsSeen = 0;
     this.retryReason = retryReason;
   }
 
   @Override
-  public boolean applicable() {
-    return !currentExecutionIsLast() || suppressFailures();
+  public void registerExecution(TestStatus status, long durationMillis) {
+    lastStatus = status;
+    ++executions;
+    if (status != TestStatus.fail) {
+      ++successfulExecutionsSeen;
+    }
+    int maxExecutionsForGivenDuration = getExecutions(durationMillis);
+    maxExecutions = Math.min(maxExecutions, maxExecutionsForGivenDuration);
   }
 
-  private boolean currentExecutionIsLast() {
-    return executions == maxExecutions - 1;
+  @Override
+  public boolean wasLastExecution() {
+    // skipped tests (either by the framework or DD) should not be retried
+    return lastStatus == TestStatus.skip || executions >= maxExecutions;
+  }
+
+  @Override
+  public boolean applicable() {
+    // executions must always be registered, therefore consider it applicable as long as there are
+    // retries left
+    return !wasLastExecution();
   }
 
   @Override
@@ -53,17 +68,6 @@ private int getExecutions(long durationMillis) {
     return 0;
   }
 
-  @Override
-  public boolean retry(boolean successful, long durationMillis) {
-    ++totalExecutionsSeen;
-    if (successful) {
-      ++successfulExecutionsSeen;
-    }
-    int maxExecutionsForGivenDuration = getExecutions(durationMillis);
-    maxExecutions = Math.min(maxExecutions, maxExecutionsForGivenDuration);
-    return ++executions < maxExecutions;
-  }
-
   @Nullable
   @Override
   public RetryReason currentExecutionRetryReason() {
@@ -76,11 +80,11 @@ private boolean currentExecutionIsRetry() {
 
   @Override
   public boolean hasFailedAllRetries() {
-    return currentExecutionIsLast() && successfulExecutionsSeen == 0;
+    return wasLastExecution() && successfulExecutionsSeen == 0;
   }
 
   @Override
   public boolean hasSucceededAllRetries() {
-    return currentExecutionIsLast() && successfulExecutionsSeen == totalExecutionsSeen;
+    return wasLastExecution() && successfulExecutionsSeen == executions;
   }
 }