http-netty: let RetryingHttpRequesterFilter return responses on failure

servicetalk-http-netty/src/main/java/io/servicetalk/http/netty/RetryingHttpRequesterFilter.java

@@ -46,6 +46,10 @@
 import io.servicetalk.transport.api.ExecutionContext;
 import io.servicetalk.transport.api.ExecutionStrategyInfluencer;
 import io.servicetalk.transport.api.RetryableException;
+import io.servicetalk.utils.internal.ThrowableUtils;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
 import java.time.Duration;
@@ -89,17 +93,21 @@
  */
 public final class RetryingHttpRequesterFilter
         implements StreamingHttpClientFilterFactory, ExecutionStrategyInfluencer<HttpExecutionStrategy> {
+
+    private static final Logger LOGGER = LoggerFactory.getLogger(RetryingHttpRequesterFilter.class);
+
     static final int DEFAULT_MAX_TOTAL_RETRIES = 4;
     private static final RetryingHttpRequesterFilter DISABLE_AUTO_RETRIES =
-            new RetryingHttpRequesterFilter(true, false, false, 1, null,
+            new RetryingHttpRequesterFilter(true, false, false, false, 1, null,
                     (__, ___) -> NO_RETRIES, null);
     private static final RetryingHttpRequesterFilter DISABLE_ALL_RETRIES =
-            new RetryingHttpRequesterFilter(false, true, false, 0, null,
+            new RetryingHttpRequesterFilter(false, true, false, false, 0, null,
                     (__, ___) -> NO_RETRIES, null);
 
     private final boolean waitForLb;
     private final boolean ignoreSdErrors;
     private final boolean mayReplayRequestPayload;
+    private final boolean returnFailedResponses;
     private final int maxTotalRetries;
     @Nullable
     private final Function<HttpResponseMetaData, HttpResponseException> responseMapper;
@@ -109,13 +117,14 @@ public final class RetryingHttpRequesterFilter
 
     RetryingHttpRequesterFilter(
             final boolean waitForLb, final boolean ignoreSdErrors, final boolean mayReplayRequestPayload,
-            final int maxTotalRetries,
+            final boolean returnFailedResponses, final int maxTotalRetries,
             @Nullable final Function<HttpResponseMetaData, HttpResponseException> responseMapper,
             final BiFunction<HttpRequestMetaData, Throwable, BackOffPolicy> retryFor,
             @Nullable final RetryCallbacks onRequestRetry) {
         this.waitForLb = waitForLb;
         this.ignoreSdErrors = ignoreSdErrors;
         this.mayReplayRequestPayload = mayReplayRequestPayload;
+        this.returnFailedResponses = returnFailedResponses;
         this.maxTotalRetries = maxTotalRetries;
         this.responseMapper = responseMapper;
         this.retryFor = retryFor;
@@ -210,8 +219,32 @@ public Completable apply(final int count, final Throwable t) {
             }
 
             Completable applyRetryCallbacks(final Completable completable, final int retryCount, final Throwable t) {
-                return retryCallbacks == null ? completable :
-                        completable.beforeOnComplete(() -> retryCallbacks.beforeRetry(retryCount, requestMetaData, t));
+                Completable result = (retryCallbacks == null ? completable :
+                        completable.beforeOnComplete(() -> retryCallbacks.beforeRetry(retryCount, requestMetaData, t)));
+                if (returnFailedResponses) {
+                    if (t instanceof HttpResponseException &&
+                            ((HttpResponseException) t).metaData() instanceof StreamingHttpResponse) {
+                        StreamingHttpResponse response = (StreamingHttpResponse) ((HttpResponseException) t).metaData();
+                        // If we succeed, we need to drain the response body before we continue. If we fail we want to
+                        // surface the original exception and don't worry about draining since it will be returned to
+                        // the user.
+                        result = result.onErrorMap(backoffError -> ThrowableUtils.addSuppressed(t, backoffError))
+                                // If we get cancelled we also need to drain the message body as there is no guarantee
+                                // we'll ever receive a completion event, error or success.
+                                .beforeCancel(() -> drain(response).subscribe())
+                                .concat(drain(response));
+                    } else if (LOGGER.isDebugEnabled()) {
+                        if (!(t instanceof HttpResponseException)) {
+                            LOGGER.debug("Couldn't unpack response due to unexpected dynamic types. Required " +
+                                    "exception of type HttpResponseException, found {}", t.getClass());
+                        } else {
+                            LOGGER.debug("Couldn't unpack response due to unexpected dynamic types. Required " +
+                                    "meta-data of type StreamingHttpResponse, found {}",
+                                    ((HttpResponseException) t).metaData().getClass());
+                        }
+                    }
+                }
+                return result;
             }
         }
 
@@ -258,19 +291,31 @@ protected Single<StreamingHttpResponse> request(final StreamingHttpRequester del
             if (responseMapper != null) {
                 single = single.flatMap(resp -> {
                     final HttpResponseException exception = responseMapper.apply(resp);
-                    return (exception != null ?
-                            // Drain response payload body before discarding it:
-                            resp.payloadBody().ignoreElements().onErrorComplete()
-                                    .concat(Single.<StreamingHttpResponse>failed(exception)) :
-                            Single.succeeded(resp))
-                            .shareContextOnSubscribe();
+                    Single<StreamingHttpResponse> response;
+                    if (exception == null) {
+                        response = Single.succeeded(resp);
+                    } else {
+                        response = Single.failed(exception);
+                        if (!returnFailedResponses) {
+                            response = drain(resp).concat(response);
+                        }
+                    }
+                    return response.shareContextOnSubscribe();
                 });
             }
 
             // 1. Metadata is shared across retries
             // 2. Publisher state is restored to original state for each retry
             // duplicatedRequest isn't used below because retryWhen must be applied outside the defer operator for (2).
-            return single.retryWhen(retryStrategy(request, executionContext(), true));
+            single = single.retryWhen(retryStrategy(request, executionContext(), true));
+            if (returnFailedResponses) {
+                single = single.onErrorResume(HttpResponseException.class, t -> {
+                    HttpResponseMetaData metaData = t.metaData();
+                    return (metaData instanceof StreamingHttpResponse ?
+                            Single.succeeded((StreamingHttpResponse) metaData) : Single.failed(t));
+                });
+            }
+            return single;
         }
     }
 
@@ -719,6 +764,7 @@ public static final class Builder {
 
         private int maxTotalRetries = DEFAULT_MAX_TOTAL_RETRIES;
         private boolean retryExpectationFailed;
+        private boolean returnFailedResponses;
 
         private BiFunction<HttpRequestMetaData, RetryableException, BackOffPolicy>
                 retryRetryableExceptions = (requestMetaData, e) -> BackOffPolicy.ofImmediateBounded();
@@ -801,6 +847,23 @@ public Builder maxTotalRetries(final int maxRetries) {
          * @return {@code this}
          */
         public Builder responseMapper(final Function<HttpResponseMetaData, HttpResponseException> mapper) {
+            return responseMapper(mapper, false);
+        }
+
+        /**
+         * Selectively map a {@link HttpResponseMetaData response} to an {@link HttpResponseException} that can match a
+         * retry behaviour through {@link #retryResponses(BiFunction)}.
+         *
+         * @param mapper a {@link Function} that maps a {@link HttpResponseMetaData} to an
+         * {@link HttpResponseException} or returns {@code null} if there is no mapping for response meta-data. The
+         * mapper should return {@code null} if no retry is needed or if it cannot be determined that a retry is needed.
+         * @param returnFailedResponses whether to unwrap the response defined by the {@link HttpResponseException}
+         * meta-data in the case that the response is not retried.
+         * @return {@code this}
+         */
+        public Builder responseMapper(final Function<HttpResponseMetaData, HttpResponseException> mapper,
+                                      boolean returnFailedResponses) {
+            this.returnFailedResponses = returnFailedResponses;
             this.responseMapper = requireNonNull(mapper);
             return this;
         }
@@ -1054,7 +1117,11 @@ public RetryingHttpRequesterFilter build() {
                         return NO_RETRIES;
                     };
             return new RetryingHttpRequesterFilter(waitForLb, ignoreSdErrors, mayReplayRequestPayload,
-                    maxTotalRetries, responseMapper, allPredicate, onRequestRetry);
+                    returnFailedResponses, maxTotalRetries, responseMapper, allPredicate, onRequestRetry);
         }
     }
+
+    private static Completable drain(StreamingHttpResponse response) {
+        return response.payloadBody().ignoreElements().onErrorComplete();
+    }
 }

servicetalk-http-netty/src/test/java/io/servicetalk/http/netty/RetryingHttpRequesterFilterTest.java

@@ -251,21 +251,30 @@ private void assertRequestRetryingPred(final BlockingHttpClient client) {
         assertThat("Unexpected calls to select.", (double) lbSelectInvoked.get(), closeTo(5.0, 1.0));
     }
 
-    @Test
-    void testResponseMapper() {
+    @ParameterizedTest
+    @ValueSource(booleans = {true, false})
+    void testResponseMapper(final boolean returnFailedResponses) throws Exception {
         AtomicInteger newConnectionCreated = new AtomicInteger();
         AtomicInteger responseDrained = new AtomicInteger();
         AtomicInteger onRequestRetryCounter = new AtomicInteger();
         final int maxTotalRetries = 4;
+        final String retryMessage = "Retryable header";
         normalClient = normalClientBuilder
                 .appendClientFilter(new Builder()
                         .maxTotalRetries(maxTotalRetries)
                         .responseMapper(metaData -> metaData.headers().contains(RETRYABLE_HEADER) ?
-                                    new HttpResponseException("Retryable header", metaData) : null)
+                                    new HttpResponseException(retryMessage, metaData) : null, returnFailedResponses)
                         // Disable request retrying
                         .retryRetryableExceptions((requestMetaData, e) -> ofNoRetries())
                         // Retry only responses marked so
-                        .retryResponses((requestMetaData, throwable) -> ofImmediate(maxTotalRetries - 1))
+                        .retryResponses((requestMetaData, throwable) -> {
+                            if (throwable instanceof HttpResponseException &&
+                                    retryMessage.equals(throwable.getMessage())) {
+                                return ofImmediate(maxTotalRetries - 1);
+                            } else {
+                                throw new RuntimeException("Unexpected exception");
+                            }
+                        })
                         .onRequestRetry((count, req, t) ->
                                 assertThat(onRequestRetryCounter.incrementAndGet(), is(count)))
                         .build())
@@ -281,9 +290,14 @@ public Single<StreamingHttpResponse> request(final StreamingHttpRequest request)
                     };
                 })
                 .buildBlocking();
-        HttpResponseException e = assertThrows(HttpResponseException.class,
-                () -> normalClient.request(normalClient.get("/")));
-        assertThat("Unexpected exception.", e, instanceOf(HttpResponseException.class));
+        if (returnFailedResponses) {
+            HttpResponse response = normalClient.request(normalClient.get("/"));
+            assertThat(response.status(), is(HttpResponseStatus.OK));
+        } else {
+            HttpResponseException e = assertThrows(HttpResponseException.class,
+                    () -> normalClient.request(normalClient.get("/")));
+            assertThat("Unexpected exception.", e, instanceOf(HttpResponseException.class));
+        }
         // The load balancer is allowed to be not ready one time, which is counted against total retry attempts but not
         // against actual requests being issued.
         assertThat("Unexpected calls to select.", lbSelectInvoked.get(), allOf(greaterThanOrEqualTo(maxTotalRetries),