mimirpb: Associate metadata to series when possible during RW2 conversion (#12342)

#### What this PR does

RW2 attaches metric metadata to series, where RW1 carries metadata
separately.
Previously, we took a fast path, where we always generate new series
just to carry each metadata.
This PR changes the conversion to metadata into the corresponding
series, if both are present in the same request.

The main benefit is on requests that are already RW2 (and tend to carry
a mixture of series and metadata). We squeeze some additional bytes out
of the new format by not repeating the label name references two times
per metric family, plus dropping extra `repeated TimeSeriesRW2`
separators in protobuf. It does not change data-only or metadata-only
pushes.

We now do extra work in the distributor to match these elements up.
Previously we theorized this wasn't worth it, but now after analyzing
typical distributor profile data, I expect the extra work here to not be
as impactful on perf or TCO as once thought. The upside is we keep the
record size absolutely minimized, as clients shift over to send RW2 push
requests.

#### Which issue(s) this PR fixes or relates to

contrib: grafana/mimir-squad#2253

#### Checklist

- [x] Tests updated.
- [ ] Documentation added.
- [ ] `CHANGELOG.md` updated - the order of entries should be
`[CHANGE]`, `[FEATURE]`, `[ENHANCEMENT]`, `[BUGFIX]`. If changelog entry
is not needed, please add the `changelog-not-needed` label to the PR.
- [ ]
[`about-versioning.md`](https://github.com/grafana/mimir/blob/main/docs/sources/mimir/configure/about-versioning.md)
updated with experimental features.

Author: alexweav

pkg/mimirpb/prealloc_rw2.go

@@ -2,7 +2,11 @@
 
 package mimirpb
 
-import "fmt"
+import (
+	"fmt"
+
+	"github.com/prometheus/prometheus/model/labels"
+)
 
 func ReuseRW2(req *WriteRequest) {
 	reuseSymbolsSlice(req.SymbolsRW2)
@@ -30,27 +34,53 @@ func FromWriteRequestToRW2Request(rw1 *WriteRequest, commonSymbols []string, off
 	defer reuseSymbolsTable(symbols)
 	symbols.ConfigureCommonSymbols(offset, commonSymbols)
 
+	var metadataMap map[string]*MetricMetadata
+	if len(rw1.Metadata) > 0 && len(rw1.Timeseries) > 0 {
+		metadataMap = make(map[string]*MetricMetadata, len(rw1.Metadata))
+		for _, meta := range rw1.Metadata {
+			metadataMap[meta.MetricFamilyName] = meta
+		}
+	}
+
 	rw2Timeseries := make([]TimeSeriesRW2, 0, len(rw1.Timeseries)+len(rw1.Metadata)) // TODO: Pool-ify this allocation
 	for _, ts := range rw1.Timeseries {
 		refs := make([]uint32, 0, len(ts.Labels)*2) // TODO: Pool-ify this allocation
+		metricName := ""
 		for i := range ts.Labels {
+			if ts.Labels[i].Name == labels.MetricName {
+				metricName = ts.Labels[i].Value
+			}
 			refs = append(refs, symbols.Symbolize(ts.Labels[i].Name), symbols.Symbolize(ts.Labels[i].Value))
 		}
 
+		var metadata MetadataRW2
+		if len(metadataMap) > 0 {
+			if rw1Metadata, ok := metadataMap[metricName]; ok {
+				metadata = FromMetricMetadataToMetadataRW2(rw1Metadata, symbols)
+				delete(metadataMap, metricName)
+			}
+		}
+
 		rw2Timeseries = append(rw2Timeseries, TimeSeriesRW2{
 			LabelsRefs:       refs,
 			Samples:          ts.Samples,
 			Histograms:       ts.Histograms,
 			Exemplars:        FromExemplarsToExemplarsRW2(ts.Exemplars, symbols),
-			Metadata:         MetadataRW2{},
+			Metadata:         metadata,
 			CreatedTimestamp: ts.CreatedTimestamp,
 		})
 	}
 
-	// Represent metadata as extra, empty timeseries rather than attaching it to existing series.
-	// We never replicate metadata if there are multiple series with the same name, and it removes the requirement to match up the metadata to the right series.
+	// If there are extra metadata not associated with any timeseries, we fabricate an empty timeseries to carry it.
 	for _, meta := range rw1.Metadata {
-		labelsRefs := []uint32{symbols.Symbolize("__name__"), symbols.Symbolize(meta.MetricFamilyName)}
+		// If we are matching metadata to series, and we already matched this metadata, we can skip it.
+		if metadataMap != nil {
+			if _, ok := metadataMap[meta.MetricFamilyName]; !ok {
+				continue
+			}
+		}
+
+		labelsRefs := []uint32{symbols.Symbolize(labels.MetricName), symbols.Symbolize(meta.MetricFamilyName)}
 		rw2meta := FromMetricMetadataToMetadataRW2(meta, symbols)
 		rw2Timeseries = append(rw2Timeseries, TimeSeriesRW2{
 			LabelsRefs: labelsRefs,

pkg/mimirpb/prealloc_rw2_test.go

@@ -237,7 +237,7 @@ func TestWriteRequestRW2Conversion(t *testing.T) {
 		require.Equal(t, expTimeseries, rw2.TimeseriesRW2)
 	})
 
-	t.Run("metadata", func(t *testing.T) {
+	t.Run("metadata matching existing series", func(t *testing.T) {
 		req := &WriteRequest{
 			Timeseries: []PreallocTimeseries{
 				{
@@ -264,15 +264,67 @@ func TestWriteRequestRW2Conversion(t *testing.T) {
 			{
 				LabelsRefs: []uint32{1, 2, 3, 4},
 				Samples:    []Sample{{Value: 123, TimestampMs: 1234567890}, {Value: 456, TimestampMs: 1234567900}},
+				Metadata: MetadataRW2{
+					Type:    METRIC_TYPE_COUNTER,
+					HelpRef: 5,
+					UnitRef: 6,
+				},
+			},
+		}
+		require.NoError(t, err)
+		require.Nil(t, rw2.Timeseries)
+		require.Nil(t, rw2.Metadata)
+		require.Equal(t, expSymbols, rw2.SymbolsRW2)
+		require.Equal(t, expTimeseries, rw2.TimeseriesRW2)
+	})
+
+	t.Run("metadata without corresponding series", func(t *testing.T) {
+		req := &WriteRequest{
+			Timeseries: []PreallocTimeseries{
+				{
+					TimeSeries: &TimeSeries{
+						Labels:  FromLabelsToLabelAdapters(labels.FromMap(map[string]string{"__name__": "my_cool_series", "job": "foo/bar"})),
+						Samples: []Sample{{Value: 123, TimestampMs: 1234567890}, {Value: 456, TimestampMs: 1234567900}},
+					},
+				},
+			},
+			Metadata: []*MetricMetadata{
+				{
+					Type:             COUNTER,
+					MetricFamilyName: "my_cool_series",
+					Help:             "It's a cool series.",
+					Unit:             "megawatts",
+				},
+				{
+					Type:             COUNTER,
+					MetricFamilyName: "my_cool_series_2",
+					Help:             "It's a second cool series.",
+					Unit:             "gigawatts",
+				},
 			},
+		}
+
+		rw2, err := FromWriteRequestToRW2Request(req, nil, 0)
+
+		expSymbols := []string{"", "__name__", "my_cool_series", "job", "foo/bar", "It's a cool series.", "megawatts", "my_cool_series_2", "It's a second cool series.", "gigawatts"}
+		expTimeseries := []TimeSeriesRW2{
 			{
-				LabelsRefs: []uint32{1, 2},
+				LabelsRefs: []uint32{1, 2, 3, 4},
+				Samples:    []Sample{{Value: 123, TimestampMs: 1234567890}, {Value: 456, TimestampMs: 1234567900}},
 				Metadata: MetadataRW2{
 					Type:    METRIC_TYPE_COUNTER,
 					HelpRef: 5,
 					UnitRef: 6,
 				},
 			},
+			{
+				LabelsRefs: []uint32{1, 7},
+				Metadata: MetadataRW2{
+					Type:    METRIC_TYPE_COUNTER,
+					HelpRef: 8,
+					UnitRef: 9,
+				},
+			},
 		}
 		require.NoError(t, err)
 		require.Nil(t, rw2.Timeseries)

pkg/mimirpb/split_test.go

@@ -388,7 +388,7 @@ func TestSplitWriteRequestByMaxMarshalSize_Fuzzy(t *testing.T) {
 				merged.Metadata = append(merged.Metadata, partial.Metadata...)
 			}
 
-			assert.Equal(t, req, merged)
+			require.Equal(t, req, merged)
 		}
 	})
 
@@ -406,14 +406,14 @@ func TestSplitWriteRequestByMaxMarshalSize_Fuzzy(t *testing.T) {
 			// Remember the original request, to compare against.
 			req := generateWriteRequestRW2(numSeries, numLabelsPerSeries, numSamplesPerSeries, numMetadata)
 			reqCpy := generateWriteRequestRW2(numSeries, numLabelsPerSeries, numSamplesPerSeries, numMetadata)
-			assert.Equal(t, req, reqCpy)
+			require.Equal(t, req, reqCpy)
 
 			maxSize := req.Size() / (1 + rnd.Intn(10))
 			partials := SplitWriteRequestByMaxMarshalSizeRW2(req, req.Size(), maxSize, 0, nil)
 
 			// Ensure the merge of all partial requests is equal to the original one.
 			merged := mergeRW2s(partials)
-			assert.Equal(t, reqCpy, merged)
+			require.Equal(t, reqCpy, merged)
 		}
 	})
 }
@@ -720,6 +720,11 @@ func mergeRW2s(partials []*WriteRequest) *WriteRequest {
 				newLbls[i] = st.Symbolize(strVal)
 			}
 
+			helpTxt := partial.SymbolsRW2[ts.Metadata.HelpRef]
+			helpRef := st.Symbolize(helpTxt)
+			unitTxt := partial.SymbolsRW2[ts.Metadata.UnitRef]
+			unitRef := st.Symbolize(unitTxt)
+
 			newExemplars := make([]ExemplarRW2, 0, len(ts.Exemplars))
 			for _, ex := range ts.Exemplars {
 				newLbls := make([]uint32, len(ex.LabelsRefs))
@@ -737,11 +742,6 @@ func mergeRW2s(partials []*WriteRequest) *WriteRequest {
 				newExemplars = nil
 			}
 
-			helpTxt := partial.SymbolsRW2[ts.Metadata.HelpRef]
-			helpRef := st.Symbolize(helpTxt)
-			unitTxt := partial.SymbolsRW2[ts.Metadata.UnitRef]
-			unitRef := st.Symbolize(unitTxt)
-
 			newTS := TimeSeriesRW2{
 				LabelsRefs:       newLbls,
 				Samples:          ts.Samples,

pkg/storage/ingest/version_test.go

@@ -243,30 +243,19 @@ func TestRecordSerializer(t *testing.T) {
 		err = DeserializeRecordContent(record.Value, resultReq, 2)
 		require.NoError(t, err)
 
-		require.Len(t, resultReq.Timeseries, 5)
+		require.Len(t, resultReq.Timeseries, 3)
 		require.Equal(t, req.Timeseries, resultReq.Timeseries[0:2])
 
 		// The only way to carry a metadata in RW2.0 is attached to a timeseries.
 		// Metadata not attached to any series in the request must fabricate extra timeseries to house it.
-		// In format V2 we also avoid the metadata-to-series matching operation to avoid ambiguity in the event multiple series share a family name.
 		expMetadataSeries := []mimirpb.PreallocTimeseries{
-			{TimeSeries: &mimirpb.TimeSeries{
-				Labels:    mimirpb.FromLabelsToLabelAdapters(labels.FromStrings(labels.MetricName, "series_1")),
-				Samples:   []mimirpb.Sample{},
-				Exemplars: []mimirpb.Exemplar{},
-			}},
-			{TimeSeries: &mimirpb.TimeSeries{
-				Labels:    mimirpb.FromLabelsToLabelAdapters(labels.FromStrings(labels.MetricName, "series_2")),
-				Samples:   []mimirpb.Sample{},
-				Exemplars: []mimirpb.Exemplar{},
-			}},
 			{TimeSeries: &mimirpb.TimeSeries{
 				Labels:    mimirpb.FromLabelsToLabelAdapters(labels.FromStrings(labels.MetricName, "series_3")),
 				Samples:   []mimirpb.Sample{},
 				Exemplars: []mimirpb.Exemplar{},
 			}},
 		}
-		require.Equal(t, expMetadataSeries, resultReq.Timeseries[2:5])
+		require.Equal(t, expMetadataSeries, resultReq.Timeseries[2:])
 
 		require.Nil(t, resultReq.SymbolsRW2)
 		require.Nil(t, resultReq.TimeseriesRW2)

pkg/storage/ingest/writer_test.go

@@ -989,27 +989,6 @@ func TestMarshalWriteRequestToRecords(t *testing.T) {
 					Exemplars: []mimirpb.Exemplar{},
 				},
 			},
-			{
-				TimeSeries: &mimirpb.TimeSeries{
-					Labels:    mimirpb.FromLabelsToLabelAdapters(labels.FromStrings("__name__", "series_1")),
-					Samples:   []mimirpb.Sample{},
-					Exemplars: []mimirpb.Exemplar{},
-				},
-			},
-			{
-				TimeSeries: &mimirpb.TimeSeries{
-					Labels:    mimirpb.FromLabelsToLabelAdapters(labels.FromStrings("__name__", "series_2")),
-					Samples:   []mimirpb.Sample{},
-					Exemplars: []mimirpb.Exemplar{},
-				},
-			},
-			{
-				TimeSeries: &mimirpb.TimeSeries{
-					Labels:    mimirpb.FromLabelsToLabelAdapters(labels.FromStrings("__name__", "series_3")),
-					Samples:   []mimirpb.Sample{},
-					Exemplars: []mimirpb.Exemplar{},
-				},
-			},
 		}
 		assert.Equal(t, expTimeseries, actual.Timeseries)
 	})
@@ -1064,7 +1043,7 @@ func TestMarshalWriteRequestToRecords(t *testing.T) {
 
 		records, err := marshalWriteRequestToRecords(1, "user-1", req, limit, splitRequestVersionTwo)
 		require.NoError(t, err)
-		require.Len(t, records, 4)
+		require.Len(t, records, 3)
 
 		// Assert each record, and decode all partial WriteRequests.
 		partials := make([]*mimirpb.WriteRequest, 0, len(records))
@@ -1097,33 +1076,6 @@ func TestMarshalWriteRequestToRecords(t *testing.T) {
 							Exemplars: []mimirpb.Exemplar{},
 						},
 					},
-					{
-						TimeSeries: &mimirpb.TimeSeries{
-							Labels:    mimirpb.FromLabelsToLabelAdapters(labels.FromStrings("__name__", "series_2")),
-							Samples:   []mimirpb.Sample{{TimestampMs: 1, Value: 2.0}},
-							Exemplars: []mimirpb.Exemplar{},
-						},
-					},
-				},
-				Metadata: []*mimirpb.MetricMetadata{},
-			}, {
-				Source:              mimirpb.RULE,
-				SkipLabelValidation: true,
-				Timeseries: []mimirpb.PreallocTimeseries{
-					{
-						TimeSeries: &mimirpb.TimeSeries{
-							Labels:    mimirpb.FromLabelsToLabelAdapters(labels.FromStrings("__name__", "series_3")),
-							Samples:   []mimirpb.Sample{{TimestampMs: 1, Value: 2.0}},
-							Exemplars: []mimirpb.Exemplar{},
-						},
-					},
-					{
-						TimeSeries: &mimirpb.TimeSeries{
-							Labels:    mimirpb.FromLabelsToLabelAdapters(labels.FromStrings("__name__", "series_1")),
-							Samples:   []mimirpb.Sample{},
-							Exemplars: []mimirpb.Exemplar{},
-						},
-					},
 				},
 				Metadata: []*mimirpb.MetricMetadata{
 					{
@@ -1139,7 +1091,7 @@ func TestMarshalWriteRequestToRecords(t *testing.T) {
 					{
 						TimeSeries: &mimirpb.TimeSeries{
 							Labels:    mimirpb.FromLabelsToLabelAdapters(labels.FromStrings("__name__", "series_2")),
-							Samples:   []mimirpb.Sample{},
+							Samples:   []mimirpb.Sample{{TimestampMs: 1, Value: 2.0}},
 							Exemplars: []mimirpb.Exemplar{},
 						},
 					},
@@ -1158,7 +1110,7 @@ func TestMarshalWriteRequestToRecords(t *testing.T) {
 					{
 						TimeSeries: &mimirpb.TimeSeries{
 							Labels:    mimirpb.FromLabelsToLabelAdapters(labels.FromStrings("__name__", "series_3")),
-							Samples:   []mimirpb.Sample{},
+							Samples:   []mimirpb.Sample{{TimestampMs: 1, Value: 2.0}},
 							Exemplars: []mimirpb.Exemplar{},
 						},
 					},