[chore][processor/tailsampling] Refactor tests

processor/tailsamplingprocessor/processor.go

@@ -70,65 +70,99 @@ var (
 	attrSampledTrue     = metric.WithAttributes(attribute.String("sampled", "true"))
 	attrSampledFalse    = metric.WithAttributes(attribute.String("sampled", "false"))
 	decisionToAttribute = map[sampling.Decision]metric.MeasurementOption{
-		sampling.Sampled:    attrSampledTrue,
-		sampling.NotSampled: attrSampledFalse,
+		sampling.Sampled:          attrSampledTrue,
+		sampling.NotSampled:       attrSampledFalse,
+		sampling.InvertNotSampled: attrSampledFalse,
+		sampling.InvertSampled:    attrSampledTrue,
 	}
 )
 
+type Option func(*tailSamplingSpanProcessor)
+
 // newTracesProcessor returns a processor.TracesProcessor that will perform tail sampling according to the given
 // configuration.
-func newTracesProcessor(ctx context.Context, settings component.TelemetrySettings, nextConsumer consumer.Traces, cfg Config) (processor.Traces, error) {
-	policyNames := map[string]bool{}
-	policies := make([]*policy, len(cfg.PolicyCfgs))
-	for i := range cfg.PolicyCfgs {
-		policyCfg := &cfg.PolicyCfgs[i]
-
-		if policyNames[policyCfg.Name] {
-			return nil, fmt.Errorf("duplicate policy name %q", policyCfg.Name)
+func newTracesProcessor(ctx context.Context, settings component.TelemetrySettings, nextConsumer consumer.Traces, cfg Config, opts ...Option) (processor.Traces, error) {
+	telemetry, err := metadata.NewTelemetryBuilder(settings)
+	if err != nil {
+		return nil, err
+	}
+
+	tsp := &tailSamplingSpanProcessor{
+		ctx:            ctx,
+		telemetry:      telemetry,
+		nextConsumer:   nextConsumer,
+		maxNumTraces:   cfg.NumTraces,
+		logger:         settings.Logger,
+		numTracesOnMap: &atomic.Uint64{},
+		deleteChan:     make(chan pcommon.TraceID, cfg.NumTraces),
+	}
+	tsp.policyTicker = &timeutils.PolicyTicker{OnTickFunc: tsp.samplingPolicyOnTick}
+
+	for _, opt := range opts {
+		opt(tsp)
+	}
+
+	if tsp.tickerFrequency == 0 {
+		tsp.tickerFrequency = time.Second
+	}
+
+	if tsp.policies == nil {
+		policyNames := map[string]bool{}
+		tsp.policies = make([]*policy, len(cfg.PolicyCfgs))
+		for i := range cfg.PolicyCfgs {
+			policyCfg := &cfg.PolicyCfgs[i]
+
+			if policyNames[policyCfg.Name] {
+				return nil, fmt.Errorf("duplicate policy name %q", policyCfg.Name)
+			}
+			policyNames[policyCfg.Name] = true
+
+			eval, err := getPolicyEvaluator(settings, policyCfg)
+			if err != nil {
+				return nil, err
+			}
+			p := &policy{
+				name:      policyCfg.Name,
+				evaluator: eval,
+				attribute: metric.WithAttributes(attribute.String("policy", policyCfg.Name)),
+			}
+			tsp.policies[i] = p
 		}
-		policyNames[policyCfg.Name] = true
+	}
 
-		eval, err := getPolicyEvaluator(settings, policyCfg)
+	if tsp.decisionBatcher == nil {
+		// this will start a goroutine in the background, so we run it only if everything went
+		// well in creating the policies
+		numDecisionBatches := math.Max(1, cfg.DecisionWait.Seconds())
+		inBatcher, err := idbatcher.New(uint64(numDecisionBatches), cfg.ExpectedNewTracesPerSec, uint64(2*runtime.NumCPU()))
 		if err != nil {
 			return nil, err
 		}
-		p := &policy{
-			name:      policyCfg.Name,
-			evaluator: eval,
-			attribute: metric.WithAttributes(attribute.String("policy", policyCfg.Name)),
-		}
-		policies[i] = p
+		tsp.decisionBatcher = inBatcher
 	}
 
-	// this will start a goroutine in the background, so we run it only if everything went
-	// well in creating the policies
-	numDecisionBatches := math.Max(1, cfg.DecisionWait.Seconds())
-	inBatcher, err := idbatcher.New(uint64(numDecisionBatches), cfg.ExpectedNewTracesPerSec, uint64(2*runtime.NumCPU()))
-	if err != nil {
-		return nil, err
-	}
+	return tsp, nil
+}
 
-	telemetry, err := metadata.NewTelemetryBuilder(settings)
-	if err != nil {
-		return nil, err
+// withDecisionBatcher sets the batcher used to batch trace IDs for policy evaluation.
+func withDecisionBatcher(batcher idbatcher.Batcher) Option {
+	return func(tsp *tailSamplingSpanProcessor) {
+		tsp.decisionBatcher = batcher
 	}
+}
 
-	tsp := &tailSamplingSpanProcessor{
-		ctx:             ctx,
-		telemetry:       telemetry,
-		nextConsumer:    nextConsumer,
-		maxNumTraces:    cfg.NumTraces,
-		logger:          settings.Logger,
-		decisionBatcher: inBatcher,
-		policies:        policies,
-		tickerFrequency: time.Second,
-		numTracesOnMap:  &atomic.Uint64{},
+// withPolicies sets the sampling policies to be used by the processor.
+func withPolicies(policies []*policy) Option {
+	return func(tsp *tailSamplingSpanProcessor) {
+		tsp.policies = policies
 	}
+}
 
-	tsp.policyTicker = &timeutils.PolicyTicker{OnTickFunc: tsp.samplingPolicyOnTick}
-	tsp.deleteChan = make(chan pcommon.TraceID, cfg.NumTraces)
-
-	return tsp, nil
+// withTickerFrequency sets the frequency at which the processor will evaluate the sampling policies.
+func withTickerFrequency(frequency time.Duration) Option {
+	return func(tsp *tailSamplingSpanProcessor) {
+		tsp.tickerFrequency = frequency
+	}
 }
 
 func getPolicyEvaluator(settings component.TelemetrySettings, cfg *PolicyCfg) (sampling.PolicyEvaluator, error) {

processor/tailsamplingprocessor/processor_benchmarks_test.go

@@ -0,0 +1,50 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package tailsamplingprocessor
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/collector/component/componenttest"
+	"go.opentelemetry.io/collector/consumer/consumertest"
+	"go.opentelemetry.io/collector/pdata/ptrace"
+
+	"github.com/open-telemetry/opentelemetry-collector-contrib/processor/tailsamplingprocessor/internal/sampling"
+)
+
+func BenchmarkSampling(b *testing.B) {
+	traceIDs, batches := generateIDsAndBatches(128)
+	cfg := Config{
+		DecisionWait:            defaultTestDecisionWait,
+		NumTraces:               uint64(2 * len(traceIDs)),
+		ExpectedNewTracesPerSec: 64,
+		PolicyCfgs:              testPolicy,
+	}
+
+	sp, _ := newTracesProcessor(context.Background(), componenttest.NewNopTelemetrySettings(), consumertest.NewNop(), cfg)
+	tsp := sp.(*tailSamplingSpanProcessor)
+	require.NoError(b, tsp.Start(context.Background(), componenttest.NewNopHost()))
+	defer func() {
+		require.NoError(b, tsp.Shutdown(context.Background()))
+	}()
+	metrics := &policyMetrics{}
+	sampleBatches := make([]*sampling.TraceData, 0, len(batches))
+
+	for i := 0; i < len(batches); i++ {
+		sampleBatches = append(sampleBatches, &sampling.TraceData{
+			ArrivalTime: time.Now(),
+			//SpanCount:       spanCount,
+			ReceivedBatches: ptrace.NewTraces(),
+		})
+	}
+
+	for i := 0; i < b.N; i++ {
+		for i, id := range traceIDs {
+			_ = tsp.makeDecision(id, sampleBatches[i], metrics)
+		}
+	}
+}