1 files changed, 36 insertions, 47 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 2e5e445d63a..d467402f79f 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
@@ -9,12 +9,11 @@ import org.apache.hc.core5.http.ClassicHttpResponse;
 import org.apache.hc.core5.http.ContentType;
 import org.apache.hc.core5.http.HttpEntity;
 import org.apache.hc.core5.http.Method;
-import org.apache.hc.core5.http.ParseException;
-import org.apache.hc.core5.http.io.entity.EntityUtils;
 import org.apache.hc.core5.http.io.entity.HttpEntities;
 import org.apache.hc.core5.net.URIBuilder;
 
 import java.io.IOException;
+import java.io.InputStream;
 import java.io.UncheckedIOException;
 import java.net.URI;
 import java.net.URISyntaxException;
@@ -22,6 +21,7 @@ import java.time.Duration;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
+import java.util.function.BiConsumer;
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -104,7 +104,7 @@ public abstract class AbstractConfigServerClient implements ConfigServerClient {
     }
 
     @Override
-    public ConfigServerClient.RequestBuilder send(HostStrategy hosts, Method method) {
+    public RequestBuilder send(HostStrategy hosts, Method method) {
         return new RequestBuilder(hosts, method);
     }
 
@@ -114,11 +114,10 @@ public abstract class AbstractConfigServerClient implements ConfigServerClient {
         private final Method method;
         private final HostStrategy hosts;
         private final URIBuilder uriBuilder = new URIBuilder();
-        private final List<String> pathSegments = new ArrayList<>();
         private HttpEntity entity;
         private RequestConfig config = ConfigServerClient.defaultRequestConfig;
-        private ResponseVerifier verifier = ConfigServerClient.throwOnError;
-        private Consumer<IOException> catcher = ConfigServerClient.retryAll;
+        private BiConsumer<ClassicHttpResponse, ClassicHttpRequest> handler = ConfigServerClient::throwOnError;
+        private Consumer<IOException> catcher = ConfigServerClient::retryAll;
 
         private RequestBuilder(HostStrategy hosts, Method method) {
             if ( ! hosts.iterator().hasNext())
@@ -130,7 +129,7 @@ public abstract class AbstractConfigServerClient implements ConfigServerClient {
 
         @Override
         public RequestBuilder at(List<String> pathSegments) {
-            this.pathSegments.addAll(pathSegments);
+            uriBuilder.setPathSegments(requireNonNull(pathSegments));
             return this;
         }
 
@@ -146,23 +145,12 @@ public abstract class AbstractConfigServerClient implements ConfigServerClient {
         }
 
         @Override
-        public ConfigServerClient.RequestBuilder emptyParameters(List<String> keys) {
-            for (String key : keys)
-                uriBuilder.setParameter(key, null);
-
-            return this;
-        }
-
-        @Override
         public RequestBuilder parameters(List<String> pairs) {
             if (pairs.size() % 2 != 0)
                 throw new IllegalArgumentException("Must supply parameter key/values in pairs");
 
-            for (int i = 0; i < pairs.size(); ) {
-                String key = pairs.get(i++), value = pairs.get(i++);
-                if (value != null)
-                    uriBuilder.setParameter(key, value);
-            }
+            for (int i = 0; i < pairs.size(); )
+                uriBuilder.setParameter(pairs.get(i++), pairs.get(i++));
 
             return this;
         }
@@ -187,54 +175,56 @@ public abstract class AbstractConfigServerClient implements ConfigServerClient {
         }
 
         @Override
-        public RequestBuilder throwing(ResponseVerifier verifier) {
-            this.verifier = requireNonNull(verifier);
+        public RequestBuilder handling(BiConsumer<ClassicHttpResponse, ClassicHttpRequest> handler) {
+            this.handler = requireNonNull(handler);
             return this;
         }
 
         @Override
-        public String read() {
-            return handle((response, __) -> {
-                try (response) {
-                    return response.getEntity() == null ? "" : EntityUtils.toString(response.getEntity());
-                }
-                catch (ParseException e) {
-                    throw new IllegalStateException(e); // This isn't actually thrown by apache >_<
-                }
-            });
-        }
-
-        @Override
         public <T> T read(Function<byte[], T> mapper) {
-            return handle((response, __) -> {
-                try (response) {
-                    return mapper.apply(response.getEntity() == null ? new byte[0] : EntityUtils.toByteArray(response.getEntity()));
+            return mapIfSuccess(input -> {
+                try (input) {
+                    return mapper.apply(input.readAllBytes());
+                }
+                catch (IOException e) {
+                    throw new RetryException(e);
                 }
             });
         }
 
         @Override
         public void discard() throws UncheckedIOException, ResponseException {
-            handle((response, __) -> {
-                try (response) {
+            mapIfSuccess(input -> {
+                try (input) {
                     return null;
                 }
+                catch (IOException e) {
+                    throw new RetryException(e);
+                }
             });
         }
 
         @Override
-        public HttpInputStream stream() throws UncheckedIOException, ResponseException {
-            return handle((response, __) -> new HttpInputStream(response));
+        public InputStream stream() throws UncheckedIOException, ResponseException {
+            return mapIfSuccess(input -> input);
         }
 
-        @Override
-        public <T> T handle(ResponseHandler<T> handler) {
-            uriBuilder.setPathSegments(pathSegments);
+        /** 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 execute(this,
                            (response, request) -> {
                                try {
-                                   verifier.verify(response, request); // This throws on unacceptable responses.
-                                   return handler.handle(response, request);
+                                   handler.accept(response, request); // This throws on unacceptable responses.
+
+                                   InputStream body = response.getEntity() != null ? response.getEntity().getContent()
+                                                                                   : InputStream.nullInputStream();
+                                   return mapper.apply(new ForwardingInputStream(body) {
+                                       @Override
+                                       public void close() throws IOException {
+                                           super.close();
+                                           response.close();
+                                       }
+                                   });
                                }
                                catch (IOException | RuntimeException | Error e) {
                                    try {
@@ -257,7 +247,6 @@ public abstract class AbstractConfigServerClient implements ConfigServerClient {
 
     }
 
-
     @SuppressWarnings("unchecked")
     private static <T extends Throwable> void sneakyThrow(Throwable t) throws T {
         throw (T) t;