fix: reload span_name_sanitization overrides during runtime (#6435)

* fix: reload MetricsGeneratorSpanNameSanitization overrides during runtime

before this fix, we were just loading the MetricsGeneratorSpanNameSanitization of the overrides at startup

if the overrides were updated (which they can without restart), we won't reload them and the MetricsGeneratorSpanNameSanitization
config change in the overrides won't take effact without restart

* Add CHANGELOG.md

* cleanup noopSanitizer because we don't need it anymore

* Add a test to ensure toggle works

* replace assert with require

Author: electron0zero

CHANGELOG.md

@@ -26,6 +26,7 @@
 * [BUGFIX] Fix query_end_cutoff param for query range [#6360](https://github.com/grafana/tempo/pull/6360) (@ruslan-mikhailov)
 * [BUGFIX] generator: fix dimension_mappings and target_info_excluded_dimensions being unconditionally overwritten even when overrides were nil [#6390](https://github.com/grafana/tempo/pull/6390) (@carles-grafana)
 * [BUGFIX] generator: fix panic when `write_relabel_configs` is configured on remote write endpoints [#6396](https://github.com/grafana/tempo/pull/6396) (@carles-grafana)
+* [BUGFIX] fix: reload span_name_sanitization overrides during runtime [#6435](https://github.com/grafana/tempo/pull/6435) (@electron0zero)
 
 ### 3.0 Cleanup
 

modules/generator/registry/builder_test.go

@@ -8,7 +8,7 @@ import (
 )
 
 func TestLabelBuilder(t *testing.T) {
-	builder := NewLabelBuilder(0, 0, &noopSanitizer{})
+	builder := NewLabelBuilder(0, 0, newTestDrainSanitizer(SpanNameSanitizationDisabled))
 	builder.Add("name", "value")
 	lbls, ok := builder.CloseAndBuildLabels()
 
@@ -23,7 +23,7 @@ func TestLabelBuilder(t *testing.T) {
 }
 
 func TestLabelBuilder_MaxLabelNameLength(t *testing.T) {
-	builder := NewLabelBuilder(10, 10, &noopSanitizer{})
+	builder := NewLabelBuilder(10, 10, newTestDrainSanitizer(SpanNameSanitizationDisabled))
 	builder.Add("name", "very_long_value")
 	builder.Add("very_long_name", "value")
 
@@ -34,7 +34,7 @@ func TestLabelBuilder_MaxLabelNameLength(t *testing.T) {
 }
 
 func TestLabelBuilder_InvalidUTF8(t *testing.T) {
-	builder := NewLabelBuilder(0, 0, &noopSanitizer{})
+	builder := NewLabelBuilder(0, 0, newTestDrainSanitizer(SpanNameSanitizationDisabled))
 	builder.Add("name", "svc-\xc3\x28") // Invalid UTF-8
 
 	_, ok := builder.CloseAndBuildLabels()

modules/generator/registry/drain_sanitizer.go

@@ -25,12 +25,17 @@ var (
 	}, []string{"tenant"})
 )
 
+// sanitizeModeFunc returns the current span name sanitization mode for the tenant.
+type sanitizeModeFunc func(tenant string) string
+
 type DrainSanitizer struct {
 	mtx    sync.Mutex
 	drain  *drain.Drain
-	dryRun bool
 	demand *Cardinality
 
+	tenant        string
+	sanitizeModeF sanitizeModeFunc
+
 	metricTotalSpansSanitized prometheus.Counter
 	demandGauge               prometheus.Gauge
 
@@ -41,10 +46,11 @@ type DrainSanitizer struct {
 	pruneChan        <-chan time.Time
 }
 
-func NewDrainSanitizer(tenant string, dryRun bool, staleDuration time.Duration) *DrainSanitizer {
+func NewDrainSanitizer(tenant string, sanitizeModeF sanitizeModeFunc, staleDuration time.Duration) *DrainSanitizer {
 	return &DrainSanitizer{
 		drain:                     drain.New(tenant, drain.DefaultConfig()),
-		dryRun:                    dryRun,
+		tenant:                    tenant,
+		sanitizeModeF:             sanitizeModeF,
 		metricTotalSpansSanitized: metricTotalSpansSanitized.WithLabelValues(tenant),
 		demand:                    NewCardinality(staleDuration, removeStaleSeriesInterval),
 		demandGauge:               metricPostSanitizationDemand.WithLabelValues(tenant),
@@ -54,6 +60,13 @@ func NewDrainSanitizer(tenant string, dryRun bool, staleDuration time.Duration)
 }
 
 func (s *DrainSanitizer) Sanitize(lbls labels.Labels) labels.Labels {
+	// Check the override at runtime so that changes to the sanitization mode override
+	// take effect without restarting the generator.
+	mode := s.sanitizeModeF(s.tenant)
+	if mode == SpanNameSanitizationDisabled {
+		return lbls
+	}
+
 	s.doPeriodicMaintenance()
 
 	s.mtx.Lock()
@@ -84,7 +97,8 @@ func (s *DrainSanitizer) Sanitize(lbls labels.Labels) labels.Labels {
 	newLabelHash := newLbls.Hash()
 	s.demand.Insert(newLabelHash)
 
-	if s.dryRun {
+	// in dry-run mode, return the labels without modifying but capture metrics
+	if mode == SpanNameSanitizationDryRun {
 		return lbls
 	}
 

modules/generator/registry/drain_sanitizer_test.go

@@ -6,13 +6,17 @@ import (
 	"time"
 
 	"github.com/prometheus/prometheus/model/labels"
-	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
+func newTestDrainSanitizer(mode string) *DrainSanitizer {
+	return NewDrainSanitizer("test-tenant", func(string) string { return mode }, 15*time.Minute)
+}
+
 func TestDrainSanitizer_PatternDetection(t *testing.T) {
 	t.Parallel()
 
-	sanitizer := NewDrainSanitizer("test-tenant", false, 15*time.Minute)
+	sanitizer := newTestDrainSanitizer(SpanNameSanitizationEnabled)
 
 	// Train with similar span names that should form a pattern
 	lbls1 := labels.FromStrings("span_name", "GET /api/users/123", "service", "api")
@@ -21,25 +25,25 @@ func TestDrainSanitizer_PatternDetection(t *testing.T) {
 
 	// First call should return original (no pattern yet)
 	result1 := sanitizer.Sanitize(lbls1)
-	assert.Equal(t, "GET /api/users/123", result1.Get("span_name"))
+	require.Equal(t, "GET /api/users/123", result1.Get("span_name"))
 
 	// After training, subsequent similar spans should be sanitized
 	result2 := sanitizer.Sanitize(lbls2)
 	result3 := sanitizer.Sanitize(lbls3)
 
 	// All should have the same sanitized span_name pattern
-	assert.Equal(t, result2.Get("span_name"), result3.Get("span_name"))
+	require.Equal(t, result2.Get("span_name"), result3.Get("span_name"))
 	// Pattern should contain the parameter marker
-	assert.Contains(t, result2.Get("span_name"), "<_>")
+	require.Contains(t, result2.Get("span_name"), "<_>")
 	// Original labels should be preserved
-	assert.Equal(t, "api", result2.Get("service"))
-	assert.Equal(t, "api", result3.Get("service"))
+	require.Equal(t, "api", result2.Get("service"))
+	require.Equal(t, "api", result3.Get("service"))
 }
 
 func TestDrainSanitizer_DryRunMode(t *testing.T) {
 	t.Parallel()
 
-	sanitizer := NewDrainSanitizer("test-tenant", true, 15*time.Minute) // dryRun = true
+	sanitizer := newTestDrainSanitizer(SpanNameSanitizationDryRun)
 
 	lbls1 := labels.FromStrings("span_name", "GET /api/users/123", "service", "api")
 	lbls2 := labels.FromStrings("span_name", "GET /api/users/456", "service", "api")
@@ -49,29 +53,75 @@ func TestDrainSanitizer_DryRunMode(t *testing.T) {
 
 	// In dry-run mode, even if pattern is detected, return original labels
 	result := sanitizer.Sanitize(lbls2)
-	assert.Equal(t, "GET /api/users/456", result.Get("span_name"))
-	assert.Equal(t, lbls2, result)
+	require.Equal(t, "GET /api/users/456", result.Get("span_name"))
+	require.Equal(t, lbls2, result)
+}
+
+func TestDrainSanitizer_RuntimeModeToggle(t *testing.T) {
+	t.Parallel()
+
+	mode := SpanNameSanitizationEnabled
+	sanitizer := NewDrainSanitizer("test-tenant", func(string) string { return mode }, 15*time.Minute)
+
+	// Train the drain tree with similar span names to establish a pattern
+	sanitizer.Sanitize(labels.FromStrings("span_name", "GET /api/users/123"))
+	sanitizer.Sanitize(labels.FromStrings("span_name", "GET /api/users/456"))
+
+	// With enabled mode, should sanitize
+	result := sanitizer.Sanitize(labels.FromStrings("span_name", "GET /api/users/789"))
+	require.Contains(t, result.Get("span_name"), "<_>")
+
+	// Toggle to disabled at runtime - should pass through
+	mode = SpanNameSanitizationDisabled
+	result = sanitizer.Sanitize(labels.FromStrings("span_name", "GET /api/users/999"))
+	require.Equal(t, "GET /api/users/999", result.Get("span_name"))
+
+	// Toggle to dry-run at runtime - should pass through but still train
+	mode = SpanNameSanitizationDryRun
+	result = sanitizer.Sanitize(labels.FromStrings("span_name", "GET /api/users/111"))
+	require.Equal(t, "GET /api/users/111", result.Get("span_name"))
+
+	// Toggle back to enabled - should sanitize again
+	mode = SpanNameSanitizationEnabled
+	result = sanitizer.Sanitize(labels.FromStrings("span_name", "GET /api/users/222"))
+	require.Contains(t, result.Get("span_name"), "<_>")
+}
+
+func TestDrainSanitizer_DisabledMode(t *testing.T) {
+	t.Parallel()
+
+	sanitizer := newTestDrainSanitizer(SpanNameSanitizationDisabled)
+
+	lbls1 := labels.FromStrings("span_name", "GET /api/users/123", "service", "api")
+	lbls2 := labels.FromStrings("span_name", "GET /api/users/456", "service", "api")
+
+	sanitizer.Sanitize(lbls1)
+
+	// When disabled, should always return original labels
+	result := sanitizer.Sanitize(lbls2)
+	require.Equal(t, "GET /api/users/456", result.Get("span_name"))
+	require.Equal(t, lbls2, result)
 }
 
 func TestDrainSanitizer_NilClusterHandling(t *testing.T) {
 	t.Parallel()
 
-	sanitizer := NewDrainSanitizer("test-tenant", false, 15*time.Minute)
+	sanitizer := newTestDrainSanitizer(SpanNameSanitizationEnabled)
 
 	// Span name with too few tokens (less than MinTokens=3)
 	// Tokenizer will produce tokens like ["a", "<END>"] which is < 3
 	lbls := labels.FromStrings("span_name", "a", "service", "api")
 	result := sanitizer.Sanitize(lbls)
 
 	// Should return original labels when cluster is nil
-	assert.Equal(t, "a", result.Get("span_name"))
-	assert.Equal(t, lbls, result)
+	require.Equal(t, "a", result.Get("span_name"))
+	require.Equal(t, lbls, result)
 }
 
 func TestDrainSanitizer_ConcurrentAccess(t *testing.T) {
 	t.Parallel()
 
-	sanitizer := NewDrainSanitizer("test-tenant", false, 15*time.Minute)
+	sanitizer := newTestDrainSanitizer(SpanNameSanitizationEnabled)
 
 	var wg sync.WaitGroup
 	numGoroutines := 10
@@ -85,8 +135,8 @@ func TestDrainSanitizer_ConcurrentAccess(t *testing.T) {
 				lbls := labels.FromStrings("span_name", "GET /api/users/123", "id", string(rune(id*1000+j)))
 				result := sanitizer.Sanitize(lbls)
 				// Should always return valid labels
-				assert.NotNil(t, result)
-				assert.NotEmpty(t, result.Get("span_name"))
+				require.NotNil(t, result)
+				require.NotEmpty(t, result.Get("span_name"))
 			}
 		}(i)
 	}
@@ -98,7 +148,7 @@ func TestDrainSanitizer_ConcurrentAccess(t *testing.T) {
 func TestDrainSanitizer_DemandTracking(t *testing.T) {
 	t.Parallel()
 
-	sanitizer := NewDrainSanitizer("test-tenant", false, 15*time.Minute)
+	sanitizer := newTestDrainSanitizer(SpanNameSanitizationEnabled)
 
 	// Create labels with different span names
 	lbls1 := labels.FromStrings("span_name", "GET /api/users/123")
@@ -114,42 +164,42 @@ func TestDrainSanitizer_DemandTracking(t *testing.T) {
 	// but we can verify the sanitizer doesn't crash and processes all labels)
 	// The demand gauge will be updated periodically via doPeriodicMaintenance
 	demandEstimate := sanitizer.demand.Estimate()
-	assert.GreaterOrEqual(t, demandEstimate, uint64(1))
+	require.GreaterOrEqual(t, demandEstimate, uint64(1))
 }
 
 func TestDrainSanitizer_NoSpanNameLabel(t *testing.T) {
 	t.Parallel()
 
-	sanitizer := NewDrainSanitizer("test-tenant", false, 15*time.Minute)
+	sanitizer := newTestDrainSanitizer(SpanNameSanitizationEnabled)
 
 	// Labels without span_name
 	lbls := labels.FromStrings("service", "api", "method", "GET")
 	result := sanitizer.Sanitize(lbls)
 
 	// Should return original labels (span_name is empty string, which drain will reject)
-	assert.Equal(t, lbls, result)
+	require.Equal(t, lbls, result)
 }
 
 func TestDrainSanitizer_PatternBeforeSanitization(t *testing.T) {
 	t.Parallel()
 
-	sanitizer := NewDrainSanitizer("test-tenant", false, 15*time.Minute)
+	sanitizer := newTestDrainSanitizer(SpanNameSanitizationEnabled)
 
 	// First span name - no pattern yet, so returns original
 	lbls1 := labels.FromStrings("span_name", "GET /api/users/123")
 	result1 := sanitizer.Sanitize(lbls1)
-	assert.Equal(t, "GET /api/users/123", result1.Get("span_name"))
+	require.Equal(t, "GET /api/users/123", result1.Get("span_name"))
 
 	// Same span name again - still no pattern (only one instance)
 	result2 := sanitizer.Sanitize(lbls1)
-	assert.Equal(t, "GET /api/users/123", result2.Get("span_name"))
+	require.Equal(t, "GET /api/users/123", result2.Get("span_name"))
 
 	// Different but similar span name - now pattern should emerge
 	lbls2 := labels.FromStrings("span_name", "GET /api/users/456")
 	result3 := sanitizer.Sanitize(lbls2)
 	// After pattern detection, should return sanitized version
-	assert.NotEqual(t, "GET /api/users/456", result3.Get("span_name"))
-	assert.Contains(t, result3.Get("span_name"), "<_>")
+	require.NotEqual(t, "GET /api/users/456", result3.Get("span_name"))
+	require.Contains(t, result3.Get("span_name"), "<_>")
 }
 
 func TestDrainSanitizer_FullSanitizedOutput(t *testing.T) {
@@ -185,7 +235,7 @@ func TestDrainSanitizer_FullSanitizedOutput(t *testing.T) {
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
-			sanitizer := NewDrainSanitizer("test-tenant", false, 15*time.Minute)
+			sanitizer := newTestDrainSanitizer(SpanNameSanitizationEnabled)
 
 			// Train with first inputs
 			for _, input := range tc.inputs[:len(tc.inputs)-1] {
@@ -195,7 +245,7 @@ func TestDrainSanitizer_FullSanitizedOutput(t *testing.T) {
 			// Last input should produce the expected sanitized output
 			lastInput := tc.inputs[len(tc.inputs)-1]
 			result := sanitizer.Sanitize(labels.FromStrings("span_name", lastInput))
-			assert.Equal(t, tc.expectedOutput, result.Get("span_name"))
+			require.Equal(t, tc.expectedOutput, result.Get("span_name"))
 		})
 	}
 }

modules/generator/registry/registry.go

@@ -98,19 +98,11 @@ type Limiter interface {
 	OnDelete(labelHash uint64, seriesCount uint32)
 }
 
-// Sanitizer is applies a transformation to all non-constant labels.
+// Sanitizer applies a transformation to all non-constant labels.
 type Sanitizer interface {
 	Sanitize(lbls labels.Labels) labels.Labels
 }
 
-type noopSanitizer struct{}
-
-var _ Sanitizer = (*noopSanitizer)(nil)
-
-func (s *noopSanitizer) Sanitize(lbls labels.Labels) labels.Labels {
-	return lbls
-}
-
 // New creates a ManagedRegistry. This Registry will scrape itself, write samples into an appender
 // and remove stale series.
 func New(cfg *Config, overrides Overrides, tenant string, appendable storage.Appendable, logger log.Logger, limiter Limiter) *ManagedRegistry {
@@ -127,10 +119,7 @@ func New(cfg *Config, overrides Overrides, tenant string, appendable storage.App
 		externalLabels[cfg.InjectTenantIDAs] = tenant
 	}
 
-	var sanitizer Sanitizer = &noopSanitizer{}
-	if sanitizationMode := overrides.MetricsGeneratorSpanNameSanitization(tenant); sanitizationMode != SpanNameSanitizationDisabled {
-		sanitizer = NewDrainSanitizer(tenant, sanitizationMode == SpanNameSanitizationDryRun, cfg.StaleDuration)
-	}
+	sanitizer := NewDrainSanitizer(tenant, overrides.MetricsGeneratorSpanNameSanitization, cfg.StaleDuration)
 
 	r := &ManagedRegistry{
 		onShutdown: cancel,

modules/generator/registry/registry_test.go

@@ -57,7 +57,7 @@ func (m *mockLimiter) OnPruneStaleSeries() {
 }
 
 func buildTestLabels(names []string, values []string) labels.Labels {
-	builder := NewLabelBuilder(0, 0, &noopSanitizer{})
+	builder := NewLabelBuilder(0, 0, newTestDrainSanitizer(SpanNameSanitizationDisabled))
 	for i := range names {
 		builder.Add(names[i], values[i])
 	}

modules/generator/registry/test.go

@@ -39,7 +39,9 @@ func (t *TestRegistry) NewGauge(name string) Gauge {
 }
 
 func (t *TestRegistry) NewLabelBuilder() LabelBuilder {
-	return NewLabelBuilder(0, 0, &noopSanitizer{})
+	nds := NewDrainSanitizer("test", func(string) string { return SpanNameSanitizationDisabled }, 0)
+
+	return NewLabelBuilder(0, 0, nds)
 }
 
 func (t *TestRegistry) NewHistogram(name string, buckets []float64, histogramOverrides HistogramMode) Histogram {