Use directly metrics.TimeSeries

Replace the internal time series data model with the new defined by k6.

Author: codebien

pkg/remotewrite/prometheus.go

@@ -4,21 +4,26 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/mstoykov/atlas"
 	prompb "go.buf.build/grpc/go/prometheus/prometheus"
 	"go.k6.io/k6/metrics"
 )
 
-func MapTagSet(t *metrics.SampleTags) []*prompb.Label {
-	tags := t.CloneTags()
-
-	labels := make([]*prompb.Label, 0, len(tags))
-	for k, v := range tags {
-		labels = append(labels, &prompb.Label{Name: k, Value: v})
+func MapTagSet(t *metrics.TagSet) []*prompb.Label {
+	n := (*atlas.Node)(t)
+	if n.Len() < 1 {
+		return nil
+	}
+	labels := make([]*prompb.Label, 0, n.Len())
+	for !n.IsRoot() {
+		prev, key, value := n.Data()
+		labels = append(labels, &prompb.Label{Name: key, Value: value})
+		n = prev
 	}
 	return labels
 }
 
-func MapSeries(ts TimeSeries) prompb.TimeSeries {
+func MapSeries(ts metrics.TimeSeries) prompb.TimeSeries {
 	return prompb.TimeSeries{
 		Labels: append(MapTagSet(ts.Tags), &prompb.Label{
 			Name:  "__name__",
@@ -27,7 +32,7 @@ func MapSeries(ts TimeSeries) prompb.TimeSeries {
 	}
 }
 
-func MapTrend(series TimeSeries, t time.Time, sink *trendSink) []*prompb.TimeSeries {
+func MapTrend(series metrics.TimeSeries, t time.Time, sink *trendSink) []*prompb.TimeSeries {
 	// Prometheus metric system does not support Trend so this mapping will
 	// store a counter for the number of reported values and gauges to keep
 	// track of aggregated values. Also store a sum of the values to allow

pkg/remotewrite/prometheus_test.go

@@ -55,11 +55,12 @@ func assertTimeSeriesEqual(t *testing.T, expected *prompb.TimeSeries, actual *pr
 
 // sortTimeSeries sorts an array of TimeSeries by name
 func sortTimeSeries(ts []*prompb.TimeSeries) []*prompb.TimeSeries {
-	sorted := make([]*prompb.TimeSeries, 0, len(ts))
+	sorted := make([]*prompb.TimeSeries, len(ts))
 	copy(sorted, ts)
 	sort.Slice(sorted, func(i int, j int) bool {
 		return getTimeSeriesName(sorted[i]) < getTimeSeriesName(sorted[j])
 	})
+
 	return sorted
 }
 
@@ -83,6 +84,7 @@ func TestMapTrend(t *testing.T) {
 	t.Parallel()
 
 	now := time.Now()
+	r := metrics.NewRegistry()
 
 	testCases := []struct {
 		sample   metrics.Sample
@@ -91,11 +93,13 @@ func TestMapTrend(t *testing.T) {
 	}{
 		{
 			sample: metrics.Sample{
-				Metric: &metrics.Metric{
-					Name: "test",
-					Type: metrics.Trend,
+				TimeSeries: metrics.TimeSeries{
+					Metric: &metrics.Metric{
+						Name: "test",
+						Type: metrics.Trend,
+					},
+					Tags: r.RootTagSet().With("tagk1", "tagv1"),
 				},
-				Tags:  metrics.NewSampleTags(map[string]string{"tagk1": "tagv1"}),
 				Value: 1.0,
 				Time:  now,
 			},
@@ -116,7 +120,7 @@ func TestMapTrend(t *testing.T) {
 		st := &trendSink{}
 		st.Add(tc.sample)
 
-		ts := MapTrend(TimeSeries{tc.sample.Metric, tc.sample.Tags}, tc.sample.Time, st)
+		ts := MapTrend(tc.sample.TimeSeries, tc.sample.Time, st)
 		assertTimeSeriesMatch(t, tc.expected, ts)
 	}
 }

pkg/remotewrite/remotewrite.go

@@ -3,8 +3,6 @@ package remotewrite
 import (
 	"context"
 	"fmt"
-	"sort"
-	"strings"
 	"time"
 
 	"github.com/grafana/xk6-output-prometheus-remote/pkg/remote"
@@ -24,7 +22,7 @@ type Output struct {
 	config          Config
 	logger          logrus.FieldLogger
 	periodicFlusher *output.PeriodicFlusher
-	tsdb            map[string]*seriesWithMeasure
+	tsdb            map[metrics.TimeSeries]*seriesWithMeasure
 
 	// TODO: copy the prometheus/remote.WriteClient interface and depend on it
 	client *remote.WriteClient
@@ -45,14 +43,14 @@ func New(params output.Params) (*Output, error) {
 
 	wc, err := remote.NewWriteClient(config.URL.String, clientConfig)
 	if err != nil {
-		return nil, fmt.Errorf("failed to initialize the client Prometheus remote write client: %w", err)
+		return nil, fmt.Errorf("failed to initialize the Prometheus remote write client: %w", err)
 	}
 
 	return &Output{
 		client: wc,
 		config: config,
 		logger: logger,
-		tsdb:   make(map[string]*seriesWithMeasure),
+		tsdb:   make(map[metrics.TimeSeries]*seriesWithMeasure),
 	}, nil
 }
 
@@ -131,33 +129,29 @@ func (o *Output) convertToPbSeries(samplesContainers []metrics.SampleContainer)
 	// so we need to aggregate all the samples in the same time bucket.
 	// More context can be found in the issue
 	// https://github.com/grafana/xk6-output-prometheus-remote/issues/11
-	seen := make(map[string]struct{})
+	seen := make(map[metrics.TimeSeries]struct{})
 
 	for _, samplesContainer := range samplesContainers {
 		samples := samplesContainer.GetSamples()
 
 		for _, sample := range samples {
 			truncTime := sample.Time.Truncate(time.Millisecond)
-			timeSeriesKey := timeSeriesKey(sample.Metric, sample.Tags)
-			swm, ok := o.tsdb[timeSeriesKey]
+			swm, ok := o.tsdb[sample.TimeSeries]
 			if !ok {
 				swm = &seriesWithMeasure{
-					TimeSeries: TimeSeries{
-						Metric: sample.Metric,
-						Tags:   sample.Tags,
-					},
-					Measure: sinkByType(sample.Metric.Type),
-					Latest:  truncTime,
+					TimeSeries: sample.TimeSeries,
+					Measure:    sinkByType(sample.Metric.Type),
+					Latest:     truncTime,
 				}
-				o.tsdb[timeSeriesKey] = swm
-				seen[timeSeriesKey] = struct{}{}
+				o.tsdb[sample.TimeSeries] = swm
+				seen[sample.TimeSeries] = struct{}{}
 			} else {
 				// save as a seen item only when the samples have a time greater than
 				// the previous saved, otherwise some implementations
 				// could see it as a duplicate and generate warnings (e.g. Mimir)
 				if truncTime.After(swm.Latest) {
 					swm.Latest = truncTime
-					seen[timeSeriesKey] = struct{}{}
+					seen[sample.TimeSeries] = struct{}{}
 				}
 
 				// If current == previous:
@@ -191,6 +185,7 @@ func (o *Output) convertToPbSeries(samplesContainers []metrics.SampleContainer)
 }
 
 type seriesWithMeasure struct {
+	metrics.TimeSeries
 	Measure metrics.Sink
 
 	// Latest tracks the latest time
@@ -200,17 +195,13 @@ type seriesWithMeasure struct {
 	// in a method in struct
 	Latest time.Time
 
-	// TimeSeries will be replaced with the native k6 version
-	// when it will be available.
-	TimeSeries
-
 	// TODO: maybe add some caching for the mapping?
 }
 
 func (swm seriesWithMeasure) MapPrompb() []*prompb.TimeSeries {
 	var newts []*prompb.TimeSeries
 
-	mapMonoSeries := func(s TimeSeries, t time.Time) prompb.TimeSeries {
+	mapMonoSeries := func(s metrics.TimeSeries, t time.Time) prompb.TimeSeries {
 		return prompb.TimeSeries{
 			Labels: append(MapTagSet(swm.Tags), &prompb.Label{
 				Name:  "__name__",
@@ -266,36 +257,3 @@ func sinkByType(mt metrics.MetricType) metrics.Sink {
 	}
 	return sink
 }
-
-// the code below will be removed
-// when TimeSeries will be a native k6's concept.
-
-type TimeSeries struct {
-	Metric *metrics.Metric
-	Tags   *metrics.SampleTags
-}
-
-var bytesep = []byte{0xff}
-
-func timeSeriesKey(m *metrics.Metric, sampleTags *metrics.SampleTags) string {
-	if sampleTags.IsEmpty() {
-		return m.Name
-	}
-
-	tmap := sampleTags.CloneTags()
-	keys := make([]string, 0, len(tmap))
-	for k := range tmap {
-		keys = append(keys, k)
-	}
-	sort.Strings(keys)
-
-	var b strings.Builder
-	b.WriteString(m.Name)
-	for i := 0; i < len(keys); i++ {
-		b.Write(bytesep)
-		b.WriteString(keys[i])
-		b.Write(bytesep)
-		b.WriteString(tmap[keys[i]])
-	}
-	return b.String()
-}

pkg/remotewrite/remotewrite_test.go

@@ -1,7 +1,6 @@
 package remotewrite
 
 import (
-	"strconv"
 	"testing"
 	"time"
 
@@ -28,31 +27,37 @@ func TestOutputConvertToPbSeries(t *testing.T) {
 
 	registry := metrics.NewRegistry()
 	metric1 := registry.MustNewMetric("metric1", metrics.Counter)
-	tagset := metrics.NewSampleTags(map[string]string{"tagk1": "tagv1"})
+	tagset := registry.RootTagSet().With("tagk1", "tagv1")
 
 	samples := []metrics.SampleContainer{
 		metrics.Sample{
-			Metric: metric1,
-			Tags:   tagset,
-			Time:   time.Date(2022, time.September, 1, 0, 0, 0, 0, time.UTC),
-			Value:  3,
+			TimeSeries: metrics.TimeSeries{
+				Metric: metric1,
+				Tags:   tagset,
+			},
+			Time:  time.Date(2022, time.September, 1, 0, 0, 0, 0, time.UTC),
+			Value: 3,
 		},
 		metrics.Sample{
-			Metric: metric1,
-			Tags:   tagset,
-			Time:   time.Date(2022, time.August, 31, 0, 0, 0, 0, time.UTC),
-			Value:  4,
+			TimeSeries: metrics.TimeSeries{
+				Metric: metric1,
+				Tags:   tagset,
+			},
+			Time:  time.Date(2022, time.August, 31, 0, 0, 0, 0, time.UTC),
+			Value: 4,
 		},
 		metrics.Sample{
-			Metric: registry.MustNewMetric("metric2", metrics.Counter),
-			Tags:   tagset,
-			Time:   time.Date(2022, time.September, 1, 0, 0, 0, 0, time.UTC),
-			Value:  2,
+			TimeSeries: metrics.TimeSeries{
+				Metric: registry.MustNewMetric("metric2", metrics.Counter),
+				Tags:   tagset,
+			},
+			Time:  time.Date(2022, time.September, 1, 0, 0, 0, 0, time.UTC),
+			Value: 2,
 		},
 	}
 
 	o := Output{
-		tsdb: make(map[string]*seriesWithMeasure),
+		tsdb: make(map[metrics.TimeSeries]*seriesWithMeasure),
 	}
 
 	pbseries := o.convertToPbSeries(samples)
@@ -89,11 +94,11 @@ func TestOutputConvertToPbSeries_WithPreviousState(t *testing.T) {
 
 	registry := metrics.NewRegistry()
 	metric1 := registry.MustNewMetric("metric1", metrics.Counter)
-	tagset := metrics.NewSampleTags(map[string]string{"tagk1": "tagv1"})
+	tagset := registry.RootTagSet().With("tagk1", "tagv1")
 	t0 := time.Date(2022, time.September, 1, 0, 0, 0, 0, time.UTC).Add(10 * time.Millisecond)
 
 	swm := &seriesWithMeasure{
-		TimeSeries: TimeSeries{
+		TimeSeries: metrics.TimeSeries{
 			Metric: metric1,
 			Tags:   tagset,
 		},
@@ -103,8 +108,8 @@ func TestOutputConvertToPbSeries_WithPreviousState(t *testing.T) {
 	}
 
 	o := Output{
-		tsdb: map[string]*seriesWithMeasure{
-			timeSeriesKey(metric1, tagset): swm,
+		tsdb: map[metrics.TimeSeries]*seriesWithMeasure{
+			swm.TimeSeries: swm,
 		},
 	}
 
@@ -142,10 +147,12 @@ func TestOutputConvertToPbSeries_WithPreviousState(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			pbseries := o.convertToPbSeries([]metrics.SampleContainer{
 				metrics.Sample{
-					Metric: metric1,
-					Tags:   tagset,
-					Value:  1,
-					Time:   tc.time,
+					TimeSeries: metrics.TimeSeries{
+						Metric: metric1,
+						Tags:   tagset,
+					},
+					Value: 1,
+					Time:  tc.time,
 				},
 			})
 			require.Len(t, o.tsdb, 1)
@@ -155,28 +162,3 @@ func TestOutputConvertToPbSeries_WithPreviousState(t *testing.T) {
 		})
 	}
 }
-
-func TestTimeSeriesKey(t *testing.T) {
-	t.Parallel()
-
-	registry := metrics.NewRegistry()
-	metric1 := registry.MustNewMetric("metric1", metrics.Counter)
-
-	tagsmap := make(map[string]string)
-	for i := 0; i < 8; i++ {
-		is := strconv.Itoa(i)
-		tagsmap["tagk"+is] = "tagv" + is
-	}
-	tagset := metrics.NewSampleTags(tagsmap)
-
-	key := timeSeriesKey(metric1, tagset)
-
-	expected := "metric1"
-	sbytesep := string(bytesep)
-	for i := 0; i < 8; i++ {
-		is := strconv.Itoa(i)
-		expected += sbytesep + "tagk" + is + sbytesep + "tagv" + is
-	}
-
-	assert.Equal(t, expected, key)
-}