feat(rw2): unmarshal RW1 request from RW2

Implement unmarshal of Promtheus remote write request 1.0 data
from Remote Write 2.0 input by converting the model.

The conversion is lossy as RW1 does not have metadata per series, only
per metric family.

Signed-off-by: György Krajcsovits <[email protected]>
Signed-off-by: Arve Knudsen <[email protected]>

Author: krajorama

integration/distributor_test.go

@@ -374,14 +374,26 @@ overrides:
 }
 
 func TestDistributorRemoteWrite2(t *testing.T) {
+	t.Run("caching_unmarshal_data_enabled", func(t *testing.T) {
+		testDistributorRemoteWrite2(t, true)
+	})
+
+	t.Run("caching_unmarshal_data_disabled", func(t *testing.T) {
+		testDistributorRemoteWrite2(t, false)
+	})
+}
+
+func testDistributorRemoteWrite2(t *testing.T, cachingUnmarshalDataEnabled bool) {
 	queryEnd := time.Now().Round(time.Second)
 	queryStart := queryEnd.Add(-1 * time.Hour)
+	queryStep := 10 * time.Minute
 
 	testCases := map[string]struct {
 		inRemoteWrite   []*promRW2.Request
 		runtimeConfig   string
 		queries         map[string]model.Matrix
 		exemplarQueries map[string][]promv1.ExemplarQueryResult
+		expectedStatus  int
 	}{
 		"no special features": {
 			inRemoteWrite: []*promRW2.Request{
@@ -401,6 +413,7 @@ func TestDistributorRemoteWrite2(t *testing.T) {
 					Values: []model.SamplePair{{Timestamp: model.Time(queryStart.UnixMilli()), Value: model.SampleValue(100)}},
 				}},
 			},
+			expectedStatus: http.StatusOK,
 		},
 	}
 
@@ -424,7 +437,7 @@ func TestDistributorRemoteWrite2(t *testing.T) {
 		"-distributor.ha-tracker.store":                      "consul",
 		"-distributor.ha-tracker.consul.hostname":            consul.NetworkHTTPEndpoint(),
 		"-distributor.ha-tracker.prefix":                     "prom_ha/",
-		"-timeseries-unmarshal-caching-optimization-enabled": strconv.FormatBool(false),
+		"-timeseries-unmarshal-caching-optimization-enabled": strconv.FormatBool(cachingUnmarshalDataEnabled),
 	}
 
 	flags := mergeFlags(
@@ -490,10 +503,22 @@ func TestDistributorRemoteWrite2(t *testing.T) {
 
 				res, err := client.PushRW2(ser)
 				require.NoError(t, err)
-				require.Equal(t, http.StatusUnsupportedMediaType, res.StatusCode)
+				require.Equal(t, tc.expectedStatus, res.StatusCode)
+			}
+
+			for q, res := range tc.queries {
+				result, err := client.QueryRange(q, queryStart, queryEnd, queryStep)
+				require.NoError(t, err)
+
+				require.Equal(t, res.String(), result.String())
 			}
 
-			// Placeholder for actual query tests.
+			for q, expResult := range tc.exemplarQueries {
+				result, err := client.QueryExemplars(q, queryStart, queryEnd)
+				require.NoError(t, err)
+
+				require.Equal(t, expResult, result)
+			}
 		})
 	}
 }

pkg/distributor/distributor.go

@@ -1645,7 +1645,7 @@ func (d *Distributor) push(ctx context.Context, pushReq *Request) error {
 		return err
 	}
 
-	d.updateReceivedMetrics(req, userID)
+	d.updateReceivedMetrics(ctx, req, userID)
 
 	if len(req.Timeseries) == 0 && len(req.Metadata) == 0 {
 		return nil
@@ -1876,18 +1876,21 @@ func tokenForMetadata(userID string, metricName string) uint32 {
 	return mimirpb.ShardByMetricName(userID, metricName)
 }
 
-func (d *Distributor) updateReceivedMetrics(req *mimirpb.WriteRequest, userID string) {
-	var receivedSamples, receivedExemplars, receivedMetadata int
+func (d *Distributor) updateReceivedMetrics(ctx context.Context, req *mimirpb.WriteRequest, userID string) {
+	var receivedSamples, receivedHistograms, receivedExemplars, receivedMetadata int
 	for _, ts := range req.Timeseries {
-		receivedSamples += len(ts.Samples) + len(ts.Histograms)
+		receivedSamples += len(ts.Samples)
+		receivedHistograms += len(ts.Histograms)
 		receivedExemplars += len(ts.Exemplars)
 	}
 	d.costAttributionMgr.SampleTracker(userID).IncrementReceivedSamples(req, mtime.Now())
 	receivedMetadata = len(req.Metadata)
 
-	d.receivedSamples.WithLabelValues(userID).Add(float64(receivedSamples))
+	d.receivedSamples.WithLabelValues(userID).Add(float64(receivedSamples + receivedHistograms))
 	d.receivedExemplars.WithLabelValues(userID).Add(float64(receivedExemplars))
 	d.receivedMetadata.WithLabelValues(userID).Add(float64(receivedMetadata))
+
+	updateWriteResponseStatsCtx(ctx, receivedSamples, receivedHistograms, receivedExemplars)
 }
 
 // forReplicationSets runs f, in parallel, for all ingesters in the input replicationSets.

pkg/distributor/push.go

@@ -25,6 +25,7 @@ import (
 	"github.com/grafana/dskit/user"
 	"github.com/opentracing/opentracing-go"
 	"github.com/pkg/errors"
+	promRemote "github.com/prometheus/prometheus/storage/remote"
 
 	"github.com/grafana/mimir/pkg/mimirpb"
 	"github.com/grafana/mimir/pkg/util"
@@ -36,6 +37,13 @@ import (
 // PushFunc defines the type of the push. It is similar to http.HandlerFunc.
 type PushFunc func(ctx context.Context, req *Request) error
 
+// The PushFunc might store promRemote.WriteResponseStats in the context.
+type pushResponseStatsContextMarker struct{}
+
+var (
+	PushResponseStatsContextKey = &pushResponseStatsContextMarker{}
+)
+
 // parserFunc defines how to read the body the request from an HTTP request. It takes an optional RequestBuffers.
 type parserFunc func(ctx context.Context, r *http.Request, maxSize int, buffers *util.RequestBuffers, req *mimirpb.PreallocWriteRequest, logger log.Logger) error
 
@@ -151,65 +159,68 @@ func handler(
 			}
 		}
 
-		var supplier supplierFunc
 		isRW2, err := isRemoteWrite2(r)
 		if err != nil {
 			http.Error(w, err.Error(), http.StatusBadRequest)
 		}
-		if isRW2 {
-			supplier = func() (*mimirpb.WriteRequest, func(), error) {
-				// Return 415 Unsupported Media Type for remote-write v2 requests for now. This is not retryable
-				// unless the client switches to remote-write v1.
-				return nil, nil, httpgrpc.Error(http.StatusUnsupportedMediaType, "remote-write v2 is not supported")
-			}
-		} else {
-			supplier = func() (*mimirpb.WriteRequest, func(), error) {
-				rb := util.NewRequestBuffers(requestBufferPool)
-				var req mimirpb.PreallocWriteRequest
+		supplier := func() (*mimirpb.WriteRequest, func(), error) {
+			rb := util.NewRequestBuffers(requestBufferPool)
+			var req mimirpb.PreallocWriteRequest
 
-				userID, err := tenant.TenantID(ctx)
-				if err != nil && !errors.Is(err, user.ErrNoOrgID) { // ignore user.ErrNoOrgID
-					return nil, nil, errors.Wrap(err, "failed to get tenant ID")
-				}
+			req.UnmarshalFromRW2 = isRW2
 
-				// userID might be empty if none was in the ctx, in this case just use the default setting.
-				if limits.MaxGlobalExemplarsPerUser(userID) == 0 {
-					// The user is not allowed to send exemplars, so there is no need to unmarshal them.
-					// Optimization to avoid the allocations required for unmarshaling exemplars.
-					req.SkipUnmarshalingExemplars = true
-				}
+			userID, err := tenant.TenantID(ctx)
+			if err != nil && !errors.Is(err, user.ErrNoOrgID) { // ignore user.ErrNoOrgID
+				return nil, nil, errors.Wrap(err, "failed to get tenant ID")
+			}
 
-				if err := parser(ctx, r, maxRecvMsgSize, rb, &req, logger); err != nil {
-					// Check for httpgrpc error, default to client error if parsing failed
-					if _, ok := httpgrpc.HTTPResponseFromError(err); !ok {
-						err = httpgrpc.Error(http.StatusBadRequest, err.Error())
-					}
+			// userID might be empty if none was in the ctx, in this case just use the default setting.
+			if limits.MaxGlobalExemplarsPerUser(userID) == 0 {
+				// The user is not allowed to send exemplars, so there is no need to unmarshal them.
+				// Optimization to avoid the allocations required for unmarshaling exemplars.
+				req.SkipUnmarshalingExemplars = true
+			}
 
-					rb.CleanUp()
-					return nil, nil, err
+			if err := parser(ctx, r, maxRecvMsgSize, rb, &req, logger); err != nil {
+				// Check for httpgrpc error, default to client error if parsing failed
+				if _, ok := httpgrpc.HTTPResponseFromError(err); !ok {
+					err = httpgrpc.Error(http.StatusBadRequest, err.Error())
 				}
 
-				if allowSkipLabelNameValidation {
-					req.SkipLabelValidation = req.SkipLabelValidation && r.Header.Get(SkipLabelNameValidationHeader) == "true"
-				} else {
-					req.SkipLabelValidation = false
-				}
+				rb.CleanUp()
+				return nil, nil, err
+			}
 
-				if allowSkipLabelCountValidation {
-					req.SkipLabelCountValidation = req.SkipLabelCountValidation && r.Header.Get(SkipLabelCountValidationHeader) == "true"
-				} else {
-					req.SkipLabelCountValidation = false
-				}
+			if allowSkipLabelNameValidation {
+				req.SkipLabelValidation = req.SkipLabelValidation && r.Header.Get(SkipLabelNameValidationHeader) == "true"
+			} else {
+				req.SkipLabelValidation = false
+			}
 
-				cleanup := func() {
-					mimirpb.ReuseSlice(req.Timeseries)
-					rb.CleanUp()
-				}
-				return &req.WriteRequest, cleanup, nil
+			if allowSkipLabelCountValidation {
+				req.SkipLabelCountValidation = req.SkipLabelCountValidation && r.Header.Get(SkipLabelCountValidationHeader) == "true"
+			} else {
+				req.SkipLabelCountValidation = false
+			}
+
+			cleanup := func() {
+				mimirpb.ReuseSlice(req.Timeseries)
+				rb.CleanUp()
 			}
+			return &req.WriteRequest, cleanup, nil
 		}
 		req := newRequest(supplier)
-		if err := push(ctx, req); err != nil {
+		ctx = contextWithWriteResponseStats(ctx)
+		err = push(ctx, req)
+		rsValue := ctx.Value(PushResponseStatsContextKey)
+		if rsValue != nil {
+			rs := rsValue.(*promRemote.WriteResponseStats)
+			addWriteResponseStats(w, rs)
+		} else {
+			// This should not happen, but if it does, we should not panic.
+			addWriteResponseStats(w, &promRemote.WriteResponseStats{})
+		}
+		if err != nil {
 			if errors.Is(err, context.Canceled) {
 				http.Error(w, err.Error(), statusClientClosedRequest)
 				level.Warn(logger).Log("msg", "push request canceled", "err", err)
@@ -277,6 +288,35 @@ func isRemoteWrite2(r *http.Request) (bool, error) {
 	return false, nil
 }
 
+// Consts from https://github.com/prometheus/prometheus/blob/main/storage/remote/stats.go
+const (
+	rw20WrittenSamplesHeader    = "X-Prometheus-Remote-Write-Samples-Written"
+	rw20WrittenHistogramsHeader = "X-Prometheus-Remote-Write-Histograms-Written"
+	rw20WrittenExemplarsHeader  = "X-Prometheus-Remote-Write-Exemplars-Written"
+)
+
+func contextWithWriteResponseStats(ctx context.Context) context.Context {
+	return context.WithValue(ctx, PushResponseStatsContextKey, &promRemote.WriteResponseStats{})
+}
+
+func addWriteResponseStats(w http.ResponseWriter, rs *promRemote.WriteResponseStats) {
+	headers := w.Header()
+	headers.Set(rw20WrittenSamplesHeader, strconv.Itoa(rs.Samples))
+	headers.Set(rw20WrittenHistogramsHeader, strconv.Itoa(rs.Histograms))
+	headers.Set(rw20WrittenExemplarsHeader, strconv.Itoa(rs.Exemplars))
+}
+
+func updateWriteResponseStatsCtx(ctx context.Context, samples, histograms, exemplars int) {
+	prs := ctx.Value(PushResponseStatsContextKey)
+	if prs == nil {
+		// Should not happen, but we should not panic anyway.
+		return
+	}
+	prs.(*promRemote.WriteResponseStats).Samples += samples
+	prs.(*promRemote.WriteResponseStats).Histograms += histograms
+	prs.(*promRemote.WriteResponseStats).Exemplars += exemplars
+}
+
 func calculateRetryAfter(retryAttemptHeader string, minBackoff, maxBackoff time.Duration) string {
 	const jitterFactor = 0.5
 

pkg/mimirpb/compat_rw2.go

@@ -0,0 +1,69 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+
+package mimirpb
+
+import (
+	"errors"
+	"fmt"
+	"sync"
+)
+
+// Remote Write 2.0 related variables and functions.
+var (
+	errorUnexpectedRW1Timeseries      = errors.New("proto: Remote Write 1.0 field Timeseries in non-Remote Write 1.0 message")
+	errorUnexpectedRW1Metadata        = errors.New("proto: Remote Write 1.0 field Metadata in non-Remote Write 1.0 message")
+	errorUnexpectedRW2Timeseries      = errors.New("proto: Remote Write 2.0 field Timeseries in non-Remote Write 2.0 message")
+	errorUnexpectedRW2Symbols         = errors.New("proto: Remote Write 2.0 field Symbols in non-Remote Write 2.0 message")
+	errorOddNumberOfLabelRefs         = errors.New("proto: Remote Write 2.0 odd number of label references")
+	errorOddNumberOfExemplarLabelRefs = errors.New("proto: Remote Write 2.0 odd number of exemplar label references")
+	errorInvalidLabelRef              = errors.New("proto: Remote Write 2.0 invalid label reference")
+	errorInvalidExemplarLabelRef      = errors.New("proto: Remote Write 2.0 invalid exemplar label reference")
+	errorInternalRW2                  = errors.New("proto: Remote Write 2.0 internal error")
+	errorInvalidHelpRef               = errors.New("proto: Remote Write 2.0 invalid help reference")
+	errorInvalidUnitRef               = errors.New("proto: Remote Write 2.0 invalid unit reference")
+)
+
+// rw2SymbolPageSize is the size of each page in bits.
+const rw2SymbolPageSize = 16
+
+// rw2PagedSymbols is a structure that holds symbols in pages.
+// The problem this solves is that protobuf doesn't tell us
+// how many symbols there are in advance. Without this paging
+// mechanism, we would have to allocate a large amount of memory
+// or do reallocation. This is a compromise between the two.
+type rw2PagedSymbols struct {
+	count uint32
+	pages [][]string
+}
+
+func (ps *rw2PagedSymbols) append(symbol string) {
+	nextPage := ps.count >> rw2SymbolPageSize
+	if int(nextPage) >= len(ps.pages) {
+		ps.pages = append(ps.pages, rw2PagedSymbolsPool.Get().([]string))
+	}
+	ps.pages[nextPage] = append(ps.pages[nextPage], symbol)
+	ps.count++
+}
+
+func (ps *rw2PagedSymbols) releasePages() {
+	for _, page := range ps.pages {
+		page = page[:0]
+		rw2PagedSymbolsPool.Put(page) //nolint:staticcheck
+	}
+}
+
+func (ps *rw2PagedSymbols) get(ref uint32) (string, error) {
+	if ref < ps.count {
+		page := ps.pages[ref>>rw2SymbolPageSize]
+		return page[ref&((1<<rw2SymbolPageSize)-1)], nil
+	}
+	return "", fmt.Errorf("symbol reference %d is out of bounds", ref)
+}
+
+var (
+	rw2PagedSymbolsPool = sync.Pool{
+		New: func() interface{} {
+			return make([]string, 0, 1<<rw2SymbolPageSize)
+		},
+	}
+)