refactor; unit tests; update README

Author: yorugac

README.md

@@ -1,40 +1,37 @@
 # xk6-output-prometheus-remote
 
-K6 extension for Prometheus remote-write output. **WIP!**
+> ### ⚠️ Beta version
+
+k6 extension for Prometheus remote-write output.
 
 *Distinguish from [Prometheus remote write **client** extension](https://github.com/grafana/xk6-client-prometheus-remote) :)*
 
-According to [Prometheus API Stability Guarantees](https://prometheus.io/docs/prometheus/latest/stability/) remote write is an **experimental feature**, thus it is unstable and is subject to change as of now. There are many options for remote-write compatible agents, the official list can be found [here](https://prometheus.io/docs/operating/integrations/). The exact details of how metrics will be processed or stored depends on the underlying agent used.
+According to [Prometheus API Stability Guarantees](https://prometheus.io/docs/prometheus/latest/stability/) remote write is an **experimental feature**, thus it is unstable and is subject to change. There are many options for remote-write compatible agents, the official list can be found [here](https://prometheus.io/docs/operating/integrations/). The exact details of how metrics will be processed or stored depends on the underlying agent used.
 
 Key points to know:
 
 - remote write format does not contain explicit definition of any metric types while metadata definition is still in flux and can have different implementation depending on the remote-write compatible agent
 - remote read is a separate interface and it is much less defined. For example, remote read may not work without precise queries; see [here](https://prometheus.io/docs/prometheus/latest/storage/#remote-storage-integrations) and [here](https://github.com/timescale/promscale/issues/64) for details
 - some remote-write compatible agents may support additional formats for remote write, like JSON, but it is not part of official Prometheus remote write specification and therefore absent here
-- exemplars are not yet officially supported by format, see [here](https://github.com/prometheus/prometheus/issues/9317)
 
 ### Usage
 
 To build k6 binary with the Prometheus remote write output extension use:
-
 ```
-xk6 build --with github.com/grafana/xk6-output-prometheus-remote=. 
+xk6 build --with github.com/grafana/xk6-output-prometheus-remote@latest 
 ```
 
 Then run new k6 binary with:
-
 ```
 K6_PROMETHEUS_REMOTE_URL=http://localhost:9090/api/v1/write ./k6 run script.js -o output-prometheus-remote
 ```
 
 Add TLS and HTTP basic authentication:
-
 ```
 K6_PROMETHEUS_REMOTE_URL=https://localhost:9090/api/v1/write K6_PROMETHEUS_INSECURE_SKIP_TLS_VERIFY=false K6_CA_CERT_FILE=example/tls.crt K6_PROMETHEUS_USER=foo K6_PROMETHEUS_PASSWORD=bar ./k6 run script.js -o output-prometheus-remote
 ```
 
-Different remote storage agents are supported with mapping option. The default is Prometheus itself but there is a simpler raw mapping that can be used as a starting point for other storages:
-
+Different remote storage agents are supported with mapping option. The default is Prometheus itself but there is a simpler raw mapping that can be used as a starting point for other remote agents:
 ```
 K6_PROMETHEUS_MAPPING=raw K6_PROMETHEUS_REMOTE_URL=http://localhost:9090/api/v1/write ./k6 run script.js -o output-prometheus-remote
 ```
@@ -43,25 +40,12 @@ Note: Prometheus remote client relies on a snappy library for serialization whic
 
 ### On sample rate
 
-K6 processes its outputs once per second and that is also a default flush period in this extension. The number of K6 builtin metrics is 26 as of now and they are collected at the rate of 50ms. In practice it means that there will be around 1000-1500 samples on average per each flush period in case of raw mapping. If custom metrics are configured, that estimate will have to be adjusted.
+k6 processes its outputs once per second and that is also a default flush period in this extension. The number of k6 builtin metrics is 26 and they are collected at the rate of 50ms. In practice it means that there will be around 1000-1500 samples on average per each flush period in case of raw mapping. If custom metrics are configured, that estimate will have to be adjusted.
 
-Depending on exact Prometheus setup, it may be necessary to configure Prometheus and / or remote-write agent to handle the load. For example, see [`queue_config` parameter](https://prometheus.io/docs/practices/remote_write/) of Prometheus.
+Depending on exact setup, it may be necessary to configure Prometheus and / or remote-write agent to handle the load. For example, see [`queue_config` parameter](https://prometheus.io/docs/practices/remote_write/) of Prometheus.
+
+If remote endpoint responds too slowly or the k6 test run generates too many metrics, extension may start discarding samples in order to continue to adhere to the flush period.
 
 ### Prometheus as remote-write agent
 
 To enable remote write in Prometheus 2.x use `--enable-feature=remote-write-receiver` option. See docker-compose samples in `example/`. Options for remote write storage can be found [here](https://prometheus.io/docs/operating/integrations/). 
-
-### Next steps
-
-*Note:* this list is meant to keep up to date with current roadmap and status of the extension and is very dynamic.
-
-- [X] support of raw format of Prometheus remote write specification
-- [X] support of Prometheus as remote-write agent
-   - [X] mapping of k6 metrics to Prometheus metrics
-      - [X] Counter
-      - [X] Gauge
-      - [X] Rate
-      - [X] Trend (draft)
-   - [X] ability to switch Prometheus as remote agent on as a configurable option
-   - [ ] decide on submetrics processing
-   - [ ] investigate: the processing loop may need an overhaul with long running tests and cloud setups and / or look into more complex aggregations

pkg/remotewrite/config_test.go

@@ -16,6 +16,8 @@ import (
 )
 
 func TestApply(t *testing.T) {
+	t.Parallel()
+
 	fullConfig := Config{
 		Url:                   null.StringFrom("some-url"),
 		InsecureSkipTLSVerify: null.BoolFrom(false),
@@ -48,6 +50,8 @@ func TestApply(t *testing.T) {
 }
 
 func TestConfigParseArg(t *testing.T) {
+	t.Parallel()
+
 	c, err := ParseArg("url=http://prometheus.remote:3412/write")
 	assert.Nil(t, err)
 	assert.Equal(t, null.StringFrom("http://prometheus.remote:3412/write"), c.Url)

pkg/remotewrite/labels.go

@@ -5,7 +5,7 @@ import (
 	"go.k6.io/k6/stats"
 )
 
-func tagsToPrometheusLabels(tags *stats.SampleTags, config Config) ([]prompb.Label, error) {
+func tagsToLabels(tags *stats.SampleTags, config Config) ([]prompb.Label, error) {
 	if !config.KeepTags.Bool {
 		return []prompb.Label{}, nil
 	}
@@ -22,11 +22,6 @@ func tagsToPrometheusLabels(tags *stats.SampleTags, config Config) ([]prompb.Lab
 			continue
 		}
 
-		// TODO add checks:
-		// - reserved underscore
-		// - sorting
-		// - duplicates?
-
 		labelPairs = append(labelPairs, prompb.Label{
 			Name:  name,
 			Value: value,
@@ -37,7 +32,3 @@ func tagsToPrometheusLabels(tags *stats.SampleTags, config Config) ([]prompb.Lab
 
 	return labelPairs[:len(labelPairs):len(labelPairs)], nil
 }
-
-// func (l labels) Len() int           { return len(l) }
-// func (l labels) Less(i, j int) bool { return l[i].Name < l[j].Name }
-// func (l labels) Swap(i, j int)      { l[i], l[j] = l[j], l[i] }

pkg/remotewrite/labels_test.go

@@ -0,0 +1,88 @@
+package remotewrite
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/prometheus/prometheus/prompb"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.k6.io/k6/stats"
+	"gopkg.in/guregu/null.v3"
+)
+
+func TestTagsToLabels(t *testing.T) {
+	t.Parallel()
+
+	testCases := map[string]struct {
+		tags   *stats.SampleTags
+		config Config
+		labels []prompb.Label
+	}{
+		"empty-tags": {
+			tags: &stats.SampleTags{},
+			config: Config{
+				KeepTags:    null.BoolFrom(true),
+				KeepNameTag: null.BoolFrom(false),
+			},
+			labels: []prompb.Label{},
+		},
+		"name-tag-discard": {
+			tags: stats.NewSampleTags(map[string]string{"foo": "bar", "name": "nnn"}),
+			config: Config{
+				KeepTags:    null.BoolFrom(true),
+				KeepNameTag: null.BoolFrom(false),
+			},
+			labels: []prompb.Label{
+				{Name: "foo", Value: "bar"},
+			},
+		},
+		"name-tag-keep": {
+			tags: stats.NewSampleTags(map[string]string{"foo": "bar", "name": "nnn"}),
+			config: Config{
+				KeepTags:    null.BoolFrom(true),
+				KeepNameTag: null.BoolFrom(true),
+			},
+			labels: []prompb.Label{
+				{Name: "foo", Value: "bar"},
+				{Name: "name", Value: "nnn"},
+			},
+		},
+		"discard-tags": {
+			tags: stats.NewSampleTags(map[string]string{"foo": "bar", "name": "nnn"}),
+			config: Config{
+				KeepTags: null.BoolFrom(false),
+			},
+			labels: []prompb.Label{},
+		},
+	}
+
+	for name, testCase := range testCases {
+		testCase := testCase
+		t.Run(name, func(t *testing.T) {
+			labels, err := tagsToLabels(testCase.tags, testCase.config)
+			require.NoError(t, err)
+
+			assert.Equal(t, len(testCase.labels), len(labels))
+
+			for i := range testCase.labels {
+				var found bool
+
+				// order is not guaranteed ATM
+				for j := range labels {
+					if labels[j].Name == testCase.labels[i].Name {
+						assert.Equal(t, testCase.labels[i].Value, labels[j].Value)
+						found = true
+						break
+					}
+
+				}
+				if !found {
+					assert.Fail(t, fmt.Sprintf("Not found label %s: \n"+
+						"expected: %v\n"+
+						"actual  : %v", testCase.labels[i].Name, testCase.labels, labels))
+				}
+			}
+		})
+	}
+}

pkg/remotewrite/metrics.go

@@ -10,6 +10,8 @@ import (
 
 // Note: k6 Registry is not used here since Output is getting
 // samples only from k6 engine, hence we assume they are already vetted.
+
+// metricsStorage is an in-memory gather point for metrics
 type metricsStorage struct {
 	m map[string]stats.Sample
 }
@@ -21,7 +23,7 @@ func newMetricsStorage() *metricsStorage {
 }
 
 // update modifies metricsStorage and returns updated sample
-// so that they hold the same value for the given metric
+// so that the stored metric and the returned metric hold the same value
 func (ms *metricsStorage) update(sample stats.Sample, add func(current, s stats.Sample) stats.Sample) stats.Sample {
 	if current, ok := ms.m[sample.Metric.Name]; ok {
 		if add == nil {

pkg/remotewrite/prometheus.go

@@ -73,39 +73,7 @@ func (pm *PrometheusMapping) MapRate(ms *metricsStorage, sample stats.Sample, la
 }
 
 func (pm *PrometheusMapping) MapTrend(ms *metricsStorage, sample stats.Sample, labels []prompb.Label) []prompb.TimeSeries {
-	sample = ms.update(sample, func(current, s stats.Sample) stats.Sample {
-		// this is an attempt to optimize what happens in TrendSink;
-		// partial copy-paste from k6/stats
-		t := current.Metric.Sink.(*stats.TrendSink)
-
-		index := sort.Search(len(t.Values), func(i int) bool {
-			return t.Values[i] > s.Value
-		})
-		t.Values = append(t.Values, 0)
-		copy(t.Values[index+1:], t.Values[index:])
-		t.Values[index] = s.Value
-
-		t.Count += 1
-		t.Sum += s.Value
-		t.Avg = t.Sum / float64(t.Count)
-
-		if s.Value > t.Max {
-			t.Max = s.Value
-		}
-		if s.Value < t.Min || t.Count == 1 {
-			t.Min = s.Value
-		}
-
-		// The median of an even number of values is the average of the middle two.
-		if (t.Count & 0x01) == 0 {
-			t.Med = (t.Values[(t.Count/2)-1] + t.Values[(t.Count/2)]) / 2
-		} else {
-			t.Med = t.Values[t.Count/2]
-		}
-
-		current.Metric.Sink = t
-		return current
-	})
+	sample = ms.update(sample, trendAdd)
 
 	s := sample.Metric.Sink.(*stats.TrendSink)
 	aggr := map[string]float64{
@@ -197,7 +165,42 @@ func (pm *PrometheusMapping) MapTrend(ms *metricsStorage, sample stats.Sample, l
 	}
 }
 
-// Note: this is a copy-paste from k6/stats/sink.go but without additional call to Calc
+// The following functions are an attempt to add ad-hoc optimization to TrendSink,
+// and are a partial copy-paste from k6/stats.
+// TODO: re-write & refactor this once metrics refactoring progresses in k6.
+
+func trendAdd(current, s stats.Sample) stats.Sample {
+	t := current.Metric.Sink.(*stats.TrendSink)
+
+	// insert into sorted array instead of sorting anew on each addition
+	index := sort.Search(len(t.Values), func(i int) bool {
+		return t.Values[i] > s.Value
+	})
+	t.Values = append(t.Values, 0)
+	copy(t.Values[index+1:], t.Values[index:])
+	t.Values[index] = s.Value
+
+	t.Count += 1
+	t.Sum += s.Value
+	t.Avg = t.Sum / float64(t.Count)
+
+	if s.Value > t.Max {
+		t.Max = s.Value
+	}
+	if s.Value < t.Min || t.Count == 1 {
+		t.Min = s.Value
+	}
+
+	if (t.Count & 0x01) == 0 {
+		t.Med = (t.Values[(t.Count/2)-1] + t.Values[(t.Count/2)]) / 2
+	} else {
+		t.Med = t.Values[t.Count/2]
+	}
+
+	current.Metric.Sink = t
+	return current
+}
+
 func p(t *stats.TrendSink, pct float64) float64 {
 	switch t.Count {
 	case 0: