diff --git a/CHANGELOG.md b/CHANGELOG.md index 56a5f900ec4..cb282b40e50 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/docs/blocks-storage/querier.md b/docs/blocks-storage/querier.md index 664c7b22f2d..99c1fe2e521 100644 --- a/docs/blocks-storage/querier.md +++ b/docs/blocks-storage/querier.md @@ -127,7 +127,7 @@ querier: [per_step_stats_enabled: | 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: | default = "gzip"] diff --git a/docs/configuration/config-file-reference.md b/docs/configuration/config-file-reference.md index a3861529ff6..90499301d4b 100644 --- a/docs/configuration/config-file-reference.md +++ b/docs/configuration/config-file-reference.md @@ -4283,7 +4283,7 @@ The `querier_config` configures the Cortex querier. [per_step_stats_enabled: | 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: | default = "gzip"] diff --git a/integration/query_frontend_test.go b/integration/query_frontend_test.go index 6d7b0651d7a..b77bfa64756 100644 --- a/integration/query_frontend_test.go +++ b/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, diff --git a/pkg/api/queryapi/compression.go b/pkg/api/queryapi/compression.go new file mode 100644 index 00000000000..7dd6fcbacab --- /dev/null +++ b/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() +} diff --git a/pkg/api/queryapi/compression_test.go b/pkg/api/queryapi/compression_test.go new file mode 100644 index 00000000000..bcd36a3728c --- /dev/null +++ b/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) +} diff --git a/pkg/api/queryapi/query_api.go b/pkg/api/queryapi/query_api.go index e3793ef5bee..5dd125a6c39 100644 --- a/pkg/api/queryapi/query_api.go +++ b/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) diff --git a/pkg/frontend/transport/handler.go b/pkg/frontend/transport/handler.go index 9001560b524..bba985ea1c0 100644 --- a/pkg/frontend/transport/handler.go +++ b/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)) } diff --git a/pkg/querier/querier.go b/pkg/querier/querier.go index 2020a160b47..78548030fba 100644 --- a/pkg/querier/querier.go +++ b/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 } diff --git a/pkg/querier/tripperware/instantquery/instant_query.go b/pkg/querier/tripperware/instantquery/instant_query.go index a3977207199..c8b41b165e1 100644 --- a/pkg/querier/tripperware/instantquery/instant_query.go +++ b/pkg/querier/tripperware/instantquery/instant_query.go @@ -47,8 +47,15 @@ type instantQueryCodec struct { func NewInstantQueryCodec(compressionStr string, defaultCodecTypeStr string) instantQueryCodec { compression := tripperware.NonCompression // default - if compressionStr == string(tripperware.GzipCompression) { + switch compressionStr { + case string(tripperware.GzipCompression): compression = tripperware.GzipCompression + + case string(tripperware.SnappyCompression): + compression = tripperware.SnappyCompression + + case string(tripperware.ZstdCompression): + compression = tripperware.ZstdCompression } defaultCodecType := tripperware.JsonCodecType // default @@ -102,13 +109,31 @@ func (c instantQueryCodec) DecodeResponse(ctx context.Context, r *http.Response, return nil, err } + responseSizeHeader := r.Header.Get("X-Uncompressed-Length") responseSizeLimiter := limiter.ResponseSizeLimiterFromContextWithFallback(ctx) - body, err := tripperware.BodyBytes(r, responseSizeLimiter, log) + responseSize, hasSizeHeader, err := tripperware.ParseResponseSizeHeader(responseSizeHeader) + if err != nil { + log.Error(err) + return nil, err + } + if hasSizeHeader { + if err := responseSizeLimiter.AddResponseBytes(responseSize); err != nil { + return nil, httpgrpc.Errorf(http.StatusUnprocessableEntity, "%s", err.Error()) + } + } + + body, err := tripperware.BodyBytes(r, log) if err != nil { log.Error(err) return nil, err } + if !hasSizeHeader { + if err := responseSizeLimiter.AddResponseBytes(len(body)); err != nil { + return nil, httpgrpc.Errorf(http.StatusUnprocessableEntity, "%s", err.Error()) + } + } + if r.StatusCode/100 != 2 { return nil, httpgrpc.Errorf(r.StatusCode, "%s", string(body)) } diff --git a/pkg/querier/tripperware/query.go b/pkg/querier/tripperware/query.go index 42e2d9eebf0..180ce1c27d0 100644 --- a/pkg/querier/tripperware/query.go +++ b/pkg/querier/tripperware/query.go @@ -4,7 +4,6 @@ import ( "bytes" "compress/gzip" "context" - "encoding/binary" "fmt" "io" "net/http" @@ -16,6 +15,8 @@ import ( "github.com/go-kit/log" "github.com/gogo/protobuf/proto" jsoniter "github.com/json-iterator/go" + "github.com/klauspost/compress/snappy" + "github.com/klauspost/compress/zstd" "github.com/opentracing/opentracing-go" otlog "github.com/opentracing/opentracing-go/log" "github.com/pkg/errors" @@ -27,7 +28,6 @@ import ( "github.com/cortexproject/cortex/pkg/chunk" "github.com/cortexproject/cortex/pkg/cortexpb" - "github.com/cortexproject/cortex/pkg/util/limiter" "github.com/cortexproject/cortex/pkg/util/runutil" "github.com/thanos-io/promql-engine/logicalplan" @@ -46,6 +46,8 @@ type Compression string const ( GzipCompression Compression = "gzip" + ZstdCompression Compression = "zstd" + SnappyCompression Compression = "snappy" NonCompression Compression = "" JsonCodecType CodecType = "json" ProtobufCodecType CodecType = "protobuf" @@ -446,7 +448,7 @@ type Buffer interface { Bytes() []byte } -func BodyBytes(res *http.Response, responseSizeLimiter *limiter.ResponseSizeLimiter, logger log.Logger) ([]byte, error) { +func BodyBytes(res *http.Response, logger log.Logger) ([]byte, error) { var buf *bytes.Buffer // Attempt to cast the response body to a Buffer and use it if possible. @@ -464,13 +466,26 @@ func BodyBytes(res *http.Response, responseSizeLimiter *limiter.ResponseSizeLimi } } - responseSize := getResponseSize(res, buf) - if err := responseSizeLimiter.AddResponseBytes(responseSize); err != nil { - return nil, httpgrpc.Errorf(http.StatusUnprocessableEntity, "%s", err.Error()) + // Handle decoding response if it was compressed + encoding := res.Header.Get("Content-Encoding") + return decode(buf, encoding, logger) +} + +func BodyBytesFromHTTPGRPCResponse(res *httpgrpc.HTTPResponse, logger log.Logger) ([]byte, error) { + headers := http.Header{} + for _, h := range res.Headers { + headers[h.Key] = h.Values } + // Handle decoding response if it was compressed + encoding := headers.Get("Content-Encoding") + buf := bytes.NewBuffer(res.Body) + return decode(buf, encoding, logger) +} + +func decode(buf *bytes.Buffer, encoding string, logger log.Logger) ([]byte, error) { // if the response is gzipped, lets unzip it here - if strings.EqualFold(res.Header.Get("Content-Encoding"), "gzip") { + if strings.EqualFold(encoding, "gzip") { gReader, err := gzip.NewReader(buf) if err != nil { return nil, err @@ -480,35 +495,24 @@ func BodyBytes(res *http.Response, responseSizeLimiter *limiter.ResponseSizeLimi return io.ReadAll(gReader) } - return buf.Bytes(), nil -} - -func BodyBytesFromHTTPGRPCResponse(res *httpgrpc.HTTPResponse, logger log.Logger) ([]byte, error) { - // if the response is gzipped, lets unzip it here - headers := http.Header{} - for _, h := range res.Headers { - headers[h.Key] = h.Values + // if the response is snappy compressed, decode it here + if strings.EqualFold(encoding, "snappy") { + sReader := snappy.NewReader(buf) + return io.ReadAll(sReader) } - if strings.EqualFold(headers.Get("Content-Encoding"), "gzip") { - gReader, err := gzip.NewReader(bytes.NewBuffer(res.Body)) + + // if the response is zstd compressed, decode it here + if strings.EqualFold(encoding, "zstd") { + zReader, err := zstd.NewReader(buf) if err != nil { return nil, err } - defer runutil.CloseWithLogOnErr(logger, gReader, "close gzip reader") + defer runutil.CloseWithLogOnErr(logger, zReader.IOReadCloser(), "close zstd decoder") - return io.ReadAll(gReader) + return io.ReadAll(zReader) } - return res.Body, nil -} - -func getResponseSize(res *http.Response, buf *bytes.Buffer) int { - if strings.EqualFold(res.Header.Get("Content-Encoding"), "gzip") && len(buf.Bytes()) >= 4 { - // GZIP body contains the size of the original (uncompressed) input data - // modulo 2^32 in the last 4 bytes (https://www.ietf.org/rfc/rfc1952.txt). - return int(binary.LittleEndian.Uint32(buf.Bytes()[len(buf.Bytes())-4:])) - } - return len(buf.Bytes()) + return buf.Bytes(), nil } // UnmarshalJSON implements json.Unmarshaler. @@ -767,9 +771,17 @@ func (s *PrometheusResponseStats) MarshalJSON() ([]byte, error) { } func SetRequestHeaders(h http.Header, defaultCodecType CodecType, compression Compression) { - if compression == GzipCompression { + switch compression { + case GzipCompression: h.Set("Accept-Encoding", string(GzipCompression)) + + case SnappyCompression: + h.Set("Accept-Encoding", string(SnappyCompression)) + + case ZstdCompression: + h.Set("Accept-Encoding", string(ZstdCompression)) } + if defaultCodecType == ProtobufCodecType { h.Set("Accept", ApplicationProtobuf+", "+ApplicationJson) } else { @@ -777,6 +789,17 @@ func SetRequestHeaders(h http.Header, defaultCodecType CodecType, compression Co } } +func ParseResponseSizeHeader(header string) (int, bool, error) { + if header == "" { + return 0, false, nil + } + size, err := strconv.Atoi(header) + if err != nil { + return 0, false, err + } + return size, true, nil +} + func UnmarshalResponse(r *http.Response, buf []byte, resp *PrometheusResponse) error { if r.Header == nil { return json.Unmarshal(buf, resp) diff --git a/pkg/querier/tripperware/query_test.go b/pkg/querier/tripperware/query_test.go index 04606df99e6..08f149f43b0 100644 --- a/pkg/querier/tripperware/query_test.go +++ b/pkg/querier/tripperware/query_test.go @@ -1,10 +1,7 @@ package tripperware import ( - "bytes" - "compress/gzip" "math" - "net/http" "strconv" "testing" "time" @@ -196,50 +193,3 @@ func generateData(timeseries, datapoints int) (floatMatrix, histogramMatrix []*S } return } - -func Test_getResponseSize(t *testing.T) { - tests := []struct { - body []byte - useGzip bool - }{ - { - body: []byte(`foo`), - useGzip: false, - }, - { - body: []byte(`foo`), - useGzip: true, - }, - { - body: []byte(`{"status":"success","data":{"resultType":"vector","result":[]}}`), - useGzip: false, - }, - { - body: []byte(`{"status":"success","data":{"resultType":"vector","result":[]}}`), - useGzip: true, - }, - } - - for i, test := range tests { - t.Run(strconv.Itoa(i), func(t *testing.T) { - expectedBodyLength := len(test.body) - buf := &bytes.Buffer{} - response := &http.Response{} - - if test.useGzip { - response = &http.Response{ - Header: http.Header{"Content-Encoding": []string{"gzip"}}, - } - w := gzip.NewWriter(buf) - _, err := w.Write(test.body) - require.NoError(t, err) - w.Close() - } else { - buf = bytes.NewBuffer(test.body) - } - - bodyLength := getResponseSize(response, buf) - require.Equal(t, expectedBodyLength, bodyLength) - }) - } -} diff --git a/pkg/querier/tripperware/queryrange/query_range.go b/pkg/querier/tripperware/queryrange/query_range.go index df721146f66..f0b11db6121 100644 --- a/pkg/querier/tripperware/queryrange/query_range.go +++ b/pkg/querier/tripperware/queryrange/query_range.go @@ -63,8 +63,15 @@ type prometheusCodec struct { func NewPrometheusCodec(sharded bool, compressionStr string, defaultCodecTypeStr string) *prometheusCodec { //nolint:revive compression := tripperware.NonCompression // default - if compressionStr == string(tripperware.GzipCompression) { + switch compressionStr { + case string(tripperware.GzipCompression): compression = tripperware.GzipCompression + + case string(tripperware.SnappyCompression): + compression = tripperware.SnappyCompression + + case string(tripperware.ZstdCompression): + compression = tripperware.ZstdCompression } defaultCodecType := tripperware.JsonCodecType // default @@ -218,13 +225,31 @@ func (c prometheusCodec) DecodeResponse(ctx context.Context, r *http.Response, _ return nil, err } + responseSizeHeader := r.Header.Get("X-Uncompressed-Length") responseSizeLimiter := limiter.ResponseSizeLimiterFromContextWithFallback(ctx) - body, err := tripperware.BodyBytes(r, responseSizeLimiter, log) + responseSize, hasSizeHeader, err := tripperware.ParseResponseSizeHeader(responseSizeHeader) + if err != nil { + log.Error(err) + return nil, err + } + if hasSizeHeader { + if err := responseSizeLimiter.AddResponseBytes(responseSize); err != nil { + return nil, httpgrpc.Errorf(http.StatusUnprocessableEntity, "%s", err.Error()) + } + } + + body, err := tripperware.BodyBytes(r, log) if err != nil { log.Error(err) return nil, err } + if !hasSizeHeader { + if err := responseSizeLimiter.AddResponseBytes(len(body)); err != nil { + return nil, httpgrpc.Errorf(http.StatusUnprocessableEntity, "%s", err.Error()) + } + } + if r.StatusCode/100 != 2 { return nil, httpgrpc.Errorf(r.StatusCode, "%s", string(body)) }