Generalize test retry logic

nikita-tkachenko-datadog · nikita-tkachenko-datadog · commit 3e58f244e535 · 2025-01-28T17:51:50.000+01:00
diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/domain/TestImpl.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/domain/TestImpl.java
@@ -15,13 +15,13 @@
 import datadog.trace.api.civisibility.domain.TestContext;
 import datadog.trace.api.civisibility.telemetry.CiVisibilityCountMetric;
 import datadog.trace.api.civisibility.telemetry.CiVisibilityMetricCollector;
+import datadog.trace.api.civisibility.telemetry.TagValue;
 import datadog.trace.api.civisibility.telemetry.tag.BrowserDriver;
 import datadog.trace.api.civisibility.telemetry.tag.EventType;
 import datadog.trace.api.civisibility.telemetry.tag.IsModified;
 import datadog.trace.api.civisibility.telemetry.tag.IsNew;
 import datadog.trace.api.civisibility.telemetry.tag.IsRetry;
 import datadog.trace.api.civisibility.telemetry.tag.IsRum;
-import datadog.trace.api.civisibility.telemetry.tag.RetryReason;
 import datadog.trace.api.civisibility.telemetry.tag.SkipReason;
 import datadog.trace.api.civisibility.telemetry.tag.TestFrameworkInstrumentation;
 import datadog.trace.api.gateway.RequestContextSlot;
@@ -263,6 +263,7 @@ public void end(@Nullable Long endTime) {
       span.finish();
     }
 
+    Object retryReason = span.getTag(Tags.TEST_RETRY_REASON);
     metricCollector.add(
         CiVisibilityCountMetric.EVENT_FINISHED,
         1,
@@ -271,25 +272,13 @@ public void end(@Nullable Long endTime) {
         span.getTag(Tags.TEST_IS_NEW) != null ? IsNew.TRUE : null,
         span.getTag(Tags.TEST_IS_MODIFIED) != null ? IsModified.TRUE : null,
         span.getTag(Tags.TEST_IS_RETRY) != null ? IsRetry.TRUE : null,
-        getRetryReason(),
+        retryReason instanceof TagValue ? (TagValue) retryReason : null,
         span.getTag(Tags.TEST_IS_RUM_ACTIVE) != null ? IsRum.TRUE : null,
         CIConstants.SELENIUM_BROWSER_DRIVER.equals(span.getTag(Tags.TEST_BROWSER_DRIVER))
             ? BrowserDriver.SELENIUM
             : null);
   }
 
-  private RetryReason getRetryReason() {
-    String retryReason = (String) span.getTag(Tags.TEST_RETRY_REASON);
-    if (retryReason != null) {
-      try {
-        return RetryReason.valueOf(retryReason.toUpperCase());
-      } catch (IllegalArgumentException e) {
-        log.debug("Non-standard retry-reason: {}", retryReason);
-      }
-    }
-    return null;
-  }
-
   /**
    * Tests often perform operations that involve APM instrumentations: sending an HTTP request,
    * executing a database query, etc. APM instrumentations create spans that correspond to those
diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/events/NoOpTestEventsHandler.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/events/NoOpTestEventsHandler.java
@@ -6,6 +6,7 @@
 import datadog.trace.api.civisibility.config.TestSourceData;
 import datadog.trace.api.civisibility.events.TestEventsHandler;
 import datadog.trace.api.civisibility.retry.TestRetryPolicy;
+import datadog.trace.api.civisibility.telemetry.tag.RetryReason;
 import datadog.trace.api.civisibility.telemetry.tag.SkipReason;
 import datadog.trace.api.civisibility.telemetry.tag.TestFrameworkInstrumentation;
 import datadog.trace.bootstrap.ContextStore;
@@ -57,7 +58,7 @@ public void onTestStart(
       @Nullable String testParameters,
       @Nullable Collection<String> categories,
       @Nonnull TestSourceData testSourceData,
-      String retryReason,
+      RetryReason retryReason,
       @Nullable Long startTime) {
     // do nothing
   }
diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/events/TestEventsHandlerImpl.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/events/TestEventsHandlerImpl.java
@@ -12,6 +12,7 @@
 import datadog.trace.api.civisibility.telemetry.CiVisibilityCountMetric;
 import datadog.trace.api.civisibility.telemetry.CiVisibilityMetricCollector;
 import datadog.trace.api.civisibility.telemetry.tag.EventType;
+import datadog.trace.api.civisibility.telemetry.tag.RetryReason;
 import datadog.trace.api.civisibility.telemetry.tag.SkipReason;
 import datadog.trace.api.civisibility.telemetry.tag.TestFrameworkInstrumentation;
 import datadog.trace.bootstrap.ContextStore;
@@ -138,7 +139,7 @@ public void onTestStart(
       final @Nullable String testParameters,
       final @Nullable Collection<String> categories,
       final @Nonnull TestSourceData testSourceData,
-      final @Nullable String retryReason,
+      final @Nullable RetryReason retryReason,
       final @Nullable Long startTime) {
     if (skipTrace(testSourceData.getTestClass())) {
       return;
diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/retry/NeverRetry.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/retry/NeverRetry.java
@@ -1,6 +1,7 @@
 package datadog.trace.civisibility.retry;
 
 import datadog.trace.api.civisibility.retry.TestRetryPolicy;
+import datadog.trace.api.civisibility.telemetry.tag.RetryReason;
 import org.jetbrains.annotations.Nullable;
 
 public class NeverRetry implements TestRetryPolicy {
@@ -20,7 +21,7 @@ public boolean suppressFailures() {
   }
 
   @Override
-  public boolean retry(boolean successful, long duration) {
+  public boolean retry(boolean successful, long durationMillis) {
     return false;
   }
 
@@ -31,7 +32,7 @@ public boolean currentExecutionIsRetry() {
 
   @Nullable
   @Override
-  public String currentExecutionRetryReason() {
+  public RetryReason currentExecutionRetryReason() {
     return null;
   }
 }
diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/retry/RetryIfFailed.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/retry/RetryIfFailed.java
@@ -1,14 +1,13 @@
 package datadog.trace.civisibility.retry;
 
 import datadog.trace.api.civisibility.retry.TestRetryPolicy;
+import datadog.trace.api.civisibility.telemetry.tag.RetryReason;
 import java.util.concurrent.atomic.AtomicInteger;
 import org.jetbrains.annotations.Nullable;
 
 /** Retries a test case if it failed, up to a maximum number of times. */
 public class RetryIfFailed implements TestRetryPolicy {
 
-  private static final String AUTO_TEST_RETRIES = "atr";
-
   private final int maxExecutions;
   private int executions;
 
@@ -32,7 +31,7 @@ public boolean suppressFailures() {
   }
 
   @Override
-  public boolean retry(boolean successful, long duration) {
+  public boolean retry(boolean successful, long durationMillis) {
     if (!successful && ++executions < maxExecutions) {
       totalExecutions.incrementAndGet();
       return true;
@@ -48,7 +47,7 @@ public boolean currentExecutionIsRetry() {
 
   @Nullable
   @Override
-  public String currentExecutionRetryReason() {
-    return currentExecutionIsRetry() ? AUTO_TEST_RETRIES : null;
+  public RetryReason currentExecutionRetryReason() {
+    return currentExecutionIsRetry() ? RetryReason.atr : null;
   }
 }
diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/retry/RetryNTimes.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/retry/RetryNTimes.java
@@ -1,14 +1,13 @@
 package datadog.trace.civisibility.retry;
 
 import datadog.trace.api.civisibility.retry.TestRetryPolicy;
+import datadog.trace.api.civisibility.telemetry.tag.RetryReason;
 import datadog.trace.civisibility.config.EarlyFlakeDetectionSettings;
 import org.jetbrains.annotations.Nullable;
 
 /** Retries a test case N times (N depends on test duration) regardless of success or failure. */
 public class RetryNTimes implements TestRetryPolicy {
 
-  private static final String EARLY_FLAKINESS_DETECTION = "efd";
-
   private final EarlyFlakeDetectionSettings earlyFlakeDetectionSettings;
   private int executions;
   private int maxExecutions;
@@ -30,9 +29,9 @@ public boolean suppressFailures() {
   }
 
   @Override
-  public boolean retry(boolean successful, long duration) {
+  public boolean retry(boolean successful, long durationMillis) {
     // adjust maximum retries based on the now known test duration
-    int maxExecutionsForGivenDuration = earlyFlakeDetectionSettings.getExecutions(duration);
+    int maxExecutionsForGivenDuration = earlyFlakeDetectionSettings.getExecutions(durationMillis);
     maxExecutions = Math.min(maxExecutions, maxExecutionsForGivenDuration);
     return ++executions < maxExecutions;
   }
@@ -44,7 +43,7 @@ public boolean currentExecutionIsRetry() {
 
   @Nullable
   @Override
-  public String currentExecutionRetryReason() {
-    return currentExecutionIsRetry() ? EARLY_FLAKINESS_DETECTION : null;
+  public RetryReason currentExecutionRetryReason() {
+    return currentExecutionIsRetry() ? RetryReason.efd : null;
   }
 }
diff --git a/dd-java-agent/instrumentation/junit-5.3/src/main/java/datadog/trace/instrumentation/junit5/JUnitPlatformUtils.java b/dd-java-agent/instrumentation/junit-5.3/src/main/java/datadog/trace/instrumentation/junit5/JUnitPlatformUtils.java
@@ -4,6 +4,7 @@
 
 import datadog.trace.api.civisibility.config.TestIdentifier;
 import datadog.trace.api.civisibility.config.TestSourceData;
+import datadog.trace.api.civisibility.telemetry.tag.RetryReason;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
@@ -185,8 +186,9 @@ public static boolean isParameterizedTest(TestDescriptor testDescriptor) {
     return "test-template".equals(lastSegment.getType());
   }
 
-  public static String retryReason(TestDescriptor testDescriptor) {
-    return getIDSegmentValue(testDescriptor, RETRY_DESCRIPTOR_REASON_SUFFIX);
+  public static RetryReason retryReason(TestDescriptor testDescriptor) {
+    String retryReasonSegment = getIDSegmentValue(testDescriptor, RETRY_DESCRIPTOR_REASON_SUFFIX);
+    return retryReasonSegment != null ? RetryReason.valueOf(retryReasonSegment) : null;
   }
 
   public static boolean isRetry(TestDescriptor testDescriptor) {
diff --git a/dd-java-agent/instrumentation/junit-5.3/src/main/java/datadog/trace/instrumentation/junit5/retry/TestDescriptorHandle.java b/dd-java-agent/instrumentation/junit-5.3/src/main/java/datadog/trace/instrumentation/junit5/retry/TestDescriptorHandle.java
@@ -43,10 +43,12 @@ public TestDescriptorHandle(TestDescriptor testDescriptor) {
   }
 
   public TestDescriptor withIdSuffix(
-      String segmentName, String segmentValue, String otherSegmentName, String otherSegmentValue) {
+      String segmentName, Object segmentValue, String otherSegmentName, Object otherSegmentValue) {
     UniqueId uniqueId = testDescriptor.getUniqueId();
     UniqueId updatedId =
-        uniqueId.append(segmentName, segmentValue).append(otherSegmentName, otherSegmentValue);
+        uniqueId
+            .append(segmentName, String.valueOf(segmentValue))
+            .append(otherSegmentName, String.valueOf(otherSegmentValue));
     TestDescriptor descriptorClone = UnsafeUtils.tryShallowClone(testDescriptor);
     METHOD_HANDLES.invoke(UNIQUE_ID_SETTER, descriptorClone, updatedId);
     return descriptorClone;
diff --git a/dd-java-agent/instrumentation/karate/src/main/java/datadog/trace/instrumentation/karate/KarateTracingHook.java b/dd-java-agent/instrumentation/karate/src/main/java/datadog/trace/instrumentation/karate/KarateTracingHook.java
@@ -19,6 +19,7 @@
 import datadog.trace.api.civisibility.config.TestSourceData;
 import datadog.trace.api.civisibility.events.TestDescriptor;
 import datadog.trace.api.civisibility.events.TestSuiteDescriptor;
+import datadog.trace.api.civisibility.telemetry.tag.RetryReason;
 import datadog.trace.api.civisibility.telemetry.tag.SkipReason;
 import datadog.trace.api.civisibility.telemetry.tag.TestFrameworkInstrumentation;
 import datadog.trace.bootstrap.ContextStore;
@@ -143,7 +144,7 @@ public boolean beforeScenario(ScenarioRuntime sr) {
         parameters,
         categories,
         TestSourceData.UNKNOWN,
-        (String) sr.magicVariables.get(KarateUtils.RETRY_MAGIC_VARIABLE),
+        (RetryReason) sr.magicVariables.get(KarateUtils.RETRY_MAGIC_VARIABLE),
         null);
     return true;
   }
diff --git a/dd-java-agent/instrumentation/testng/src/main/java/datadog/trace/instrumentation/testng/TracingListener.java b/dd-java-agent/instrumentation/testng/src/main/java/datadog/trace/instrumentation/testng/TracingListener.java
@@ -2,6 +2,7 @@
 
 import datadog.trace.api.civisibility.config.TestSourceData;
 import datadog.trace.api.civisibility.events.TestSuiteDescriptor;
+import datadog.trace.api.civisibility.telemetry.tag.RetryReason;
 import datadog.trace.api.civisibility.telemetry.tag.TestFrameworkInstrumentation;
 import datadog.trace.instrumentation.testng.retry.RetryAnalyzer;
 import java.util.List;
@@ -94,7 +95,7 @@ public void onTestStart(final ITestResult result) {
         null);
   }
 
-  private String retryReason(final ITestResult result) {
+  private RetryReason retryReason(final ITestResult result) {
     IRetryAnalyzer retryAnalyzer = TestNGUtils.getRetryAnalyzer(result);
     if (retryAnalyzer instanceof RetryAnalyzer) {
       RetryAnalyzer datadogAnalyzer = (RetryAnalyzer) retryAnalyzer;
diff --git a/dd-java-agent/instrumentation/testng/src/main/java/datadog/trace/instrumentation/testng/retry/RetryAnalyzer.java b/dd-java-agent/instrumentation/testng/src/main/java/datadog/trace/instrumentation/testng/retry/RetryAnalyzer.java
@@ -3,6 +3,7 @@
 import datadog.trace.api.civisibility.config.TestIdentifier;
 import datadog.trace.api.civisibility.config.TestSourceData;
 import datadog.trace.api.civisibility.retry.TestRetryPolicy;
+import datadog.trace.api.civisibility.telemetry.tag.RetryReason;
 import datadog.trace.instrumentation.testng.TestEventsHandlerHolder;
 import datadog.trace.instrumentation.testng.TestNGUtils;
 import org.testng.IRetryAnalyzer;
@@ -31,7 +32,7 @@ public boolean retry(ITestResult result) {
     return retryPolicy.retry(result.isSuccess(), result.getEndMillis() - result.getStartMillis());
   }
 
-  public String currentExecutionRetryReason() {
+  public RetryReason currentExecutionRetryReason() {
     return retryPolicy != null ? retryPolicy.currentExecutionRetryReason() : null;
   }
 }
diff --git a/dd-java-agent/instrumentation/testng/testng-7/src/main/java/datadog/trace/instrumentation/testng7/TestNGClassListenerInstrumentation.java b/dd-java-agent/instrumentation/testng/testng-7/src/main/java/datadog/trace/instrumentation/testng7/TestNGClassListenerInstrumentation.java
@@ -60,6 +60,7 @@ public String[] helperClassNames() {
   }
 
   public static class InvokeBeforeClassAdvice {
+    @SuppressWarnings("bytebuddy-exception-suppression")
     @Advice.OnMethodEnter
     public static void invokeBeforeClass(
         @Advice.FieldValue("m_testContext") final ITestContext testContext,
@@ -81,6 +82,7 @@ public static String muzzleCheck(final CustomAttribute customAttribute) {
   }
 
   public static class InvokeAfterClassAdvice {
+    @SuppressWarnings("bytebuddy-exception-suppression")
     @Advice.OnMethodExit
     public static void invokeAfterClass(
         @Advice.FieldValue("m_testContext") final ITestContext testContext,
diff --git a/dd-java-agent/instrumentation/testng/testng-7/src/main/java/datadog/trace/instrumentation/testng7/TestNGRetryInstrumentation.java b/dd-java-agent/instrumentation/testng/testng-7/src/main/java/datadog/trace/instrumentation/testng7/TestNGRetryInstrumentation.java
@@ -55,6 +55,7 @@ public String[] helperClassNames() {
   }
 
   public static class RetryAdvice {
+    @SuppressWarnings("bytebuddy-exception-suppression")
     @SuppressFBWarnings("NP_BOOLEAN_RETURN_NULL")
     @Advice.OnMethodExit
     public static void shouldRetryTestMethod(
diff --git a/dd-java-agent/instrumentation/weaver/src/main/java/datadog/trace/instrumentation/weaver/DatadogWeaverReporter.java b/dd-java-agent/instrumentation/weaver/src/main/java/datadog/trace/instrumentation/weaver/DatadogWeaverReporter.java
@@ -5,6 +5,7 @@
 import datadog.trace.api.civisibility.events.TestDescriptor;
 import datadog.trace.api.civisibility.events.TestEventsHandler;
 import datadog.trace.api.civisibility.events.TestSuiteDescriptor;
+import datadog.trace.api.civisibility.telemetry.tag.RetryReason;
 import datadog.trace.api.civisibility.telemetry.tag.TestFrameworkInstrumentation;
 import datadog.trace.api.time.SystemTimeSource;
 import datadog.trace.util.AgentThreadFactory;
@@ -91,7 +92,7 @@ public static void onTestFinished(TestFinished event, TaskDef taskDef) {
         new TestDescriptor(testSuiteName, testClass, testName, testParameters, testQualifier);
     String testMethodName = null;
     Method testMethod = null;
-    String retryReason = null;
+    RetryReason retryReason = null;
 
     // Only test finish is reported, so fake test start timestamp
     long endMicros = SystemTimeSource.INSTANCE.getCurrentTimeMicros();
diff --git a/internal-api/src/main/java/datadog/trace/api/civisibility/events/TestEventsHandler.java b/internal-api/src/main/java/datadog/trace/api/civisibility/events/TestEventsHandler.java
@@ -5,6 +5,7 @@
 import datadog.trace.api.civisibility.config.TestIdentifier;
 import datadog.trace.api.civisibility.config.TestSourceData;
 import datadog.trace.api.civisibility.retry.TestRetryPolicy;
+import datadog.trace.api.civisibility.telemetry.tag.RetryReason;
 import datadog.trace.api.civisibility.telemetry.tag.SkipReason;
 import datadog.trace.api.civisibility.telemetry.tag.TestFrameworkInstrumentation;
 import datadog.trace.bootstrap.ContextStore;
@@ -64,7 +65,7 @@ void onTestStart(
       @Nullable String testParameters,
       @Nullable Collection<String> categories,
       @Nonnull TestSourceData testSourceData,
-      @Nullable String retryReason,
+      @Nullable RetryReason retryReason,
       @Nullable Long startTime);
 
   void onTestSkip(TestKey descriptor, @Nullable String reason);
diff --git a/internal-api/src/main/java/datadog/trace/api/civisibility/retry/TestRetryPolicy.java b/internal-api/src/main/java/datadog/trace/api/civisibility/retry/TestRetryPolicy.java
@@ -1,20 +1,32 @@
 package datadog.trace.api.civisibility.retry;
 
+import datadog.trace.api.civisibility.telemetry.tag.RetryReason;
 import javax.annotation.Nullable;
 
 public interface TestRetryPolicy {
-  boolean retriesLeft();
 
+  /** @return {@code true} if failure of this test should not affect the build result */
   boolean suppressFailures();
 
-  boolean retry(boolean successful, long duration);
+  /**
+   * @return {@code true} if this test can be retried. Current execution is NOT taken into account
+   */
+  boolean retriesLeft();
+
+  /**
+   * @param successful {@code true} if test passed or was skipped, {@code false} otherwise
+   * @param durationMillis test duration in milliseconds
+   * @return {@code true} if this test should be retried
+   */
+  boolean retry(boolean successful, long durationMillis);
 
+  /** @return {@code true} if current execution is a retry */
   boolean currentExecutionIsRetry();
 
   /**
    * Returns retry reason for current execution (will be {@code null} if current execution is not a
    * retry)
    */
   @Nullable
-  String currentExecutionRetryReason();
+  RetryReason currentExecutionRetryReason();
 }
diff --git a/internal-api/src/main/java/datadog/trace/api/civisibility/telemetry/tag/RetryReason.java b/internal-api/src/main/java/datadog/trace/api/civisibility/telemetry/tag/RetryReason.java
@@ -3,13 +3,19 @@
 import datadog.trace.api.civisibility.telemetry.TagValue;
 
 public enum RetryReason implements TagValue {
-  ATR,
-  EFD;
+  atr("Auto Test Retries"),
+  efd("Early Flakiness Detection");
 
   private final String s;
+  private final String description;
 
-  RetryReason() {
-    s = "retry_reason:" + name().toLowerCase();
+  RetryReason(String description) {
+    this.description = description;
+    this.s = "retry_reason:" + name();
+  }
+
+  public String getDescription() {
+    return description;
   }
 
   @Override

Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,7 @@`
`1`	`1`	`package datadog.trace.civisibility.retry;`
`2`	`2`
`3`	`3`	`import datadog.trace.api.civisibility.retry.TestRetryPolicy;`
	`4`	`+import datadog.trace.api.civisibility.telemetry.tag.RetryReason;`
`4`	`5`	`import org.jetbrains.annotations.Nullable;`
`5`	`6`
`6`	`7`	`public class NeverRetry implements TestRetryPolicy {`
`@@ -20,7 +21,7 @@ public boolean suppressFailures() {`
`20`	`21`	`}`
`21`	`22`
`22`	`23`	`@Override`
`23`		`- public boolean retry(boolean successful, long duration) {`
	`24`	`+ public boolean retry(boolean successful, long durationMillis) {`
`24`	`25`	`return false;`
`25`	`26`	`}`
`26`	`27`
`@@ -31,7 +32,7 @@ public boolean currentExecutionIsRetry() {`
`31`	`32`
`32`	`33`	`@Nullable`
`33`	`34`	`@Override`
`34`		`- public String currentExecutionRetryReason() {`
	`35`	`+ public RetryReason currentExecutionRetryReason() {`
`35`	`36`	`return null;`
`36`	`37`	`}`
`37`	`38`	`}`