feat(ci_visibility): dynamic atr retries

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

@@ -276,7 +276,8 @@ private static TestFrameworkSession.Factory childTestFrameworkSessionFactory(
               services.config,
               executionSettings,
               repoServices.sourcePathResolver,
-              services.linesResolver);
+              services.linesResolver,
+              services.metricCollector);
 
       // only add report upload capability for children sessions,
       // because report upload is only supported when the build system is instrumented
@@ -319,7 +320,8 @@ private static TestFrameworkSession.Factory headlessTestFrameworkSessionFactory(
               services.config,
               executionSettings,
               repoServices.sourcePathResolver,
-              services.linesResolver);
+              services.linesResolver,
+              services.metricCollector);
       return new HeadlessTestSession(
           projectName,
           startTime,

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/config/EarlyFlakeDetectionSettings.java

@@ -40,6 +40,40 @@ public List<ExecutionsByDuration> getExecutionsByDuration() {
     return executionsByDuration;
   }
 
+  /**
+   * Returns the EFD retry-bucket index for an initial test duration. Bucket boundaries match the EFD
+   * duration buckets: {@code <=5s -> 0, <=10s -> 1, <=30s -> 2, <=5m -> 3, >5m -> 4}.
+   */
+  public int retryBucketIndexForDuration(long durationMillis) {
+    if (durationMillis <= 5_000) {
+      return 0;
+    }
+    if (durationMillis <= 10_000) {
+      return 1;
+    }
+    if (durationMillis <= 30_000) {
+      return 2;
+    }
+    if (durationMillis <= 300_000) {
+      return 3;
+    }
+    return 4;
+  }
+
+  /**
+   * Returns the configured retry budget for an initial test duration, based on the backend-provided
+   * {@code executionsByDuration} list. Returns 0 if no bucket matches (e.g. test ran longer than the
+   * longest configured duration).
+   */
+  public int retriesForDuration(long durationMillis) {
+    for (ExecutionsByDuration e : executionsByDuration) {
+      if (durationMillis <= e.getDurationMillis()) {
+        return e.getExecutions();
+      }
+    }
+    return 0;
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) {

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

@@ -17,7 +17,7 @@
         "TestExecutionPolicy instances are confined to a single thread and are not meant to be thread-safe")
 public class AutoTestRetry implements TestExecutionPolicy {
 
-  private final int maxExecutions;
+  protected int maxExecutions;
   private final boolean suppressFailures;
   private int executions;
   private ExecutionAggregation results;

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

@@ -0,0 +1,49 @@
+package datadog.trace.civisibility.execution;
+
+import datadog.trace.api.civisibility.execution.TestStatus;
+import datadog.trace.api.civisibility.telemetry.tag.RetryReason;
+import datadog.trace.civisibility.config.EarlyFlakeDetectionSettings;
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * Execution policy for dynamic Auto Test Retries (ATR). Instead of a flat per-test retry limit, the
+ * number of retries is determined by the duration of the initial attempt, using the same duration
+ * buckets as Early Flake Detection. When custom buckets are provided they override the EFD retry
+ * settings; otherwise the EFD settings from the backend are used.
+ */
+public class DynamicAutoTestRetry extends AutoTestRetry {
+
+  private final EarlyFlakeDetectionSettings efdSettings;
+  private final int[] customBuckets; // null = use EFD settings
+  private boolean maxExecutionsDetermined = false;
+
+  public DynamicAutoTestRetry(
+      EarlyFlakeDetectionSettings efdSettings,
+      int[] customBuckets,
+      boolean suppressFailures,
+      AtomicInteger totalRetryCount) {
+    super(Integer.MAX_VALUE, suppressFailures, totalRetryCount);
+    this.efdSettings = efdSettings;
+    this.customBuckets = customBuckets;
+  }
+
+  @Override
+  public ExecutionOutcome registerExecution(TestStatus status, long durationMillis) {
+    if (!maxExecutionsDetermined) {
+      maxExecutions = computeMaxExecutions(durationMillis);
+      maxExecutionsDetermined = true;
+    }
+    return super.registerExecution(status, durationMillis);
+  }
+
+  private int computeMaxExecutions(long durationMillis) {
+    int retries;
+    if (customBuckets != null) {
+      int index = efdSettings.retryBucketIndexForDuration(durationMillis);
+      retries = customBuckets[index];
+    } else {
+      retries = efdSettings.retriesForDuration(durationMillis);
+    }
+    return Math.max(1, retries) + 1; // +1 for the initial attempt
+  }
+}

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

@@ -4,8 +4,7 @@
 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;
+import datadog.trace.civisibility.config.EarlyFlakeDetectionSettings;
 
 /**
  * Execution policy for Early Flake Detection. Runs a new or modified test case multiple times to
@@ -15,18 +14,18 @@
 public class EarlyFlakeDetection implements TestExecutionPolicy {
 
   private final boolean suppressFailures;
-  private final List<ExecutionsByDuration> executionsByDuration;
+  private final EarlyFlakeDetectionSettings efdSettings;
   private int executions;
   private int maxExecutions;
   private ExecutionAggregation results;
   private TestStatus lastStatus;
 
   public EarlyFlakeDetection(
-      List<ExecutionsByDuration> executionsByDuration, boolean suppressFailures) {
+      EarlyFlakeDetectionSettings efdSettings, boolean suppressFailures) {
     this.suppressFailures = suppressFailures;
-    this.executionsByDuration = executionsByDuration;
+    this.efdSettings = efdSettings;
     this.executions = 0;
-    this.maxExecutions = getExecutions(0);
+    this.maxExecutions = efdSettings.retriesForDuration(0);
     this.results = ExecutionAggregation.NONE;
   }
 
@@ -35,7 +34,7 @@ public ExecutionOutcome registerExecution(TestStatus status, long durationMillis
     lastStatus = status;
     ++executions;
     results = results.withExecution(status);
-    int maxExecutionsForGivenDuration = getExecutions(durationMillis);
+    int maxExecutionsForGivenDuration = efdSettings.retriesForDuration(durationMillis);
     maxExecutions = Math.min(maxExecutions, maxExecutionsForGivenDuration);
 
     boolean lastExecution = !retriesLeft();
@@ -80,15 +79,6 @@ public boolean propagateFailure() {
     return !suppressFailures && results == ExecutionAggregation.MIXED;
   }
 
-  private int getExecutions(long durationMillis) {
-    for (ExecutionsByDuration e : executionsByDuration) {
-      if (durationMillis <= e.getDurationMillis()) {
-        return e.getExecutions();
-      }
-    }
-    return 0;
-  }
-
   @Override
   public boolean failedTestReplayApplicable() {
     return false;

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/test/ExecutionStrategy.java

@@ -6,13 +6,17 @@
 import datadog.trace.api.civisibility.config.TestMetadata;
 import datadog.trace.api.civisibility.config.TestSourceData;
 import datadog.trace.api.civisibility.execution.TestExecutionPolicy;
+import datadog.trace.api.civisibility.telemetry.CiVisibilityMetricCollector;
+import datadog.trace.api.civisibility.telemetry.CiVisibilityCountMetric;
+import datadog.trace.api.civisibility.telemetry.tag.HasCustomBuckets;
 import datadog.trace.api.civisibility.telemetry.tag.SkipReason;
 import datadog.trace.civisibility.config.EarlyFlakeDetectionSettings;
 import datadog.trace.civisibility.config.ExecutionSettings;
 import datadog.trace.civisibility.config.TestManagementSettings;
 import datadog.trace.civisibility.config.TestSetting;
 import datadog.trace.civisibility.execution.AttemptToFix;
 import datadog.trace.civisibility.execution.AutoTestRetry;
+import datadog.trace.civisibility.execution.DynamicAutoTestRetry;
 import datadog.trace.civisibility.execution.EarlyFlakeDetection;
 import datadog.trace.civisibility.execution.Quarantine;
 import datadog.trace.civisibility.execution.Regular;
@@ -38,16 +42,29 @@ public class ExecutionStrategy {
   @Nonnull private final ExecutionSettings executionSettings;
   @Nonnull private final SourcePathResolver sourcePathResolver;
   @Nonnull private final LinesResolver linesResolver;
+  @Nonnull private final CiVisibilityMetricCollector metricCollector;
+  private final int[] dynamicAtrBuckets;
 
   public ExecutionStrategy(
       @Nonnull Config config,
       @Nonnull ExecutionSettings executionSettings,
       @Nonnull SourcePathResolver sourcePathResolver,
-      @Nonnull LinesResolver linesResolver) {
+      @Nonnull LinesResolver linesResolver,
+      @Nonnull CiVisibilityMetricCollector metricCollector) {
     this.config = config;
     this.executionSettings = executionSettings;
     this.sourcePathResolver = sourcePathResolver;
     this.linesResolver = linesResolver;
+    this.metricCollector = metricCollector;
+    this.dynamicAtrBuckets = parseDynamicAtrBuckets(config.getCiVisibilityDynamicAtrBuckets());
+
+    if (config.isCiVisibilityDynamicAtrEnabled()
+        && executionSettings.isFlakyTestRetriesEnabled()) {
+      metricCollector.add(
+          CiVisibilityCountMetric.DYNAMIC_ATR_RETRIES_ENABLED,
+          1,
+          dynamicAtrBuckets != null ? HasCustomBuckets.TRUE : null);
+    }
   }
 
   @Nonnull
@@ -135,13 +152,20 @@ public TestExecutionPolicy executionPolicy(
       // but we don't care if we go "a bit" over the limit, it does not have to be precise
       earlyFlakeDetectionsUsed.incrementAndGet();
       return new EarlyFlakeDetection(
-          executionSettings.getEarlyFlakeDetectionSettings().getExecutionsByDuration(),
+          executionSettings.getEarlyFlakeDetectionSettings(),
           isQuarantined(test));
     }
 
     if (isAutoRetryApplicable(test)) {
       // check-then-act with "autoRetriesUsed" is not atomic here,
       // but we don't care if we go "a bit" over the limit, it does not have to be precise
+      if (config.isCiVisibilityDynamicAtrEnabled()) {
+        return new DynamicAutoTestRetry(
+            executionSettings.getEarlyFlakeDetectionSettings(),
+            dynamicAtrBuckets,
+            isQuarantined(test),
+            autoRetriesUsed);
+      }
       return new AutoTestRetry(
           config.getCiVisibilityFlakyRetryCount(), isQuarantined(test), autoRetriesUsed);
     }
@@ -247,4 +271,50 @@ public int executionPriority(@Nullable TestIdentifier test, @Nonnull TestSourceD
     // then the rest
     return 0;
   }
+
+  private static final int RETRY_BUCKET_COUNT = 5;
+  private static final int MAX_RETRIES_PER_BUCKET = 20;
+
+  /**
+   * Parses the {@code DD_CIVISIBILITY_DYNAMIC_ATR_BUCKETS} env var into five positive integers in
+   * [1, 20]. Returns {@code null} if the value is unset/empty or invalid (wrong count,
+   * non-integer, out of range) — in which case the EFD retry settings are used as fallback.
+   */
+  private static int[] parseDynamicAtrBuckets(String rawBuckets) {
+    if (rawBuckets == null || rawBuckets.isEmpty()) {
+      return null;
+    }
+    String[] parts = rawBuckets.split(",", -1);
+    if (parts.length != RETRY_BUCKET_COUNT) {
+      LOGGER.warn(
+          "Invalid {} value '{}'; expected five comma-separated integers in [1, {}]",
+          "DD_CIVISIBILITY_DYNAMIC_ATR_BUCKETS",
+          rawBuckets,
+          MAX_RETRIES_PER_BUCKET);
+      return null;
+    }
+    int[] buckets = new int[RETRY_BUCKET_COUNT];
+    try {
+      for (int i = 0; i < RETRY_BUCKET_COUNT; i++) {
+        int value = Integer.parseInt(parts[i].trim());
+        if (value < 1 || value > MAX_RETRIES_PER_BUCKET) {
+          LOGGER.warn(
+              "Invalid {} value '{}'; expected five comma-separated integers in [1, {}]",
+              "DD_CIVISIBILITY_DYNAMIC_ATR_BUCKETS",
+              rawBuckets,
+              MAX_RETRIES_PER_BUCKET);
+          return null;
+        }
+        buckets[i] = value;
+      }
+    } catch (NumberFormatException e) {
+      LOGGER.warn(
+          "Invalid {} value '{}'; expected five comma-separated integers in [1, {}]",
+          "DD_CIVISIBILITY_DYNAMIC_ATR_BUCKETS",
+          rawBuckets,
+          MAX_RETRIES_PER_BUCKET);
+      return null;
+    }
+    return buckets;
+  }
 }

dd-java-agent/agent-ci-visibility/src/test/groovy/datadog/trace/civisibility/domain/buildsystem/ProxyTestModuleTest.groovy

@@ -9,6 +9,7 @@ import datadog.trace.civisibility.config.EarlyFlakeDetectionSettings
 import datadog.trace.api.civisibility.config.TestIdentifier
 import datadog.trace.api.civisibility.coverage.CoverageStore
 import datadog.trace.api.civisibility.telemetry.CiVisibilityMetricCollector
+import datadog.trace.api.civisibility.telemetry.NoOpMetricCollector
 import datadog.trace.civisibility.codeowners.Codeowners
 import datadog.trace.civisibility.config.ExecutionSettings
 import datadog.trace.civisibility.coverage.report.child.ChildProcessCoverageReporter
@@ -30,7 +31,7 @@ class ProxyTestModuleTest extends DDSpecification {
     config.getCiVisibilityFlakyRetryCount() >> 2 // this counts all executions of a test case (first attempt is counted too)
     config.getCiVisibilityTotalFlakyRetryCount() >> 2 // this counts retries across all tests (first attempt is not a retry, so it is not counted)
 
-    def executionStrategy = new ExecutionStrategy(config, executionSettings, Stub(SourcePathResolver), Stub(LinesResolver))
+    def executionStrategy = new ExecutionStrategy(config, executionSettings, Stub(SourcePathResolver), Stub(LinesResolver), NoOpMetricCollector.INSTANCE)
 
     def traceId = Stub(DDTraceId)
     traceId.toLong() >> 123

dd-java-agent/agent-ci-visibility/src/test/groovy/datadog/trace/civisibility/domain/headless/HeadlessTestModuleTest.groovy

@@ -6,6 +6,7 @@ import datadog.trace.api.civisibility.config.TestSourceData
 import datadog.trace.api.civisibility.coverage.CoverageStore
 import datadog.trace.api.civisibility.execution.TestStatus
 import datadog.trace.api.civisibility.telemetry.CiVisibilityMetricCollector
+import datadog.trace.api.civisibility.telemetry.NoOpMetricCollector
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext
 import datadog.trace.civisibility.codeowners.Codeowners
 import datadog.trace.civisibility.config.EarlyFlakeDetectionSettings
@@ -58,7 +59,7 @@ class HeadlessTestModuleTest extends SpanWriterTest {
     config.getCiVisibilityTotalFlakyRetryCount() >> 2
     // this counts retries across all tests (first attempt is not a retry, so it is not counted)
 
-    def executionStrategy = new ExecutionStrategy(config, executionSettings, Stub(SourcePathResolver), Stub(LinesResolver))
+    def executionStrategy = new ExecutionStrategy(config, executionSettings, Stub(SourcePathResolver), Stub(LinesResolver), NoOpMetricCollector.INSTANCE)
 
     new HeadlessTestModule(
     Stub(AgentSpanContext),

dd-java-agent/agent-ci-visibility/src/test/groovy/datadog/trace/civisibility/domain/headless/HeadlessTestSessionTest.groovy

@@ -5,6 +5,7 @@ import datadog.trace.api.Config
 import datadog.trace.api.DDSpanTypes
 import datadog.trace.api.civisibility.coverage.CoverageStore
 import datadog.trace.api.civisibility.telemetry.CiVisibilityMetricCollector
+import datadog.trace.api.civisibility.telemetry.NoOpMetricCollector
 import datadog.trace.api.civisibility.telemetry.tag.Provider
 import datadog.trace.bootstrap.instrumentation.api.Tags
 import datadog.trace.civisibility.codeowners.Codeowners
@@ -47,7 +48,7 @@ class HeadlessTestSessionTest extends SpanWriterTest {
     def executionSettings = Stub(ExecutionSettings)
     executionSettings.getTestManagementSettings() >> new TestManagementSettings(true, 10)
 
-    def executionStrategy = new ExecutionStrategy(Stub(Config), executionSettings, Stub(SourcePathResolver), Stub(LinesResolver))
+    def executionStrategy = new ExecutionStrategy(Stub(Config), executionSettings, Stub(SourcePathResolver), Stub(LinesResolver), NoOpMetricCollector.INSTANCE)
 
     new HeadlessTestSession(
       "project-name",