Start implementing direct unmarshal from RW2 into mimirpb.Timeseries

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

Author: krajorama

pkg/distributor/distributor_test.go

@@ -8543,7 +8543,7 @@ func cloneTimeseries(orig *mimirpb.TimeSeries) (*mimirpb.TimeSeries, error) {
 	}
 
 	cloned := &mimirpb.TimeSeries{}
-	err = cloned.Unmarshal(data)
+	err = cloned.Unmarshal(data, nil)
 	return cloned, err
 }
 

pkg/ingester/client/ingester.pb.go

@@ -0,0 +1,51 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+
+package mimirpb
+
+import (
+	"testing"
+
+	rw2 "github.com/prometheus/prometheus/prompb/io/prometheus/write/v2"
+	"github.com/stretchr/testify/require"
+
+	rw2util "github.com/grafana/mimir/pkg/util/test"
+)
+
+func TestXxx(t *testing.T) {
+	syms := rw2util.NewSymbolTableBuilder(nil)
+	
+	// Create a new WriteRequest with some sample data.
+	writeRequest := &rw2.Request{
+		Timeseries: []rw2.TimeSeries{
+			{
+				LabelsRefs: []uint32{syms.GetSymbol("__name__"), syms.GetSymbol("test_metric"), syms.GetSymbol("job"), syms.GetSymbol("test_job")},
+				Samples: []rw2.Sample{
+					{
+						Value:     123.456,
+						Timestamp: 1234567890,
+					},
+				},
+			},
+		},
+	}
+	writeRequest.Symbols = syms.GetSymbols()
+	data, err := writeRequest.Marshal()
+	require.NoError(t, err)
+
+	// Unmarshal the data back into Mimir's WriteRequest.
+	received := PreallocWriteRequest{}
+	received.UnmarshalFromRW2 = true
+	err = received.Unmarshal(data)
+	require.NoError(t, err)
+
+	// Check that the unmarshalled data matches the original data.
+	require.Len(t,received.Timeseries, 1)
+	require.Len(t, received.Timeseries[0].TimeSeries.Labels, 2)
+	require.Equal(t, received.Timeseries[0].TimeSeries.Labels[0].Name, "__name__")
+	require.Equal(t, received.Timeseries[0].TimeSeries.Labels[0].Value, "test_metric")
+	require.Equal(t, received.Timeseries[0].TimeSeries.Labels[1].Name, "job")
+	require.Equal(t, received.Timeseries[0].TimeSeries.Labels[1].Value, "test_job")
+	require.Len(t, received.Timeseries[0].TimeSeries.Samples, 1)
+	require.Equal(t, received.Timeseries[0].TimeSeries.Samples[0].Value, 123.456)
+	require.Equal(t, received.Timeseries[0].TimeSeries.Samples[0].TimestampMs, 1234567890)
+}

pkg/mimirpb/compat_rw2_test.go

@@ -74,92 +74,11 @@ type PreallocWriteRequest struct {
 func (p *PreallocWriteRequest) Unmarshal(dAtA []byte) error {
 	p.Timeseries = PreallocTimeseriesSliceFromPool()
 	p.WriteRequest.skipUnmarshalingExemplars = p.SkipUnmarshalingExemplars
-
-	if p.UnmarshalFromRW2 {
-		return p.unmarshalRW2(dAtA)
-	}
+	p.WriteRequest.UnmarshalFromRW2 = p.UnmarshalFromRW2
 
 	return p.WriteRequest.Unmarshal(dAtA)
 }
 
-// UnmarshalRW2 unmarshals the given remote write 2.0 data and converts it to a WriteRequest.
-func (p *PreallocWriteRequest) unmarshalRW2(data []byte) error {
-	rw2req := &WriteRequestRW2{}
-	rw2req.skipUnmarshalingExemplars = p.SkipUnmarshalingExemplars
-	if err := rw2req.Unmarshal(data); err != nil {
-		return err
-	}
-
-	metricFamilies := map[string]*MetricMetadata{}
-
-	for _, ts := range rw2req.Timeseries {
-		p.Timeseries = append(p.Timeseries, PreallocTimeseries{})
-		p.Timeseries[len(p.Timeseries)-1].TimeSeries = TimeseriesFromPool()
-		p.Timeseries[len(p.Timeseries)-1].TimeSeries.CreatedTimestamp = ts.CreatedTimestamp
-		var err error
-		p.Timeseries[len(p.Timeseries)-1].TimeSeries.Labels, err = labelRefsToLabelAdapter(ts.LabelsRefs, rw2req.Symbols)
-		if err != nil {
-			return err
-		}
-		p.Timeseries[len(p.Timeseries)-1].TimeSeries.Samples = ts.Samples
-		p.Timeseries[len(p.Timeseries)-1].TimeSeries.Histograms = ts.Histograms
-		if len(ts.Exemplars) > 0 {
-			p.Timeseries[len(p.Timeseries)-1].TimeSeries.Exemplars = make([]Exemplar, 0, len(ts.Exemplars))
-			for i := range ts.Exemplars {
-				lbls, err := labelRefsToLabelAdapter(ts.Exemplars[i].LabelsRefs, rw2req.Symbols)
-				if err != nil {
-					return err
-				}
-				p.Timeseries[len(p.Timeseries)-1].TimeSeries.Exemplars = append(p.Timeseries[len(p.Timeseries)-1].TimeSeries.Exemplars, Exemplar{
-					Labels:      lbls,
-					Value:       ts.Exemplars[i].Value,
-					TimestampMs: ts.Exemplars[i].Timestamp,
-				})
-			}
-		}
-
-		// Convert RW2 metadata to RW1 metadata.
-		seriesName := getSeriesName(p.Timeseries[len(p.Timeseries)-1].TimeSeries.Labels)
-		if seriesName == "" {
-			continue
-		}
-		metricFamily, _ := getMetricName(seriesName, ts.Metadata.Type)
-		if metricFamily == "" {
-			continue
-		}
-		help, _ := getSymbol(ts.Metadata.HelpRef, rw2req.Symbols)
-		unit, _ := getSymbol(ts.Metadata.UnitRef, rw2req.Symbols)
-
-		if ts.Metadata.Type == METRIC_TYPE_UNSPECIFIED && help == "" && unit == "" {
-			// Nothing to do here.
-			continue
-		}
-
-		metricFamilies[metricFamily] = &MetricMetadata{
-			Type:             MetricMetadata_MetricType(ts.Metadata.Type),
-			MetricFamilyName: metricFamily,
-			Help:             help,
-			Unit:             unit,
-		}
-	}
-
-	// Fill the metadata
-	p.Metadata = make([]*MetricMetadata, 0, len(metricFamilies))
-	for _, metadata := range metricFamilies {
-		p.Metadata = append(p.Metadata, metadata)
-	}
-
-	return nil
-}
-
-// getSymbol resolves the symbol reference to a string.
-func getSymbol(ref uint32, symbols []string) (string, error) {
-	if ref < uint32(len(symbols)) {
-		return symbols[ref], nil
-	}
-	return "", fmt.Errorf("symbol reference %d is out of bounds", ref)
-}
-
 // labelRefsToLabelAdapter converts a slice of label references to a slice
 // of LabelAdapter.
 func labelRefsToLabelAdapter(refs []uint32, symbols []string) ([]LabelAdapter, error) {
@@ -168,11 +87,11 @@ func labelRefsToLabelAdapter(refs []uint32, symbols []string) ([]LabelAdapter, e
 	}
 	labels := make([]LabelAdapter, 0, len(refs)/2)
 	for i := 0; i < len(refs); i += 2 {
-		name, err := getSymbol(refs[i], symbols)
+		name, err := getRW2Symbol(refs[i], symbols)
 		if err != nil {
 			return nil, err
 		}
-		value, err := getSymbol(refs[i+1], symbols)
+		value, err := getRW2Symbol(refs[i+1], symbols)
 		if err != nil {
 			return nil, err
 		}
@@ -254,6 +173,7 @@ type PreallocTimeseries struct {
 	marshalledData []byte
 
 	skipUnmarshalingExemplars bool
+	unmarshalFromRW2 	      bool
 }
 
 // RemoveLabel removes the label labelName from this timeseries, if it exists.
@@ -376,13 +296,15 @@ var TimeseriesUnmarshalCachingEnabled = true
 // Unmarshal implements proto.Message. Input data slice is retained.
 // Copied from the protobuf generated code, the only change is that in case 3 the exemplars don't get unmarshaled
 // if p.skipUnmarshalingExemplars is false.
-func (p *PreallocTimeseries) Unmarshal(dAtA []byte) error {
-	if TimeseriesUnmarshalCachingEnabled {
+func (p *PreallocTimeseries) Unmarshal(dAtA []byte, symbols []string) error {
+	if TimeseriesUnmarshalCachingEnabled && !p.unmarshalFromRW2 {
+		// TODO(krajorama): check if it makes sense for RW2 as well.
 		p.marshalledData = dAtA
 	}
 	p.TimeSeries = TimeseriesFromPool()
 	p.TimeSeries.SkipUnmarshalingExemplars = p.skipUnmarshalingExemplars
-	return p.TimeSeries.Unmarshal(dAtA)
+	p.TimeSeries.UnmarshalFromRW2 = p.unmarshalFromRW2
+	return p.TimeSeries.Unmarshal(dAtA, symbols)
 }
 
 func (p *PreallocTimeseries) Size() int {