feat: Support zstd encoding

This allows endpoints to respond with zstd compressed metric data, if
the requester supports it. For backwards compatibility, gzip compression
will take precedence.

Signed-off-by: Manuel Rüger <manuel@rueg.eu>

go.mod

@@ -7,6 +7,7 @@ require (
 	github.com/cespare/xxhash/v2 v2.2.0
 	github.com/davecgh/go-spew v1.1.1
 	github.com/json-iterator/go v1.1.12
+	github.com/klauspost/compress v1.17.8
 	github.com/prometheus/client_model v0.6.0
 	github.com/prometheus/common v0.48.0
 	github.com/prometheus/procfs v0.13.0

go.sum

@@ -17,6 +17,8 @@ github.com/jpillora/backoff v1.0.0 h1:uvFg412JmmHBHw7iwprIxkPMI+sGQ4kzOWsMeHnm2E
 github.com/jpillora/backoff v1.0.0/go.mod h1:J/6gKK9jxlEcS3zixgDgUAsiuZ7yrSoa/FX5e0EB2j4=
 github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
 github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=
+github.com/klauspost/compress v1.17.8 h1:YcnTYrq7MikUT7k0Yb5eceMmALQPYBW/Xltxn0NAMnU=
+github.com/klauspost/compress v1.17.8/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ibi7bDH9ttBbw=
 github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
 github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=

prometheus/promhttp/http.go

@@ -42,6 +42,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/klauspost/compress/zstd"
 	"github.com/prometheus/common/expfmt"
 
 	"github.com/prometheus/client_golang/prometheus"
@@ -169,15 +170,31 @@ func HandlerForTransactional(reg prometheus.TransactionalGatherer, opts HandlerO
 		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)
+		if !opts.DisableCompression {
+			// Gzip takes precedence over zstd
+			// TODO(mrueg): Replace klauspost/compress with stdlib implementation once https://github.com/golang/go/issues/62513 is implemented.
+			if encodingAccepted(req.Header, "zstd") {
+				header.Set(contentEncodingHeader, "zstd")
+				z, err := zstd.NewWriter(rsp, zstd.WithEncoderLevel(zstd.SpeedFastest))
+				if err != nil {
+					return
+				}
+				z.Reset(w)
+				defer z.Close()
+
+				w = z
+			}
+			if encodingAccepted(req.Header, "gzip") {
+				header.Set(contentEncodingHeader, "gzip")
+				gz := gzipPool.Get().(*gzip.Writer)
+				defer gzipPool.Put(gz)
+
+				gz.Reset(w)
+				defer gz.Close()
 
-			gz.Reset(w)
-			defer gz.Close()
+				w = gz
+			}
 
-			w = gz
 		}
 
 		enc := expfmt.NewEncoder(w, contentType)
@@ -381,13 +398,13 @@ type HandlerOpts struct {
 	ProcessStartTime time.Time
 }
 
-// gzipAccepted returns whether the client will accept gzip-encoded content.
-func gzipAccepted(header http.Header) bool {
+// encodingAccepted returns whether the client will accept encoded content.
+func encodingAccepted(header http.Header, encoding string) bool {
 	a := header.Get(acceptEncodingHeader)
 	parts := strings.Split(a, ",")
 	for _, part := range parts {
 		part = strings.TrimSpace(part)
-		if part == "gzip" || strings.HasPrefix(part, "gzip;") {
+		if part == encoding || strings.HasPrefix(part, encoding+";") {
 			return true
 		}
 	}

prometheus/promhttp/http_test.go

@@ -331,3 +331,95 @@ func TestHandlerTimeout(t *testing.T) {
 
 	close(c.Block) // To not leak a goroutine.
 }
+
+func BenchmarkEncoding(b *testing.B) {
+	benchmarks := []struct {
+		name         string
+		encodingType string
+	}{
+		{
+			name:         "test with gzip encoding",
+			encodingType: "gzip",
+		},
+		{
+			name:         "test with zstd encoding",
+			encodingType: "zstd",
+		},
+		{
+			name:         "test with no encoding",
+			encodingType: "identity",
+		},
+	}
+	sizes := []struct {
+		name         string
+		metricCount  int
+		labelCount   int
+		labelLength  int
+		metricLength int
+	}{
+		{
+			name:         "small",
+			metricCount:  50,
+			labelCount:   5,
+			labelLength:  5,
+			metricLength: 5,
+		},
+		{
+			name:         "medium",
+			metricCount:  500,
+			labelCount:   10,
+			labelLength:  5,
+			metricLength: 10,
+		},
+		{
+			name:         "large",
+			metricCount:  5000,
+			labelCount:   10,
+			labelLength:  5,
+			metricLength: 10,
+		},
+		{
+			name:         "extra-large",
+			metricCount:  50000,
+			labelCount:   20,
+			labelLength:  5,
+			metricLength: 10,
+		},
+	}
+
+	for _, size := range sizes {
+		reg := prometheus.NewRegistry()
+		handler := HandlerFor(reg, HandlerOpts{})
+
+		// Generate Metrics
+		// Original source: https://github.com/prometheus-community/avalanche/blob/main/metrics/serve.go
+		labelKeys := make([]string, size.labelCount)
+		for idx := 0; idx < size.labelCount; idx++ {
+			labelKeys[idx] = fmt.Sprintf("label_key_%s_%v", strings.Repeat("k", size.labelLength), idx)
+		}
+		labelValues := make([]string, size.labelCount)
+		for idx := 0; idx < size.labelCount; idx++ {
+			labelValues[idx] = fmt.Sprintf("label_val_%s_%v", strings.Repeat("v", size.labelLength), idx)
+		}
+		metrics := make([]*prometheus.GaugeVec, size.metricCount)
+		for idx := 0; idx < size.metricCount; idx++ {
+			gauge := prometheus.NewGaugeVec(prometheus.GaugeOpts{
+				Name: fmt.Sprintf("avalanche_metric_%s_%v_%v", strings.Repeat("m", size.metricLength), 0, idx),
+				Help: "A tasty metric morsel",
+			}, append([]string{"series_id", "cycle_id"}, labelKeys...))
+			reg.MustRegister(gauge)
+			metrics[idx] = gauge
+		}
+
+		for _, benchmark := range benchmarks {
+			b.Run(benchmark.name+"_"+size.name, func(b *testing.B) {
+				for i := 0; i < b.N; i++ {
+					writer := httptest.NewRecorder()
+					request, _ := http.NewRequest("GET", "/", nil)
+					request.Header.Add("Accept-Encoding", benchmark.encodingType)
+					handler.ServeHTTP(writer, request)
+				}
+			})
+		}
+	}
+}