Allow RetryOptions to configure HttpResponse and Throwable retry logic (

#38469)

Allow for RetryOptions to configure HttpResponse and Throwable retry behaviors

sdk/core/azure-core/src/main/java/com/azure/core/http/policy/ExponentialBackoff.java

@@ -3,6 +3,8 @@
 
 package com.azure.core.http.policy;
 
+import com.azure.core.http.HttpResponse;
+import com.azure.core.implementation.accesshelpers.ExponentialBackoffAccessHelper;
 import com.azure.core.implementation.util.ObjectsUtil;
 import com.azure.core.util.Configuration;
 import com.azure.core.util.CoreUtils;
@@ -12,6 +14,7 @@
 import java.time.Duration;
 import java.util.Objects;
 import java.util.concurrent.ThreadLocalRandom;
+import java.util.function.Predicate;
 
 import static com.azure.core.util.Configuration.PROPERTY_AZURE_REQUEST_RETRY_COUNT;
 
@@ -44,11 +47,15 @@ public class ExponentialBackoff implements RetryStrategy {
         }
 
         DEFAULT_MAX_RETRIES = defaultMaxRetries;
+
+        ExponentialBackoffAccessHelper.setAccessor(ExponentialBackoff::new);
     }
 
     private final int maxRetries;
     private final long baseDelayNanos;
     private final long maxDelayNanos;
+    private final Predicate<HttpResponse> shouldRetry;
+    private final Predicate<Throwable> shouldRetryException;
 
     /**
      * Creates an instance of {@link ExponentialBackoff} with a maximum number of retry attempts configured by the
@@ -68,18 +75,25 @@ public ExponentialBackoff() {
      */
     public ExponentialBackoff(ExponentialBackoffOptions options) {
         this(
-            ObjectsUtil.requireNonNullElse(
-                Objects.requireNonNull(options, "'options' cannot be null.").getMaxRetries(),
+            ObjectsUtil.requireNonNullElse(Objects.requireNonNull(options, "'options' cannot be null.").getMaxRetries(),
                 DEFAULT_MAX_RETRIES),
-            ObjectsUtil.requireNonNullElse(
-                Objects.requireNonNull(options, "'options' cannot be null.").getBaseDelay(),
+            ObjectsUtil.requireNonNullElse(Objects.requireNonNull(options, "'options' cannot be null.").getBaseDelay(),
                 DEFAULT_BASE_DELAY),
-            ObjectsUtil.requireNonNullElse(
-                Objects.requireNonNull(options, "'options' cannot be null.").getMaxDelay(),
+            ObjectsUtil.requireNonNullElse(Objects.requireNonNull(options, "'options' cannot be null.").getMaxDelay(),
                 DEFAULT_MAX_DELAY)
         );
     }
 
+    private ExponentialBackoff(ExponentialBackoffOptions options, Predicate<HttpResponse> shouldRetry,
+        Predicate<Throwable> shouldRetryException) {
+        this(ObjectsUtil.requireNonNullElse(
+            Objects.requireNonNull(options, "'options' cannot be null.").getMaxRetries(), DEFAULT_MAX_RETRIES),
+            ObjectsUtil.requireNonNullElse(Objects.requireNonNull(options, "'options' cannot be null.").getBaseDelay(),
+                DEFAULT_BASE_DELAY),
+            ObjectsUtil.requireNonNullElse(Objects.requireNonNull(options, "'options' cannot be null.").getMaxDelay(),
+                DEFAULT_MAX_DELAY), shouldRetry, shouldRetryException);
+    }
+
     /**
      * Creates an instance of {@link ExponentialBackoff}.
      *
@@ -90,6 +104,11 @@ public ExponentialBackoff(ExponentialBackoffOptions options) {
      * to 0 or {@code maxDelay} is less than {@code baseDelay}.
      */
     public ExponentialBackoff(int maxRetries, Duration baseDelay, Duration maxDelay) {
+        this(maxRetries, baseDelay, maxDelay, null, null);
+    }
+
+    private ExponentialBackoff(int maxRetries, Duration baseDelay, Duration maxDelay,
+        Predicate<HttpResponse> shouldRetry, Predicate<Throwable> shouldRetryException) {
         if (maxRetries < 0) {
             throw LOGGER.logExceptionAsError(new IllegalArgumentException("Max retries cannot be less than 0."));
         }
@@ -107,6 +126,8 @@ public ExponentialBackoff(int maxRetries, Duration baseDelay, Duration maxDelay)
         this.maxRetries = maxRetries;
         this.baseDelayNanos = baseDelay.toNanos();
         this.maxDelayNanos = maxDelay.toNanos();
+        this.shouldRetry = shouldRetry;
+        this.shouldRetryException = shouldRetryException;
     }
 
     @Override
@@ -121,4 +142,18 @@ public Duration calculateRetryDelay(int retryAttempts) {
             .nextLong((long) (baseDelayNanos * (1 - JITTER_FACTOR)), (long) (baseDelayNanos * (1 + JITTER_FACTOR)));
         return Duration.ofNanos(Math.min((1L << retryAttempts) * delayWithJitterInNanos, maxDelayNanos));
     }
+
+    @Override
+    public boolean shouldRetry(HttpResponse httpResponse) {
+        return (shouldRetry == null)
+            ? RetryStrategy.super.shouldRetry(httpResponse)
+            : shouldRetry.test(httpResponse);
+    }
+
+    @Override
+    public boolean shouldRetryException(Throwable throwable) {
+        return (shouldRetryException == null)
+            ? RetryStrategy.super.shouldRetryException(throwable)
+            : shouldRetryException.test(throwable);
+    }
 }

sdk/core/azure-core/src/main/java/com/azure/core/http/policy/FixedDelay.java

@@ -3,18 +3,27 @@
 
 package com.azure.core.http.policy;
 
+import com.azure.core.http.HttpResponse;
+import com.azure.core.implementation.accesshelpers.FixedDelayAccessHelper;
 import com.azure.core.util.logging.ClientLogger;
 import java.time.Duration;
 import java.util.Objects;
+import java.util.function.Predicate;
 
 /**
  * A fixed-delay implementation of {@link RetryStrategy} that has a fixed delay duration between each retry attempt.
  */
 public class FixedDelay implements RetryStrategy {
     private static final ClientLogger LOGGER = new ClientLogger(FixedDelay.class);
 
+    static {
+        FixedDelayAccessHelper.setAccessor(FixedDelay::new);
+    }
+
     private final int maxRetries;
     private final Duration delay;
+    private final Predicate<HttpResponse> shouldRetry;
+    private final Predicate<Throwable> shouldRetryException;
 
     /**
      * Creates an instance of {@link FixedDelay}.
@@ -25,11 +34,7 @@ public class FixedDelay implements RetryStrategy {
      * @throws NullPointerException If {@code delay} is {@code null}.
      */
     public FixedDelay(int maxRetries, Duration delay) {
-        if (maxRetries < 0) {
-            throw LOGGER.logExceptionAsError(new IllegalArgumentException("Max retries cannot be less than 0."));
-        }
-        this.maxRetries = maxRetries;
-        this.delay = Objects.requireNonNull(delay, "'delay' cannot be null.");
+        this(maxRetries, delay, null, null);
     }
 
     /**
@@ -38,10 +43,26 @@ public FixedDelay(int maxRetries, Duration delay) {
      * @param fixedDelayOptions The {@link FixedDelayOptions}.
      */
     public FixedDelay(FixedDelayOptions fixedDelayOptions) {
-        this(
-            Objects.requireNonNull(fixedDelayOptions, "'fixedDelayOptions' cannot be null.").getMaxRetries(),
-            Objects.requireNonNull(fixedDelayOptions, "'fixedDelayOptions' cannot be null.").getDelay()
-        );
+        this(Objects.requireNonNull(fixedDelayOptions, "'fixedDelayOptions' cannot be null.").getMaxRetries(),
+            Objects.requireNonNull(fixedDelayOptions, "'fixedDelayOptions' cannot be null.").getDelay());
+    }
+
+    private FixedDelay(FixedDelayOptions fixedDelayOptions, Predicate<HttpResponse> shouldRetry,
+        Predicate<Throwable> shouldRetryException) {
+        this(Objects.requireNonNull(fixedDelayOptions, "'fixedDelayOptions' cannot be null.").getMaxRetries(),
+            Objects.requireNonNull(fixedDelayOptions, "'fixedDelayOptions' cannot be null.").getDelay(), shouldRetry,
+            shouldRetryException);
+    }
+
+    private FixedDelay(int maxRetries, Duration delay, Predicate<HttpResponse> shouldRetry,
+        Predicate<Throwable> shouldRetryException) {
+        if (maxRetries < 0) {
+            throw LOGGER.logExceptionAsError(new IllegalArgumentException("Max retries cannot be less than 0."));
+        }
+        this.maxRetries = maxRetries;
+        this.delay = Objects.requireNonNull(delay, "'delay' cannot be null.");
+        this.shouldRetry = shouldRetry;
+        this.shouldRetryException = shouldRetryException;
     }
 
     @Override
@@ -53,4 +74,18 @@ public int getMaxRetries() {
     public Duration calculateRetryDelay(int retryAttempts) {
         return delay;
     }
+
+    @Override
+    public boolean shouldRetry(HttpResponse httpResponse) {
+        return (shouldRetry == null)
+            ? RetryStrategy.super.shouldRetry(httpResponse)
+            : shouldRetry.test(httpResponse);
+    }
+
+    @Override
+    public boolean shouldRetryException(Throwable throwable) {
+        return (shouldRetryException == null)
+            ? RetryStrategy.super.shouldRetryException(throwable)
+            : shouldRetryException.test(throwable);
+    }
 }

sdk/core/azure-core/src/main/java/com/azure/core/http/policy/RetryOptions.java

@@ -3,7 +3,10 @@
 
 package com.azure.core.http.policy;
 
+import com.azure.core.http.HttpResponse;
+
 import java.util.Objects;
+import java.util.function.Predicate;
 
 /**
  * The configuration for retries.
@@ -12,6 +15,9 @@ public class RetryOptions {
     private final ExponentialBackoffOptions exponentialBackoffOptions;
     private final FixedDelayOptions fixedDelayOptions;
 
+    private Predicate<HttpResponse> shouldRetry;
+    private Predicate<Throwable> shouldRetryException;
+
     /**
      * Creates a new instance that uses {@link ExponentialBackoffOptions}.
      *
@@ -51,4 +57,56 @@ public ExponentialBackoffOptions getExponentialBackoffOptions() {
     public FixedDelayOptions getFixedDelayOptions() {
         return fixedDelayOptions;
     }
+
+    /**
+     * Gets the predicate that determines if a retry should be attempted for the given {@link HttpResponse}.
+     * <p>
+     * If null, the default behavior is to retry HTTP responses with status codes 408, 429, and any 500 status code that
+     * isn't 501 or 505.
+     *
+     * @return The predicate that determines if a retry should be attempted for the given {@link HttpResponse}.
+     */
+    public Predicate<HttpResponse> getShouldRetry() {
+        return shouldRetry;
+    }
+
+    /**
+     * Sets the predicate that determines if a retry should be attempted for the given {@link HttpResponse}.
+     * <p>
+     * If null, the default behavior is to retry HTTP responses with status codes 408, 429, and any 500 status code that
+     * isn't 501 or 505.
+     *
+     * @param shouldRetry The predicate that determines if a retry should be attempted for the given
+     * {@link HttpResponse}.
+     * @return The updated {@link RetryOptions} object.
+     */
+    public RetryOptions setShouldRetry(Predicate<HttpResponse> shouldRetry) {
+        this.shouldRetry = shouldRetry;
+        return this;
+    }
+
+    /**
+     * Gets the predicate that determines if a retry should be attempted for the given {@link Throwable}.
+     * <p>
+     * If null, the default behavior is to retry any {@link Exception}.
+     *
+     * @return The predicate that determines if a retry should be attempted for the given {@link Throwable}.
+     */
+    public Predicate<Throwable> getShouldRetryException() {
+        return shouldRetryException;
+    }
+
+    /**
+     * Sets the predicate that determines if a retry should be attempted for the given {@link Throwable}.
+     * <p>
+     * If null, the default behavior is to retry any {@link Exception}.
+     *
+     * @param shouldRetryException The predicate that determines if a retry should be attempted for the given
+     * {@link Throwable}.
+     * @return The updated {@link RetryOptions} object.
+     */
+    public RetryOptions setShouldRetryException(Predicate<Throwable> shouldRetryException) {
+        this.shouldRetryException = shouldRetryException;
+        return this;
+    }
 }

sdk/core/azure-core/src/main/java/com/azure/core/implementation/ImplUtils.java

@@ -5,10 +5,10 @@
 
 import com.azure.core.http.HttpHeaderName;
 import com.azure.core.http.HttpHeaders;
-import com.azure.core.http.policy.ExponentialBackoff;
-import com.azure.core.http.policy.FixedDelay;
 import com.azure.core.http.policy.RetryOptions;
 import com.azure.core.http.policy.RetryStrategy;
+import com.azure.core.implementation.accesshelpers.ExponentialBackoffAccessHelper;
+import com.azure.core.implementation.accesshelpers.FixedDelayAccessHelper;
 import com.azure.core.util.CoreUtils;
 import com.azure.core.util.DateTimeRfc1123;
 import com.azure.core.util.FluxUtil;
@@ -406,9 +406,11 @@ public static RetryStrategy getRetryStrategyFromOptions(RetryOptions retryOption
         Objects.requireNonNull(retryOptions, "'retryOptions' cannot be null.");
 
         if (retryOptions.getExponentialBackoffOptions() != null) {
-            return new ExponentialBackoff(retryOptions.getExponentialBackoffOptions());
+            return ExponentialBackoffAccessHelper.create(retryOptions.getExponentialBackoffOptions(),
+                retryOptions.getShouldRetry(), retryOptions.getShouldRetryException());
         } else if (retryOptions.getFixedDelayOptions() != null) {
-            return new FixedDelay(retryOptions.getFixedDelayOptions());
+            return FixedDelayAccessHelper.create(retryOptions.getFixedDelayOptions(), retryOptions.getShouldRetry(),
+                retryOptions.getShouldRetryException());
         } else {
             // This should never happen.
             throw new IllegalArgumentException("'retryOptions' didn't define any retry strategy options");

sdk/core/azure-core/src/main/java/com/azure/core/implementation/accesshelpers/ExponentialBackoffAccessHelper.java

@@ -0,0 +1,64 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.core.implementation.accesshelpers;
+
+import com.azure.core.http.HttpResponse;
+import com.azure.core.http.policy.ExponentialBackoff;
+import com.azure.core.http.policy.ExponentialBackoffOptions;
+
+import java.util.function.Predicate;
+
+/**
+ * Class containing helper methods for accessing private members of {@link ExponentialBackoff}.
+ */
+public final class ExponentialBackoffAccessHelper {
+    private static ExponentialBackoffAccessor accessor;
+
+    /**
+     * Type defining the methods to set the non-public properties of an {@link ExponentialBackoff} instance.
+     */
+    public interface ExponentialBackoffAccessor {
+        /**
+         * Creates an {@link ExponentialBackoff} instance with the passed {@code exponentialBackoffOptions},
+         * {@code shouldRetry} and {@code shouldRetryException}.
+         *
+         * @param exponentialBackoffOptions The {@link ExponentialBackoffOptions}.
+         * @param shouldRetry The {@link Predicate} to determine if a response should be retried.
+         * @param shouldRetryException The {@link Predicate} to determine if a {@link Throwable} should be retried.
+         * @return The created {@link ExponentialBackoff} instance.
+         */
+        ExponentialBackoff create(ExponentialBackoffOptions exponentialBackoffOptions,
+            Predicate<HttpResponse> shouldRetry, Predicate<Throwable> shouldRetryException);
+    }
+
+    /**
+     * The method called from {@link ExponentialBackoff} to set it's accessor.
+     *
+     * @param exponentialBackoffAccessor The accessor.
+     */
+    public static void setAccessor(final ExponentialBackoffAccessor exponentialBackoffAccessor) {
+        accessor = exponentialBackoffAccessor;
+    }
+
+    /**
+     * Creates an {@link ExponentialBackoff} instance with the passed {@code exponentialBackoffOptions},
+     * {@code shouldRetry} and {@code shouldRetryException}.
+     *
+     * @param exponentialBackoffOptions The {@link ExponentialBackoffOptions}.
+     * @param shouldRetry The {@link Predicate} to determine if a response should be retried.
+     * @param shouldRetryException The {@link Predicate} to determine if a {@link Throwable} should be retried.
+     * @return The created {@link ExponentialBackoff} instance.
+     */
+    public static ExponentialBackoff create(ExponentialBackoffOptions exponentialBackoffOptions,
+        Predicate<HttpResponse> shouldRetry, Predicate<Throwable> shouldRetryException) {
+        if (accessor == null) {
+            new ExponentialBackoff();
+        }
+
+        assert accessor != null;
+        return accessor.create(exponentialBackoffOptions, shouldRetry, shouldRetryException);
+    }
+
+    private ExponentialBackoffAccessHelper() {
+    }
+}