Separate handling of responses and exceptions

2 files changed, 40 insertions, 39 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 0ee9e320259..31707579aa6 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
@@ -52,7 +52,7 @@ public abstract class AbstractConfigServerClient implements ConfigServerClient {
     protected abstract ClassicHttpResponse execute(ClassicHttpRequest request, HttpClientContext context) throws IOException;
 
     /** Executes the given request with response/error handling and retries. */
-    private <T> T execute(RequestBuilder builder, BiFunction<ClassicHttpResponse, IOException, T> handler) {
+    private <T> T execute(RequestBuilder builder, Function<ClassicHttpResponse, T> handler, Function<IOException, T> catcher) {
         HttpClientContext context = HttpClientContext.create();
         context.setRequestConfig(builder.config);
 
@@ -62,10 +62,10 @@ public abstract class AbstractConfigServerClient implements ConfigServerClient {
             request.setEntity(builder.entity);
             try {
                 try {
-                    return handler.apply(execute(request, context), null);
+                    return handler.apply(execute(request, context));
                 }
                 catch (IOException e) {
-                    return handler.apply(null, e);
+                    return catcher.apply(e);
                 }
             }
             catch (RetryException e) {
@@ -162,7 +162,7 @@ public abstract class AbstractConfigServerClient implements ConfigServerClient {
 
         @Override
         public RequestBuilder timeout(Duration timeout) {
-            return config(RequestConfig.copy(defaultRequestConfig)
+            return config(RequestConfig.copy(config)
                                        .setResponseTimeout(timeout.toMillis(), TimeUnit.MILLISECONDS)
                                        .build());
         }
@@ -175,8 +175,8 @@ public abstract class AbstractConfigServerClient implements ConfigServerClient {
         }
 
         @Override
-        public <T> T handle(BiFunction<ClassicHttpResponse, IOException, T> handler) throws UncheckedIOException {
-            return execute(this, requireNonNull(handler));
+        public <T> T handle(Function<ClassicHttpResponse, T> handler, Function<IOException, T> catcher) throws UncheckedIOException {
+            return execute(this, requireNonNull(handler), requireNonNull(catcher));
         }
 
         @Override
@@ -210,37 +210,38 @@ public abstract class AbstractConfigServerClient implements ConfigServerClient {
 
         /** Returns the mapped body, if successful, retrying any IOException. The caller must close the body stream. */
         private <T> T mapIfSuccess(Function<InputStream, T> mapper) {
-            return handle((response, ioException) -> {
-                if (response != null) {
-                    try {
-                        InputStream body = response.getEntity() != null ? response.getEntity().getContent()
-                                                                        : InputStream.nullInputStream();
-                        if (response.getCode() >= HttpStatus.SC_REDIRECTION)
-                            throw readException(body.readAllBytes());
-
-                        return mapper.apply(new ForwardingInputStream(body) {
-                            @Override
-                            public void close() throws IOException {
-                                super.close();
-                                response.close();
+            return handle(response -> {
+                            try {
+                                InputStream body = response.getEntity() != null ? response.getEntity().getContent()
+                                                                                : InputStream.nullInputStream();
+                                if (response.getCode() >= HttpStatus.SC_REDIRECTION)
+                                    throw readException(body.readAllBytes());
+
+                                return mapper.apply(new ForwardingInputStream(body) {
+                                    @Override
+                                    public void close() throws IOException {
+                                        super.close();
+                                        response.close();
+                                    }
+                                });
                             }
-                        });
-                    }
-                    catch (IOException | RuntimeException | Error e) {
-                        try {
-                            response.close();
-                        }
-                        catch (IOException f) {
-                            e.addSuppressed(f);
-                        }
-                        if (e instanceof IOException)
-                            ioException = (IOException) e;
-                        else
-                            sneakyThrow(e);
-                    }
-                }
-                throw new RetryException(ioException);
-            });
+                            catch (IOException | RuntimeException | Error e) {
+                                try {
+                                    response.close();
+                                }
+                                catch (IOException f) {
+                                    e.addSuppressed(f);
+                                }
+                                if (e instanceof IOException)
+                                    throw new RetryException((IOException) e);
+                                else
+                                    sneakyThrow(e);
+                                throw new AssertionError("Should not happen");
+                            }
+                          },
+                          ioException -> {
+                              throw new RetryException(ioException);
+                          });
         }
 
     }
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 234dbe9ee06..482e8d2eb40 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
@@ -47,19 +47,19 @@ public interface ConfigServerClient extends AutoCloseable {
         /** Overrides the default socket read timeout of the request. {@code Duration.ZERO} gives infinite timeout. */
         RequestBuilder timeout(Duration timeout);
 
-        /** Overrides the default socket read timeout of the request. {@code null} allows infinite timeout. */
+        /** Overrides the default request config of the request. */
         RequestBuilder config(RequestConfig config);
 
         /**
          * Sets custom retry/failure logic for this.
          * <p>
-         * Exactly one of the arguments (response, exception) are non-null.
+         * Exactly one of the callbacks will be invoked, with a non-null argument.
          * Return a value to have that returned to the caller;
          * throw a {@link RetryException} to have the request retried; or
          * throw any other unchecked exception to have this propagate out to the caller.
          * The caller must close the provided response, if any.
          */
-        <T> T handle(BiFunction<ClassicHttpResponse, IOException, T> handler) throws UncheckedIOException;
+        <T> T handle(Function<ClassicHttpResponse, T> handler, Function<IOException, T> catcher) throws UncheckedIOException;
 
         /** Sets the response body mapper for this, for successful requests. */
         <T> T read(Function<byte[], T> mapper) throws UncheckedIOException, ConfigServerException;