Introduce response level retry strategy

2 files changed, 254 insertions, 0 deletions

diff --git a/http-utils/src/main/java/ai/vespa/util/http/retry/DelayedResponseLevelRetryHandler.java b/http-utils/src/main/java/ai/vespa/util/http/retry/DelayedResponseLevelRetryHandler.java
new file mode 100644
index 00000000000..041b501809c
--- /dev/null
+++ b/http-utils/src/main/java/ai/vespa/util/http/retry/DelayedResponseLevelRetryHandler.java
@@ -0,0 +1,125 @@
+// Copyright 2020 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package ai.vespa.util.http.retry;
+
+import org.apache.http.HttpResponse;
+import org.apache.http.annotation.Contract;
+import org.apache.http.annotation.ThreadingBehavior;
+import org.apache.http.client.ServiceUnavailableRetryStrategy;
+import org.apache.http.client.protocol.HttpClientContext;
+import org.apache.http.protocol.HttpContext;
+
+import java.time.Duration;
+import java.util.List;
+import java.util.function.Predicate;
+import java.util.logging.Logger;
+
+/**
+ * A {@link ServiceUnavailableRetryStrategy} that supports delayed retries on any response types.
+ *
+ * @author bjorncs
+ */
+@Contract(threading = ThreadingBehavior.IMMUTABLE)
+public class DelayedResponseLevelRetryHandler implements ServiceUnavailableRetryStrategy {
+
+    private static final Logger log = Logger.getLogger(DelayedResponseLevelRetryHandler.class.getName());
+
+    private final DelaySupplier delaySupplier;
+    private final int maxRetries;
+    private final RetryPredicate<HttpResponse> predicate;
+    private final RetryConsumer<HttpResponse> retryConsumer;
+    private final RetryFailedConsumer<HttpResponse> retryFailedConsumer;
+    private final ThreadLocal<Long> retryInterval = ThreadLocal.withInitial(() -> 0L);
+
+    private DelayedResponseLevelRetryHandler(
+            DelaySupplier delaySupplier,
+            int maxRetries,
+            RetryPredicate<HttpResponse> predicate,
+            RetryConsumer<HttpResponse> retryConsumer,
+            RetryFailedConsumer<HttpResponse> retryFailedConsumer) {
+
+        this.delaySupplier = delaySupplier;
+        this.maxRetries = maxRetries;
+        this.predicate = predicate;
+        this.retryConsumer = retryConsumer;
+        this.retryFailedConsumer = retryFailedConsumer;
+    }
+
+    @Override
+    public boolean retryRequest(HttpResponse response, int executionCount, HttpContext ctx) {
+        log.fine(() -> String.format("retryRequest(responseCode='%s', executionCount='%d', ctx='%s'",
+                                     response.getStatusLine().getStatusCode(), executionCount, ctx));
+        HttpClientContext clientCtx = HttpClientContext.adapt(ctx);
+        if (!predicate.test(response, clientCtx)) {
+            log.fine(() -> String.format("Not retrying for '%s'", ctx));
+            return false;
+        }
+        if (executionCount > maxRetries) {
+            log.fine(() -> String.format("Max retries exceeded for '%s'", ctx));
+            retryFailedConsumer.onRetryFailed(response, executionCount, clientCtx);
+            return false;
+        }
+        Duration delay = delaySupplier.getDelay(executionCount);
+        log.fine(() -> String.format("Retrying after %s for '%s'", delay, ctx));
+        retryInterval.set(delay.toMillis());
+        retryConsumer.onRetry(response, delay, executionCount, clientCtx);
+        return true;
+    }
+
+    @Override
+    public long getRetryInterval() {
+        // Calls to getRetryInterval are always guarded by a call to retryRequest (using the same thread).
+        // A thread local allows this retry handler to be thread safe and support dynamic retry intervals
+        return retryInterval.get();
+    }
+
+    public static class Builder {
+
+        private final DelaySupplier delaySupplier;
+        private final int maxRetries;
+        private RetryPredicate<HttpResponse> predicate = (response, ctx) -> true;
+        private RetryConsumer<HttpResponse> retryConsumer = (response, delay, count, ctx) -> {};
+        private RetryFailedConsumer<HttpResponse> retryFailedConsumer = (response, count, ctx) -> {};
+
+        private Builder(DelaySupplier delaySupplier, int maxRetries) {
+            this.delaySupplier = delaySupplier;
+            this.maxRetries = maxRetries;
+        }
+
+        public static Builder withFixedDelay(Duration delay, int maxRetries) {
+            return new Builder(new DelaySupplier.Fixed(delay), maxRetries);
+        }
+
+        public static Builder withExponentialBackoff(Duration startDelay, Duration maxDelay, int maxRetries) {
+            return new Builder(new DelaySupplier.Exponential(startDelay, maxDelay), maxRetries);
+        }
+
+        public Builder retryForStatusCodes(List<Integer> statusCodes) {
+            this.predicate = (response, ctx) -> statusCodes.contains(response.getStatusLine().getStatusCode());
+            return this;
+        }
+
+        public Builder retryForResponses(Predicate<HttpResponse> predicate) {
+            this.predicate = (response, ctx) -> predicate.test(response);
+            return this;
+        }
+
+        public Builder retryFor(RetryPredicate<HttpResponse> predicate) {
+            this.predicate = predicate;
+            return this;
+        }
+
+        public Builder onRetry(RetryConsumer<HttpResponse> consumer) {
+            this.retryConsumer = consumer;
+            return this;
+        }
+
+        public Builder onRetryFailed(RetryFailedConsumer<HttpResponse> consumer) {
+            this.retryFailedConsumer = consumer;
+            return this;
+        }
+
+        public DelayedResponseLevelRetryHandler build() {
+            return new DelayedResponseLevelRetryHandler(delaySupplier, maxRetries, predicate, retryConsumer, retryFailedConsumer);
+        }
+    }
+}
diff --git a/http-utils/src/test/java/ai/vespa/util/http/retry/DelayedResponseLevelRetryHandlerTest.java b/http-utils/src/test/java/ai/vespa/util/http/retry/DelayedResponseLevelRetryHandlerTest.java
new file mode 100644
index 00000000000..7be1e143078
--- /dev/null
+++ b/http-utils/src/test/java/ai/vespa/util/http/retry/DelayedResponseLevelRetryHandlerTest.java
@@ -0,0 +1,129 @@
+package ai.vespa.util.http.retry;// Copyright 2020 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+import org.apache.http.HttpResponse;
+import org.apache.http.HttpStatus;
+import org.apache.http.HttpVersion;
+import org.apache.http.client.protocol.HttpClientContext;
+import org.apache.http.message.BasicHttpResponse;
+import org.apache.http.message.BasicStatusLine;
+import org.junit.Test;
+
+import java.time.Duration;
+import java.util.List;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+
+/**
+ * @author bjorncs
+ */
+public class DelayedResponseLevelRetryHandlerTest {
+
+    @Test
+    @SuppressWarnings("unchecked")
+    public void retry_consumers_are_invoked() {
+        RetryConsumer<HttpResponse> retryConsumer = mock(RetryConsumer.class);
+        RetryFailedConsumer<HttpResponse> retryFailedConsumer = mock(RetryFailedConsumer.class);
+
+        Duration delay = Duration.ofSeconds(10);
+        int maxRetries = 5;
+
+        DelayedResponseLevelRetryHandler handler = DelayedResponseLevelRetryHandler.Builder
+                .withFixedDelay(delay, maxRetries)
+                .onRetry(retryConsumer)
+                .onRetryFailed(retryFailedConsumer)
+                .build();
+
+        HttpResponse response = createResponse(HttpStatus.SC_SERVICE_UNAVAILABLE);
+        HttpClientContext ctx = new HttpClientContext();
+        int lastExecutionCount = maxRetries + 1;
+        for (int i = 1; i <= lastExecutionCount; i++) {
+            handler.retryRequest(response, i, ctx);
+        }
+
+        verify(retryFailedConsumer).onRetryFailed(response, lastExecutionCount, ctx);
+        for (int i = 1; i < lastExecutionCount; i++) {
+            verify(retryConsumer).onRetry(response, delay, i, ctx);
+        }
+    }
+
+    @Test
+    public void retry_with_fixed_delay_sleeps_for_expected_duration() {
+        Duration delay = Duration.ofSeconds(2);
+        int maxRetries = 2;
+
+        DelayedResponseLevelRetryHandler handler = DelayedResponseLevelRetryHandler.Builder
+                .withFixedDelay(delay, maxRetries)
+                .build();
+
+        HttpResponse response = createResponse(HttpStatus.SC_SERVICE_UNAVAILABLE);
+        HttpClientContext ctx = new HttpClientContext();
+        int lastExecutionCount = maxRetries + 1;
+        for (int i = 1; i <= lastExecutionCount; i++) {
+            handler.retryRequest(response, i, ctx);
+            assertThat(handler.getRetryInterval()).isEqualTo(delay.toMillis());
+        }
+    }
+
+    @Test
+    public void retry_with_fixed_backoff_sleeps_for_expected_durations() {
+        Duration startDelay = Duration.ofMillis(500);
+        Duration maxDelay = Duration.ofSeconds(5);
+        int maxRetries = 10;
+
+        DelayedResponseLevelRetryHandler handler = DelayedResponseLevelRetryHandler.Builder
+                .withExponentialBackoff(startDelay, maxDelay, maxRetries)
+                .build();
+
+        HttpResponse response = createResponse(HttpStatus.SC_SERVICE_UNAVAILABLE);
+        HttpClientContext ctx = new HttpClientContext();
+        int lastExecutionCount = maxRetries + 1;
+        List<Duration> expectedIntervals =
+                com.yahoo.vespa.jdk8compat.List.of(
+                        startDelay, Duration.ofSeconds(1), Duration.ofSeconds(2), Duration.ofSeconds(4),
+                        Duration.ofSeconds(5), Duration.ofSeconds(5), Duration.ofSeconds(5), Duration.ofSeconds(5),
+                        Duration.ofSeconds(5), Duration.ofSeconds(5), Duration.ofSeconds(5));
+        for (int i = 1; i <= lastExecutionCount; i++) {
+            handler.retryRequest(response, i, ctx);
+            assertThat(handler.getRetryInterval()).isEqualTo(expectedIntervals.get(i-1).toMillis());
+        }
+    }
+
+    @Test
+    public void retries_for_listed_exceptions_until_max_retries_exceeded() {
+        int maxRetries = 2;
+
+        DelayedResponseLevelRetryHandler handler = DelayedResponseLevelRetryHandler.Builder
+                .withFixedDelay(Duration.ofSeconds(2), maxRetries)
+                .retryForStatusCodes(com.yahoo.vespa.jdk8compat.List.of(HttpStatus.SC_SERVICE_UNAVAILABLE, HttpStatus.SC_BAD_GATEWAY))
+                .build();
+
+        HttpResponse response = createResponse(HttpStatus.SC_SERVICE_UNAVAILABLE);
+        HttpClientContext ctx = new HttpClientContext();
+        int lastExecutionCount = maxRetries + 1;
+        for (int i = 1; i < lastExecutionCount; i++) {
+            assertTrue(handler.retryRequest(response, i, ctx));
+        }
+        assertFalse(handler.retryRequest(response, lastExecutionCount, ctx));
+    }
+
+    @Test
+    public void does_not_retry_for_non_listed_exception() {
+        DelayedResponseLevelRetryHandler handler = DelayedResponseLevelRetryHandler.Builder
+                .withFixedDelay(Duration.ofSeconds(2), 2)
+                .retryForStatusCodes(com.yahoo.vespa.jdk8compat.List.of(HttpStatus.SC_SERVICE_UNAVAILABLE, HttpStatus.SC_BAD_GATEWAY))
+                .build();
+
+        HttpResponse response = createResponse(HttpStatus.SC_OK);
+        HttpClientContext ctx = new HttpClientContext();
+        assertFalse(handler.retryRequest(response, 1, ctx));
+    }
+
+    private static HttpResponse createResponse(int statusCode) {
+        return new BasicHttpResponse(new BasicStatusLine(HttpVersion.HTTP_1_1, statusCode, "reason phrase"));
+    }
+
+}
\ No newline at end of file