Merge branch 'use-gzip-pool' of github.com:glefloch/client_golang into beorn7/http

Author: beorn7

prometheus/promhttp/http.go

@@ -53,19 +53,16 @@ const (
 	acceptEncodingHeader  = "Accept-Encoding"
 )
 
-var bufPool sync.Pool
-
-func getBuf() *bytes.Buffer {
-	buf := bufPool.Get()
-	if buf == nil {
-		return &bytes.Buffer{}
-	}
-	return buf.(*bytes.Buffer)
+var gzipPool = sync.Pool{
+	New: func() interface{} {
+		return gzip.NewWriter(nil)
+	},
 }
 
-func giveBuf(buf *bytes.Buffer) {
-	buf.Reset()
-	bufPool.Put(buf)
+var bufPool = sync.Pool{
+	New: func() interface{} {
+		return &bytes.Buffer{}
+	},
 }
 
 // Handler returns an http.Handler for the prometheus.DefaultGatherer, using
@@ -112,7 +109,6 @@ func HandlerFor(reg prometheus.Gatherer, opts HandlerOpts) http.Handler {
 				return
 			}
 		}
-
 		mfs, err := reg.Gather()
 		if err != nil {
 			if opts.ErrorLog != nil {
@@ -133,10 +129,12 @@ func HandlerFor(reg prometheus.Gatherer, opts HandlerOpts) http.Handler {
 		}
 
 		contentType := expfmt.Negotiate(req.Header)
-		buf := getBuf()
-		defer giveBuf(buf)
-		writer, encoding := decorateWriter(req, buf, opts.DisableCompression)
-		enc := expfmt.NewEncoder(writer, contentType)
+		buf := bufPool.Get().(*bytes.Buffer)
+		buf.Reset()
+		enc := expfmt.NewEncoder(buf, contentType)
+
+		defer bufPool.Put(buf)
+
 		var lastErr error
 		for _, mf := range mfs {
 			if err := enc.Encode(mf); err != nil {
@@ -155,22 +153,28 @@ func HandlerFor(reg prometheus.Gatherer, opts HandlerOpts) http.Handler {
 				}
 			}
 		}
-		if closer, ok := writer.(io.Closer); ok {
-			closer.Close()
-		}
+
 		if lastErr != nil && buf.Len() == 0 {
 			http.Error(w, "No metrics encoded, last error:\n\n"+lastErr.Error(), http.StatusInternalServerError)
 			return
 		}
 		header := w.Header()
 		header.Set(contentTypeHeader, string(contentType))
 		header.Set(contentLengthHeader, fmt.Sprint(buf.Len()))
-		if encoding != "" {
-			header.Set(contentEncodingHeader, encoding)
-		}
-		if _, err := w.Write(buf.Bytes()); err != nil && opts.ErrorLog != nil {
-			opts.ErrorLog.Println("error while sending encoded metrics:", err)
+
+		if !opts.DisableCompression && gzipAccepted(req.Header) {
+			header.Set(contentEncodingHeader, "gzip")
+			gz := gzipPool.Get().(*gzip.Writer)
+			defer gzipPool.Put(gz)
+
+			gz.Reset(w)
+			defer gz.Close()
+
+			zipWriter := gzipResponseWriter{gz, w}
+			writeResult(zipWriter, buf, opts)
+			return
 		}
+		writeResult(w, buf, opts)
 		// TODO(beorn7): Consider streaming serving of metrics.
 	})
 
@@ -292,20 +296,36 @@ type HandlerOpts struct {
 	Timeout time.Duration
 }
 
-// decorateWriter wraps a writer to handle gzip compression if requested.  It
-// returns the decorated writer and the appropriate "Content-Encoding" header
-// (which is empty if no compression is enabled).
-func decorateWriter(request *http.Request, writer io.Writer, compressionDisabled bool) (io.Writer, string) {
-	if compressionDisabled {
-		return writer, ""
+// gzipResponseWriter in charge of wrapping io.Writer and http.ReponseWriter
+// together, allowing to get a single struct which implements both interface.
+type gzipResponseWriter struct {
+	io.Writer
+	http.ResponseWriter
+}
+
+func (w gzipResponseWriter) Write(b []byte) (int, error) {
+	return w.Writer.Write(b)
+}
+
+// writeResult to buf using http.ResponseWriter.
+// If ErrorLog is enabled, err is logged in.
+func writeResult(w http.ResponseWriter, buf *bytes.Buffer, opts HandlerOpts) {
+	if _, err := w.Write(buf.Bytes()); err != nil && opts.ErrorLog != nil {
+		opts.ErrorLog.Println("error while sending encoded metrics:", err)
 	}
-	header := request.Header.Get(acceptEncodingHeader)
-	parts := strings.Split(header, ",")
+}
+
+// gzipHandler return a http.HandlerFunc in charge of compressing the content
+// of the given http.HandlerFunc
+func gzipAccepted(header http.Header) bool {
+
+	a := header.Get(acceptEncodingHeader)
+	parts := strings.Split(a, ",")
 	for _, part := range parts {
 		part = strings.TrimSpace(part)
 		if part == "gzip" || strings.HasPrefix(part, "gzip;") {
-			return gzip.NewWriter(writer), "gzip"
+			return true
 		}
 	}
-	return writer, ""
+	return false
 }