Improve decompression within ParseProtoReader. (#3682)

* Improve decompression within ParseProtoReader.

This uses the last weaveworks, which include a fix for the content-type size to be set correctly which allows to allocate buffers correctly instead of let it grow naturally.

I've also implemented an alternative decompression from a `bytes.Buffer` which avoids allocating and reading a new buffer when using httpgrpc since the underlaying reader is already a `bytes.Buffer`.

I took the liberty to also improve the validation when using RawSnappy, I'm checking the length using the snappy header. This will avoid wasting CPU when the decompressed size is too big since we will know in advance.

I've added also missing  tests that covers the old and new code.

Signed-off-by: Cyril Tovena <[email protected]>

* missing go.sum update.

Signed-off-by: Cyril Tovena <[email protected]>

* Update changelog.

Signed-off-by: Cyril Tovena <[email protected]>

* Removes FramedSnappy encoding support.

Signed-off-by: Cyril Tovena <[email protected]>

* update go.mod

Signed-off-by: Cyril Tovena <[email protected]>

* Simplify implementation.

From Review feedback from @pracucci

Signed-off-by: Cyril Tovena <[email protected]>

* Nits & limitReader in both cases when reading from a buffer.

Signed-off-by: Cyril Tovena <[email protected]>

* Moves comment.

Signed-off-by: Cyril Tovena <[email protected]>

Co-authored-by: Marco Pracucci <[email protected]>

Author: cyriltovena

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## master / unreleased
 
+* [CHANGE] FramedSnappy encoding support has been removed from Push and Remote Read APIs. This means Prometheus 1.6 support has been removed and the oldest Prometheus version supported in the remote write is 1.7. #3682
 * [CHANGE] Ruler: removed the flag `-ruler.evaluation-delay-duration-deprecated` which was deprecated in 1.4.0. Please use the `ruler_evaluation_delay_duration` per-tenant limit instead. #3693
 * [CHANGE] Removed the flags `-<prefix>.grpc-use-gzip-compression` which were deprecated in 1.3.0: #3693
   * `-query-scheduler.grpc-client-config.grpc-use-gzip-compression`: use `-query-scheduler.grpc-client-config.grpc-compression` instead
@@ -23,6 +24,7 @@
   * `-cluster.peer-timeout` in favor of `-alertmanager.cluster.peer-timeout`
 * [CHANGE] Blocks storage: the default value of `-blocks-storage.bucket-store.sync-interval` has been changed from `5m` to `15m`. #3724
 * [FEATURE] Querier: Queries can be federated across multiple tenants. The tenants IDs involved need to be specified separated by a `|` character in the `X-Scope-OrgID` request header. This is an experimental feature, which can be enabled by setting `-tenant-federation.enabled=true` on all Cortex services. #3250
+* [ENHANCEMENT] Allow specifying JAEGER_ENDPOINT instead of sampling server or local agent port. #3682
 * [FEATURE] Alertmanager: introduced the experimental option `-alertmanager.sharding-enabled` to shard tenants across multiple Alertmanager instances. This feature is still under heavy development and its usage is discouraged. The following new metrics are exported by the Alertmanager: #3664
   * `cortex_alertmanager_ring_check_errors_total`
   * `cortex_alertmanager_sync_configs_total`

go.mod

@@ -13,7 +13,6 @@ require (
 	github.com/alicebob/miniredis v2.5.0+incompatible
 	github.com/armon/go-metrics v0.3.3
 	github.com/aws/aws-sdk-go v1.35.31
-	github.com/blang/semver v3.5.0+incompatible
 	github.com/bradfitz/gomemcache v0.0.0-20190913173617-a41fca850d0b
 	github.com/cespare/xxhash v1.1.0
 	github.com/dustin/go-humanize v1.0.0
@@ -56,7 +55,7 @@ require (
 	github.com/stretchr/testify v1.6.1
 	github.com/thanos-io/thanos v0.13.1-0.20210108102609-f85e4003ba51
 	github.com/uber/jaeger-client-go v2.25.0+incompatible
-	github.com/weaveworks/common v0.0.0-20201119133501-0619918236ec
+	github.com/weaveworks/common v0.0.0-20210112142934-23c8d7fa6120
 	go.etcd.io/bbolt v1.3.5-0.20200615073812-232d8fc87f50
 	go.etcd.io/etcd v0.5.0-alpha.5.0.20200520232829-54ba9589114f
 	go.uber.org/atomic v1.7.0

go.sum

@@ -1179,8 +1179,8 @@ github.com/weaveworks/common v0.0.0-20200206153930-760e36ae819a/go.mod h1:6enWAq
 github.com/weaveworks/common v0.0.0-20200625145055-4b1847531bc9/go.mod h1:c98fKi5B9u8OsKGiWHLRKus6ToQ1Tubeow44ECO1uxY=
 github.com/weaveworks/common v0.0.0-20200914083218-61ffdd448099 h1:MS5M2antM8wzMUqVxIfAi+yb6yjXvDINRFvLnmNXeIw=
 github.com/weaveworks/common v0.0.0-20200914083218-61ffdd448099/go.mod h1:hz10LOsAdzC3K/iXaKoFxOKTDRgxJl+BTGX1GY+TzO4=
-github.com/weaveworks/common v0.0.0-20201119133501-0619918236ec h1:5JmevdpzK10Z2ua0VDToj7Kg2+/t0FzdYBjsurYRE8k=
-github.com/weaveworks/common v0.0.0-20201119133501-0619918236ec/go.mod h1:ykzWac1LtVfOxdCK+jD754at1Ws9dKCwFeUzkFBffPs=
+github.com/weaveworks/common v0.0.0-20210112142934-23c8d7fa6120 h1:zQtcwREXYNvW116ipgc0bRDg1avD2b6QP0RGPLlPWkc=
+github.com/weaveworks/common v0.0.0-20210112142934-23c8d7fa6120/go.mod h1:ykzWac1LtVfOxdCK+jD754at1Ws9dKCwFeUzkFBffPs=
 github.com/weaveworks/promrus v1.2.0 h1:jOLf6pe6/vss4qGHjXmGz4oDJQA+AOCqEL3FvvZGz7M=
 github.com/weaveworks/promrus v1.2.0/go.mod h1:SaE82+OJ91yqjrE1rsvBWVzNZKcHYFtMUyS1+Ogs/KA=
 github.com/willf/bitset v1.1.3/go.mod h1:RjeCKbqT1RxIR/KWY6phxZiaY1IyutSBfGjNPySAYV4=

pkg/querier/remote_read.go

@@ -16,13 +16,11 @@ const maxRemoteReadQuerySize = 1024 * 1024
 // RemoteReadHandler handles Prometheus remote read requests.
 func RemoteReadHandler(q storage.Queryable) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		compressionType := util.CompressionTypeFor(r.Header.Get("X-Prometheus-Remote-Read-Version"))
-
 		ctx := r.Context()
 		var req client.ReadRequest
 		logger := util.WithContext(r.Context(), util.Logger)
-		if err := util.ParseProtoReader(ctx, r.Body, int(r.ContentLength), maxRemoteReadQuerySize, &req, compressionType); err != nil {
-			level.Error(logger).Log("err", err.Error())
+		if err := util.ParseProtoReader(ctx, r.Body, int(r.ContentLength), maxRemoteReadQuerySize, &req, util.RawSnappy); err != nil {
+			level.Error(logger).Log("msg", "failed to parse proto", "err", err.Error())
 			http.Error(w, err.Error(), http.StatusBadRequest)
 			return
 		}
@@ -68,7 +66,7 @@ func RemoteReadHandler(q storage.Queryable) http.Handler {
 			return
 		}
 		w.Header().Add("Content-Type", "application/x-protobuf")
-		if err := util.SerializeProtoResponse(w, &resp, compressionType); err != nil {
+		if err := util.SerializeProtoResponse(w, &resp, util.RawSnappy); err != nil {
 			level.Error(logger).Log("msg", "error sending remote read response", "err", err)
 		}
 	})

pkg/util/http.go

@@ -10,14 +10,15 @@ import (
 	"net/http"
 	"strings"
 
-	"github.com/blang/semver"
 	"github.com/gogo/protobuf/proto"
 	"github.com/golang/snappy"
 	"github.com/opentracing/opentracing-go"
 	otlog "github.com/opentracing/opentracing-go/log"
 	"gopkg.in/yaml.v2"
 )
 
+const messageSizeLargerErrFmt = "received message larger than max (%d vs %d)"
+
 // WriteJSONResponse writes some JSON as a HTTP response.
 func WriteJSONResponse(w http.ResponseWriter, v interface{}) {
 	w.Header().Set("Content-Type", "application/json")
@@ -81,71 +82,22 @@ type CompressionType int
 // Values for CompressionType
 const (
 	NoCompression CompressionType = iota
-	FramedSnappy
 	RawSnappy
 )
 
-var rawSnappyFromVersion = semver.MustParse("0.1.0")
-
-// CompressionTypeFor a given version of the Prometheus remote storage protocol.
-// See https://github.com/prometheus/prometheus/issues/2692.
-func CompressionTypeFor(version string) CompressionType {
-	ver, err := semver.Make(version)
-	if err != nil {
-		return FramedSnappy
-	}
-
-	if ver.GTE(rawSnappyFromVersion) {
-		return RawSnappy
-	}
-	return FramedSnappy
-}
-
 // ParseProtoReader parses a compressed proto from an io.Reader.
 func ParseProtoReader(ctx context.Context, reader io.Reader, expectedSize, maxSize int, req proto.Message, compression CompressionType) error {
-	var body []byte
-	var err error
 	sp := opentracing.SpanFromContext(ctx)
 	if sp != nil {
 		sp.LogFields(otlog.String("event", "util.ParseProtoRequest[start reading]"))
 	}
-	var buf bytes.Buffer
-	if expectedSize > 0 {
-		if expectedSize > maxSize {
-			return fmt.Errorf("message expected size larger than max (%d vs %d)", expectedSize, maxSize)
-		}
-		buf.Grow(expectedSize + bytes.MinRead) // extra space guarantees no reallocation
-	}
-	switch compression {
-	case NoCompression:
-		// Read from LimitReader with limit max+1. So if the underlying
-		// reader is over limit, the result will be bigger than max.
-		_, err = buf.ReadFrom(io.LimitReader(reader, int64(maxSize)+1))
-		body = buf.Bytes()
-	case FramedSnappy:
-		_, err = buf.ReadFrom(io.LimitReader(snappy.NewReader(reader), int64(maxSize)+1))
-		body = buf.Bytes()
-	case RawSnappy:
-		_, err = buf.ReadFrom(reader)
-		body = buf.Bytes()
-		if sp != nil {
-			sp.LogFields(otlog.String("event", "util.ParseProtoRequest[decompress]"),
-				otlog.Int("size", len(body)))
-		}
-		if err == nil && len(body) <= maxSize {
-			body, err = snappy.Decode(nil, body)
-		}
-	}
+	body, err := decompressRequest(reader, expectedSize, maxSize, compression, sp)
 	if err != nil {
 		return err
 	}
-	if len(body) > maxSize {
-		return fmt.Errorf("received message larger than max (%d vs %d)", len(body), maxSize)
-	}
 
 	if sp != nil {
-		sp.LogFields(otlog.String("event", "util.ParseProtoRequest[unmarshal]"),
-			otlog.Int("size", len(body)))
+		sp.LogFields(otlog.String("event", "util.ParseProtoRequest[unmarshal]"), otlog.Int("size", len(body)))
 	}
 
 	// We re-implement proto.Unmarshal here as it calls XXX_Unmarshal first,
@@ -163,6 +115,89 @@ func ParseProtoReader(ctx context.Context, reader io.Reader, expectedSize, maxSi
 	return nil
 }
 
+func decompressRequest(reader io.Reader, expectedSize, maxSize int, compression CompressionType, sp opentracing.Span) (body []byte, err error) {
+	defer func() {
+		if err != nil && len(body) > maxSize {
+			err = fmt.Errorf(messageSizeLargerErrFmt, len(body), maxSize)
+		}
+	}()
+	if expectedSize > maxSize {
+		return nil, fmt.Errorf(messageSizeLargerErrFmt, expectedSize, maxSize)
+	}
+	buffer, ok := tryBufferFromReader(reader)
+	if ok {
+		body, err = decompressFromBuffer(buffer, maxSize, compression, sp)
+		return
+	}
+	body, err = decompressFromReader(reader, expectedSize, maxSize, compression, sp)
+	return
+}
+
+func decompressFromReader(reader io.Reader, expectedSize, maxSize int, compression CompressionType, sp opentracing.Span) ([]byte, error) {
+	var (
+		buf  bytes.Buffer
+		body []byte
+		err  error
+	)
+	if expectedSize > 0 {
+		buf.Grow(expectedSize + bytes.MinRead) // extra space guarantees no reallocation
+	}
+	// Read from LimitReader with limit max+1. So if the underlying
+	// reader is over limit, the result will be bigger than max.
+	reader = io.LimitReader(reader, int64(maxSize)+1)
+	switch compression {
+	case NoCompression:
+		_, err = buf.ReadFrom(reader)
+		body = buf.Bytes()
+	case RawSnappy:
+		_, err = buf.ReadFrom(reader)
+		if err != nil {
+			return nil, err
+		}
+		body, err = decompressFromBuffer(&buf, maxSize, RawSnappy, sp)
+	}
+	return body, err
+}
+
+func decompressFromBuffer(buffer *bytes.Buffer, maxSize int, compression CompressionType, sp opentracing.Span) ([]byte, error) {
+	if len(buffer.Bytes()) > maxSize {
+		return nil, fmt.Errorf(messageSizeLargerErrFmt, len(buffer.Bytes()), maxSize)
+	}
+	switch compression {
+	case NoCompression:
+		return buffer.Bytes(), nil
+	case RawSnappy:
+		if sp != nil {
+			sp.LogFields(otlog.String("event", "util.ParseProtoRequest[decompress]"),
+				otlog.Int("size", len(buffer.Bytes())))
+		}
+		size, err := snappy.DecodedLen(buffer.Bytes())
+		if err != nil {
+			return nil, err
+		}
+		if size > maxSize {
+			return nil, fmt.Errorf(messageSizeLargerErrFmt, size, maxSize)
+		}
+		body, err := snappy.Decode(nil, buffer.Bytes())
+		if err != nil {
+			return nil, err
+		}
+		return body, nil
+	}
+	return nil, nil
+}
+
+// tryBufferFromReader attempts to cast the reader to a `*bytes.Buffer` this is possible when using httpgrpc.
+// If it fails it will return nil and false.
+func tryBufferFromReader(reader io.Reader) (*bytes.Buffer, bool) {
+	if bufReader, ok := reader.(interface {
+		BytesBuffer() *bytes.Buffer
+	}); ok && bufReader != nil {
+		return bufReader.BytesBuffer(), true
+	}
+	return nil, false
+}
+
 // SerializeProtoResponse serializes a protobuf response into an HTTP response.
 func SerializeProtoResponse(w http.ResponseWriter, resp proto.Message, compression CompressionType) error {
 	data, err := proto.Marshal(resp)
@@ -173,14 +208,6 @@ func SerializeProtoResponse(w http.ResponseWriter, resp proto.Message, compressi
 
 	switch compression {
 	case NoCompression:
-	case FramedSnappy:
-		buf := bytes.Buffer{}
-		writer := snappy.NewBufferedWriter(&buf)
-		if _, err := writer.Write(data); err != nil {
-			return err
-		}
-		writer.Close()
-		data = buf.Bytes()
 	case RawSnappy:
 		data = snappy.Encode(nil, data)
 	}

pkg/util/http_test.go

@@ -1,11 +1,16 @@
-package util
+package util_test
 
 import (
+	"bytes"
+	"context"
 	"html/template"
 	"net/http/httptest"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+
+	"github.com/cortexproject/cortex/pkg/ingester/client"
+	"github.com/cortexproject/cortex/pkg/util"
 )
 
 func TestRenderHTTPResponse(t *testing.T) {
@@ -58,7 +63,7 @@ func TestRenderHTTPResponse(t *testing.T) {
 				request.Header.Add(k, v)
 			}
 
-			RenderHTTPResponse(writer, tt.value, tmpl, request)
+			util.RenderHTTPResponse(writer, tt.value, tmpl, request)
 
 			assert.Equal(t, tt.expectedContentType, writer.Header().Get("Content-Type"))
 			assert.Equal(t, 200, writer.Code)
@@ -70,9 +75,85 @@ func TestRenderHTTPResponse(t *testing.T) {
 func TestWriteTextResponse(t *testing.T) {
 	w := httptest.NewRecorder()
 
-	WriteTextResponse(w, "hello world")
+	util.WriteTextResponse(w, "hello world")
 
 	assert.Equal(t, 200, w.Code)
 	assert.Equal(t, "hello world", w.Body.String())
 	assert.Equal(t, "text/plain", w.Header().Get("Content-Type"))
 }
+
+func TestParseProtoReader(t *testing.T) {
+	// 47 bytes compressed and 53 uncompressed
+	req := &client.PreallocWriteRequest{
+		WriteRequest: client.WriteRequest{
+			Timeseries: []client.PreallocTimeseries{
+				{
+					TimeSeries: &client.TimeSeries{
+						Labels: []client.LabelAdapter{
+							{Name: "foo", Value: "bar"},
+						},
+						Samples: []client.Sample{
+							{Value: 10, TimestampMs: 1},
+							{Value: 20, TimestampMs: 2},
+							{Value: 30, TimestampMs: 3},
+						},
+					},
+				},
+			},
+		},
+	}
+
+	for _, tt := range []struct {
+		name           string
+		compression    util.CompressionType
+		maxSize        int
+		expectErr      bool
+		useBytesBuffer bool
+	}{
+		{"rawSnappy", util.RawSnappy, 53, false, false},
+		{"noCompression", util.NoCompression, 53, false, false},
+		{"too big rawSnappy", util.RawSnappy, 10, true, false},
+		{"too big decoded rawSnappy", util.RawSnappy, 50, true, false},
+		{"too big noCompression", util.NoCompression, 10, true, false},
+
+		{"bytesbuffer rawSnappy", util.RawSnappy, 53, false, true},
+		{"bytesbuffer noCompression", util.NoCompression, 53, false, true},
+		{"bytesbuffer too big rawSnappy", util.RawSnappy, 10, true, true},
+		{"bytesbuffer too big decoded rawSnappy", util.RawSnappy, 50, true, true},
+		{"bytesbuffer too big noCompression", util.NoCompression, 10, true, true},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			w := httptest.NewRecorder()
+			assert.Nil(t, util.SerializeProtoResponse(w, req, tt.compression))
+			var fromWire client.PreallocWriteRequest
+
+			reader := w.Result().Body
+			if tt.useBytesBuffer {
+				buf := bytes.Buffer{}
+				_, err := buf.ReadFrom(reader)
+				assert.Nil(t, err)
+				reader = bytesBuffered{Buffer: &buf}
+			}
+
+			err := util.ParseProtoReader(context.Background(), reader, 0, tt.maxSize, &fromWire, tt.compression)
+			if tt.expectErr {
+				assert.NotNil(t, err)
+				return
+			}
+			assert.Nil(t, err)
+			assert.Equal(t, req, &fromWire)
+		})
+	}
+}
+
+type bytesBuffered struct {
+	*bytes.Buffer
+}
+
+func (b bytesBuffered) Close() error {
+	return nil
+}
+
+func (b bytesBuffered) BytesBuffer() *bytes.Buffer {
+	return b.Buffer
+}

pkg/util/push/push.go

@@ -25,9 +25,8 @@ func Handler(cfg distributor.Config, sourceIPs *middleware.SourceIPExtractor, pu
 				logger = util.WithSourceIPs(source, logger)
 			}
 		}
-		compressionType := util.CompressionTypeFor(r.Header.Get("X-Prometheus-Remote-Write-Version"))
 		var req client.PreallocWriteRequest
-		err := util.ParseProtoReader(ctx, r.Body, int(r.ContentLength), cfg.MaxRecvMsgSize, &req, compressionType)
+		err := util.ParseProtoReader(ctx, r.Body, int(r.ContentLength), cfg.MaxRecvMsgSize, &req, util.RawSnappy)
 		if err != nil {
 			level.Error(logger).Log("err", err.Error())
 			http.Error(w, err.Error(), http.StatusBadRequest)