From 1aa8f071b9375a19f9bb5b4def69d08fdefe5f78 Mon Sep 17 00:00:00 2001
From: Julius Volz <julius.volz@gmail.com>
Date: Fri, 28 Jun 2013 10:19:16 +0200
Subject: [PATCH] Add content compression support to API HTTP responses.

---
 web/compression.go | 92 ++++++++++++++++++++++++++++++++++++++++++++++
 web/web.go         | 15 +++++---
 2 files changed, 101 insertions(+), 6 deletions(-)
 create mode 100644 web/compression.go

diff --git a/web/compression.go b/web/compression.go
new file mode 100644
index 000000000..92f226a9a
--- /dev/null
+++ b/web/compression.go
@@ -0,0 +1,92 @@
+// Copyright 2013 Prometheus Team
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package web
+
+import (
+	"compress/gzip"
+	"compress/zlib"
+	"io"
+	"net/http"
+	"strings"
+)
+
+const (
+	acceptEncodingHeader  = "Accept-Encoding"
+	contentEncodingHeader = "Content-Encoding"
+	gzipEncoding          = "gzip"
+	deflateEncoding       = "deflate"
+)
+
+// Wrapper around http.Handler which adds suitable response compression based
+// on the client's Accept-Encoding headers.
+type compressedResponseWriter struct {
+	http.ResponseWriter
+	writer io.Writer
+}
+
+// Writes HTTP response content data.
+func (c *compressedResponseWriter) Write(p []byte) (int, error) {
+	return c.writer.Write(p)
+}
+
+// Closes the compressedResponseWriter and ensures to flush all data before.
+func (c *compressedResponseWriter) Close() {
+	if zlibWriter, ok := c.writer.(*zlib.Writer); ok {
+		zlibWriter.Flush()
+	}
+	if gzipWriter, ok := c.writer.(*gzip.Writer); ok {
+		gzipWriter.Flush()
+	}
+	if closer, ok := c.writer.(io.Closer); ok {
+		defer closer.Close()
+	}
+}
+
+// Constructs a new compressedResponseWriter based on client request headers.
+func newCompressedResponseWriter(writer http.ResponseWriter, req *http.Request) *compressedResponseWriter {
+	encodings := strings.Split(req.Header.Get(acceptEncodingHeader), ",")
+	for _, encoding := range encodings {
+		switch strings.TrimSpace(encoding) {
+		case gzipEncoding:
+			writer.Header().Set(contentEncodingHeader, gzipEncoding)
+			return &compressedResponseWriter{
+				ResponseWriter: writer,
+				writer:         gzip.NewWriter(writer),
+			}
+		case deflateEncoding:
+			writer.Header().Set(contentEncodingHeader, deflateEncoding)
+			return &compressedResponseWriter{
+				ResponseWriter: writer,
+				writer:         zlib.NewWriter(writer),
+			}
+		}
+	}
+	return &compressedResponseWriter{
+		ResponseWriter: writer,
+		writer:         writer,
+	}
+}
+
+// Wrapper around http.Handler which adds suitable response compression based
+// on the client's Accept-Encoding headers.
+type compressionHandler struct {
+	handler http.Handler
+}
+
+// Adds compression to the original http.Handler's ServeHTTP() method.
+func (c compressionHandler) ServeHTTP(writer http.ResponseWriter, req *http.Request) {
+	compWriter := newCompressedResponseWriter(writer, req)
+	c.handler.ServeHTTP(compWriter, req)
+	compWriter.Close()
+}
diff --git a/web/web.go b/web/web.go
index 754d3a85d..f09830489 100644
--- a/web/web.go
+++ b/web/web.go
@@ -14,17 +14,20 @@
 package web
 
 import (
-	"code.google.com/p/gorest"
 	"flag"
 	"fmt"
-	"github.com/prometheus/client_golang/prometheus"
-	"github.com/prometheus/client_golang/prometheus/exp"
-	"github.com/prometheus/prometheus/web/api"
-	"github.com/prometheus/prometheus/web/blob"
 	"html/template"
 	"log"
 	"net/http"
 	"net/http/pprof"
+
+	"code.google.com/p/gorest"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/exp"
+
+	"github.com/prometheus/prometheus/web/api"
+	"github.com/prometheus/prometheus/web/blob"
 )
 
 // Commandline flags.
@@ -60,7 +63,7 @@ func (w WebService) ServeForever() error {
 	exp.Handle("/alerts", w.AlertsHandler)
 	exp.HandleFunc("/graph", graphHandler)
 
-	exp.Handle("/api/", gorest.Handle())
+	exp.Handle("/api/", compressionHandler{handler: gorest.Handle()})
 	exp.Handle("/metrics.json", prometheus.DefaultHandler)
 	if *useLocalAssets {
 		exp.Handle("/static/", http.StripPrefix("/static/", http.FileServer(http.Dir("web/static"))))