test(loki.process): Add benchmarks to test performance with a large number of stage.match components

internal/component/loki/process/stages/pipeline_test.go

@@ -4,6 +4,9 @@ import (
 	"fmt"
 	"io"
 	"log/slog"
+	"maps"
+	"runtime"
+	"strings"
 	"sync"
 	"testing"
 	"time"
@@ -282,6 +285,152 @@ func BenchmarkPipeline(b *testing.B) {
 	}
 }
 
+// genRulesConfig builds numRules independent stage.match blocks, each gated
+// on a distinct label value and containing one small nested stage. This
+// mirrors a pipeline assembled from many separately configured rules (for
+// example, one rule per tenant policy) rather than a handful of stages
+// written by hand, so it can be scaled up to see how per-entry cost scales
+// with rule count.
+func genRulesConfig(numRules int) string {
+	var sb strings.Builder
+	for i := 0; i < numRules; i++ {
+		fmt.Fprintf(&sb, `
+stage.match {
+	selector = "{rule_id=\"rule-%d\"}"
+	stage.static_labels {
+		values = {
+			rule_matched = "%d",
+		}
+	}
+}
+`, i, i)
+	}
+	return sb.String()
+}
+
+// BenchmarkPipelineManyRules measures how per-entry latency scales with the
+// number of sequential stages in a pipeline. Every entry pays the full
+// traversal cost of checking each rule's selector regardless of whether it
+// ends up matching one — a matching and a non-matching entry were measured
+// separately and found equivalent, so this only exercises the non-matching
+// case.
+func BenchmarkPipelineManyRules(b *testing.B) {
+	ruleCounts := []int{1, 10, 50, 100, 500, 1000}
+
+	for _, numRules := range ruleCounts {
+		stgs := loadConfig(genRulesConfig(numRules))
+
+		b.Run(fmt.Sprintf("rules=%d", numRules), func(b *testing.B) {
+			lb := model.LabelSet{"rule_id": "no-match"}
+			benchmarkPipelineEntries(b, stgs, lb, rawTestLine)
+		})
+	}
+}
+
+// genRulesConfigWithOneMultiline is genRulesConfig, but the rule at index
+// multilineAt nests a stage.multiline instead of stage.static_labels.
+// multiline needs its own goroutine — it must flush a stale block even when
+// no new entry arrives — so exactly that one rule falls back to its own
+// channel (see newMatcherStage); every other rule still fuses into a plain
+// function call. See BenchmarkPipelineOneMultilineAmongManyRules.
+func genRulesConfigWithOneMultiline(numRules, multilineAt int) string {
+	var sb strings.Builder
+	for i := 0; i < numRules; i++ {
+		if i == multilineAt {
+			fmt.Fprintf(&sb, `
+stage.match {
+	selector = "{rule_id=\"rule-%d\"}"
+	stage.multiline {
+		firstline     = "^NEVER_MATCHES"
+		max_wait_time = "3s"
+	}
+}
+`, i)
+			continue
+		}
+		fmt.Fprintf(&sb, `
+stage.match {
+	selector = "{rule_id=\"rule-%d\"}"
+	stage.static_labels {
+		values = {
+			rule_matched = "%d",
+		}
+	}
+}
+`, i, i)
+	}
+	return sb.String()
+}
+
+// BenchmarkPipelineOneMultilineAmongManyRules checks that one rule needing
+// its own channel (because it nests stage.multiline) doesn't drag the rest
+// of a large rule set back into the old per-stage-channel behavior: only the
+// fusion immediately around that one rule is affected, not the whole
+// pipeline. "all_sync" (no multiline anywhere) is the reference point.
+//
+// "in_middle" tends to measure faster than "all_sync" under a small
+// GOMAXPROCS: splitting ~1000 fused rules into two ~500-rule halves turns
+// them into two independent goroutines that can pipeline different entries
+// across separate cores at once, an incidental win from an even split
+// rather than something this design specifically provides — it disappears
+// under GOMAXPROCS=1, where all four variants converge. "at_start" and
+// "at_end" don't get this, because one side of the split is nearly the
+// entire rule set and the other is nearly empty, so there's nothing to
+// overlap. The number that matters here is that none of these are anywhere
+// near the cost of giving every one of the 1000 rules its own channel.
+func BenchmarkPipelineOneMultilineAmongManyRules(b *testing.B) {
+	const n = 1000
+	variants := []struct {
+		name string
+		cfg  string
+	}{
+		{"all_sync", genRulesConfig(n)},
+		{"one_multiline_at_start", genRulesConfigWithOneMultiline(n, 0)},
+		{"one_multiline_in_middle", genRulesConfigWithOneMultiline(n, n/2)},
+		{"one_multiline_at_end", genRulesConfigWithOneMultiline(n, n-1)},
+	}
+
+	for _, v := range variants {
+		b.Run(v.name, func(b *testing.B) {
+			stgs := loadConfig(v.cfg)
+			lb := model.LabelSet{"rule_id": "no-match"}
+			benchmarkPipelineEntries(b, stgs, lb, rawTestLine)
+		})
+	}
+}
+
+func benchmarkPipelineEntries(b *testing.B, stgs []StageConfig, lb model.LabelSet, entry string) {
+	pl, err := NewPipeline(logging.NewSlogNop(), stgs, prometheus.NewRegistry(), featuregate.StabilityGenerallyAvailable)
+	if err != nil {
+		panic(err)
+	}
+	ts := time.Now()
+
+	in := make(chan Entry)
+	out := pl.Run(in)
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		for range out {
+		}
+	}()
+
+	b.ResetTimer()
+	for b.Loop() {
+		// The pipeline may still be processing an earlier entry when this
+		// send returns, so every entry needs its own label map: a matching
+		// rule mutates Labels in place, and sharing one map across entries
+		// that are processed concurrently races.
+		in <- newEntry(nil, maps.Clone(lb), entry, ts)
+	}
+	close(in)
+	// b.Loop already stopped the timer, so waiting for the last entries to
+	// finish draining here doesn't affect ns/op — it just makes sure this
+	// sub-benchmark's goroutines are gone before the next one starts,
+	// instead of bleeding CPU time into its measurement.
+	<-done
+}
+
 func TestPipeline_Wrap(t *testing.T) {
 	now := time.Now()
 	p, err := NewPipeline(logging.NewSlogNop(), loadConfig(testMultiStageAlloy), prometheus.DefaultRegisterer, featuregate.StabilityGenerallyAvailable)
@@ -411,3 +560,61 @@ stage.match {
 	out.Stop()
 	t.Log(out.Received())
 }
+
+// BenchmarkPipelineManyStreamsSaturated models many concurrent log streams
+// actually competing for the CPU, as opposed to BenchmarkPipelineManyRules'
+// single stream with idle cores to spare. The two regimes favor opposite
+// designs: with idle cores available, spreading tiny amounts of relay work
+// across them can make a single stream's latency look fine even if it's
+// wasteful, since there's nothing else for those cores to do. Once enough
+// streams are genuinely competing for the same cores, that per-stage
+// goroutine overhead has nowhere to hide and becomes a direct tax on
+// aggregate throughput, which this benchmark's reported entries/sec makes
+// visible.
+func BenchmarkPipelineManyStreamsSaturated(b *testing.B) {
+	const (
+		numRules      = 1000
+		entriesPerRun = 200
+	)
+	streams := runtime.GOMAXPROCS(0) * 4
+
+	stgs := loadConfig(genRulesConfig(numRules))
+	pipelines := make([]*Pipeline, streams)
+	for i := range pipelines {
+		pl, err := NewPipeline(logging.NewSlogNop(), stgs, prometheus.NewRegistry(), featuregate.StabilityGenerallyAvailable)
+		if err != nil {
+			b.Fatal(err)
+		}
+		pipelines[i] = pl
+	}
+	lb := model.LabelSet{"rule_id": "no-match"}
+	ts := time.Now()
+
+	b.ResetTimer()
+	for b.Loop() {
+		var wg sync.WaitGroup
+		wg.Add(streams)
+		for _, pl := range pipelines {
+			go func(pl *Pipeline) {
+				defer wg.Done()
+
+				in := make(chan Entry)
+				out := pl.Run(in)
+				done := make(chan struct{})
+				go func() {
+					defer close(done)
+					for range out {
+					}
+				}()
+
+				for i := 0; i < entriesPerRun; i++ {
+					in <- newEntry(nil, maps.Clone(lb), rawTestLine, ts)
+				}
+				close(in)
+				<-done
+			}(pl)
+		}
+		wg.Wait()
+	}
+	b.ReportMetric(float64(streams*entriesPerRun*b.N)/b.Elapsed().Seconds(), "entries/sec")
+}