diff options
Diffstat (limited to 'vespa-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() { |