Add --gzip-requests to compress request bodies in feed client

3 files changed, 49 insertions, 3 deletions

diff --git a/vespa-feed-client/pom.xml b/vespa-feed-client/pom.xml
index 1cc2f2adee1..01b9b00b8a0 100644
--- a/vespa-feed-client/pom.xml
+++ b/vespa-feed-client/pom.xml
@@ -46,6 +46,11 @@
       <artifactId>junit-jupiter</artifactId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>com.github.tomakehurst</groupId>
+      <artifactId>wiremock-jre8-standalone</artifactId>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 
   <build>
diff --git a/vespa-feed-client/src/main/java/ai/vespa/feed/client/impl/ApacheCluster.java b/vespa-feed-client/src/main/java/ai/vespa/feed/client/impl/ApacheCluster.java
index 3ffbaf136f2..e2678503c68 100644
--- a/vespa-feed-client/src/main/java/ai/vespa/feed/client/impl/ApacheCluster.java
+++ b/vespa-feed-client/src/main/java/ai/vespa/feed/client/impl/ApacheCluster.java
@@ -2,23 +2,39 @@
 package ai.vespa.feed.client.impl;
 
 import ai.vespa.feed.client.HttpResponse;
+import org.apache.hc.client5.http.async.methods.SimpleBody;
 import org.apache.hc.client5.http.async.methods.SimpleHttpRequest;
 import org.apache.hc.client5.http.async.methods.SimpleHttpResponse;
 import org.apache.hc.client5.http.config.RequestConfig;
+import org.apache.hc.client5.http.entity.GzipCompressingEntity;
 import org.apache.hc.client5.http.impl.async.CloseableHttpAsyncClient;
 import org.apache.hc.client5.http.impl.async.HttpAsyncClients;
 import org.apache.hc.client5.http.ssl.ClientTlsStrategyBuilder;
 import org.apache.hc.core5.concurrent.FutureCallback;
 import org.apache.hc.core5.http.ContentType;
+import org.apache.hc.core5.http.Header;
+import org.apache.hc.core5.http.HttpHeaders;
 import org.apache.hc.core5.http.HttpHost;
+import org.apache.hc.core5.http.io.entity.BasicHttpEntity;
+import org.apache.hc.core5.http.io.entity.ByteArrayEntity;
+import org.apache.hc.core5.http.io.entity.StringEntity;
 import org.apache.hc.core5.http.message.BasicHeader;
+import org.apache.hc.core5.http.message.BasicHttpRequest;
+import org.apache.hc.core5.http.nio.AsyncEntityProducer;
+import org.apache.hc.core5.http.nio.AsyncRequestProducer;
+import org.apache.hc.core5.http.nio.entity.AsyncEntityProducers;
+import org.apache.hc.core5.http.nio.entity.BasicAsyncEntityProducer;
+import org.apache.hc.core5.http.nio.support.AsyncRequestBuilder;
+import org.apache.hc.core5.http.nio.support.BasicRequestProducer;
 import org.apache.hc.core5.http2.config.H2Config;
 import org.apache.hc.core5.net.URIAuthority;
 import org.apache.hc.core5.reactor.IOReactorConfig;
 import org.apache.hc.core5.util.Timeout;
 
 import javax.net.ssl.SSLContext;
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.io.UncheckedIOException;
 import java.net.URI;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -29,6 +45,7 @@ import java.util.concurrent.Future;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.zip.GZIPOutputStream;
 
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.apache.hc.core5.http.ssl.TlsCiphers.excludeH2Blacklisted;
@@ -40,9 +57,11 @@ import static org.apache.hc.core5.http.ssl.TlsCiphers.excludeWeak;
 class ApacheCluster implements Cluster {
 
     private final List<Endpoint> endpoints = new ArrayList<>();
-    private final List<BasicHeader> defaultHeaders = Arrays.asList(new BasicHeader("User-Agent", String.format("vespa-feed-client/%s", Vespa.VERSION)),
+    private final List<BasicHeader> defaultHeaders = Arrays.asList(new BasicHeader(HttpHeaders.USER_AGENT, String.format("vespa-feed-client/%s", Vespa.VERSION)),
                                                                    new BasicHeader("Vespa-Client-Version", Vespa.VERSION));
+    private final Header gzipEncodingHeader = new BasicHeader(HttpHeaders.CONTENT_ENCODING, "gzip");
     private final RequestConfig requestConfig;
+    private final boolean gzip;
     private int someNumber = 0;
 
     private final ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor(t -> new Thread(t, "request-timeout-thread"));
@@ -52,6 +71,7 @@ class ApacheCluster implements Cluster {
             for (URI endpoint : builder.endpoints)
                 endpoints.add(new Endpoint(createHttpClient(builder), endpoint));
         this.requestConfig = createRequestConfig(builder);
+        this.gzip = builder.gzipRequests;
     }
 
     @Override
@@ -77,8 +97,14 @@ class ApacheCluster implements Cluster {
             request.setConfig(requestConfig);
             defaultHeaders.forEach(request::setHeader);
             wrapped.headers().forEach((name, value) -> request.setHeader(name, value.get()));
-            if (wrapped.body() != null)
-                request.setBody(wrapped.body(), ContentType.APPLICATION_JSON);
+            if (wrapped.body() != null) {
+                byte[] body = wrapped.body();
+                if (gzip) {
+                    request.setHeader(gzipEncodingHeader);
+                    body = gzipped(body);
+                }
+                request.setBody(body, ContentType.APPLICATION_JSON);
+            }
 
             Future<?> future = endpoint.client.execute(request,
                                                        new FutureCallback<SimpleHttpResponse>() {
@@ -96,6 +122,14 @@ class ApacheCluster implements Cluster {
         vessel.whenComplete((__, ___) -> endpoint.inflight.decrementAndGet());
     }
 
+    private byte[] gzipped(byte[] content) throws IOException{
+        ByteArrayOutputStream buffer = new ByteArrayOutputStream(1 << 10);
+        try (GZIPOutputStream zip = new GZIPOutputStream(buffer)) {
+            zip.write(content);
+        }
+        return buffer.toByteArray();
+    }
+
     @Override
     public void close() {
         Throwable thrown = null;
diff --git a/vespa-feed-client/src/main/java/ai/vespa/feed/client/impl/FeedClientBuilderImpl.java b/vespa-feed-client/src/main/java/ai/vespa/feed/client/impl/FeedClientBuilderImpl.java
index 1960991792f..13556b60a28 100644
--- a/vespa-feed-client/src/main/java/ai/vespa/feed/client/impl/FeedClientBuilderImpl.java
+++ b/vespa-feed-client/src/main/java/ai/vespa/feed/client/impl/FeedClientBuilderImpl.java
@@ -50,6 +50,7 @@ public class FeedClientBuilderImpl implements FeedClientBuilder {
     boolean benchmark = true;
     boolean dryrun = false;
     boolean speedTest = false;
+    boolean gzipRequests = false;
     URI proxy;
 
 
@@ -200,6 +201,12 @@ public class FeedClientBuilderImpl implements FeedClientBuilder {
         return this;
     }
 
+    @Override
+    public FeedClientBuilderImpl setGzipRequests(boolean gzip) {
+        this.gzipRequests = gzip;
+        return this;
+    }
+
     /** Constructs instance of {@link ai.vespa.feed.client.FeedClient} from builder configuration */
     @Override
     public FeedClient build() {