traces/modules/generator/generator_test.go

package generator

import (
	"context"
	"flag"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"testing"
	"time"

	"github.com/go-kit/log"
	"github.com/grafana/dskit/services"
	"github.com/grafana/tempo/modules/generator/processor"
	"github.com/grafana/tempo/modules/generator/processor/spanmetrics"
	"github.com/grafana/tempo/modules/generator/storage"
	"github.com/grafana/tempo/modules/overrides"
	"github.com/grafana/tempo/modules/overrides/histograms"
	"github.com/grafana/tempo/pkg/tempopb"
	common_v1 "github.com/grafana/tempo/pkg/tempopb/common/v1"
	trace_v1 "github.com/grafana/tempo/pkg/tempopb/trace/v1"
	"github.com/grafana/tempo/pkg/util/test"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/model"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

const (
	user1 = "user1"
	user2 = "user2"
)

func TestGeneratorSpanMetrics_subprocessorConcurrency(t *testing.T) {
	overridesFile := filepath.Join(t.TempDir(), "Overrides.yaml")
	overridesConfig := overrides.Config{
		Defaults: overrides.Overrides{
			MetricsGenerator: overrides.MetricsGeneratorOverrides{
				Processors: map[string]struct{}{
					processor.SpanMetricsName: {},
				},
				CollectionInterval: 2 * time.Second,
			},
		},
		PerTenantOverrideConfig: overridesFile,
		PerTenantOverridePeriod: model.Duration(time.Second),
	}

	require.NoError(t, os.WriteFile(overridesFile, []byte(fmt.Sprintf(`
overrides:
  %s:
    metrics_generator:
      collection_interval: 1s
      processors:
        - %s
`, user1, processor.SpanMetricsName)), 0o700))

	o, err := overrides.NewOverrides(overridesConfig, nil, prometheus.NewRegistry())
	require.NoError(t, err)
	require.NoError(t, services.StartAndAwaitRunning(context.Background(), o))

	generatorConfig := &Config{}
	generatorConfig.RegisterFlagsAndApplyDefaults("", &flag.FlagSet{})
	generatorConfig.Storage.Path = t.TempDir()
	generatorConfig.Ring.KVStore.Store = "inmemory"
	ifaces, err := net.Interfaces()
	require.NoError(t, err)
	netWorkInteraces := make([]string, len(ifaces))
	for i, iface := range ifaces {
		netWorkInteraces[i] = iface.Name
	}
	generatorConfig.Ring.InstanceInterfaceNames = netWorkInteraces
	g, err := New(generatorConfig, o, prometheus.NewRegistry(), nil, nil, newTestLogger(t))
	require.NoError(t, err)
	require.NoError(t, services.StartAndAwaitRunning(context.Background(), g))

	t.Cleanup(func() {
		require.NoError(t, services.StopAndAwaitTerminated(context.Background(), o))
		require.NoError(t, services.StopAndAwaitTerminated(context.Background(), g))
	})

	allSubprocessors := map[spanmetrics.Subprocessor]bool{spanmetrics.Count: true, spanmetrics.Latency: true, spanmetrics.Size: true}

	// All subprocessors should be enabled for user1
	instance1, err := g.getOrCreateInstance(user1)
	require.NoError(t, err)
	verifySubprocessors(t, instance1, allSubprocessors)

	// All subprocessors should be enabled for user2
	instance2, err := g.getOrCreateInstance(user2)
	require.NoError(t, err)
	verifySubprocessors(t, instance2, allSubprocessors)

	// Change overrides for user1
	require.NoError(t, os.WriteFile(overridesFile, []byte(fmt.Sprintf(`
overrides:
  %s:
    metrics_generator:
      collection_interval: 1s
      processors:
        - %s
`, user1, spanmetrics.Count.String())), 0o700))
	time.Sleep(15 * time.Second) // Wait for overrides to be applied. Reload is hardcoded to 10s :(

	// Only Count should be enabled for user1
	instance1, err = g.getOrCreateInstance(user1)
	require.NoError(t, err)
	verifySubprocessors(t, instance1, map[spanmetrics.Subprocessor]bool{spanmetrics.Count: true, spanmetrics.Latency: false, spanmetrics.Size: false})

	// All subprocessors should be enabled for user2
	instance2, err = g.getOrCreateInstance(user2)
	require.NoError(t, err)
	verifySubprocessors(t, instance2, allSubprocessors)
}

func verifySubprocessors(t *testing.T, instance *instance, expected map[spanmetrics.Subprocessor]bool) {
	instance.processorsMtx.RLock()
	defer instance.processorsMtx.RUnlock()

	require.Len(t, instance.processors, 1)

	processor, ok := instance.processors[processor.SpanMetricsName]
	require.True(t, ok)

	require.Equal(t, len(processor.(*spanmetrics.Processor).Cfg.Subprocessors), len(expected))

	cfg := processor.(*spanmetrics.Processor).Cfg
	for sub, enabled := range expected {
		assert.Equal(t, enabled, cfg.Subprocessors[sub])
	}
}

var _ log.Logger = (*testLogger)(nil)

type testLogger struct {
	t *testing.T
}

func newTestLogger(t *testing.T) log.Logger {
	return testLogger{t: t}
}

func (l testLogger) Log(keyvals ...interface{}) error {
	l.t.Log(keyvals...)
	return nil
}

func BenchmarkPushSpans(b *testing.B) {
	var (
		tenant = "test-tenant"
		reg    = prometheus.NewRegistry()
		ctx    = context.Background()
		log    = log.NewNopLogger()
		cfg    = &Config{}

		walcfg = &storage.Config{
			Path: b.TempDir(),
		}

		o = &mockOverrides{
			processors: map[string]struct{}{
				"span-metrics":   {},
				"service-graphs": {},
			},
			spanMetricsEnableTargetInfo:             boolPtr(true),
			spanMetricsTargetInfoExcludedDimensions: []string{"excluded}"},
		}
	)

	cfg.RegisterFlagsAndApplyDefaults("", &flag.FlagSet{})

	wal, err := storage.New(walcfg, o, tenant, reg, log)
	require.NoError(b, err)

	inst, err := newInstance(cfg, tenant, o, wal, log, nil, nil, nil)
	require.NoError(b, err)
	defer inst.shutdown()

	req := &tempopb.PushSpansRequest{
		Batches: []*trace_v1.ResourceSpans{
			test.MakeBatch(100, nil),
			test.MakeBatch(100, nil),
			test.MakeBatch(100, nil),
			test.MakeBatch(100, nil),
			test.MakeBatch(100, nil),
			test.MakeBatch(100, nil),
			test.MakeBatch(100, nil),
			test.MakeBatch(100, nil),
			test.MakeBatch(100, nil),
			test.MakeBatch(100, nil),
		},
	}

	// Add more resource attributes to get closer to real data
	// Add integer to increase cardinality.
	// Currently this is about 80 active series
	// TODO - Get more series
	for i, b := range req.Batches {
		b.Resource.Attributes = append(b.Resource.Attributes, []*common_v1.KeyValue{
			{Key: "k8s.cluster.name", Value: &common_v1.AnyValue{Value: &common_v1.AnyValue_StringValue{StringValue: "test" + strconv.Itoa(i)}}},
			{Key: "k8s.namespace.name", Value: &common_v1.AnyValue{Value: &common_v1.AnyValue_StringValue{StringValue: "test" + strconv.Itoa(i)}}},
			{Key: "k8s.node.name", Value: &common_v1.AnyValue{Value: &common_v1.AnyValue_StringValue{StringValue: "test" + strconv.Itoa(i)}}},
			{Key: "k8s.pod.ip", Value: &common_v1.AnyValue{Value: &common_v1.AnyValue_StringValue{StringValue: "test" + strconv.Itoa(i)}}},
			{Key: "k8s.pod.name", Value: &common_v1.AnyValue{Value: &common_v1.AnyValue_StringValue{StringValue: "test" + strconv.Itoa(i)}}},
			{Key: "service.instance.id", Value: &common_v1.AnyValue{Value: &common_v1.AnyValue_StringValue{StringValue: "test" + strconv.Itoa(i)}}},
			{Key: "excluded", Value: &common_v1.AnyValue{Value: &common_v1.AnyValue_StringValue{StringValue: "test" + strconv.Itoa(i)}}},
		}...)
	}

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		inst.pushSpans(ctx, req)
	}

	b.StopTimer()
	runtime.GC()
	mem := runtime.MemStats{}
	runtime.ReadMemStats(&mem)
	b.ReportMetric(float64(mem.HeapInuse), "heap_in_use")
	b.ReportMetric(float64(mem.HeapAlloc), "heap_alloc")
}

func TestGetOrCreateInstance_FailureCaching(t *testing.T) {
	cfg := &Config{}
	cfg.RegisterFlagsAndApplyDefaults("", &flag.FlagSet{})
	cfg.Storage.Path = t.TempDir()

	g := &Generator{
		cfg:             cfg,
		overrides:       &mockOverrides{processors: map[string]struct{}{"invalid-processor": {}}}, // will fail
		instances:       map[string]*instance{},
		failedInstances: map[string]time.Time{},
		logger:          newTestLogger(t),
		reg:             prometheus.NewRegistry(),
	}

	// First call should fail and cache the failure
	_, err := g.getOrCreateInstance("tenant-1")
	require.Error(t, err)
	require.NotErrorIs(t, err, errInstanceCreationBackoff)

	// Second call should return cached error
	_, err = g.getOrCreateInstance("tenant-1")
	require.ErrorIs(t, err, errInstanceCreationBackoff)

	// Different tenant should also fail (not cached)
	_, err = g.getOrCreateInstance("tenant-2")
	require.Error(t, err)
	require.NotErrorIs(t, err, errInstanceCreationBackoff)
}

func TestGetOrCreateInstance_FailureCachingExpiry(t *testing.T) {
	cfg := &Config{}
	cfg.RegisterFlagsAndApplyDefaults("", &flag.FlagSet{})
	cfg.Storage.Path = t.TempDir()

	g := &Generator{
		cfg:             cfg,
		overrides:       &mockOverrides{processors: map[string]struct{}{"invalid-processor": {}}}, // will fail
		instances:       map[string]*instance{},
		failedInstances: map[string]time.Time{},
		logger:          newTestLogger(t),
		reg:             prometheus.NewRegistry(),
	}

	// First call should fail and cache the failure
	_, err := g.getOrCreateInstance("tenant-1")
	require.Error(t, err)
	require.NotErrorIs(t, err, errInstanceCreationBackoff)

	// Simulate backoff expiry by setting the failure time in the past
	g.instancesMtx.Lock()
	g.failedInstances["tenant-1"] = time.Now().Add(-failureBackoff - time.Second)
	g.instancesMtx.Unlock()

	// After backoff expires, should retry (and fail again with actual error, not backoff error)
	_, err = g.getOrCreateInstance("tenant-1")
	require.Error(t, err)
	require.NotErrorIs(t, err, errInstanceCreationBackoff)
}

func BenchmarkCollect(b *testing.B) {
	var (
		tenant = "test-tenant"
		reg    = prometheus.NewRegistry()
		ctx    = context.Background()
		log    = log.NewNopLogger()
		cfg    = &Config{}

		walcfg = &storage.Config{
			Path: b.TempDir(),
		}

		o = &mockOverrides{
			processors: map[string]struct{}{
				"span-metrics":   {},
				"service-graphs": {},
			},
			spanMetricsDimensions:                   []string{"k8s.cluster.name", "k8s.namespace.name"},
			spanMetricsEnableTargetInfo:             boolPtr(true),
			spanMetricsTargetInfoExcludedDimensions: []string{"excluded}"},
			nativeHistograms:                        histograms.HistogramMethodBoth,
		}
	)

	cfg.RegisterFlagsAndApplyDefaults("", &flag.FlagSet{})

	wal, err := storage.New(walcfg, o, tenant, reg, log)
	require.NoError(b, err)

	inst, err := newInstance(cfg, tenant, o, wal, log, nil, nil, nil)
	require.NoError(b, err)
	defer inst.shutdown()

	req := &tempopb.PushSpansRequest{
		Batches: []*trace_v1.ResourceSpans{
			test.MakeBatch(100, nil),
			test.MakeBatch(100, nil),
			test.MakeBatch(100, nil),
			test.MakeBatch(100, nil),
		},
	}

	// Add more resource attributes to get closer to real data
	// Add integer to increase cardinality.
	// Currently this is about 80 active series
	// TODO - Get more series
	for i, b := range req.Batches {
		b.Resource.Attributes = append(b.Resource.Attributes, []*common_v1.KeyValue{
			{Key: "k8s.cluster.name", Value: &common_v1.AnyValue{Value: &common_v1.AnyValue_StringValue{StringValue: "test" + strconv.Itoa(i)}}},
			{Key: "k8s.namespace.name", Value: &common_v1.AnyValue{Value: &common_v1.AnyValue_StringValue{StringValue: "test" + strconv.Itoa(i)}}},
			{Key: "k8s.node.name", Value: &common_v1.AnyValue{Value: &common_v1.AnyValue_StringValue{StringValue: "test" + strconv.Itoa(i)}}},
			{Key: "k8s.pod.ip", Value: &common_v1.AnyValue{Value: &common_v1.AnyValue_StringValue{StringValue: "test" + strconv.Itoa(i)}}},
			{Key: "k8s.pod.name", Value: &common_v1.AnyValue{Value: &common_v1.AnyValue_StringValue{StringValue: "test" + strconv.Itoa(i)}}},
			{Key: "excluded", Value: &common_v1.AnyValue{Value: &common_v1.AnyValue_StringValue{StringValue: "test" + strconv.Itoa(i)}}},
		}...)
	}
	inst.pushSpans(ctx, req)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		inst.registry.CollectMetrics(ctx)
	}

	b.StopTimer()
	runtime.GC()
	mem := runtime.MemStats{}
	runtime.ReadMemStats(&mem)
	b.ReportMetric(float64(mem.HeapInuse), "heap_in_use")
}