Close entities, and support retries with supplier<http-entity> which is non-repeatable

2 files changed, 14 insertions, 9 deletions

diff --git a/http-client/src/main/java/ai/vespa/hosted/client/AbstractHttpClient.java b/http-client/src/main/java/ai/vespa/hosted/client/AbstractHttpClient.java
index 68741d6d509..57cb14e3c66 100644
--- a/http-client/src/main/java/ai/vespa/hosted/client/AbstractHttpClient.java
+++ b/http-client/src/main/java/ai/vespa/hosted/client/AbstractHttpClient.java
@@ -79,9 +79,9 @@ public abstract class AbstractHttpClient implements HttpClient {
                                                                              .asURI())
                                                               .build();
             builder.headers.forEach((name, values) -> values.forEach(value -> request.setHeader(name, value)));
-            request.setEntity(builder.entity);
 
-            try {
+            try (HttpEntity entity = builder.entity.get()) {
+                request.setEntity(entity);
                 try {
                     return handler.apply(execute(request, contextWithTimeout(builder)), request);
                 }
@@ -90,16 +90,15 @@ public abstract class AbstractHttpClient implements HttpClient {
                     throw RetryException.wrap(e, request);
                 }
             }
+            catch (IOException e) {
+                throw new UncheckedIOException("failed closing request entity", e);
+            }
             catch (RetryException e) {
                 if (thrown == null)
                     thrown = e.getCause();
                 else
                     thrown.addSuppressed(e.getCause());
 
-                if (builder.entity != null && ! builder.entity.isRepeatable()) {
-                    log.log(WARNING, "Cannot retry " + request + " as entity is not repeatable");
-                    break;
-                }
                 log.log(FINE, e.getCause(), () -> request + " failed; will retry");
             }
         }
@@ -152,7 +151,7 @@ public abstract class AbstractHttpClient implements HttpClient {
         private HttpURL.Query query = Query.empty();
         private List<Supplier<Query>> dynamicQuery = new ArrayList<>();
         private Map<String, List<String>> headers = new LinkedHashMap<>();
-        private HttpEntity entity;
+        private Supplier<HttpEntity> entity;
         private RequestConfig config = HttpClient.defaultRequestConfig;
         private ResponseVerifier verifier = HttpClient.throwOnError;
         private ExceptionHandler catcher = HttpClient.retryAll;
@@ -178,7 +177,7 @@ public abstract class AbstractHttpClient implements HttpClient {
         }
 
         @Override
-        public RequestBuilder body(HttpEntity entity) {
+        public RequestBuilder body(Supplier<HttpEntity> entity) {
             this.entity = requireNonNull(entity);
             return this;
         }
diff --git a/http-client/src/main/java/ai/vespa/hosted/client/HttpClient.java b/http-client/src/main/java/ai/vespa/hosted/client/HttpClient.java
index b41b16c25be..16a419bf324 100644
--- a/http-client/src/main/java/ai/vespa/hosted/client/HttpClient.java
+++ b/http-client/src/main/java/ai/vespa/hosted/client/HttpClient.java
@@ -78,7 +78,13 @@ public interface HttpClient extends Closeable {
         RequestBuilder body(byte[] json);
 
         /** Sets the request body. */
-        RequestBuilder body(HttpEntity entity);
+        default RequestBuilder body(HttpEntity entity) {
+            if (entity.isRepeatable()) return body(() -> entity);
+            throw new IllegalArgumentException("entitiy must be repeatable, or a supplier must be used");
+        }
+
+        /** Sets the request body. */
+        RequestBuilder body(Supplier<HttpEntity> entity);
 
         /** Sets query parameters without a value, like {@code ?debug&recursive}. */
         default RequestBuilder emptyParameters(String... keys) {