HTTP Server: Add support for Gzip Compression + Reuse ByteArray (prometheus#291)

erebe · brian-brazil · commit f1950039de5b · 2017-09-13T12:26:49.000+01:00
diff --git a/simpleclient_httpserver/src/main/java/io/prometheus/client/exporter/HTTPServer.java b/simpleclient_httpserver/src/main/java/io/prometheus/client/exporter/HTTPServer.java
@@ -6,11 +6,14 @@
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.OutputStreamWriter;
+import java.net.HttpURLConnection;
 import java.net.InetSocketAddress;
 import java.net.URLDecoder;
+import java.util.List;
 import java.util.Set;
 import java.util.HashSet;
 import java.util.concurrent.Executors;
+import java.util.zip.GZIPOutputStream;
 
 import com.sun.net.httpserver.HttpHandler;
 import com.sun.net.httpserver.HttpServer;
@@ -27,8 +30,16 @@
  * </pre>
  * */
 public class HTTPServer {
+    private static class LocalByteArray extends ThreadLocal<ByteArrayOutputStream> {
+        protected ByteArrayOutputStream initialValue()
+        {
+            return new ByteArrayOutputStream(1 << 20);
+        }
+    }
+
     static class HTTPMetricHandler implements HttpHandler {
         private CollectorRegistry registry;
+        private final LocalByteArray response = new LocalByteArray();
 
         HTTPMetricHandler(CollectorRegistry registry) {
           this.registry = registry;
@@ -38,7 +49,8 @@ static class HTTPMetricHandler implements HttpHandler {
         public void handle(HttpExchange t) throws IOException {
             String query = t.getRequestURI().getRawQuery();
 
-            ByteArrayOutputStream response = new ByteArrayOutputStream(1 << 20);
+            ByteArrayOutputStream response = this.response.get();
+            response.reset();
             OutputStreamWriter osw = new OutputStreamWriter(response);
             TextFormat.write004(osw,
                     registry.filteredMetricFamilySamples(parseQuery(query)));
@@ -51,13 +63,36 @@ public void handle(HttpExchange t) throws IOException {
                     TextFormat.CONTENT_TYPE_004);
             t.getResponseHeaders().set("Content-Length",
                     String.valueOf(response.size()));
-            t.sendResponseHeaders(200, response.size());
-            response.writeTo(t.getResponseBody());
+            if (shouldUseCompression(t)) {
+                t.getResponseHeaders().set("Content-Encoding", "gzip");
+                t.sendResponseHeaders(HttpURLConnection.HTTP_OK, 0);
+                final GZIPOutputStream os = new GZIPOutputStream(t.getResponseBody());
+                response.writeTo(os);
+                os.finish();
+            } else {
+                t.sendResponseHeaders(HttpURLConnection.HTTP_OK, response.size());
+                response.writeTo(t.getResponseBody());
+            }
             t.close();
         }
 
     }
 
+    protected static boolean shouldUseCompression(HttpExchange exchange) {
+        List<String> encodingHeaders = exchange.getRequestHeaders().get("Accept-Encoding");
+        if (encodingHeaders == null) return false;
+
+        for (String encodingHeader : encodingHeaders) {
+            String[] encodings = encodingHeader.split(",");
+            for (String encoding : encodings) {
+                if (encoding.trim().toLowerCase().equals("gzip")) {
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
+
     protected static Set<String> parseQuery(String query) throws IOException {
         Set<String> names = new HashSet<String>();
         if (query != null) {
diff --git a/simpleclient_httpserver/src/test/java/io/prometheus/client/exporter/TestHTTPServer.java b/simpleclient_httpserver/src/test/java/io/prometheus/client/exporter/TestHTTPServer.java
@@ -7,6 +7,8 @@
 import java.net.URL;
 import java.net.URLConnection;
 import java.util.Scanner;
+import java.util.zip.GZIPInputStream;
+
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -40,6 +42,18 @@ String request(String suffix) throws IOException {
     return s.hasNext() ? s.next() : "";
   }
 
+  String requestWithCompression(String suffix) throws IOException {
+    String url = "http://localhost:" + s.server.getAddress().getPort() + "/metrics" + suffix;
+    URLConnection connection = new URL(url).openConnection();
+    connection.setDoOutput(true);
+    connection.setDoInput(true);
+    connection.setRequestProperty("Accept-Encoding", "gzip, deflate");
+    connection.connect();
+    GZIPInputStream gzs = new GZIPInputStream(connection.getInputStream());
+    Scanner s = new Scanner(gzs).useDelimiter("\\A");
+    return s.hasNext() ? s.next() : "";
+  }
+
   @Test
   public void testSimpleRequest() throws IOException {
     String response = request("");
@@ -80,4 +94,11 @@ public void testDecoding() throws IOException {
     assertThat(response).doesNotContain("c 0.0");
   }
 
+  @Test
+  public void testGzipCompression() throws IOException {
+    String response = requestWithCompression("");
+    assertThat(response).contains("a 0.0");
+    assertThat(response).contains("b 0.0");
+    assertThat(response).contains("c 0.0");
+  }
 }
