Consume Response Body Before Retrying Request (#16607)

* Updated retry policies to consume response body before triggering retry * Validate response body isn't null before attempting to consume
Azure · Oct 21, 2020 · 3bce404 · 3bce404
1 parent 0dabe6f
commit 3bce404
Show file tree

Hide file tree

Showing 3 changed files with 90 additions and 6 deletions.
diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/http/policy/RetryPolicy.java b/sdk/core/azure-core/src/main/java/com/azure/core/http/policy/RetryPolicy.java
@@ -8,8 +8,10 @@
 import com.azure.core.http.HttpRequest;
 import com.azure.core.http.HttpResponse;
 import com.azure.core.util.logging.ClientLogger;
+import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
+import java.nio.ByteBuffer;
 import java.time.Duration;
 import java.time.temporal.ChronoUnit;
 import java.util.Objects;
@@ -99,8 +101,17 @@ private Mono<HttpResponse> attemptAsync(final HttpPipelineCallContext context, f
                     final Duration delayDuration = determineDelayDuration(httpResponse, tryCount);
                     logger.verbose("[Retrying] Try count: {}, Delay duration in seconds: {}", tryCount,
                         delayDuration.getSeconds());
-                    return attemptAsync(context, next, originalHttpRequest, tryCount + 1)
-                        .delaySubscription(delayDuration);
+
+                    Flux<ByteBuffer> responseBody = httpResponse.getBody();
+                    if (responseBody == null) {
+                        return attemptAsync(context, next, originalHttpRequest, tryCount + 1)
+                            .delaySubscription(delayDuration);
+                    } else {
+                        return httpResponse.getBody()
+                            .ignoreElements()
+                            .then(attemptAsync(context, next, originalHttpRequest, tryCount + 1)
+                                .delaySubscription(delayDuration));
+                    }
                 } else {
                     return Mono.just(httpResponse);
                 }

diff --git a/sdk/core/azure-core/src/test/java/com/azure/core/http/policy/RetryPolicyTests.java b/sdk/core/azure-core/src/test/java/com/azure/core/http/policy/RetryPolicyTests.java
@@ -3,21 +3,30 @@
 
 package com.azure.core.http.policy;
 
+import com.azure.core.http.HttpHeaders;
 import com.azure.core.http.HttpMethod;
 import com.azure.core.http.HttpPipeline;
 import com.azure.core.http.HttpPipelineBuilder;
 import com.azure.core.http.HttpRequest;
 import com.azure.core.http.HttpResponse;
 import com.azure.core.http.MockHttpResponse;
 import com.azure.core.http.clients.NoOpHttpClient;
+import com.azure.core.util.FluxUtil;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
+import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
 import java.net.URL;
+import java.nio.ByteBuffer;
+import java.nio.charset.Charset;
+import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.time.temporal.ChronoUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
 
 public class RetryPolicyTests {
 
@@ -40,7 +49,7 @@ public Mono<HttpResponse> send(HttpRequest request) {
         HttpResponse response = pipeline.send(new HttpRequest(HttpMethod.GET,
             new URL("http://localhost/"))).block();
 
-        Assertions.assertEquals(501, response.getStatusCode());
+        assertEquals(501, response.getStatusCode());
     }
 
     @Test
@@ -62,7 +71,7 @@ public Mono<HttpResponse> send(HttpRequest request) {
         HttpResponse response = pipeline.send(new HttpRequest(HttpMethod.GET,
             new URL("http://localhost/"))).block();
 
-        Assertions.assertEquals(500, response.getStatusCode());
+        assertEquals(500, response.getStatusCode());
     }
 
     @Test
@@ -124,5 +133,59 @@ public Mono<HttpResponse> send(HttpRequest request) {
             .verifyComplete();
     }
 
+    @Test
+    public void retryConsumesBody() {
+        final AtomicInteger bodyConsumptionCount = new AtomicInteger();
+        Flux<ByteBuffer> errorBody = Flux.generate(sink -> {
+            bodyConsumptionCount.incrementAndGet();
+            sink.next(ByteBuffer.wrap("Should be consumed".getBytes(StandardCharsets.UTF_8)));
+            sink.complete();
+        });
+
+        final HttpPipeline pipeline = new HttpPipelineBuilder()
+            .policies(new RetryPolicy(new FixedDelay(2, Duration.ofMillis(1))))
+            .httpClient(request -> Mono.just(new HttpResponse(request) {
+                @Override
+                public int getStatusCode() {
+                    return 503;
+                }
+
+                @Override
+                public String getHeaderValue(String name) {
+                    return getHeaders().getValue(name);
+                }
+
+                @Override
+                public HttpHeaders getHeaders() {
+                    return new HttpHeaders();
+                }
+
+                @Override
+                public Flux<ByteBuffer> getBody() {
+                    return errorBody;
+                }
 
+                @Override
+                public Mono<byte[]> getBodyAsByteArray() {
+                    return FluxUtil.collectBytesInByteBufferStream(getBody());
+                }
+
+                @Override
+                public Mono<String> getBodyAsString() {
+                    return getBodyAsString(StandardCharsets.UTF_8);
+                }
+
+                @Override
+                public Mono<String> getBodyAsString(Charset charset) {
+                    return getBodyAsByteArray().map(bytes -> new String(bytes, charset));
+                }
+            }))
+            .build();
+
+        StepVerifier.create(pipeline.send(new HttpRequest(HttpMethod.GET, "https://example.com")))
+            .expectNextCount(1)
+            .verifyComplete();
+
+        assertEquals(2, bodyConsumptionCount.get());
+    }
 }
diff --git a/...zure-storage-common/src/main/java/com/azure/storage/common/policy/RequestRetryPolicy.java b/...zure-storage-common/src/main/java/com/azure/storage/common/policy/RequestRetryPolicy.java
@@ -140,8 +140,18 @@ we do not consider the secondary at all (considerSecondary==false)). This will
                         ensure primaryTry is correct when passed to calculate the delay.
                          */
                     int newPrimaryTry = (!tryingPrimary || !considerSecondary) ? primaryTry + 1 : primaryTry;
-                    return attemptAsync(context, next, originalRequest, newConsiderSecondary, newPrimaryTry,
-                        attempt + 1);
+
+                    Flux<ByteBuffer> responseBody = response.getBody();
+                    if (responseBody == null) {
+                        return attemptAsync(context, next, originalRequest, newConsiderSecondary, newPrimaryTry,
+                            attempt + 1);
+                    } else {
+                        return response.getBody()
+                            .ignoreElements()
+                            .then(attemptAsync(context, next, originalRequest, newConsiderSecondary, newPrimaryTry,
+                                attempt + 1));
+                    }
+
                 }
                 return Mono.just(response);
             }).onErrorResume(throwable -> {