No Content-Encoding when disabling compression for localhost (elastic#2238)

Author: felixbarny

apm-agent-core/src/main/java/co/elastic/apm/agent/report/AbstractIntakeApiHandler.java

@@ -51,6 +51,8 @@ public class AbstractIntakeApiHandler {
     protected HttpURLConnection connection;
     @Nullable
     protected OutputStream os;
+    @Nullable
+    private CountingOutputStream countingOs;
     protected int errorCount;
     protected volatile boolean shutDown;
     private volatile boolean healthy = true;
@@ -60,7 +62,7 @@ public AbstractIntakeApiHandler(ReporterConfiguration reporterConfiguration, Pay
         this.reporterConfiguration = reporterConfiguration;
         this.payloadSerializer = payloadSerializer;
         this.apmServerClient = apmServerClient;
-        this.deflater = new Deflater();
+        this.deflater = new Deflater(Deflater.BEST_SPEED);
     }
 
     /*
@@ -79,7 +81,10 @@ static long getBackoffTimeSeconds(long errorCount) {
     }
 
     protected boolean shouldEndRequest() {
-        final long written = deflater.getBytesWritten() + DslJsonSerializer.BUFFER_SIZE;
+        if (countingOs == null) {
+            return false;
+        }
+        final long written = countingOs.getCount() + payloadSerializer.getBufferSize();
         final boolean endRequest = written >= reporterConfiguration.getApiRequestSize();
         if (endRequest && logger.isDebugEnabled()) {
             logger.debug("Flushing, because request size limit exceeded {}/{}", written, reporterConfiguration.getApiRequestSize());
@@ -92,23 +97,26 @@ protected HttpURLConnection startRequest(String endpoint) throws Exception {
         payloadSerializer.blockUntilReady();
         final HttpURLConnection connection = apmServerClient.startRequest(endpoint);
         if (connection != null) {
-            if (isLocalhost(connection)) {
-                deflater.setLevel(Deflater.NO_COMPRESSION);
-            } else {
-                deflater.setLevel(Deflater.BEST_SPEED);
-            }
+            boolean useCompression = !isLocalhost(connection);
             try {
                 if (logger.isDebugEnabled()) {
                     logger.debug("Starting new request to {}", connection.getURL());
                 }
                 connection.setRequestMethod("POST");
                 connection.setDoOutput(true);
                 connection.setChunkedStreamingMode(DslJsonSerializer.BUFFER_SIZE);
-                connection.setRequestProperty("Content-Encoding", "deflate");
+                if (useCompression) {
+                    connection.setRequestProperty("Content-Encoding", "deflate");
+                }
                 connection.setRequestProperty("Content-Type", "application/x-ndjson");
                 connection.setUseCaches(false);
                 connection.connect();
-                os = new DeflaterOutputStream(connection.getOutputStream(), deflater, true);
+                countingOs = new CountingOutputStream(connection.getOutputStream());
+                if (useCompression) {
+                    os = new DeflaterOutputStream(countingOs, deflater, true);
+                } else {
+                    os = countingOs;
+                }
                 payloadSerializer.setOutputStream(os);
                 payloadSerializer.appendMetaDataNdJsonToStream();
                 payloadSerializer.flushToOutputStream();
@@ -181,6 +189,7 @@ public void endRequest() {
                 HttpUtils.consumeAndClose(connection);
                 connection = null;
                 os = null;
+                countingOs = null;
                 deflater.reset();
                 currentlyTransmitting = 0;
             }

apm-agent-core/src/main/java/co/elastic/apm/agent/report/CountingOutputStream.java

@@ -0,0 +1,41 @@
+package co.elastic.apm.agent.report;
+
+import javax.annotation.Nonnull;
+import java.io.IOException;
+import java.io.OutputStream;
+
+class CountingOutputStream extends OutputStream {
+
+    private final OutputStream out;
+    private long count;
+
+    CountingOutputStream(OutputStream out) {
+        this.out = out;
+    }
+
+    @Override
+    public void write(int b) throws IOException {
+        out.write(b);
+        count++;
+    }
+
+    @Override
+    public void write(@Nonnull byte[] b, int off, int len) throws IOException {
+        out.write(b, off, len);
+        count += len;
+    }
+
+    @Override
+    public void close() throws IOException {
+        out.close();
+    }
+
+    @Override
+    public void flush() throws IOException {
+        out.flush();
+    }
+
+    public long getCount() {
+        return count;
+    }
+}

apm-agent-core/src/test/java/co/elastic/apm/agent/report/ApmServerReporterIntegrationTest.java

@@ -44,12 +44,12 @@
 import org.junit.jupiter.api.Test;
 import org.stagemonitor.configuration.ConfigurationRegistry;
 
+import java.io.InputStream;
 import java.net.InetSocketAddress;
 import java.net.URL;
 import java.util.Collections;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
-import java.util.zip.InflaterInputStream;
 
 import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
 import static org.awaitility.Awaitility.await;
@@ -92,7 +92,7 @@ void setUp() throws Exception {
         handler = new BlockingHandler(exchange -> {
             if (statusCode < 300 && exchange.getRequestPath().equals("/intake/v2/events")) {
                 receivedIntakeApiCalls.incrementAndGet();
-                InflaterInputStream in = new InflaterInputStream(exchange.getInputStream());
+                InputStream in = exchange.getInputStream();
                 try (in) {
                     for (int n = 0; -1 != n; n = in.read()) {
                         if (n == '\n') {

apm-agent-core/src/test/java/co/elastic/apm/agent/report/IntakeV2ReportingEventHandlerTest.java

@@ -55,7 +55,6 @@
 import java.util.List;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
-import java.util.zip.InflaterInputStream;
 
 import static co.elastic.apm.agent.report.IntakeV2ReportingEventHandler.INTAKE_V2_URL;
 import static com.github.tomakehurst.wiremock.client.WireMock.ok;
@@ -262,7 +261,7 @@ private List<JsonNode> getNdJsonNodes() {
         return Stream.of(mockApmServer1, mockApmServer2)
             .flatMap(apmServer -> apmServer.findAll(postRequestedFor(urlEqualTo(INTAKE_V2_URL))).stream())
             .findFirst()
-            .map(request -> new BufferedReader(new InputStreamReader(new InflaterInputStream(new ByteArrayInputStream(request.getBody()))))
+            .map(request -> new BufferedReader(new InputStreamReader(new ByteArrayInputStream(request.getBody())))
                 .lines()
                 .map(IntakeV2ReportingEventHandlerTest::getReadTree)
                 .collect(Collectors.toList()))

apm-agent-plugins/apm-log-shipper-plugin/src/test/java/co/elastic/apm/agent/log/shipper/ApmServerLogShipperTest.java

@@ -51,7 +51,6 @@
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
-import java.util.zip.InflaterInputStream;
 
 import static com.github.tomakehurst.wiremock.client.WireMock.get;
 import static com.github.tomakehurst.wiremock.client.WireMock.ok;
@@ -143,7 +142,7 @@ void testSendLogsAfterServerUrlsSet() throws Exception {
     private List<String> getEvents() {
         return mockApmServer.findAll(postRequestedFor(urlEqualTo(ApmServerLogShipper.LOGS_ENDPOINT)))
             .stream()
-            .flatMap(request -> new BufferedReader(new InputStreamReader(new InflaterInputStream(new ByteArrayInputStream(request.getBody())))).lines())
+            .flatMap(request -> new BufferedReader(new InputStreamReader(new ByteArrayInputStream(request.getBody()))).lines())
             .collect(Collectors.toList());
     }
 }