Retry Requests Based on Causal Exception Too (Azure#32648)

Retry Requests Based on Causal Exception Too

Author: alzimmermsft

sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/policy/RequestRetryPolicy.java

@@ -110,6 +110,7 @@ private Mono<HttpResponse> attemptAsync(final HttpPipelineCallContext context, H
             Flux<ByteBuffer> bufferedBody = context.getHttpRequest().getBody().map(ByteBuffer::duplicate);
             context.getHttpRequest().setBody(bufferedBody);
         }
+
         if (!tryingPrimary) {
             UrlBuilder builder = UrlBuilder.parse(context.getHttpRequest().getUrl());
             builder.setHost(this.requestRetryOptions.getSecondaryHost());
@@ -119,23 +120,18 @@ private Mono<HttpResponse> attemptAsync(final HttpPipelineCallContext context, H
                 return Mono.error(e);
             }
         }
-        /*
-        Update the RETRY_COUNT_CONTEXT to log retries.
-         */
+
+        // Update the RETRY_COUNT_CONTEXT to log retries.
         context.setData(HttpLoggingPolicy.RETRY_COUNT_CONTEXT, attempt);
 
-        /*
-        Reset progress if progress is tracked.
-         */
+        // Reset progress if progress is tracked.
         ProgressReporter progressReporter = Contexts.with(context.getContext()).getHttpRequestProgressReporter();
         if (progressReporter != null) {
             progressReporter.reset();
         }
 
-        /*
-         We want to send the request with a given timeout, but we don't want to kickoff that timeout-bound operation
-         until after the retry backoff delay, so we call delaySubscription.
-         */
+        // We want to send the request with a given timeout, but we don't want to kick off that timeout-bound operation
+        // until after the retry backoff delay, so we call delaySubscription.
         Mono<HttpResponse> responseMono = next.clone().process();
 
         // Default try timeout is Integer.MAX_VALUE seconds, if it's that don't set a timeout as that's about 68 years
@@ -152,19 +148,11 @@ private Mono<HttpResponse> attemptAsync(final HttpPipelineCallContext context, H
 
         return responseMono.flatMap(response -> {
             boolean newConsiderSecondary = considerSecondary;
-            boolean retry = false;
             int statusCode = response.getStatusCode();
 
-            /*
-             * If attempt was against the secondary & it returned a StatusNotFound (404), then the resource was not
-             * found. This may be due to replication delay. So, in this case, we'll never try the secondary again for
-             * this operation.
-             */
+            boolean retry = shouldStatusCodeBeRetried(statusCode, tryingPrimary);
             if (!tryingPrimary && statusCode == 404) {
                 newConsiderSecondary = false;
-                retry = true;
-            } else if (statusCode == 503 || statusCode == 500) {
-                retry = true;
             }
 
             if (retry && attempt < requestRetryOptions.getMaxTries()) {
@@ -208,15 +196,10 @@ private Mono<HttpResponse> attemptAsync(final HttpPipelineCallContext context, H
              * better to optimistically retry instead of failing too soon. A Timeout Exception is a client-side timeout
              * coming from Rx.
              */
-            boolean retry = false;
-            Throwable unwrappedThrowable = Exceptions.unwrap(throwable);
-            if (unwrappedThrowable instanceof IOException) {
-                retry = true;
-            } else if (unwrappedThrowable instanceof TimeoutException) {
-                retry = true;
-            }
+            ExceptionRetryStatus exceptionRetryStatus = shouldErrorBeRetried(throwable, attempt,
+                requestRetryOptions.getMaxTries());
 
-            if (retry && attempt < requestRetryOptions.getMaxTries()) {
+            if (exceptionRetryStatus.canBeRetried) {
                 /*
                  * We increment primaryTry if we are about to try the primary again (which is when we consider the
                  * secondary and tried the secondary this time (tryingPrimary==false) or we do not consider the
@@ -225,14 +208,63 @@ private Mono<HttpResponse> attemptAsync(final HttpPipelineCallContext context, H
                  */
                 int newPrimaryTry = (!tryingPrimary || !considerSecondary) ? primaryTry + 1 : primaryTry;
                 List<Throwable> suppressedLocal = suppressed == null ? new LinkedList<>() : suppressed;
-                suppressedLocal.add(unwrappedThrowable);
+                suppressedLocal.add(exceptionRetryStatus.unwrappedThrowable);
                 return attemptAsync(context, next, originalRequest, considerSecondary, newPrimaryTry, attempt + 1,
                     suppressedLocal);
             }
+
             if (suppressed != null) {
                 suppressed.forEach(throwable::addSuppressed);
             }
+
             return Mono.error(throwable);
         });
     }
+
+    static ExceptionRetryStatus shouldErrorBeRetried(Throwable error, int attempt, int maxAttempts) {
+        Throwable unwrappedThrowable = Exceptions.unwrap(error);
+
+        // Check if there are any attempts remaining.
+        if (attempt >= maxAttempts) {
+            return new ExceptionRetryStatus(false, unwrappedThrowable);
+        }
+
+        // Check if the unwrapped error is an IOException or TimeoutException.
+        if (unwrappedThrowable instanceof IOException || unwrappedThrowable instanceof TimeoutException) {
+            return new ExceptionRetryStatus(true, unwrappedThrowable);
+        }
+
+        // Check the causal exception chain for this exception being caused by an IOException or TimeoutException.
+        Throwable causalException = unwrappedThrowable.getCause();
+        while (causalException != null) {
+            if (causalException instanceof IOException || causalException instanceof TimeoutException) {
+                return new ExceptionRetryStatus(true, unwrappedThrowable);
+            }
+
+            causalException = causalException.getCause();
+        }
+
+        // Finally all exceptions have been checked and none can be retried.
+        return new ExceptionRetryStatus(false, unwrappedThrowable);
+    }
+
+    static boolean shouldStatusCodeBeRetried(int statusCode, boolean isPrimary) {
+        /*
+         * Retry the request if the server had an error (500), was unavailable (503), or requested a backoff (429),
+         * or if the secondary was being tried and the resources didn't exist there (404). Only the secondary can retry
+         * if the resource wasn't found as there may be a delay in replication from the primary.
+         */
+        return (statusCode == 429 || statusCode == 500 || statusCode == 503)
+            || (!isPrimary && statusCode == 404);
+    }
+
+    static final class ExceptionRetryStatus {
+        final boolean canBeRetried;
+        final Throwable unwrappedThrowable;
+
+        ExceptionRetryStatus(boolean canBeRetried, Throwable unwrappedThrowable) {
+            this.canBeRetried = canBeRetried;
+            this.unwrappedThrowable = unwrappedThrowable;
+        }
+    }
 }

sdk/storage/azure-storage-common/src/test/java/com/azure/storage/common/policy/RequestRetryPolicyTest.groovy

@@ -0,0 +1,47 @@
+package com.azure.storage.common.policy
+
+import reactor.core.Exceptions
+import spock.lang.Specification
+
+import java.util.concurrent.TimeoutException
+
+/**
+ * Tests {@link RequestRetryPolicy}
+ */
+class RequestRetryPolicyTest extends Specification {
+    def "should exception be retried"() {
+        expect:
+        RequestRetryPolicy.shouldErrorBeRetried(error, 0, 1).canBeRetried == shouldBeRetried
+
+        where:
+        error                                                       | shouldBeRetried
+        new IOException()                                           | true
+        new TimeoutException()                                      | true
+        new RuntimeException()                                      | false
+        Exceptions.propagate(new IOException())                     | true
+        Exceptions.propagate(new TimeoutException())                | true
+        Exceptions.propagate(new RuntimeException())                | false
+        new Exception(new IOException())                            | true
+        new Exception(new TimeoutException())                       | true
+        new Exception(new RuntimeException())                       | false
+        Exceptions.propagate(new Exception(new IOException()))      | true
+        Exceptions.propagate(new Exception(new TimeoutException())) | true
+        Exceptions.propagate(new Exception(new RuntimeException())) | false
+    }
+
+    def "should status code be retried"() {
+        expect:
+        RequestRetryPolicy.shouldStatusCodeBeRetried(statusCode, isPrimary) == shouldBeRetried
+
+        where:
+        statusCode | isPrimary | shouldBeRetried
+        429        | true      | true
+        429        | false     | true
+        500        | true      | true
+        500        | false     | true
+        503        | true      | true
+        503        | false     | true
+        404        | true      | false
+        404        | false     | true
+    }
+}