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,270 @@
+package prometheus_test
+
+import (
+	"context"
+	"fmt"
+	"log/slog"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/go-kit/log"
+	promclient "github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/prometheus/model/labels"
+	"github.com/prometheus/prometheus/storage"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/exp/maps"
+
+	"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/component/prometheus/scrape"
+	"github.com/grafana/alloy/internal/runtime/componenttest"
+	"github.com/grafana/alloy/internal/runtime/logging"
+	"github.com/grafana/alloy/internal/service/labelstore"
+	"github.com/grafana/alloy/internal/service/livedebugging"
+	"github.com/grafana/alloy/internal/static/metrics/wal"
+	"github.com/grafana/alloy/internal/util"
+	"github.com/grafana/alloy/internal/util/testappender"
+	"github.com/grafana/alloy/syntax"
+)
+
+// This test simulates a scrape -> remote_write pipeline, without actually scraping
+func TestPipeline(t *testing.T) {
+	pipeline, ls, destination := newDefaultPipeline(t, util.TestLogger(t))
+
+	// 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")
+	ref1 := sendMetric(t, pipeline.Appender(t.Context()), lset1, sampleTimestamp, 12)
+	lset2 := labels.FromStrings("fizz", "buzz")
+	ref2 := sendMetric(t, pipeline.Appender(t.Context()), lset2, sampleTimestamp, 34)
+
+	expect := []*testappender.MetricSample{{
+		Labels:    lset1,
+		Timestamp: sampleTimestamp,
+		Value:     12,
+	}, {
+		Labels:    lset2,
+		Timestamp: sampleTimestamp,
+		Value:     34,
+	}}
+
+	require.EventuallyWithT(t, func(t *assert.CollectT) {
+		require.Len(t, destination.CollectedSamples(), 2)
+		require.ElementsMatch(t, expect, maps.Values(destination.CollectedSamples()))
+	}, 5*time.Second, 10*time.Millisecond, "timed out waiting for metrics to be written to destination")
+
+	ref := ls.GetOrAddGlobalRefID(lset1)
+	require.NotZero(t, ref)
+	// Append result ref should match the labelstore ref
+	require.Equal(t, ref, ref1)
+	localRef := ls.GetLocalRefID("prometheus.remote_write.test", ref)
+	require.NotZero(t, localRef)
+
+	ref = ls.GetOrAddGlobalRefID(lset2)
+	require.NotZero(t, ref)
+	// Append result ref should match the labelstore ref
+	require.Equal(t, ref, ref2)
+	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) {
+	pipeline, ls, destination := newRelabelPipeline(t, util.TestLogger(t))
+
+	// 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")
+	ref1 := sendMetric(t, pipeline.Appender(t.Context()), lset1, sampleTimestamp, 12)
+	lset2 := labels.FromStrings("fizz", "buzz")
+	ref2 := sendMetric(t, pipeline.Appender(t.Context()), lset2, sampleTimestamp, 34)
+
+	expect := []*testappender.MetricSample{{
+		Labels:    labels.NewBuilder(lset1).Set("lbl", "foo").Labels(),
+		Timestamp: sampleTimestamp,
+		Value:     12,
+	}, {
+		Labels:    labels.NewBuilder(lset2).Set("lbl", "foo").Labels(),
+		Timestamp: sampleTimestamp,
+		Value:     34,
+	}}
+
+	require.EventuallyWithT(t, func(t *assert.CollectT) {
+		require.Len(t, destination.CollectedSamples(), 2)
+		require.ElementsMatch(t, expect, maps.Values(destination.CollectedSamples()))
+	}, 1*time.Minute, 100*time.Millisecond, "timed out waiting for metrics to be written to destination")
+
+	ref := ls.GetOrAddGlobalRefID(lset1)
+	require.NotZero(t, ref)
+	// Append result ref should match the labelstore ref
+	require.Equal(t, ref, ref1)
+	// 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)
+	// Append result ref should match the labelstore ref
+	require.Equal(t, ref, ref2)
+
+	// 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 BenchmarkPipelines(b *testing.B) {
+	tests := []struct {
+		name            string
+		pipelineBuilder func(t testing.TB, logger log.Logger) (storage.Appendable, labelstore.LabelStore, testappender.CollectingAppender)
+	}{
+		{"default", newDefaultPipeline},
+		{"relabel", newRelabelPipeline},
+	}
+
+	numberOfMetrics := []int{2, 10, 100, 1000}
+
+	for _, n := range numberOfMetrics {
+		for _, tt := range tests {
+			// Don't need care about the labelstore and destination for benchmarks
+			pipeline, _, _ := tt.pipelineBuilder(b, log.NewNopLogger())
+			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),
+						)
+					}
+				}
+			})
+		}
+	}
+}
+
+func newDefaultPipeline(t testing.TB, logger log.Logger) (storage.Appendable, labelstore.LabelStore, testappender.CollectingAppender) {
+	ls := labelstore.New(logger, promclient.DefaultRegisterer)
+	rwAppendable, rwDestination := newRemoteWriteComponent(t, logger, ls)
+	pipelineAppendable := prometheus.NewFanout([]storage.Appendable{rwAppendable}, promclient.DefaultRegisterer, ls)
+	scrapeInterceptor := scrape.NewInterceptor("prometheus.scrape.test", ls, livedebugging.NewLiveDebugging(), pipelineAppendable)
+
+	return scrapeInterceptor, ls, rwDestination
+}
+
+func newRelabelPipeline(t testing.TB, logger log.Logger) (storage.Appendable, labelstore.LabelStore, testappender.CollectingAppender) {
+	ls := labelstore.New(logger, promclient.DefaultRegisterer)
+	rwAppendable, rwDestination := newRemoteWriteComponent(t, logger, ls)
+	relabelAppendable := newRelabelComponent(t, logger, []storage.Appendable{rwAppendable}, ls)
+	pipelineAppendable := prometheus.NewFanout([]storage.Appendable{relabelAppendable}, promclient.DefaultRegisterer, ls)
+	scrapeInterceptor := scrape.NewInterceptor("prometheus.scrape.test", ls, livedebugging.NewLiveDebugging(), pipelineAppendable)
+
+	return scrapeInterceptor, ls, rwDestination
+}
+
+func newRemoteWriteComponent(t testing.TB, logger log.Logger, ls *labelstore.Service) (storage.Appendable, testappender.CollectingAppender) {
+	walDir := t.TempDir()
+
+	walStorage, err := wal.NewStorage(logger, promclient.NewRegistry(), walDir)
+	require.NoError(t, err)
+
+	fanoutLogger := slog.New(
+		logging.NewSlogGoKitHandler(
+			log.With(logger, "subcomponent", "fanout"),
+		),
+	)
+
+	inMemoryAppendable := testappender.ConstantAppendable{Inner: testappender.NewCollectingAppender()}
+	store := storage.NewFanout(fanoutLogger, walStorage, testStorage{inMemoryAppendable: inMemoryAppendable})
+
+	return remotewrite.NewInterceptor("prometheus.remote_write.test", &atomic.Bool{}, livedebugging.NewLiveDebugging(), ls, store), inMemoryAppendable.Inner
+}
+
+type testStorage struct {
+	// Embed Queryable/ChunkQueryable for compatibility, but don't actually implement it.
+	storage.Queryable
+	storage.ChunkQueryable
+
+	inMemoryAppendable storage.Appendable
+}
+
+func (t testStorage) Appender(ctx context.Context) storage.Appender {
+	return t.inMemoryAppendable.Appender(ctx)
+}
+
+func (t testStorage) StartTime() (int64, error) {
+	return 0, nil
+}
+
+func (t testStorage) Close() error {
+	return nil
+}
+
+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) uint64 {
+	ref, err := appender.Append(0, labels, time, value)
+	require.NoError(t, err)
+	require.NoError(t, appender.Commit())
+
+	return uint64(ref)
+}

internal/component/prometheus/receive_http/receive_http.go

@@ -65,7 +65,7 @@ func New(opts component.Options, args Arguments) (*Component, error) {
 		return nil, err
 	}
 	ls := service.(labelstore.LabelStore)
-	fanout := alloyprom.NewFanout(args.ForwardTo, opts.ID, opts.Registerer, ls)
+	fanout := alloyprom.NewFanout(args.ForwardTo, opts.Registerer, ls)
 
 	uncheckedCollector := util.NewUncheckedCollector(nil)
 	opts.Registerer.MustRegister(uncheckedCollector)

internal/component/prometheus/relabel/relabel.go

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