Test with gzip compression

Author: mrueg

prometheus/promhttp/http.go

@@ -190,7 +190,11 @@ func HandlerForTransactional(reg prometheus.TransactionalGatherer, opts HandlerO
 		}
 		rsp.Header().Set(contentTypeHeader, string(contentType))
 
-		w, encodingHeader, err := NegotiateEncodingWriter(req, rsp, opts.DisableCompression, compressions)
+		w, encodingHeader, closeWriter, err := NegotiateEncodingWriter(req, rsp, opts.DisableCompression, compressions)
+
+		if closeWriter != nil {
+			defer closeWriter()
+		}
 		if err != nil {
 			if opts.ErrorLog != nil {
 				opts.ErrorLog.Println("error getting writer", err)
@@ -432,11 +436,11 @@ func httpError(rsp http.ResponseWriter, err error) {
 // selects the right compression based on an allow-list of supported
 // compressions. It returns a writer implementing the compression and an the
 // correct value that the caller can set in the response header.
-func NegotiateEncodingWriter(r *http.Request, rw io.Writer, disableCompression bool, compressions []string) (_ io.Writer, encodingHeaderValue string, _ error) {
+func NegotiateEncodingWriter(r *http.Request, rw io.Writer, disableCompression bool, compressions []string) (_ io.Writer, encodingHeaderValue string, closeWriter func() error, _ error) {
 	w := rw
 
 	if disableCompression {
-		return w, string(Identity), nil
+		return w, string(Identity), nil, nil
 	}
 
 	// TODO(mrueg): Replace internal/github.com/gddo once https://github.com/golang/go/issues/19307 is implemented.
@@ -447,26 +451,26 @@ func NegotiateEncodingWriter(r *http.Request, rw io.Writer, disableCompression b
 		// TODO(mrueg): Replace klauspost/compress with stdlib implementation once https://github.com/golang/go/issues/62513 is implemented.
 		z, err := zstd.NewWriter(rw, zstd.WithEncoderLevel(zstd.SpeedFastest))
 		if err != nil {
-			return nil, "", err
+			return nil, "", nil, err
 		}
 
 		z.Reset(w)
-		defer z.Close()
-
 		w = z
+
+		return w, compression, z.Close, nil
 	case "gzip":
 		gz := gzipPool.Get().(*gzip.Writer)
 		defer gzipPool.Put(gz)
 
 		gz.Reset(w)
-		defer gz.Close()
 
 		w = gz
+		return w, compression, gz.Close, nil
 	case "identity":
 		// This means the content is not compressed.
+		return w, compression, nil, nil
 	default:
 		// The content encoding was not implemented yet.
-		return nil, "", fmt.Errorf("content compression format not recognized: %s. Valid formats are: %s", compression, defaultCompressionFormats)
+		return nil, "", nil, fmt.Errorf("content compression format not recognized: %s. Valid formats are: %s", compression, defaultCompressionFormats)
 	}
-	return w, compression, nil
 }

prometheus/promhttp/http_test.go

@@ -15,8 +15,10 @@ package promhttp
 
 import (
 	"bytes"
+	"compress/gzip"
 	"errors"
 	"fmt"
+	"io"
 	"log"
 	"net/http"
 	"net/http/httptest"
@@ -31,6 +33,11 @@ import (
 
 type errorCollector struct{}
 
+const (
+	acceptHeader    = "Accept"
+	acceptTextPlain = "text/plain"
+)
+
 func (e errorCollector) Describe(ch chan<- *prometheus.Desc) {
 	ch <- prometheus.NewDesc("invalid_metric", "not helpful", nil, nil)
 }
@@ -71,6 +78,16 @@ func (g *mockTransactionGatherer) Gather() (_ []*dto.MetricFamily, done func(),
 	return mfs, func() { g.doneInvoked++ }, err
 }
 
+func readGzippedBody(r io.Reader) (string, error) {
+	reader, err := gzip.NewReader(r)
+	if err != nil {
+		return "", err
+	}
+	defer reader.Close()
+	got, err := io.ReadAll(reader)
+	return string(got), err
+}
+
 func TestHandlerErrorHandling(t *testing.T) {
 	// Create a registry that collects a MetricFamily with two elements,
 	// another with one, and reports an error. Further down, we'll use the
@@ -223,7 +240,7 @@ func TestInstrumentMetricHandler(t *testing.T) {
 	InstrumentMetricHandler(reg, HandlerForTransactional(mReg, HandlerOpts{}))
 	writer := httptest.NewRecorder()
 	request, _ := http.NewRequest("GET", "/", nil)
-	request.Header.Add("Accept", "test/plain")
+	request.Header.Add(acceptHeader, acceptTextPlain)
 
 	handler.ServeHTTP(writer, request)
 	if got := mReg.gatherInvoked; got != 1 {
@@ -237,6 +254,10 @@ func TestInstrumentMetricHandler(t *testing.T) {
 		t.Errorf("got HTTP status code %d, want %d", got, want)
 	}
 
+	if got, want := writer.Header().Get(contentEncodingHeader), string(Identity); got != want {
+		t.Errorf("got HTTP content encoding header %s, want %s", got, want)
+	}
+
 	want := "promhttp_metric_handler_requests_in_flight 1\n"
 	if got := writer.Body.String(); !strings.Contains(got, want) {
 		t.Errorf("got body %q, does not contain %q", got, want)
@@ -278,7 +299,7 @@ func TestHandlerMaxRequestsInFlight(t *testing.T) {
 	w2 := httptest.NewRecorder()
 	w3 := httptest.NewRecorder()
 	request, _ := http.NewRequest("GET", "/", nil)
-	request.Header.Add("Accept", "test/plain")
+	request.Header.Add(acceptHeader, acceptTextPlain)
 
 	c := blockingCollector{Block: make(chan struct{}), CollectStarted: make(chan struct{}, 1)}
 	reg.MustRegister(c)
@@ -332,6 +353,69 @@ func TestHandlerTimeout(t *testing.T) {
 	close(c.Block) // To not leak a goroutine.
 }
 
+func TestInstrumentMetricHandlerWithCompression(t *testing.T) {
+	reg := prometheus.NewRegistry()
+	mReg := &mockTransactionGatherer{g: reg}
+	handler := InstrumentMetricHandler(reg, HandlerForTransactional(mReg, HandlerOpts{DisableCompression: false}))
+	writer := httptest.NewRecorder()
+	request, _ := http.NewRequest("GET", "/", nil)
+	request.Header.Add(acceptHeader, acceptTextPlain)
+	request.Header.Add(acceptEncodingHeader, string(Gzip))
+
+	handler.ServeHTTP(writer, request)
+	if got := mReg.gatherInvoked; got != 1 {
+		t.Fatalf("unexpected number of gather invokes, want 1, got %d", got)
+	}
+	if got := mReg.doneInvoked; got != 1 {
+		t.Fatalf("unexpected number of done invokes, want 1, got %d", got)
+	}
+
+	if got, want := writer.Code, http.StatusOK; got != want {
+		t.Errorf("got HTTP status code %d, want %d", got, want)
+	}
+
+	if got, want := writer.Header().Get(contentEncodingHeader), string(Gzip); got != want {
+		t.Errorf("got HTTP content encoding header %s, want %s", got, want)
+	}
+
+	body, err := readGzippedBody(writer.Body)
+	want := "promhttp_metric_handler_requests_in_flight 1\n"
+	if got := body; !strings.Contains(got, want) {
+		t.Errorf("got body %q, does not contain %q, err: %v", got, want, err)
+	}
+
+	want = "promhttp_metric_handler_requests_total{code=\"200\"} 0\n"
+	if got := body; !strings.Contains(got, want) {
+		t.Errorf("got body %q, does not contain %q, err: %v", got, want, err)
+	}
+
+	for i := 0; i < 100; i++ {
+		writer.Body.Reset()
+		handler.ServeHTTP(writer, request)
+
+		if got, want := mReg.gatherInvoked, i+2; got != want {
+			t.Fatalf("unexpected number of gather invokes, want %d, got %d", want, got)
+		}
+		if got, want := mReg.doneInvoked, i+2; got != want {
+			t.Fatalf("unexpected number of done invokes, want %d, got %d", want, got)
+		}
+		if got, want := writer.Code, http.StatusOK; got != want {
+			t.Errorf("got HTTP status code %d, want %d", got, want)
+		}
+		body, err := readGzippedBody(writer.Body)
+
+		want := "promhttp_metric_handler_requests_in_flight 1\n"
+		if got := body; !strings.Contains(got, want) {
+			t.Errorf("got body %q, does not contain %q, err: %v", got, want, err)
+		}
+
+		want = fmt.Sprintf("promhttp_metric_handler_requests_total{code=\"200\"} %d\n", i+1)
+		if got := body; !strings.Contains(got, want) {
+			t.Errorf("got body %q, does not contain %q, err: %v", got, want, err)
+		}
+	}
+}
+
 func TestNegotiateEncodingWriter(t *testing.T) {
 	var defaultCompressions []string
 
@@ -393,7 +477,7 @@ func TestNegotiateEncodingWriter(t *testing.T) {
 		request, _ := http.NewRequest("GET", "/", nil)
 		request.Header.Add(acceptEncodingHeader, test.acceptEncoding)
 		rr := httptest.NewRecorder()
-		_, encodingHeader, err := NegotiateEncodingWriter(request, rr, test.disableCompression, test.offeredCompressions)
+		_, encodingHeader, _, err := NegotiateEncodingWriter(request, rr, test.disableCompression, test.offeredCompressions)
 
 		if !errors.Is(err, test.err) {
 			t.Errorf("got error: %v, expected: %v", err, test.err)