Add duration based maxRetries option to Retry class

Author: DomGarguilo

core/src/main/java/org/apache/accumulo/core/util/Retry.java

@@ -272,6 +272,12 @@ public interface NeedsRetries {
      * @return this builder with the maximum number of retries set to the provided value
      */
     NeedsRetryDelay maxRetries(long max);
+
+    /**
+     * @return this builder with the maximum number of retries set to the number of retries that can
+     *         occur within the given duration
+     */
+    NeedsRetryDelay retriesForDuration(Duration duration);
   }
 
   public interface NeedsRetryDelay {
@@ -358,6 +364,7 @@ private static class RetryFactoryBuilder
     private Duration waitIncrement;
     private Duration logInterval;
     private double backOffFactor = 1.5;
+    private Duration retriesForDuration = null;
 
     RetryFactoryBuilder() {}
 
@@ -381,6 +388,55 @@ public NeedsRetryDelay maxRetries(long max) {
       return this;
     }
 
+    @Override
+    public NeedsRetryDelay retriesForDuration(Duration duration) {
+      checkState();
+      Preconditions.checkState(!duration.isNegative(), "Duration must not be negative");
+      this.retriesForDuration = duration;
+      return this;
+    }
+
+    private void calculateRetriesForDuration() {
+      long durationMillis = this.retriesForDuration.toMillis();
+      long totalWaitMillis = initialWait.toMillis();
+      long currentWaitMillis = totalWaitMillis;
+
+      long retries = 0L;
+      // if the initial wait already exceeds or meets the duration
+      if (totalWaitMillis >= durationMillis) {
+        this.maxRetries = retries;
+        return;
+      }
+
+      while (totalWaitMillis + currentWaitMillis <= durationMillis) {
+
+        if (backOffFactor > 1) {
+          currentWaitMillis = (long) (currentWaitMillis * backOffFactor);
+        } else {
+          currentWaitMillis += waitIncrement.toMillis();
+        }
+
+        // Ensure the wait time does not exceed maxWait
+        if (currentWaitMillis > maxWait.toMillis()) {
+          currentWaitMillis = maxWait.toMillis();
+        }
+
+        // Break if adding another wait period would exceed the duration
+        if (totalWaitMillis + currentWaitMillis > durationMillis) {
+          break;
+        }
+
+        totalWaitMillis += currentWaitMillis;
+        retries++;
+
+        if (retries > 100) {
+          break;
+        }
+      }
+
+      this.maxRetries = retries;
+    }
+
     @Override
     public NeedsTimeIncrement retryAfter(Duration duration) {
       checkState();
@@ -434,6 +490,10 @@ public RetryFactory createFactory() {
 
     @Override
     public Retry createRetry() {
+      if (retriesForDuration != null) {
+        calculateRetriesForDuration();
+      }
+      this.modifiable = false;
       return new Retry(maxRetries, initialWait, waitIncrement, maxWait, logInterval, backOffFactor);
     }
 

core/src/test/java/org/apache/accumulo/core/util/RetryTest.java

@@ -34,6 +34,7 @@
 import org.easymock.EasyMock;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -351,4 +352,92 @@ public void testInfiniteRetryWithBackoff() throws InterruptedException {
       }
     }
   }
+
+  @Timeout(30)
+  @Test
+  public void testRetriesForDuration() throws InterruptedException {
+    final Duration duration = Duration.ofSeconds(10);
+    final Duration increment = Duration.ofMillis(500);
+    Retry retry = Retry.builder().retriesForDuration(duration).retryAfter(Duration.ofMillis(0))
+        .incrementBy(increment).maxWait(increment).backOffFactor(1)
+        .logInterval(Duration.ofMinutes(3)).createRetry();
+
+    // with the backoff factor set to 1 and with max wait set to the increment value, the expected
+    // number of retries is the duration divided by the increment
+    long expectedRetries = duration.dividedBy(increment);
+
+    assertExpectedRetries(retry, duration, expectedRetries);
+  }
+
+  @Timeout(30)
+  @Test
+  public void testRetriesForDurationWithInitial() throws InterruptedException {
+    final Duration duration = Duration.ofSeconds(10);
+    final Duration increment = Duration.ofMillis(500);
+    final Duration initialWait = Duration.ofMillis(400);
+    Retry retry =
+        Retry.builder().retriesForDuration(duration).retryAfter(initialWait).incrementBy(increment)
+            .maxWait(increment).backOffFactor(1).logInterval(Duration.ofMinutes(3)).createRetry();
+
+    // with the backoff factor set to 1 and with max wait set to the increment value, the expected
+    // number of retries is the duration minus the initial wait divided by the increment
+    long expectedRetries = duration.minus(initialWait).dividedBy(increment);
+
+    assertExpectedRetries(retry, duration, expectedRetries);
+  }
+
+  @Timeout(30)
+  @Test
+  public void testRetriesForDuration2() throws InterruptedException {
+    final Duration increment = Duration.ofMillis(100);
+    final double backOffFactor = 2;
+    Duration maxWait = Duration.ofMillis(800); // Allows for 4 doublings 100- > 200 -> 400 -> 800
+    final Duration initialWait = Duration.ofMillis(50);
+
+    // Total duration accounting for the initial wait, the max wait, and the backoff factor
+    Duration duration =
+        initialWait.plus(Duration.ofMillis(100 + 200 + 400 + 800 + 800 + 800 + 800));
+
+    // count up the expected retries from the calculation above
+    long expectedRetries = 7;
+
+    Retry retry = Retry.builder().retriesForDuration(duration).retryAfter(initialWait)
+        .incrementBy(increment).maxWait(maxWait).backOffFactor(backOffFactor)
+        .logInterval(Duration.ofMinutes(3)).createRetry();
+
+    assertExpectedRetries(retry, duration, expectedRetries);
+  }
+
+  private static void assertExpectedRetries(Retry retry, Duration duration, long expectedRetries)
+      throws InterruptedException {
+    Duration totalTime = Duration.ZERO;
+
+    int iterations = 0;
+    // While the total wait time is less than the duration and there are retries left
+    while (retry.canRetry() && (totalTime.compareTo(duration) <= 0)) {
+      iterations++;
+      totalTime = totalTime.plus(retry.getCurrentWait());
+      if (totalTime.compareTo(duration) <= 0) {
+        retry.useRetry();
+        if (retry.canRetry()) {
+          log.info("Iteration {} - Total time: {}ms - Current wait: {}ms", iterations,
+              totalTime.toMillis(), retry.getCurrentWait().toMillis());
+          retry.waitForNextAttempt(log, "testRetriesForDuration");
+        }
+      }
+    }
+
+    assertEquals(expectedRetries, retry.retriesCompleted(),
+        () -> getErrorMessage(expectedRetries, retry));
+  }
+
+  static String getErrorMessage(long expectedRetries, Retry retry) {
+    return String.format(
+        "Expected %d retries, but only completed %d retries. Retry parameters: maxRetries=%dms, startWait=%sms, maxWait=%sms, waitIncrement=%sms, backOffFactor=%f, logInterval=%sms",
+        expectedRetries, retry.retriesCompleted(), retry.getMaxRetries(),
+        retry.getCurrentWait().toMillis(), retry.getMaxWait().toMillis(),
+        retry.getWaitIncrement().toMillis(), retry.getWaitFactor(),
+        retry.getLogInterval().toMillis());
+  }
+
 }