Add context to caught IOExceptions

2 files changed, 29 insertions, 16 deletions

diff --git a/configserver-client/src/main/java/ai/vespa/hosted/client/AbstractConfigServerClient.java b/configserver-client/src/main/java/ai/vespa/hosted/client/AbstractConfigServerClient.java
index 81970d47bc7..608e9fe18b4 100644
--- a/configserver-client/src/main/java/ai/vespa/hosted/client/AbstractConfigServerClient.java
+++ b/configserver-client/src/main/java/ai/vespa/hosted/client/AbstractConfigServerClient.java
@@ -23,7 +23,6 @@ import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 import java.util.function.BiFunction;
-import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.logging.Logger;
 
@@ -44,7 +43,7 @@ public abstract class AbstractConfigServerClient implements ConfigServerClient {
     /** Executes the given request with response/error handling and retries. */
     private <T> T execute(RequestBuilder builder,
                           BiFunction<ClassicHttpResponse, ClassicHttpRequest, T> handler,
-                          Consumer<IOException> catcher) {
+                          ExceptionHandler catcher) {
         HttpClientContext context = HttpClientContext.create();
         context.setRequestConfig(builder.config);
 
@@ -57,8 +56,8 @@ public abstract class AbstractConfigServerClient implements ConfigServerClient {
                     return handler.apply(execute(request, context), request);
                 }
                 catch (IOException e) {
-                    catcher.accept(e);
-                    throw new UncheckedIOException(e); // Throw unchecked if catcher doesn't throw.
+                    catcher.handle(e, request);
+                    throw RetryException.wrap(e, request);
                 }
             }
             catch (RetryException e) {
@@ -118,7 +117,7 @@ public abstract class AbstractConfigServerClient implements ConfigServerClient {
         private HttpEntity entity;
         private RequestConfig config = ConfigServerClient.defaultRequestConfig;
         private ResponseVerifier verifier = ConfigServerClient.throwOnError;
-        private Consumer<IOException> catcher = ConfigServerClient.retryAll;
+        private ExceptionHandler catcher = ConfigServerClient.retryAll;
 
         private RequestBuilder(HostStrategy hosts, Method method) {
             if ( ! hosts.iterator().hasNext())
@@ -181,7 +180,7 @@ public abstract class AbstractConfigServerClient implements ConfigServerClient {
         }
 
         @Override
-        public RequestBuilder catching(Consumer<IOException> catcher) {
+        public RequestBuilder catching(ExceptionHandler catcher) {
             this.catcher = requireNonNull(catcher);
             return this;
         }
@@ -244,8 +243,8 @@ public abstract class AbstractConfigServerClient implements ConfigServerClient {
                                        e.addSuppressed(f);
                                    }
                                    if (e instanceof IOException) {
-                                       catcher.accept((IOException) e);
-                                       throw new UncheckedIOException((IOException) e);
+                                       catcher.handle((IOException) e, request);
+                                       throw RetryException.wrap((IOException) e, request);
                                    }
                                    else
                                        sneakyThrow(e); // e is a runtime exception or an error, so this is fine.
diff --git a/configserver-client/src/main/java/ai/vespa/hosted/client/ConfigServerClient.java b/configserver-client/src/main/java/ai/vespa/hosted/client/ConfigServerClient.java
index c92acd7cd0b..d5a4153fb8d 100644
--- a/configserver-client/src/main/java/ai/vespa/hosted/client/ConfigServerClient.java
+++ b/configserver-client/src/main/java/ai/vespa/hosted/client/ConfigServerClient.java
@@ -20,7 +20,6 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
-import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.stream.IntStream;
 
@@ -39,10 +38,8 @@ public interface ConfigServerClient extends Closeable {
                                                       .setRedirectsEnabled(false)
                                                       .build();
 
-    /** Wraps with a {@link RetryException} and rethrows. */
-    Consumer<IOException> retryAll = (e) -> {
-        throw new RetryException(e);
-    };
+    /** Does nothing, letting the client wrap with a {@link RetryException} and re-throw. */
+    ExceptionHandler retryAll = (exception, request) -> { };
 
     /** Throws a a {@link RetryException} if {@code statusCode == 503}, or a {@link ResponseException} unless {@code 200 <= statusCode < 300}. */
     ResponseVerifier throwOnError = new DefaultResponseVerifier() { };
@@ -101,15 +98,15 @@ public interface ConfigServerClient extends Closeable {
          * Sets the catch clause for {@link IOException}s during execution of this.
          * The default is to wrap the IOException in a {@link RetryException} and rethrow this;
          * this makes the client retry the request, as long as there are remaining entries in the {@link HostStrategy}.
-         * If the catcher returns normally, the {@link IOException} is unchecked and thrown instead.
+         * If the catcher returns normally, the exception is wrapped and retried, as per the default.
          */
-         RequestBuilder catching(Consumer<IOException> catcher);
+        RequestBuilder catching(ExceptionHandler catcher);
 
         /**
          * Sets the (error) response handler for this request. The default is {@link #throwOnError}.
          * When the handler returns normally, the response is treated as a success, and passed on to a response mapper.
          */
-         RequestBuilder throwing(ResponseVerifier handler);
+        RequestBuilder throwing(ResponseVerifier handler);
 
         /** Reads the response as a {@link String}, or throws if unsuccessful. */
         String read();
@@ -203,6 +200,19 @@ public interface ConfigServerClient extends Closeable {
     }
 
 
+    @FunctionalInterface
+    interface ExceptionHandler {
+
+        /**
+         * Called with any IO exception that might occur when attempting to send the request.
+         * To retry, wrap the exception with a {@link RetryException} and re-throw, or exit normally.
+         * Any other thrown exception will propagate out to the caller.
+         */
+        void handle(IOException exception, ClassicHttpRequest request);
+
+    }
+
+
     /** What host(s) to try for a request, in what order. A host may be specified multiple times, for retries.  */
     @FunctionalInterface
     interface HostStrategy extends Iterable<URI> {
@@ -240,6 +250,10 @@ public interface ConfigServerClient extends Closeable {
             super(requireNonNull(cause));
         }
 
+        static RetryException wrap(IOException exception, ClassicHttpRequest request) {
+            return new RetryException(new UncheckedIOException(request + " failed (" + exception.getClass().getSimpleName() + ")", exception));
+        }
+
     }