Removed the prometheus/prometheus go library

It has been replaced with a lighter version based on buf.build

Author: codebien

README.md

@@ -32,8 +32,6 @@ K6_PROMETHEUS_PASSWORD=bar \
 ./k6 run script.js -o output-prometheus-remote
 ```
 
-Note: Prometheus remote client relies on a snappy library for serialization which could panic on [encode operations](https://github.com/golang/snappy/blob/544b4180ac705b7605231d4a4550a1acb22a19fe/encode.go#L22).
-
 ### 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 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.

go.mod

@@ -1,58 +1,36 @@
 module github.com/grafana/xk6-output-prometheus-remote
 
-go 1.17
+go 1.18
 
 require (
-	github.com/golang/protobuf v1.5.2
+	github.com/gogo/protobuf v1.3.2
 	github.com/golang/snappy v0.0.4
 	github.com/kubernetes/helm v2.17.0+incompatible
-	github.com/prometheus/common v0.32.1
-	github.com/prometheus/prometheus v1.8.2-0.20211005150130-f29caccc4255
 	github.com/sirupsen/logrus v1.8.1
 	github.com/stretchr/testify v1.8.0
+	go.buf.build/grpc/go/prometheus/prometheus v1.4.3
 	go.k6.io/k6 v0.40.0
 	gopkg.in/guregu/null.v3 v3.5.0
 )
 
 require (
-	github.com/alecthomas/units v0.0.0-20210208195552-ff826a37aa15 // indirect
-	github.com/aws/aws-sdk-go v1.40.37 // indirect
-	github.com/beorn7/perks v1.0.1 // indirect
-	github.com/cespare/xxhash/v2 v2.1.2 // indirect
-	github.com/cncf/xds/go v0.0.0-20211011173535-cb28da3451f1 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
-	github.com/dennwc/varint v1.0.0 // indirect
-	github.com/envoyproxy/go-control-plane v0.9.10-0.20210907150352-cf90f659a021 // indirect
 	github.com/fatih/color v1.13.0 // indirect
 	github.com/ghodss/yaml v1.0.0 // indirect
-	github.com/go-kit/log v0.1.0 // indirect
-	github.com/go-logfmt/logfmt v0.5.1 // indirect
-	github.com/gogo/protobuf v1.3.2 // indirect
-	github.com/jmespath/go-jmespath v0.4.0 // indirect
-	github.com/jpillora/backoff v1.0.0 // indirect
+	github.com/google/go-cmp v0.5.6 // indirect
+	github.com/kr/text v0.2.0 // indirect
 	github.com/mailru/easyjson v0.7.7 // indirect
 	github.com/mattn/go-colorable v0.1.12 // indirect
 	github.com/mattn/go-isatty v0.0.14 // indirect
-	github.com/matttproud/golang_protobuf_extensions v1.0.2-0.20181231171920-c182affec369 // indirect
-	github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f // indirect
-	github.com/opentracing-contrib/go-stdlib v1.0.0 // indirect
-	github.com/opentracing/opentracing-go v1.2.0 // indirect
 	github.com/oxtoacart/bpool v0.0.0-20190530202638-03653db5a59c // indirect
-	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
-	github.com/prometheus/client_golang v1.11.0 // indirect
-	github.com/prometheus/client_model v0.2.0 // indirect
-	github.com/prometheus/common/sigv4 v0.1.0 // indirect
-	github.com/prometheus/procfs v0.6.0 // indirect
 	github.com/spf13/afero v1.3.4 // indirect
-	go.uber.org/atomic v1.9.0 // indirect
-	golang.org/x/net v0.0.0-20220630215102-69896b714898 // indirect
-	golang.org/x/oauth2 v0.0.0-20210819190943-2bc19b11175f // indirect
+	go.buf.build/grpc/go/gogo/protobuf v1.4.9 // indirect
 	golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a // indirect
 	golang.org/x/text v0.3.7 // indirect
 	golang.org/x/time v0.0.0-20220609170525-579cf78fd858 // indirect
-	google.golang.org/appengine v1.6.7 // indirect
-	google.golang.org/protobuf v1.27.1 // indirect
+	google.golang.org/genproto v0.0.0-20210903162649-d08c68adba83 // indirect
+	google.golang.org/protobuf v1.28.1 // indirect
 	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect

go.sum

@@ -0,0 +1,113 @@
+package remote
+
+import (
+	"bytes"
+	"context"
+	"crypto/tls"
+	"fmt"
+	"io"
+	"net/http"
+	"net/url"
+	"time"
+
+	"github.com/gogo/protobuf/proto"
+	"github.com/golang/snappy"
+	prompb "go.buf.build/grpc/go/prometheus/prometheus"
+)
+
+type HTTPConfig struct {
+	Timeout   time.Duration
+	TLSConfig *tls.Config
+	BasicAuth *BasicAuth
+	Headers   http.Header
+}
+
+type BasicAuth struct {
+	Username, Password string
+}
+
+// WriteClient is a client implementation of the Prometheus remote write protocol.
+// It follows the specs defined by the official design document:
+// https://docs.google.com/document/d/1LPhVRSFkGNSuU1fBd81ulhsCPR4hkSZyyBj1SZ8fWOM
+type WriteClient struct {
+	hc  *http.Client
+	url *url.URL
+	cfg *HTTPConfig
+}
+
+func NewWriteClient(endpoint string, cfg *HTTPConfig) (*WriteClient, error) {
+	if cfg == nil {
+		cfg = &HTTPConfig{}
+	}
+	u, err := url.Parse(endpoint)
+	if err != nil {
+		return nil, err
+	}
+	wc := &WriteClient{
+		hc: &http.Client{
+			Timeout: cfg.Timeout,
+		},
+		url: u,
+		cfg: cfg,
+	}
+	if cfg.TLSConfig != nil {
+		wc.hc.Transport = &http.Transport{
+			TLSClientConfig: cfg.TLSConfig,
+		}
+	}
+	return wc, nil
+}
+
+// Store sends a batch of samples to the HTTP endpoint,
+// the request is the proto marshaled and encoded.
+func (c *WriteClient) Store(ctx context.Context, series []*prompb.TimeSeries) error {
+	b, err := newWriteRequestBody(series)
+	if err != nil {
+		return err
+	}
+	req, err := http.NewRequestWithContext(
+		ctx, http.MethodPost, c.url.String(), bytes.NewReader(b))
+	if err != nil {
+		return fmt.Errorf("create new HTTP request failed: %w", err)
+	}
+	if c.cfg.BasicAuth != nil {
+		req.SetBasicAuth(c.cfg.BasicAuth.Username, c.cfg.BasicAuth.Password)
+	}
+
+	if len(c.cfg.Headers) > 0 {
+		req.Header = c.cfg.Headers.Clone()
+	}
+
+	req.Header.Set("User-Agent", "k6-prometheus-rw-output")
+
+	// They are mostly defined by the specs
+	req.Header.Set("Content-Encoding", "snappy")
+	req.Header.Set("Content-Type", "application/x-protobuf")
+	req.Header.Set("X-Prometheus-Remote-Write-Version", "0.1.0")
+
+	resp, err := c.hc.Do(req)
+	if err != nil {
+		return fmt.Errorf("HTTP POST request failed: %w", err)
+	}
+	defer resp.Body.Close()
+	io.Copy(io.Discard, resp.Body) //nolint:errcheck
+
+	if resp.StatusCode != http.StatusNoContent {
+		return fmt.Errorf("got status code: %s instead expected: 204 No Content", resp.Status)
+	}
+	return nil
+}
+
+func newWriteRequestBody(series []*prompb.TimeSeries) ([]byte, error) {
+	b, err := proto.Marshal(&prompb.WriteRequest{
+		Timeseries: series,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("encoding series as protobuf write request failed: %w", err)
+	}
+	if snappy.MaxEncodedLen(len(b)) < 0 {
+		return nil, fmt.Errorf("the protobuf message is too large to be handled by Snappy encoder; "+
+			"size: %d, limit: %d", len(b), 0xffffffff)
+	}
+	return snappy.Encode(nil, b), nil
+}

pkg/remote/client.go

@@ -0,0 +1,172 @@
+package remote
+
+import (
+	"context"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	prompb "go.buf.build/grpc/go/prometheus/prometheus"
+)
+
+func TestNewWrtiteClient(t *testing.T) {
+	t.Parallel()
+	t.Run("DefaultConfig", func(t *testing.T) {
+		t.Parallel()
+		wc, err := NewWriteClient("http://example.com/api/v1/write", nil)
+		require.NoError(t, err)
+		require.NotNil(t, wc)
+		assert.Equal(t, wc.cfg, &HTTPConfig{})
+	})
+
+	t.Run("CustomConfig", func(t *testing.T) {
+		t.Parallel()
+		hc := &HTTPConfig{Timeout: time.Second}
+		wc, err := NewWriteClient("http://example.com/api/v1/write", hc)
+		require.NoError(t, err)
+		require.NotNil(t, wc)
+		assert.Equal(t, wc.cfg, hc)
+	})
+
+	t.Run("InvalidURL", func(t *testing.T) {
+		t.Parallel()
+		wc, err := NewWriteClient("fake://bad url", nil)
+		require.Error(t, err)
+		assert.Nil(t, wc)
+	})
+}
+
+func TestClientStore(t *testing.T) {
+	t.Parallel()
+	h := func(rw http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, r.Header.Get("Content-Encoding"), "snappy")
+		assert.Equal(t, r.Header.Get("Content-Type"), "application/x-protobuf")
+		assert.Equal(t, r.Header.Get("User-Agent"), "k6-prometheus-rw-output")
+		assert.Equal(t, r.Header.Get("X-Prometheus-Remote-Write-Version"), "0.1.0")
+		assert.NotEmpty(t, r.Header.Get("Content-Length"))
+
+		b, err := io.ReadAll(r.Body)
+		assert.NoError(t, err)
+		assert.NotEmpty(t, len(b))
+
+		rw.WriteHeader(http.StatusNoContent)
+	}
+	ts := httptest.NewServer(http.HandlerFunc(h))
+	defer ts.Close()
+
+	u, err := url.Parse(ts.URL)
+	require.NoError(t, err)
+
+	c := &WriteClient{
+		hc:  ts.Client(),
+		url: u,
+		cfg: &HTTPConfig{},
+	}
+	data := &prompb.TimeSeries{
+		Labels: []*prompb.Label{
+			{
+				Name:  "label1",
+				Value: "label1-val",
+			},
+		},
+		Samples: []*prompb.Sample{
+			{
+				Value:     8.5,
+				Timestamp: time.Now().UnixMilli(),
+			},
+		},
+	}
+	err = c.Store(context.Background(), []*prompb.TimeSeries{data})
+	assert.NoError(t, err)
+}
+
+func TestClientStoreHTTPError(t *testing.T) {
+	t.Parallel()
+	h := func(w http.ResponseWriter, r *http.Request) {
+		http.Error(w, "bad bad", http.StatusUnauthorized)
+	}
+	ts := httptest.NewServer(http.HandlerFunc(h))
+	defer ts.Close()
+
+	u, err := url.Parse(ts.URL)
+	require.NoError(t, err)
+
+	c := &WriteClient{
+		hc:  ts.Client(),
+		url: u,
+		cfg: &HTTPConfig{},
+	}
+	assert.Error(t, c.Store(context.Background(), nil))
+}
+
+func TestClientStoreHTTPBasic(t *testing.T) {
+	t.Parallel()
+	h := func(w http.ResponseWriter, r *http.Request) {
+		u, pwd, ok := r.BasicAuth()
+		require.True(t, ok)
+		assert.Equal(t, "usertest", u)
+		assert.Equal(t, "pwdtest", pwd)
+	}
+	ts := httptest.NewServer(http.HandlerFunc(h))
+	defer ts.Close()
+
+	u, err := url.Parse(ts.URL)
+	require.NoError(t, err)
+
+	c := &WriteClient{
+		hc:  ts.Client(),
+		url: u,
+		cfg: &HTTPConfig{
+			BasicAuth: &BasicAuth{
+				Username: "usertest",
+				Password: "pwdtest",
+			},
+		},
+	}
+	assert.Error(t, c.Store(context.Background(), nil))
+}
+
+func TestClientStoreHeaders(t *testing.T) {
+	t.Parallel()
+	h := func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, r.Header.Get("X-Prometheus-Remote-Write-Version"), "0.1.0")
+		assert.Equal(t, r.Header.Get("X-MY-CUSTOM-HEADER"), "fake")
+	}
+	ts := httptest.NewServer(http.HandlerFunc(h))
+	defer ts.Close()
+
+	u, err := url.Parse(ts.URL)
+	require.NoError(t, err)
+
+	c := &WriteClient{
+		hc:  ts.Client(),
+		url: u,
+		cfg: &HTTPConfig{
+			Headers: http.Header(map[string][]string{
+				"X-MY-CUSTOM-HEADER": {"fake"},
+				// If the same key, of a mandatory protocol's header
+				// is provided, it will be overwritten.
+				"X-Prometheus-Remote-Write-Version": {"fake"},
+			}),
+		},
+	}
+	assert.Error(t, c.Store(context.Background(), nil))
+}
+
+func TestNewWriteRequestBody(t *testing.T) {
+	ts := []*prompb.TimeSeries{
+		{
+			Labels:  []*prompb.Label{{Name: "label1", Value: "val1"}},
+			Samples: []*prompb.Sample{{Value: 10.1, Timestamp: time.Unix(1, 0).Unix()}},
+		},
+	}
+	b, err := newWriteRequestBody(ts)
+	require.NoError(t, err)
+	require.NotEmpty(t, string(b))
+	assert.Contains(t, string(b), `label1`)
+}