Merge pull request #482 from prometheus/beorn7/http

Use gzip from a pool and stream while encoding metrics

Author: beorn7

prometheus/http.go

@@ -15,9 +15,7 @@ package prometheus
 
 import (
 	"bufio"
-	"bytes"
 	"compress/gzip"
-	"fmt"
 	"io"
 	"net"
 	"net/http"
@@ -41,19 +39,10 @@ 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)
-}
-
-func giveBuf(buf *bytes.Buffer) {
-	buf.Reset()
-	bufPool.Put(buf)
+var gzipPool = sync.Pool{
+	New: func() interface{} {
+		return gzip.NewWriter(nil)
+	},
 }
 
 // Handler returns an HTTP handler for the DefaultGatherer. It is
@@ -71,58 +60,40 @@ func Handler() http.Handler {
 // Deprecated: Use promhttp.HandlerFor(DefaultGatherer, promhttp.HandlerOpts{})
 // instead. See there for further documentation.
 func UninstrumentedHandler() http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+	return http.HandlerFunc(func(rsp http.ResponseWriter, req *http.Request) {
 		mfs, err := DefaultGatherer.Gather()
 		if err != nil {
-			http.Error(w, "An error has occurred during metrics collection:\n\n"+err.Error(), http.StatusInternalServerError)
+			httpError(rsp, err)
 			return
 		}
 
 		contentType := expfmt.Negotiate(req.Header)
-		buf := getBuf()
-		defer giveBuf(buf)
-		writer, encoding := decorateWriter(req, buf)
-		enc := expfmt.NewEncoder(writer, contentType)
-		var lastErr error
+		header := rsp.Header()
+		header.Set(contentTypeHeader, string(contentType))
+
+		w := io.Writer(rsp)
+		if gzipAccepted(req.Header) {
+			header.Set(contentEncodingHeader, "gzip")
+			gz := gzipPool.Get().(*gzip.Writer)
+			defer gzipPool.Put(gz)
+
+			gz.Reset(w)
+			defer gz.Close()
+
+			w = gz
+		}
+
+		enc := expfmt.NewEncoder(w, contentType)
+
 		for _, mf := range mfs {
 			if err := enc.Encode(mf); err != nil {
-				lastErr = err
-				http.Error(w, "An error has occurred during metrics encoding:\n\n"+err.Error(), http.StatusInternalServerError)
+				httpError(rsp, err)
 				return
 			}
 		}
-		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)
-		}
-		w.Write(buf.Bytes())
 	})
 }
 
-// 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) (io.Writer, string) {
-	header := request.Header.Get(acceptEncodingHeader)
-	parts := strings.Split(header, ",")
-	for _, part := range parts {
-		part = strings.TrimSpace(part)
-		if part == "gzip" || strings.HasPrefix(part, "gzip;") {
-			return gzip.NewWriter(writer), "gzip"
-		}
-	}
-	return writer, ""
-}
-
 var instLabels = []string{"method", "code"}
 
 type nower interface {
@@ -503,3 +474,31 @@ func sanitizeCode(s int) string {
 		return strconv.Itoa(s)
 	}
 }
+
+// gzipAccepted returns whether the client will accept gzip-encoded content.
+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 true
+		}
+	}
+	return false
+}
+
+// httpError removes any content-encoding header and then calls http.Error with
+// the provided error and http.StatusInternalServerErrer. Error contents is
+// supposed to be uncompressed plain text. However, same as with a plain
+// http.Error, any header settings will be void if the header has already been
+// sent. The error message will still be written to the writer, but it will
+// probably be of limited use.
+func httpError(rsp http.ResponseWriter, err error) {
+	rsp.Header().Del(contentEncodingHeader)
+	http.Error(
+		rsp,
+		"An error has occurred while serving metrics:\n\n"+err.Error(),
+		http.StatusInternalServerError,
+	)
+}

prometheus/promhttp/http.go

@@ -32,7 +32,6 @@
 package promhttp
 
 import (
-	"bytes"
 	"compress/gzip"
 	"fmt"
 	"io"
@@ -53,19 +52,10 @@ 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)
-}
-
-func giveBuf(buf *bytes.Buffer) {
-	buf.Reset()
-	bufPool.Put(buf)
+var gzipPool = sync.Pool{
+	New: func() interface{} {
+		return gzip.NewWriter(nil)
+	},
 }
 
 // Handler returns an http.Handler for the prometheus.DefaultGatherer, using
@@ -100,19 +90,18 @@ func HandlerFor(reg prometheus.Gatherer, opts HandlerOpts) http.Handler {
 		inFlightSem = make(chan struct{}, opts.MaxRequestsInFlight)
 	}
 
-	h := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+	h := http.HandlerFunc(func(rsp http.ResponseWriter, req *http.Request) {
 		if inFlightSem != nil {
 			select {
 			case inFlightSem <- struct{}{}: // All good, carry on.
 				defer func() { <-inFlightSem }()
 			default:
-				http.Error(w, fmt.Sprintf(
+				http.Error(rsp, fmt.Sprintf(
 					"Limit of concurrent requests reached (%d), try again later.", opts.MaxRequestsInFlight,
 				), http.StatusServiceUnavailable)
 				return
 			}
 		}
-
 		mfs, err := reg.Gather()
 		if err != nil {
 			if opts.ErrorLog != nil {
@@ -123,55 +112,56 @@ func HandlerFor(reg prometheus.Gatherer, opts HandlerOpts) http.Handler {
 				panic(err)
 			case ContinueOnError:
 				if len(mfs) == 0 {
-					http.Error(w, "No metrics gathered, last error:\n\n"+err.Error(), http.StatusInternalServerError)
+					// Still report the error if no metrics have been gathered.
+					httpError(rsp, err)
 					return
 				}
 			case HTTPErrorOnError:
-				http.Error(w, "An error has occurred during metrics gathering:\n\n"+err.Error(), http.StatusInternalServerError)
+				httpError(rsp, err)
 				return
 			}
 		}
 
 		contentType := expfmt.Negotiate(req.Header)
-		buf := getBuf()
-		defer giveBuf(buf)
-		writer, encoding := decorateWriter(req, buf, opts.DisableCompression)
-		enc := expfmt.NewEncoder(writer, contentType)
+		header := rsp.Header()
+		header.Set(contentTypeHeader, string(contentType))
+
+		w := io.Writer(rsp)
+		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()
+
+			w = gz
+		}
+
+		enc := expfmt.NewEncoder(w, contentType)
+
 		var lastErr error
 		for _, mf := range mfs {
 			if err := enc.Encode(mf); err != nil {
 				lastErr = err
 				if opts.ErrorLog != nil {
-					opts.ErrorLog.Println("error encoding metric family:", err)
+					opts.ErrorLog.Println("error encoding and sending metric family:", err)
 				}
 				switch opts.ErrorHandling {
 				case PanicOnError:
 					panic(err)
 				case ContinueOnError:
 					// Handled later.
 				case HTTPErrorOnError:
-					http.Error(w, "An error has occurred during metrics encoding:\n\n"+err.Error(), http.StatusInternalServerError)
+					httpError(rsp, err)
 					return
 				}
 			}
 		}
-		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 lastErr != nil {
+			httpError(rsp, lastErr)
 		}
-		// TODO(beorn7): Consider streaming serving of metrics.
 	})
 
 	if opts.Timeout <= 0 {
@@ -292,20 +282,30 @@ 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, ""
-	}
-	header := request.Header.Get(acceptEncodingHeader)
-	parts := strings.Split(header, ",")
+// gzipAccepted returns whether the client will accept gzip-encoded content.
+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
+}
+
+// httpError removes any content-encoding header and then calls http.Error with
+// the provided error and http.StatusInternalServerErrer. Error contents is
+// supposed to be uncompressed plain text. However, same as with a plain
+// http.Error, any header settings will be void if the header has already been
+// sent. The error message will still be written to the writer, but it will
+// probably be of limited use.
+func httpError(rsp http.ResponseWriter, err error) {
+	rsp.Header().Del(contentEncodingHeader)
+	http.Error(
+		rsp,
+		"An error has occurred while serving metrics:\n\n"+err.Error(),
+		http.StatusInternalServerError,
+	)
 }

prometheus/promhttp/http_test.go

@@ -103,7 +103,7 @@ func TestHandlerErrorHandling(t *testing.T) {
 	})
 	wantMsg := `error gathering metrics: error collecting metric Desc{fqName: "invalid_metric", help: "not helpful", constLabels: {}, variableLabels: []}: collect error
 `
-	wantErrorBody := `An error has occurred during metrics gathering:
+	wantErrorBody := `An error has occurred while serving metrics:
 
 error collecting metric Desc{fqName: "invalid_metric", help: "not helpful", constLabels: {}, variableLabels: []}: collect error
 `

prometheus/registry_test.go

@@ -250,7 +250,7 @@ metric: <
 		},
 	}
 
-	expectedMetricFamilyInvalidLabelValueAsText := []byte(`An error has occurred during metrics gathering:
+	expectedMetricFamilyInvalidLabelValueAsText := []byte(`An error has occurred while serving metrics:
 
 collected metric "name" { label:<name:"constname" value:"\377" > label:<name:"labelname" value:"different_val" > counter:<value:42 > } has a label named "constname" whose value is not utf8: "\xff"
 `)
@@ -299,15 +299,15 @@ complex_bucket 1
 			},
 		},
 	}
-	bucketCollisionMsg := []byte(`An error has occurred during metrics gathering:
+	bucketCollisionMsg := []byte(`An error has occurred while serving metrics:
 
 collected metric named "complex_bucket" collides with previously collected histogram named "complex"
 `)
-	summaryCountCollisionMsg := []byte(`An error has occurred during metrics gathering:
+	summaryCountCollisionMsg := []byte(`An error has occurred while serving metrics:
 
 collected metric named "complex_count" collides with previously collected summary named "complex"
 `)
-	histogramCountCollisionMsg := []byte(`An error has occurred during metrics gathering:
+	histogramCountCollisionMsg := []byte(`An error has occurred while serving metrics:
 
 collected metric named "complex_count" collides with previously collected histogram named "complex"
 `)
@@ -333,7 +333,7 @@ collected metric named "complex_count" collides with previously collected histog
 			},
 		},
 	}
-	duplicateLabelMsg := []byte(`An error has occurred during metrics gathering:
+	duplicateLabelMsg := []byte(`An error has occurred while serving metrics:
 
 collected metric "broken_metric" { label:<name:"foo" value:"bar" > label:<name:"foo" value:"baz" > counter:<value:2.7 > } has two or more labels with the same name: foo
 `)