Add httpclient retry handler with delay retry

2 files changed, 302 insertions, 0 deletions

diff --git a/http-utils/src/main/java/ai/vespa/util/http/retry/DelayedHttpRequestRetryHandler.java b/http-utils/src/main/java/ai/vespa/util/http/retry/DelayedHttpRequestRetryHandler.java
new file mode 100644
index 00000000000..5922040a49a
--- /dev/null
+++ b/http-utils/src/main/java/ai/vespa/util/http/retry/DelayedHttpRequestRetryHandler.java
@@ -0,0 +1,166 @@
+// 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.client.HttpRequestRetryHandler;
+import org.apache.http.client.protocol.HttpClientContext;
+import org.apache.http.protocol.HttpContext;
+
+import java.io.IOException;
+import java.time.Duration;
+import java.util.List;
+import java.util.function.BiPredicate;
+import java.util.function.Predicate;
+import java.util.logging.Logger;
+
+/**
+ * A {@link HttpRequestRetryHandler} that supports delayed retries.
+ *
+ * @author bjorncs
+ */
+public class DelayedHttpRequestRetryHandler implements HttpRequestRetryHandler {
+
+    private static final Logger log = Logger.getLogger(HttpRequestRetryHandler.class.getName());
+
+    @FunctionalInterface
+    public interface RetryConsumer {
+        void onRetry(IOException exception, Duration delay, int executionCount, HttpClientContext context);
+    }
+
+    @FunctionalInterface
+    public interface RetryFailedConsumer {
+        void onRetryFailed(IOException exception, int executionCount, HttpClientContext context);
+    }
+
+    @FunctionalInterface
+    public interface RetryPredicate extends BiPredicate<IOException, HttpClientContext> {}
+
+    private final DelaySupplier delaySupplier;
+    private final int maxRetries;
+    private final RetryPredicate predicate;
+    private final RetryConsumer retryConsumer;
+    private final RetryFailedConsumer retryFailedConsumer;
+    private final Sleeper sleeper;
+
+    private DelayedHttpRequestRetryHandler(
+            DelaySupplier delaySupplier,
+            int maxRetries,
+            RetryPredicate predicate,
+            RetryConsumer retryConsumer,
+            RetryFailedConsumer retryFailedConsumer,
+            Sleeper sleeper) {
+        this.delaySupplier = delaySupplier;
+        this.maxRetries = maxRetries;
+        this.predicate = predicate;
+        this.retryConsumer = retryConsumer;
+        this.retryFailedConsumer = retryFailedConsumer;
+        this.sleeper = sleeper;
+    }
+
+    @Override
+    public boolean retryRequest(IOException exception, int executionCount, HttpContext ctx) {
+        log.fine(() -> String.format("retryRequest(exception='%s', executionCount='%d', ctx='%s'",
+                                     exception.getClass().getName(), executionCount, ctx));
+        HttpClientContext clientCtx = HttpClientContext.adapt(ctx);
+        if (!predicate.test(exception, 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(exception, executionCount, clientCtx);
+            return false;
+        }
+        Duration delay = delaySupplier.getDelay(executionCount);
+        log.fine(() -> String.format("Retrying after %s for '%s'", delay, ctx));
+        retryConsumer.onRetry(exception, delay, executionCount, clientCtx);
+        sleeper.sleep(delay);
+        return true;
+    }
+
+    public static class Builder {
+
+        private final DelaySupplier delaySupplier;
+        private final int maxRetries;
+        private RetryPredicate predicate = (ioException, ctx) -> true;
+        private RetryConsumer retryConsumer = (exception, delay, count, ctx) -> {};
+        private RetryFailedConsumer retryFailedConsumer = (exception, count, ctx) -> {};
+        private Sleeper sleeper = new DefaultSleeper();
+
+        private Builder(DelaySupplier delaySupplier, int maxRetries) {
+            this.delaySupplier = delaySupplier;
+            this.maxRetries = maxRetries;
+        }
+
+        public static Builder withFixedDelay(Duration delay, int maxRetries) {
+            return new Builder(executionCount -> delay, maxRetries);
+        }
+
+        public static Builder withExponentialBackoff(Duration startDelay, Duration maxDelay, int maxRetries) {
+            return new Builder(
+                    executionCount -> {
+                        Duration nextDelay = startDelay;
+                        for (int i = 1; i < executionCount; ++i) {
+                            nextDelay = nextDelay.multipliedBy(2);
+                        }
+                        return maxDelay.compareTo(nextDelay) > 0 ? nextDelay : maxDelay;
+                    },
+                    maxRetries);
+        }
+
+        public Builder retryForExceptions(List<Class<? extends IOException>> exceptionTypes) {
+            this.predicate = (ioException, ctx) -> exceptionTypes.stream().anyMatch(type -> type.isInstance(ioException));
+            return this;
+        }
+
+        public Builder retryForExceptions(Predicate<IOException> predicate) {
+            this.predicate = (ioException, ctx) -> predicate.test(ioException);
+            return this;
+        }
+
+        public Builder retryFor(RetryPredicate predicate) {
+            this.predicate = predicate;
+            return this;
+        }
+
+        public Builder onRetry(RetryConsumer consumer) {
+            this.retryConsumer = consumer;
+            return this;
+        }
+
+        public Builder onRetryFailed(RetryFailedConsumer consumer) {
+            this.retryFailedConsumer = consumer;
+            return this;
+        }
+
+        // For unit testing
+        Builder withSleeper(Sleeper sleeper) {
+            this.sleeper = sleeper;
+            return this;
+        }
+
+        public DelayedHttpRequestRetryHandler build() {
+            return new DelayedHttpRequestRetryHandler(delaySupplier, maxRetries, predicate, retryConsumer, retryFailedConsumer, sleeper);
+        }
+
+        private static class DefaultSleeper implements Sleeper {
+            @Override
+            public void sleep(Duration duration) {
+                try {
+                    Thread.sleep(duration.toMillis());
+                } catch (InterruptedException e) {
+                    throw new RuntimeException(e);
+                }
+            }
+        }
+    }
+
+    // For unit testing
+    interface Sleeper {
+        void sleep(Duration duration);
+    }
+
+    @FunctionalInterface
+    private interface DelaySupplier {
+        Duration getDelay(int executionCount);
+    }
+}
diff --git a/http-utils/src/test/java/ai/vespa/util/http/retry/DelayedHttpRequestRetryHandlerTest.java b/http-utils/src/test/java/ai/vespa/util/http/retry/DelayedHttpRequestRetryHandlerTest.java
new file mode 100644
index 00000000000..51a05f6b2a7
--- /dev/null
+++ b/http-utils/src/test/java/ai/vespa/util/http/retry/DelayedHttpRequestRetryHandlerTest.java
@@ -0,0 +1,136 @@
+// 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 ai.vespa.util.http.retry.DelayedHttpRequestRetryHandler.RetryConsumer;
+import ai.vespa.util.http.retry.DelayedHttpRequestRetryHandler.RetryFailedConsumer;
+import ai.vespa.util.http.retry.DelayedHttpRequestRetryHandler.Sleeper;
+import com.yahoo.vespa.jdk8compat.List;
+import org.apache.http.client.protocol.HttpClientContext;
+import org.junit.Test;
+
+import javax.net.ssl.SSLException;
+import java.io.IOException;
+import java.net.ConnectException;
+import java.time.Duration;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+/**
+ * @author bjorncs
+ */
+public class DelayedHttpRequestRetryHandlerTest {
+
+    @Test
+    public void retry_consumers_are_invoked() {
+        RetryConsumer retryConsumer = mock(RetryConsumer.class);
+        RetryFailedConsumer retryFailedConsumer = mock(RetryFailedConsumer.class);
+
+        Duration delay = Duration.ofSeconds(10);
+        int maxRetries = 5;
+
+        DelayedHttpRequestRetryHandler handler = DelayedHttpRequestRetryHandler.Builder
+                .withFixedDelay(delay, maxRetries)
+                .withSleeper(mock(Sleeper.class))
+                .onRetry(retryConsumer)
+                .onRetryFailed(retryFailedConsumer)
+                .build();
+
+        IOException exception = new IOException();
+        HttpClientContext ctx = new HttpClientContext();
+        int lastExecutionCount = maxRetries + 1;
+        for (int i = 1; i <= lastExecutionCount; i++) {
+            handler.retryRequest(exception, i, ctx);
+        }
+
+        verify(retryFailedConsumer).onRetryFailed(exception, lastExecutionCount, ctx);
+        for (int i = 1; i < lastExecutionCount; i++) {
+            verify(retryConsumer).onRetry(exception, delay, i, ctx);
+        }
+    }
+
+    @Test
+    public void retry_with_fixed_delay_sleeps_for_expected_duration() {
+        Sleeper sleeper = mock(Sleeper.class);
+
+        Duration delay = Duration.ofSeconds(2);
+        int maxRetries = 2;
+
+        DelayedHttpRequestRetryHandler handler = DelayedHttpRequestRetryHandler.Builder
+                .withFixedDelay(delay, maxRetries)
+                .withSleeper(sleeper)
+                .build();
+
+        IOException exception = new IOException();
+        HttpClientContext ctx = new HttpClientContext();
+        int lastExecutionCount = maxRetries + 1;
+        for (int i = 1; i <= lastExecutionCount; i++) {
+            handler.retryRequest(exception, i, ctx);
+        }
+
+        verify(sleeper, times(2)).sleep(delay);
+    }
+
+    @Test
+    public void retry_with_fixed_backoff_sleeps_for_expected_durations() {
+        Sleeper sleeper = mock(Sleeper.class);
+
+        Duration startDelay = Duration.ofMillis(500);
+        Duration maxDelay = Duration.ofSeconds(5);
+        int maxRetries = 10;
+
+        DelayedHttpRequestRetryHandler handler = DelayedHttpRequestRetryHandler.Builder
+                .withExponentialBackoff(startDelay, maxDelay, maxRetries)
+                .withSleeper(sleeper)
+                .build();
+
+        IOException exception = new IOException();
+        HttpClientContext ctx = new HttpClientContext();
+        int lastExecutionCount = maxRetries + 1;
+        for (int i = 1; i <= lastExecutionCount; i++) {
+            handler.retryRequest(exception, i, ctx);
+        }
+
+        verify(sleeper).sleep(startDelay);
+        verify(sleeper).sleep(Duration.ofSeconds(1));
+        verify(sleeper).sleep(Duration.ofSeconds(2));
+        verify(sleeper).sleep(Duration.ofSeconds(4));
+        verify(sleeper, times(6)).sleep(Duration.ofSeconds(5));
+    }
+
+    @Test
+    public void retries_for_listed_exceptions_until_max_retries_exceeded() {
+        int maxRetries = 2;
+
+        DelayedHttpRequestRetryHandler handler = DelayedHttpRequestRetryHandler.Builder
+                .withFixedDelay(Duration.ofSeconds(2), maxRetries)
+                .retryForExceptions(List.of(SSLException.class, ConnectException.class))
+                .withSleeper(mock(Sleeper.class))
+                .build();
+
+        SSLException sslException = new SSLException("ssl error");
+        HttpClientContext ctx = new HttpClientContext();
+        int lastExecutionCount = maxRetries + 1;
+        for (int i = 1; i < lastExecutionCount; i++) {
+            assertTrue(handler.retryRequest(sslException, i, ctx));
+        }
+        assertFalse(handler.retryRequest(sslException, lastExecutionCount, ctx));
+    }
+
+    @Test
+    public void does_not_retry_for_non_listed_exception() {
+        DelayedHttpRequestRetryHandler handler = DelayedHttpRequestRetryHandler.Builder
+                .withFixedDelay(Duration.ofSeconds(2), 2)
+                .retryForExceptions(List.of(SSLException.class, ConnectException.class))
+                .withSleeper(mock(Sleeper.class))
+                .build();
+
+        IOException ioException = new IOException();
+        HttpClientContext ctx = new HttpClientContext();
+        assertFalse(handler.retryRequest(ioException, 1, ctx));
+    }
+
+}
\ No newline at end of file