add zstd and snappy compression for query api

CHANGELOG.md

@@ -20,6 +20,7 @@
 * [FEATURE] Compactor: Add support for percentage based sharding for compactors. #6738
 * [FEATURE] Querier: Allow choosing PromQL engine via header. #6777
 * [FEATURE] Querier: Support for configuring query optimizers and enabling XFunctions in the Thanos engine. #6873
+* [ENHANCEMENT] Querier: Support snappy and zstd response compression for `-querier.response-compression` flag. #6848
 * [ENHANCEMENT] Tenant Federation: Add a # of query result limit logic when the `-tenant-federation.regex-matcher-enabled` is enabled. #6845
 * [ENHANCEMENT] Query Frontend: Add a `cortex_slow_queries_total` metric to track # of slow queries per user. #6859
 * [ENHANCEMENT] Query Frontend: Change to return 400 when the tenant resolving fail. #6715

docs/blocks-storage/querier.md

@@ -127,7 +127,7 @@ querier:
   [per_step_stats_enabled: <boolean> | default = false]
 
   # Use compression for metrics query API or instant and range query APIs.
-  # Supports 'gzip' and '' (disable compression)
+  # Supported compression 'gzip', 'snappy', 'zstd' and '' (disable compression)
   # CLI flag: -querier.response-compression
   [response_compression: <string> | default = "gzip"]
 

docs/configuration/config-file-reference.md

@@ -4283,7 +4283,7 @@ The `querier_config` configures the Cortex querier.
 [per_step_stats_enabled: <boolean> | default = false]
 
 # Use compression for metrics query API or instant and range query APIs.
-# Supports 'gzip' and '' (disable compression)
+# Supported compression 'gzip', 'snappy', 'zstd' and '' (disable compression)
 # CLI flag: -querier.response-compression
 [response_compression: <string> | default = "gzip"]
 

integration/query_frontend_test.go

@@ -216,14 +216,34 @@ func TestQueryFrontendProtobufCodec(t *testing.T) {
 			require.NoError(t, s.StartAndWaitReady(minio))
 
 			flags = mergeFlags(e2e.EmptyFlags(), map[string]string{
-				"-api.querier-default-codec":    "protobuf",
-				"-querier.response-compression": "gzip",
+				"-api.querier-default-codec": "protobuf",
 			})
 			return cortexConfigFile, flags
 		},
 	})
 }
 
+func TestQuerierToQueryFrontendCompression(t *testing.T) {
+	for _, compression := range []string{"gzip", "zstd", "snappy", ""} {
+		runQueryFrontendTest(t, queryFrontendTestConfig{
+			testMissingMetricName: false,
+			querySchedulerEnabled: true,
+			queryStatsEnabled:     true,
+			setup: func(t *testing.T, s *e2e.Scenario) (configFile string, flags map[string]string) {
+				require.NoError(t, writeFileToSharedDir(s, cortexConfigFile, []byte(BlocksStorageConfig)))
+
+				minio := e2edb.NewMinio(9000, BlocksStorageFlags()["-blocks-storage.s3.bucket-name"])
+				require.NoError(t, s.StartAndWaitReady(minio))
+
+				flags = mergeFlags(e2e.EmptyFlags(), map[string]string{
+					"-querier.response-compression": compression,
+				})
+				return cortexConfigFile, flags
+			},
+		})
+	}
+}
+
 func TestQueryFrontendRemoteRead(t *testing.T) {
 	runQueryFrontendTest(t, queryFrontendTestConfig{
 		remoteReadEnabled: true,

pkg/api/queryapi/compression.go

@@ -0,0 +1,90 @@
+package queryapi
+
+import (
+	"io"
+	"net/http"
+	"strings"
+
+	"github.com/klauspost/compress/gzip"
+	"github.com/klauspost/compress/snappy"
+	"github.com/klauspost/compress/zlib"
+	"github.com/klauspost/compress/zstd"
+)
+
+const (
+	acceptEncodingHeader  = "Accept-Encoding"
+	contentEncodingHeader = "Content-Encoding"
+	gzipEncoding          = "gzip"
+	deflateEncoding       = "deflate"
+	snappyEncoding        = "snappy"
+	zstdEncoding          = "zstd"
+)
+
+// 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 zstdWriter, ok := c.writer.(*zstd.Encoder); ok {
+		zstdWriter.Flush()
+	}
+	if snappyWriter, ok := c.writer.(*snappy.Writer); ok {
+		snappyWriter.Flush()
+	}
+	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 zstdEncoding:
+			encoder, err := zstd.NewWriter(writer)
+			if err == nil {
+				writer.Header().Set(contentEncodingHeader, zstdEncoding)
+				return &compressedResponseWriter{ResponseWriter: writer, writer: encoder}
+			}
+		case snappyEncoding:
+			writer.Header().Set(contentEncodingHeader, snappyEncoding)
+			return &compressedResponseWriter{ResponseWriter: writer, writer: snappy.NewBufferedWriter(writer)}
+		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}
+}
+
+// CompressionHandler is a wrapper around http.Handler which adds suitable
+// response compression based on the client's Accept-Encoding headers.
+type CompressionHandler struct {
+	Handler http.Handler
+}
+
+// ServeHTTP 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()
+}

pkg/api/queryapi/compression_test.go

@@ -0,0 +1,159 @@
+package queryapi
+
+import (
+	"bytes"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/klauspost/compress/gzip"
+	"github.com/klauspost/compress/snappy"
+	"github.com/klauspost/compress/zlib"
+	"github.com/klauspost/compress/zstd"
+	"github.com/stretchr/testify/require"
+)
+
+func decompress(t *testing.T, encoding string, b []byte) []byte {
+	t.Helper()
+
+	switch encoding {
+	case gzipEncoding:
+		r, err := gzip.NewReader(bytes.NewReader(b))
+		require.NoError(t, err)
+		defer r.Close()
+		data, err := io.ReadAll(r)
+		require.NoError(t, err)
+		return data
+	case deflateEncoding:
+		r, err := zlib.NewReader(bytes.NewReader(b))
+		require.NoError(t, err)
+		defer r.Close()
+		data, err := io.ReadAll(r)
+		require.NoError(t, err)
+		return data
+	case snappyEncoding:
+		data, err := io.ReadAll(snappy.NewReader(bytes.NewReader(b)))
+		require.NoError(t, err)
+		return data
+	case zstdEncoding:
+		r, err := zstd.NewReader(bytes.NewReader(b))
+		require.NoError(t, err)
+		defer r.Close()
+		data, err := io.ReadAll(r)
+		require.NoError(t, err)
+		return data
+	default:
+		return b
+	}
+}
+
+func TestNewCompressedResponseWriter_SupportedEncodings(t *testing.T) {
+	for _, tc := range []string{gzipEncoding, deflateEncoding, snappyEncoding, zstdEncoding} {
+		t.Run(tc, func(t *testing.T) {
+			rec := httptest.NewRecorder()
+			req := httptest.NewRequest(http.MethodGet, "/", nil)
+			req.Header.Set(acceptEncodingHeader, tc)
+
+			cw := newCompressedResponseWriter(rec, req)
+			payload := []byte("hello world")
+			_, err := cw.Write(payload)
+			require.NoError(t, err)
+			cw.Close()
+
+			require.Equal(t, tc, rec.Header().Get(contentEncodingHeader))
+
+			decompressed := decompress(t, tc, rec.Body.Bytes())
+			require.Equal(t, payload, decompressed)
+
+			switch tc {
+			case gzipEncoding:
+				_, ok := cw.writer.(*gzip.Writer)
+				require.True(t, ok)
+			case deflateEncoding:
+				_, ok := cw.writer.(*zlib.Writer)
+				require.True(t, ok)
+			case snappyEncoding:
+				_, ok := cw.writer.(*snappy.Writer)
+				require.True(t, ok)
+			case zstdEncoding:
+				_, ok := cw.writer.(*zstd.Encoder)
+				require.True(t, ok)
+			}
+		})
+	}
+}
+
+func TestNewCompressedResponseWriter_UnsupportedEncoding(t *testing.T) {
+	for _, tc := range []string{"", "br", "unknown"} {
+		t.Run(tc, func(t *testing.T) {
+			rec := httptest.NewRecorder()
+			req := httptest.NewRequest(http.MethodGet, "/", nil)
+			if tc != "" {
+				req.Header.Set(acceptEncodingHeader, tc)
+			}
+
+			cw := newCompressedResponseWriter(rec, req)
+			payload := []byte("data")
+			_, err := cw.Write(payload)
+			require.NoError(t, err)
+			cw.Close()
+
+			require.Empty(t, rec.Header().Get(contentEncodingHeader))
+			require.Equal(t, payload, rec.Body.Bytes())
+			require.Same(t, rec, cw.writer)
+		})
+	}
+}
+
+func TestNewCompressedResponseWriter_MultipleEncodings(t *testing.T) {
+	tests := []struct {
+		header     string
+		expectEnc  string
+		expectType interface{}
+	}{
+		{"snappy, gzip", snappyEncoding, &snappy.Writer{}},
+		{"unknown, gzip", gzipEncoding, &gzip.Writer{}},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.header, func(t *testing.T) {
+			rec := httptest.NewRecorder()
+			req := httptest.NewRequest(http.MethodGet, "/", nil)
+			req.Header.Set(acceptEncodingHeader, tc.header)
+
+			cw := newCompressedResponseWriter(rec, req)
+			_, err := cw.Write([]byte("payload"))
+			require.NoError(t, err)
+			cw.Close()
+
+			require.Equal(t, tc.expectEnc, rec.Header().Get(contentEncodingHeader))
+			decompressed := decompress(t, tc.expectEnc, rec.Body.Bytes())
+			require.Equal(t, []byte("payload"), decompressed)
+
+			switch tc.expectEnc {
+			case gzipEncoding:
+				require.IsType(t, &gzip.Writer{}, cw.writer)
+			case snappyEncoding:
+				require.IsType(t, &snappy.Writer{}, cw.writer)
+			}
+		})
+	}
+}
+
+func TestCompressionHandler_ServeHTTP(t *testing.T) {
+	handler := CompressionHandler{Handler: http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		_, err := w.Write([]byte("hello"))
+		require.NoError(t, err)
+	})}
+
+	rec := httptest.NewRecorder()
+	req := httptest.NewRequest(http.MethodGet, "/", nil)
+	req.Header.Set(acceptEncodingHeader, gzipEncoding)
+
+	handler.ServeHTTP(rec, req)
+
+	require.Equal(t, gzipEncoding, rec.Header().Get(contentEncodingHeader))
+	decompressed := decompress(t, gzipEncoding, rec.Body.Bytes())
+	require.Equal(t, []byte("hello"), decompressed)
+}

pkg/api/queryapi/query_api.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"strconv"
 	"time"
 
 	"github.com/go-kit/log"
@@ -208,7 +209,7 @@ func (q *QueryAPI) Wrap(f apiFunc) http.HandlerFunc {
 		w.WriteHeader(http.StatusNoContent)
 	}
 
-	return httputil.CompressionHandler{
+	return CompressionHandler{
 		Handler: http.HandlerFunc(hf),
 	}.ServeHTTP
 }
@@ -237,6 +238,7 @@ func (q *QueryAPI) respond(w http.ResponseWriter, req *http.Request, data interf
 	}
 
 	w.Header().Set("Content-Type", codec.ContentType().String())
+	w.Header().Set("X-Uncompressed-Length", strconv.Itoa(len(b)))
 	w.WriteHeader(http.StatusOK)
 	if n, err := w.Write(b); err != nil {
 		level.Error(q.logger).Log("error writing response", "url", req.URL, "bytesWritten", n, "err", err)

pkg/frontend/transport/handler.go

@@ -75,8 +75,6 @@ const (
 	limitBytesStoreGateway  = `exceeded bytes limit`
 )
 
-var noopResponseSizeLimiter = limiter.NewResponseSizeLimiter(0)
-
 // Config for a Handler.
 type HandlerConfig struct {
 	LogQueriesLongerThan      time.Duration `yaml:"log_queries_longer_than"`
@@ -332,7 +330,7 @@ func (f *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			// If the response status code is not 2xx, try to get the
 			// error message from response body.
 			if resp.StatusCode/100 != 2 {
-				body, err2 := tripperware.BodyBytes(resp, noopResponseSizeLimiter, f.log)
+				body, err2 := tripperware.BodyBytes(resp, f.log)
 				if err2 == nil {
 					err = httpgrpc.Errorf(resp.StatusCode, "%s", string(body))
 				}

pkg/querier/querier.go

@@ -102,7 +102,7 @@ var (
 	errBadLookbackConfigs                             = errors.New("bad settings, query_store_after >= query_ingesters_within which can result in queries not being sent")
 	errShuffleShardingLookbackLessThanQueryStoreAfter = errors.New("the shuffle-sharding lookback period should be greater or equal than the configured 'query store after'")
 	errEmptyTimeRange                                 = errors.New("empty time range")
-	errUnsupportedResponseCompression                 = errors.New("unsupported response compression. Supported compression 'gzip' and '' (disable compression)")
+	errUnsupportedResponseCompression                 = errors.New("unsupported response compression. Supported compression 'gzip', 'snappy', 'zstd' and '' (disable compression)")
 	errInvalidConsistencyCheckAttempts                = errors.New("store gateway consistency check max attempts should be greater or equal than 1")
 	errInvalidIngesterQueryMaxAttempts                = errors.New("ingester query max attempts should be greater or equal than 1")
 	errInvalidParquetQueryableDefaultBlockStore       = errors.New("unsupported parquet queryable default block store. Supported options are tsdb and parquet")
@@ -129,7 +129,7 @@ func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 	f.IntVar(&cfg.MaxSamples, "querier.max-samples", 50e6, "Maximum number of samples a single query can load into memory.")
 	f.DurationVar(&cfg.QueryIngestersWithin, "querier.query-ingesters-within", 0, "Maximum lookback beyond which queries are not sent to ingester. 0 means all queries are sent to ingester.")
 	f.BoolVar(&cfg.EnablePerStepStats, "querier.per-step-stats-enabled", false, "Enable returning samples stats per steps in query response.")
-	f.StringVar(&cfg.ResponseCompression, "querier.response-compression", "gzip", "Use compression for metrics query API or instant and range query APIs. Supports 'gzip' and '' (disable compression)")
+	f.StringVar(&cfg.ResponseCompression, "querier.response-compression", "gzip", "Use compression for metrics query API or instant and range query APIs. Supported compression 'gzip', 'snappy', 'zstd' and '' (disable compression)")
 	f.DurationVar(&cfg.MaxQueryIntoFuture, "querier.max-query-into-future", 10*time.Minute, "Maximum duration into the future you can query. 0 to disable.")
 	f.DurationVar(&cfg.DefaultEvaluationInterval, "querier.default-evaluation-interval", time.Minute, "The default evaluation interval or step size for subqueries.")
 	f.DurationVar(&cfg.QueryStoreAfter, "querier.query-store-after", 0, "The time after which a metric should be queried from storage and not just ingesters. 0 means all queries are sent to store. When running the blocks storage, if this option is enabled, the time range of the query sent to the store will be manipulated to ensure the query end is not more recent than 'now - query-store-after'.")
@@ -158,7 +158,7 @@ func (cfg *Config) Validate() error {
 		}
 	}
 
-	if cfg.ResponseCompression != "" && cfg.ResponseCompression != "gzip" {
+	if cfg.ResponseCompression != "" && cfg.ResponseCompression != "gzip" && cfg.ResponseCompression != "snappy" && cfg.ResponseCompression != "zstd" {
 		return errUnsupportedResponseCompression
 	}