Add some high level pipeline tests for prometheus

Author: kgeckhart

internal/component/otelcol/exporter/prometheus/prometheus.go

@@ -8,14 +8,15 @@ import (
 	"time"
 
 	"github.com/go-kit/log"
+	"github.com/prometheus/prometheus/storage"
+
 	"github.com/grafana/alloy/internal/component"
 	"github.com/grafana/alloy/internal/component/otelcol"
 	"github.com/grafana/alloy/internal/component/otelcol/exporter/prometheus/internal/convert"
 	"github.com/grafana/alloy/internal/component/otelcol/internal/lazyconsumer"
 	"github.com/grafana/alloy/internal/component/prometheus"
 	"github.com/grafana/alloy/internal/featuregate"
 	"github.com/grafana/alloy/internal/service/labelstore"
-	"github.com/prometheus/prometheus/storage"
 )
 
 func init() {
@@ -87,7 +88,7 @@ func New(o component.Options, c Arguments) (*Component, error) {
 		return nil, err
 	}
 	ls := service.(labelstore.LabelStore)
-	fanout := prometheus.NewFanout(nil, o.ID, o.Registerer, ls)
+	fanout := prometheus.NewFanout(nil, o.Registerer, ls)
 
 	converter := convert.New(o.Logger, fanout, convertArgumentsToConvertOptions(c))
 

internal/component/prometheus/enrich/enrich.go

@@ -97,7 +97,7 @@ func New(opts component.Options, args Arguments) (*Component, error) {
 		}
 	}
 
-	c.fanout = prometheus.NewFanout(args.ForwardTo, opts.ID, opts.Registerer, ls)
+	c.fanout = prometheus.NewFanout(args.ForwardTo, opts.Registerer, ls)
 	c.receiver = prometheus.NewInterceptor(
 		c.fanout,
 		ls,

internal/component/prometheus/fanout.go

@@ -26,9 +26,7 @@ var _ storage.Appendable = (*Fanout)(nil)
 type Fanout struct {
 	mut sync.RWMutex
 	// children is where to fan out.
-	children []storage.Appendable
-	// ComponentID is what component this belongs to.
-	componentID    string
+	children       []storage.Appendable
 	writeLatency   prometheus.Histogram
 	samplesCounter prometheus.Counter
 	ls             labelstore.LabelStore
@@ -39,7 +37,7 @@ type Fanout struct {
 }
 
 // NewFanout creates a fanout appendable.
-func NewFanout(children []storage.Appendable, componentID string, register prometheus.Registerer, ls labelstore.LabelStore) *Fanout {
+func NewFanout(children []storage.Appendable, register prometheus.Registerer, ls labelstore.LabelStore) *Fanout {
 	wl := prometheus.NewHistogram(prometheus.HistogramOpts{
 		Name:    "prometheus_fanout_latency",
 		Help:    "Write latency for sending to direct and indirect components",
@@ -55,7 +53,6 @@ func NewFanout(children []storage.Appendable, componentID string, register prome
 
 	return &Fanout{
 		children:       children,
-		componentID:    componentID,
 		writeLatency:   wl,
 		samplesCounter: s,
 		ls:             ls,

internal/component/prometheus/fanout_test.go

@@ -13,15 +13,15 @@ import (
 
 func TestRollback(t *testing.T) {
 	ls := labelstore.New(nil, prometheus.DefaultRegisterer)
-	fanout := NewFanout([]storage.Appendable{NewFanout(nil, "1", prometheus.DefaultRegisterer, ls)}, "", prometheus.DefaultRegisterer, ls)
+	fanout := NewFanout([]storage.Appendable{NewFanout(nil, prometheus.DefaultRegisterer, ls)}, prometheus.DefaultRegisterer, ls)
 	app := fanout.Appender(t.Context())
 	err := app.Rollback()
 	require.NoError(t, err)
 }
 
 func TestCommit(t *testing.T) {
 	ls := labelstore.New(nil, prometheus.DefaultRegisterer)
-	fanout := NewFanout([]storage.Appendable{NewFanout(nil, "1", prometheus.DefaultRegisterer, ls)}, "", prometheus.DefaultRegisterer, ls)
+	fanout := NewFanout([]storage.Appendable{NewFanout(nil, prometheus.DefaultRegisterer, ls)}, prometheus.DefaultRegisterer, ls)
 	app := fanout.Appender(t.Context())
 	err := app.Commit()
 	require.NoError(t, err)

internal/component/prometheus/operator/common/crdmanager.go

@@ -145,7 +145,7 @@ func (c *crdManager) Run(ctx context.Context) error {
 	}()
 
 	// Start prometheus scrape manager.
-	alloyAppendable := prometheus.NewFanout(c.args.ForwardTo, c.opts.ID, c.opts.Registerer, c.ls)
+	alloyAppendable := prometheus.NewFanout(c.args.ForwardTo, c.opts.Registerer, c.ls)
 	opts := &scrape.Options{}
 	c.scrapeManager, err = scrape.NewManager(opts, slog.New(logging.NewSlogGoKitHandler(c.logger)), nil, alloyAppendable, unregisterer)
 	if err != nil {

internal/component/prometheus/pipeline_test.go

@@ -0,0 +1,309 @@
+package prometheus_test
+
+import (
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/go-kit/log"
+	promclient "github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/prometheus/model/labels"
+	"github.com/prometheus/prometheus/prompb"
+	"github.com/prometheus/prometheus/storage"
+	"github.com/prometheus/prometheus/storage/remote"
+	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/alloy/internal/component"
+	"github.com/grafana/alloy/internal/component/prometheus"
+	"github.com/grafana/alloy/internal/component/prometheus/relabel"
+	"github.com/grafana/alloy/internal/component/prometheus/remotewrite"
+	"github.com/grafana/alloy/internal/runtime/componenttest"
+	"github.com/grafana/alloy/internal/service/labelstore"
+	"github.com/grafana/alloy/internal/util"
+	"github.com/grafana/alloy/syntax"
+)
+
+// This test simulates a scrape -> remote_write pipeline, without actually scraping
+func TestPipeline(t *testing.T) {
+	handler, writeResult := channelServerHandler(t)
+	pipeline, srv, ls := newDefaultPipeline(t, util.TestLogger(t), handler)
+	defer srv.Close()
+
+	// We need to use a future timestamp since remote_write will ignore any
+	// sample which is earlier than the time when it started. Adding a minute
+	// ensures that our samples will never get ignored.
+	sampleTimestamp := time.Now().Add(time.Minute).UnixMilli()
+
+	// Send metrics to our component. These will be written to the WAL and
+	// subsequently written to our HTTP server.
+	lset1 := labels.FromStrings("foo", "bar")
+	sendMetric(t, pipeline.Appender(t.Context()), lset1, sampleTimestamp, 12)
+	lset2 := labels.FromStrings("fizz", "buzz")
+	sendMetric(t, pipeline.Appender(t.Context()), lset2, sampleTimestamp, 34)
+
+	expect := []prompb.TimeSeries{{
+		Labels: []prompb.Label{
+			{Name: "cluster", Value: "local"},
+			{Name: "foo", Value: "bar"},
+		},
+		Samples: []prompb.Sample{
+			{Timestamp: sampleTimestamp, Value: 12},
+		},
+	}, {
+		Labels: []prompb.Label{
+			{Name: "cluster", Value: "local"},
+			{Name: "fizz", Value: "buzz"},
+		},
+		Samples: []prompb.Sample{
+			{Timestamp: sampleTimestamp, Value: 34},
+		},
+	}}
+
+	select {
+	case <-time.After(time.Minute):
+		require.FailNow(t, "timed out waiting for metrics")
+	case res := <-writeResult:
+		require.Equal(t, expect, res.Timeseries)
+	}
+
+	ref := ls.GetOrAddGlobalRefID(lset1)
+	require.NotZero(t, ref)
+	localRef := ls.GetLocalRefID("prometheus.remote_write.test", ref)
+	require.NotZero(t, localRef)
+
+	ref = ls.GetOrAddGlobalRefID(lset2)
+	require.NotZero(t, ref)
+	localRef = ls.GetLocalRefID("prometheus.remote_write.test", ref)
+	require.NotZero(t, localRef)
+}
+
+// This test simulates a scrape -> relabel -> remote_write pipeline, without actually scraping
+func TestRelabelPipeline(t *testing.T) {
+	handler, writeResult := channelServerHandler(t)
+	pipeline, srv, ls := newRelabelPipeline(t, util.TestLogger(t), handler)
+	defer srv.Close()
+
+	// We need to use a future timestamp since remote_write will ignore any
+	// sample which is earlier than the time when it started. Adding a minute
+	// ensures that our samples will never get ignored.
+	sampleTimestamp := time.Now().Add(time.Minute).UnixMilli()
+
+	// Send metrics to our component. These will be written to the WAL and
+	// subsequently written to our HTTP server.
+	lset1 := labels.FromStrings("foo", "bar")
+	sendMetric(t, pipeline.Appender(t.Context()), lset1, sampleTimestamp, 12)
+	lset2 := labels.FromStrings("fizz", "buzz")
+	sendMetric(t, pipeline.Appender(t.Context()), lset2, sampleTimestamp, 34)
+
+	expect := []prompb.TimeSeries{{
+		Labels: []prompb.Label{
+			{Name: "cluster", Value: "local"},
+			{Name: "foo", Value: "bar"},
+			{Name: "lbl", Value: "foo"},
+		},
+		Samples: []prompb.Sample{
+			{Timestamp: sampleTimestamp, Value: 12},
+		},
+	}, {
+		Labels: []prompb.Label{
+			{Name: "cluster", Value: "local"},
+			{Name: "fizz", Value: "buzz"},
+			{Name: "lbl", Value: "foo"},
+		},
+		Samples: []prompb.Sample{
+			{Timestamp: sampleTimestamp, Value: 34},
+		},
+	}}
+
+	select {
+	case <-time.After(time.Minute):
+		require.FailNow(t, "timed out waiting for metrics")
+	case res := <-writeResult:
+		require.Equal(t, expect, res.Timeseries)
+	}
+
+	ref := ls.GetOrAddGlobalRefID(lset1)
+	require.NotZero(t, ref)
+	// This was relabeled, so we shouldn't have a local ref
+	localRef := ls.GetLocalRefID("prometheus.remote_write.test", ref)
+	require.Zero(t, localRef)
+
+	ref = ls.GetOrAddGlobalRefID(lset2)
+	require.NotZero(t, ref)
+	// This was relabeled, so we shouldn't have a local ref
+	localRef = ls.GetLocalRefID("prometheus.remote_write.test", ref)
+	require.Zero(t, localRef)
+
+	lset1Relabeled := labels.NewBuilder(lset1).Set("lbl", "foo").Labels()
+	ref = ls.GetOrAddGlobalRefID(lset1Relabeled)
+	require.NotZero(t, ref)
+	localRef = ls.GetLocalRefID("prometheus.remote_write.test", ref)
+	require.NotZero(t, localRef)
+
+	lset2Relabeled := labels.NewBuilder(lset2).Set("lbl", "foo").Labels()
+	ref = ls.GetOrAddGlobalRefID(lset2Relabeled)
+	require.NotZero(t, ref)
+	localRef = ls.GetLocalRefID("prometheus.remote_write.test", ref)
+	require.NotZero(t, localRef)
+}
+
+func channelServerHandler(t *testing.T) (http.HandlerFunc, chan *prompb.WriteRequest) {
+	writeResult := make(chan *prompb.WriteRequest)
+	serverHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		req, err := remote.DecodeWriteRequest(r.Body)
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusBadRequest)
+			return
+		}
+
+		select {
+		case writeResult <- req:
+		default:
+			require.Fail(t, "failed to send remote_write result over channel")
+		}
+	})
+
+	return serverHandler, writeResult
+}
+
+func BenchmarkPipelines(b *testing.B) {
+	tests := []struct {
+		name            string
+		pipelineBuilder func(t testing.TB, logger log.Logger, handlerFunc http.HandlerFunc) (storage.Appendable, *httptest.Server, labelstore.LabelStore)
+	}{
+		{"default", newDefaultPipeline},
+		{"relabel", newRelabelPipeline},
+	}
+
+	numberOfMetrics := []int{2, 10, 100, 1000}
+
+	for _, n := range numberOfMetrics {
+		for _, tt := range tests {
+			// Don't need the server to do anything here
+			pipeline, srv, _ := tt.pipelineBuilder(b, log.NewNopLogger(), func(w http.ResponseWriter, r *http.Request) {})
+			b.Run(fmt.Sprintf("%s/%d-metrics", tt.name, n), func(b *testing.B) {
+				b.ReportAllocs()
+				b.ResetTimer()
+
+				for b.Loop() {
+					for i := 0; i < n; i++ {
+						sendMetric(
+							b,
+							pipeline.Appender(b.Context()),
+							labels.FromStrings(fmt.Sprintf("metric-%d", i), fmt.Sprintf("metric-%d", i)),
+							time.Now().Add(time.Minute).UnixMilli(),
+							float64(i),
+						)
+					}
+				}
+			})
+			srv.Close()
+		}
+	}
+}
+
+func newDefaultPipeline(t testing.TB, logger log.Logger, handlerFunc http.HandlerFunc) (storage.Appendable, *httptest.Server, labelstore.LabelStore) {
+	srv := httptest.NewServer(handlerFunc)
+
+	ls := labelstore.New(logger, promclient.DefaultRegisterer)
+	rwAppendable := newRemoteWriteComponent(t, logger, srv.URL, ls)
+	pipelineAppendable := prometheus.NewFanout([]storage.Appendable{rwAppendable}, promclient.DefaultRegisterer, ls)
+
+	return pipelineAppendable, srv, ls
+}
+
+func newRelabelPipeline(t testing.TB, logger log.Logger, handlerFunc http.HandlerFunc) (storage.Appendable, *httptest.Server, labelstore.LabelStore) {
+	srv := httptest.NewServer(handlerFunc)
+
+	ls := labelstore.New(logger, promclient.DefaultRegisterer)
+	rwAppendable := newRemoteWriteComponent(t, logger, srv.URL, ls)
+	relabelAppendable := newRelabelComponent(t, logger, []storage.Appendable{rwAppendable}, ls)
+	pipelineAppendable := prometheus.NewFanout([]storage.Appendable{relabelAppendable}, promclient.DefaultRegisterer, ls)
+
+	return pipelineAppendable, srv, ls
+}
+
+func newRemoteWriteComponent(t testing.TB, logger log.Logger, url string, ls *labelstore.Service) storage.Appendable {
+	// Create our component and wait for it to start running, so we can write
+	// metrics to the WAL.
+	cfg := fmt.Sprintf(`
+		external_labels = {
+			cluster = "local",
+		}
+		endpoint {
+			name           = "test-url"
+			url            = "%s/api/v1/write"
+			remote_timeout = "100ms"
+
+			queue_config {
+				batch_send_deadline = "100ms"
+			}
+		}
+	`, url)
+	var args remotewrite.Arguments
+	require.NoError(t, syntax.Unmarshal([]byte(cfg), &args))
+
+	tc, err := componenttest.NewControllerFromID(logger, "prometheus.remote_write")
+	require.NoError(t, err)
+	go func() {
+		err = tc.Run(componenttest.TestContext(t), args, func(opts component.Options) component.Options {
+			inner := opts.GetServiceData
+			opts.GetServiceData = func(name string) (interface{}, error) {
+				if name == labelstore.ServiceName {
+					return ls, nil
+				}
+				return inner(name)
+			}
+			return opts
+		})
+		require.NoError(t, err)
+	}()
+	require.NoError(t, tc.WaitRunning(5*time.Second))
+
+	return tc.Exports().(remotewrite.Exports).Receiver
+}
+
+func newRelabelComponent(t testing.TB, logger log.Logger, forwardTo []storage.Appendable, ls *labelstore.Service) storage.Appendable {
+	cfg := `forward_to = []
+			rule {
+				action       = "replace"
+				target_label = "lbl"
+				replacement  = "foo"
+			}`
+	var args relabel.Arguments
+	require.NoError(t, syntax.Unmarshal([]byte(cfg), &args))
+	args.ForwardTo = forwardTo
+
+	tc, err := componenttest.NewControllerFromID(logger, "prometheus.relabel")
+	require.NoError(t, err)
+	go func() {
+		err = tc.Run(componenttest.TestContext(t), args, func(opts component.Options) component.Options {
+			inner := opts.GetServiceData
+			opts.GetServiceData = func(name string) (interface{}, error) {
+				if name == labelstore.ServiceName {
+					return ls, nil
+				}
+				return inner(name)
+			}
+			return opts
+		})
+		require.NoError(t, err)
+	}()
+	require.NoError(t, tc.WaitRunning(5*time.Second))
+
+	return tc.Exports().(relabel.Exports).Receiver
+}
+
+func sendMetric(
+	t testing.TB,
+	appender storage.Appender,
+	labels labels.Labels,
+	time int64,
+	value float64,
+) {
+	_, err := appender.Append(0, labels, time, value)
+	require.NoError(t, err)
+	require.NoError(t, appender.Commit())
+}