Support remote write v2 by converting request

Signed-off-by: SungJin1212 <tjdwls1201@gmail.com>

Author: SungJin1212

docs/configuration/config-file-reference.md

@@ -3135,6 +3135,11 @@ ha_tracker:
 # CLI flag: -distributor.sign-write-requests
 [sign_write_requests: <boolean> | default = false]
 
+# EXPERIMENTAL: If true, accept prometheus remote write v2 protocol push
+# request.
+# CLI flag: -distributor.remote-write2-enabled
+[remote_write2_enabled: <boolean> | default = false]
+
 # EXPERIMENTAL: If enabled, distributor would use stream connection to send
 # requests to ingesters.
 # CLI flag: -distributor.use-stream-push

integration/remote_write_v2_test.go

@@ -25,8 +25,8 @@ import (
 )
 
 func TestIngesterRollingUpdate(t *testing.T) {
-	// Test ingester rolling update situation: when -distributor.remote-writev2-enabled is true, and ingester uses the v1.19.0 image.
-	// Expected: remote write 2.0 push success
+	// Test ingester rolling update situation: when -distributor.remote-write2-enabled is true, and ingester uses the v1.19.0 image.
+	// Expected: remote write 2.0 push success, but response header values are set to "0".
 	const blockRangePeriod = 5 * time.Second
 	ingesterImage := "quay.io/cortexproject/cortex:v1.19.0"
 

pkg/cortexpb/cortex.proto

@@ -43,6 +43,14 @@ message WriteResponse {
   // Exemplars represents X-Prometheus-Remote-Write-Written-Exemplars
   int64 Exemplars = 5;
 }
+message WriteResponse {
+  // Samples represents X-Prometheus-Remote-Write-Written-Samples
+  int64 Samples = 1;
+  // Histograms represents X-Prometheus-Remote-Write-Written-Histograms
+  int64 Histograms = 2;
+  // Exemplars represents X-Prometheus-Remote-Write-Written-Exemplars
+  int64 Exemplars = 3;
+}
 
 message TimeSeries {
   repeated LabelPair labels = 1 [(gogoproto.nullable) = false, (gogoproto.customtype) = "LabelAdapter"];

pkg/distributor/write_stats.go

@@ -0,0 +1,62 @@
+package distributor
+
+import (
+	"go.uber.org/atomic"
+)
+
+type WriteStats struct {
+	// Samples represents X-Prometheus-Remote-Write-Written-Samples
+	Samples atomic.Int64
+	// Histograms represents X-Prometheus-Remote-Write-Written-Histograms
+	Histograms atomic.Int64
+	// Exemplars represents X-Prometheus-Remote-Write-Written-Exemplars
+	Exemplars atomic.Int64
+}
+
+func (w *WriteStats) SetSamples(samples int64) {
+	if w == nil {
+		return
+	}
+
+	w.Samples.Store(samples)
+}
+
+func (w *WriteStats) SetHistograms(histograms int64) {
+	if w == nil {
+		return
+	}
+
+	w.Histograms.Store(histograms)
+}
+
+func (w *WriteStats) SetExemplars(exemplars int64) {
+	if w == nil {
+		return
+	}
+
+	w.Exemplars.Store(exemplars)
+}
+
+func (w *WriteStats) LoadSamples() int64 {
+	if w == nil {
+		return 0
+	}
+
+	return w.Samples.Load()
+}
+
+func (w *WriteStats) LoadHistogram() int64 {
+	if w == nil {
+		return 0
+	}
+
+	return w.Histograms.Load()
+}
+
+func (w *WriteStats) LoadExemplars() int64 {
+	if w == nil {
+		return 0
+	}
+
+	return w.Exemplars.Load()
+}

pkg/distributor/write_stats_test.go

@@ -0,0 +1,41 @@
+package distributor
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_SetAndLoad(t *testing.T) {
+	ws := &WriteStats{}
+
+	t.Run("Samples", func(t *testing.T) {
+		ws.SetSamples(3)
+		assert.Equal(t, int64(3), ws.LoadSamples())
+	})
+	t.Run("Histograms", func(t *testing.T) {
+		ws.SetHistograms(10)
+		assert.Equal(t, int64(10), ws.LoadHistogram())
+	})
+	t.Run("Exemplars", func(t *testing.T) {
+		ws.SetExemplars(2)
+		assert.Equal(t, int64(2), ws.LoadExemplars())
+	})
+}
+
+func Test_NilReceiver(t *testing.T) {
+	var ws *WriteStats
+
+	t.Run("Samples", func(t *testing.T) {
+		ws.SetSamples(3)
+		assert.Equal(t, int64(0), ws.LoadSamples())
+	})
+	t.Run("Histograms", func(t *testing.T) {
+		ws.SetHistograms(10)
+		assert.Equal(t, int64(0), ws.LoadHistogram())
+	})
+	t.Run("Exemplars", func(t *testing.T) {
+		ws.SetExemplars(2)
+		assert.Equal(t, int64(0), ws.LoadExemplars())
+	})
+}

pkg/ingester/ingester.go

@@ -1579,7 +1579,13 @@ func (i *Ingester) Push(ctx context.Context, req *cortexpb.WriteRequest) (*corte
 		return &cortexpb.WriteResponse{}, httpgrpc.Errorf(code, "%s", wrapWithUser(firstPartialErr, userID).Error())
 	}
 
-	return &cortexpb.WriteResponse{}, nil
+	writeResponse := &cortexpb.WriteResponse{
+		Samples:    int64(succeededSamplesCount),
+		Histograms: int64(succeededHistogramsCount),
+		Exemplars:  int64(succeededExemplarsCount),
+	}
+
+	return writeResponse, nil
 }
 
 func (i *Ingester) PushStream(srv client.Ingester_PushStreamServer) error {

pkg/util/push/push.go

@@ -31,6 +31,8 @@ const (
 	rw20WrittenSamplesHeader    = "X-Prometheus-Remote-Write-Samples-Written"
 	rw20WrittenHistogramsHeader = "X-Prometheus-Remote-Write-Histograms-Written"
 	rw20WrittenExemplarsHeader  = "X-Prometheus-Remote-Write-Exemplars-Written"
+
+	errMsgNotEnabledPRW2 = "Not enabled prometheus remote write v2 push request"
 )
 
 // Func defines the type of the push. It is similar to http.HandlerFunc.

pkg/util/push/push_test.go

@@ -35,30 +35,34 @@ var (
 	}
 )
 
-func makeV2ReqWithSeries(num int) *writev2.Request {
-	ts := make([]writev2.TimeSeries, 0, num)
+func makeV2ReqWithSeries(num int) *cortexpb.PreallocWriteRequestV2 {
+	ts := make([]cortexpb.PreallocTimeseriesV2, 0, num)
 	symbols := []string{"", "__name__", "test_metric1", "b", "c", "baz", "qux", "d", "e", "foo", "bar", "f", "g", "h", "i", "Test gauge for test purposes", "Maybe op/sec who knows (:", "Test counter for test purposes"}
 	for range num {
-		ts = append(ts, writev2.TimeSeries{
-			LabelsRefs: []uint32{1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
-			Metadata: writev2.Metadata{
-				Type: writev2.Metadata_METRIC_TYPE_GAUGE,
-
-				HelpRef: 15,
-				UnitRef: 16,
-			},
-			Samples:   []writev2.Sample{{Value: 1, Timestamp: 10}},
-			Exemplars: []writev2.Exemplar{{LabelsRefs: []uint32{11, 12}, Value: 1, Timestamp: 10}},
-			Histograms: []writev2.Histogram{
-				writev2.FromIntHistogram(10, &testHistogram),
-				writev2.FromFloatHistogram(20, testHistogram.ToFloat(nil)),
+		ts = append(ts, cortexpb.PreallocTimeseriesV2{
+			TimeSeriesV2: &cortexpb.TimeSeriesV2{
+				LabelsRefs: []uint32{1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
+				Metadata: cortexpb.MetadataV2{
+					Type: cortexpb.METRIC_TYPE_GAUGE,
+
+					HelpRef: 15,
+					UnitRef: 16,
+				},
+				Samples:   []cortexpb.Sample{{Value: 1, TimestampMs: 10}},
+				Exemplars: []cortexpb.ExemplarV2{{LabelsRefs: []uint32{11, 12}, Value: 1, Timestamp: 10}},
+				Histograms: []cortexpb.Histogram{
+					cortexpb.HistogramToHistogramProto(10, &testHistogram),
+					cortexpb.FloatHistogramToHistogramProto(20, testHistogram.ToFloat(nil)),
+				},
 			},
 		})
 	}
 
-	return &writev2.Request{
-		Symbols:    symbols,
-		Timeseries: ts,
+	return &cortexpb.PreallocWriteRequestV2{
+		WriteRequestV2: cortexpb.WriteRequestV2{
+			Symbols:    symbols,
+			Timeseries: ts,
+		},
 	}
 }
 
@@ -161,36 +165,44 @@ func Benchmark_convertV2RequestToV1(b *testing.B) {
 }
 
 func Test_convertV2RequestToV1(t *testing.T) {
-	var v2Req writev2.Request
+	var v2Req cortexpb.PreallocWriteRequestV2
 
 	fh := tsdbutil.GenerateTestFloatHistogram(1)
-	ph := writev2.FromFloatHistogram(4, fh)
+	ph := cortexpb.FloatHistogramToHistogramProto(4, fh)
 
 	symbols := []string{"", "__name__", "test_metric", "b", "c", "baz", "qux", "d", "e", "foo", "bar", "f", "g", "h", "i", "Test gauge for test purposes", "Maybe op/sec who knows (:", "Test counter for test purposes"}
-	timeseries := []writev2.TimeSeries{
+	timeseries := []cortexpb.PreallocTimeseriesV2{
 		{
-			LabelsRefs: []uint32{1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
-			Metadata: writev2.Metadata{
-				Type: writev2.Metadata_METRIC_TYPE_COUNTER,
+			TimeSeriesV2: &cortexpb.TimeSeriesV2{
+				LabelsRefs: []uint32{1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
+				Metadata: cortexpb.MetadataV2{
+					Type: cortexpb.METRIC_TYPE_COUNTER,
 
-				HelpRef: 15,
-				UnitRef: 16,
+					HelpRef: 15,
+					UnitRef: 16,
+				},
+				Samples:   []cortexpb.Sample{{Value: 1, TimestampMs: 1}},
+				Exemplars: []cortexpb.ExemplarV2{{LabelsRefs: []uint32{11, 12}, Value: 1, Timestamp: 1}},
 			},
-			Samples:   []writev2.Sample{{Value: 1, Timestamp: 1}},
-			Exemplars: []writev2.Exemplar{{LabelsRefs: []uint32{11, 12}, Value: 1, Timestamp: 1}},
 		},
 		{
-			LabelsRefs: []uint32{1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
-			Samples:    []writev2.Sample{{Value: 2, Timestamp: 2}},
+			TimeSeriesV2: &cortexpb.TimeSeriesV2{
+				LabelsRefs: []uint32{1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
+				Samples:    []cortexpb.Sample{{Value: 2, TimestampMs: 2}},
+			},
 		},
 		{
-			LabelsRefs: []uint32{1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
-			Samples:    []writev2.Sample{{Value: 3, Timestamp: 3}},
+			TimeSeriesV2: &cortexpb.TimeSeriesV2{
+				LabelsRefs: []uint32{1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
+				Samples:    []cortexpb.Sample{{Value: 3, TimestampMs: 3}},
+			},
 		},
 		{
-			LabelsRefs: []uint32{1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
-			Histograms: []writev2.Histogram{ph, ph},
-			Exemplars:  []writev2.Exemplar{{LabelsRefs: []uint32{11, 12}, Value: 1, Timestamp: 1}},
+			TimeSeriesV2: &cortexpb.TimeSeriesV2{
+				LabelsRefs: []uint32{1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
+				Histograms: []cortexpb.Histogram{ph, ph},
+				Exemplars:  []cortexpb.ExemplarV2{{LabelsRefs: []uint32{11, 12}, Value: 1, Timestamp: 1}},
+			},
 		},
 	}
 
@@ -391,7 +403,7 @@ func TestHandler_ignoresSkipLabelNameValidationIfSet(t *testing.T) {
 	}
 }
 
-func verifyWriteRequestHandler(t *testing.T, expectSource cortexpb.WriteRequest_SourceEnum) func(ctx context.Context, request *cortexpb.WriteRequest) (response *cortexpb.WriteResponse, err error) {
+func verifyWriteRequestHandler(t *testing.T, expectSource cortexpb.SourceEnum) func(ctx context.Context, request *cortexpb.WriteRequest) (response *cortexpb.WriteResponse, err error) {
 	t.Helper()
 	return func(ctx context.Context, request *cortexpb.WriteRequest) (response *cortexpb.WriteResponse, err error) {
 		assert.Len(t, request.Timeseries, 1)
@@ -441,7 +453,7 @@ func createRequest(t *testing.T, protobuf []byte, isV2 bool) *http.Request {
 	return req
 }
 
-func createCortexRemoteWriteV2Protobuf(t *testing.T, skipLabelNameValidation bool, source cortexpb.WriteRequest_SourceEnum) []byte {
+func createCortexRemoteWriteV2Protobuf(t *testing.T, skipLabelNameValidation bool, source cortexpb.SourceEnum) []byte {
 	t.Helper()
 	input := writev2.Request{
 		Symbols: []string{"", "__name__", "foo"},
@@ -497,7 +509,7 @@ func createPrometheusRemoteWriteProtobuf(t *testing.T) []byte {
 	require.NoError(t, err)
 	return inoutBytes
 }
-func createCortexWriteRequestProtobuf(t *testing.T, skipLabelNameValidation bool, source cortexpb.WriteRequest_SourceEnum) []byte {
+func createCortexWriteRequestProtobuf(t *testing.T, skipLabelNameValidation bool, source cortexpb.SourceEnum) []byte {
 	t.Helper()
 	ts := cortexpb.PreallocTimeseries{
 		TimeSeries: &cortexpb.TimeSeries{