[refactor] move root span handler into aggregator (#5478)

## Which problem is this PR solving?
-
#5389 (comment)

## Description of the changes
- Refactored `handleRootSpan` logic into a helper method in
aggregator.go.

Author: pushkarm029

cmd/collector/app/collector.go

@@ -30,6 +30,7 @@ import (
 	"github.com/jaegertracing/jaeger/cmd/collector/app/processor"
 	"github.com/jaegertracing/jaeger/cmd/collector/app/sampling/strategystore"
 	"github.com/jaegertracing/jaeger/cmd/collector/app/server"
+	"github.com/jaegertracing/jaeger/model"
 	"github.com/jaegertracing/jaeger/pkg/healthcheck"
 	"github.com/jaegertracing/jaeger/pkg/metrics"
 	"github.com/jaegertracing/jaeger/pkg/tenancy"
@@ -103,7 +104,9 @@ func (c *Collector) Start(options *flags.CollectorOptions) error {
 
 	var additionalProcessors []ProcessSpan
 	if c.aggregator != nil {
-		additionalProcessors = append(additionalProcessors, handleRootSpan(c.aggregator, c.logger))
+		additionalProcessors = append(additionalProcessors, func(span *model.Span, tenant string) {
+			c.aggregator.HandleRootSpan(span, c.logger)
+		})
 	}
 
 	c.spanProcessor = handlerBuilder.BuildSpanProcessor(additionalProcessors...)

cmd/collector/app/collector_test.go

@@ -17,6 +17,7 @@ package app
 import (
 	"context"
 	"io"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -47,6 +48,26 @@ func optionsForEphemeralPorts() *flags.CollectorOptions {
 	return collectorOpts
 }
 
+type mockAggregator struct {
+	callCount  atomic.Int32
+	closeCount atomic.Int32
+}
+
+func (t *mockAggregator) RecordThroughput(service, operation string, samplerType model.SamplerType, probability float64) {
+	t.callCount.Add(1)
+}
+
+func (t *mockAggregator) HandleRootSpan(span *model.Span, logger *zap.Logger) {
+	t.callCount.Add(1)
+}
+
+func (t *mockAggregator) Start() {}
+
+func (t *mockAggregator) Close() error {
+	t.closeCount.Add(1)
+	return nil
+}
+
 func TestNewCollector(t *testing.T) {
 	// prepare
 	hc := healthcheck.New()

cmd/collector/app/root_span_handler.go

@@ -18,6 +18,8 @@ import (
 	"context"
 	"io"
 
+	"go.uber.org/zap"
+
 	"github.com/jaegertracing/jaeger/model"
 	"github.com/jaegertracing/jaeger/proto-gen/api_v2"
 )
@@ -36,6 +38,10 @@ type Aggregator interface {
 	// Close() from io.Closer stops the aggregator from aggregating throughput.
 	io.Closer
 
+	// The HandleRootSpan function processes a span, checking if it's a root span.
+	// If it is, it extracts sampler parameters, then calls RecordThroughput.
+	HandleRootSpan(span *model.Span, logger *zap.Logger)
+
 	// RecordThroughput records throughput for an operation for aggregation.
 	RecordThroughput(service, operation string, samplerType model.SamplerType, probability float64)
 

cmd/collector/app/root_span_handler_test.go

@@ -144,3 +144,20 @@ func (a *aggregator) Close() error {
 	a.bgFinished.Wait()
 	return nil
 }
+
+func (a *aggregator) HandleRootSpan(span *span_model.Span, logger *zap.Logger) {
+	// simply checking parentId to determine if a span is a root span is not sufficient. However,
+	// we can be sure that only a root span will have sampler tags.
+	if span.ParentSpanID() != span_model.NewSpanID(0) {
+		return
+	}
+	service := span.Process.ServiceName
+	if service == "" || span.OperationName == "" {
+		return
+	}
+	samplerType, samplerParam := span.GetSamplerParams(logger)
+	if samplerType == span_model.SamplerTypeUnrecognized {
+		return
+	}
+	a.RecordThroughput(service, span.OperationName, samplerType, samplerParam)
+}

cmd/collector/app/sampling/strategystore/interface.go

@@ -115,3 +115,43 @@ func TestLowerboundThroughput(t *testing.T) {
 	assert.EqualValues(t, 0, a.(*aggregator).currentThroughput["A"]["GET"].Count)
 	assert.Empty(t, a.(*aggregator).currentThroughput["A"]["GET"].Probabilities["0.001000"])
 }
+
+func TestRecordThroughput(t *testing.T) {
+	metricsFactory := metricstest.NewFactory(0)
+	mockStorage := &mocks.Store{}
+	mockEP := &epmocks.ElectionParticipant{}
+	testOpts := Options{
+		CalculationInterval:   1 * time.Second,
+		AggregationBuckets:    1,
+		BucketsForCalculation: 1,
+	}
+	logger := zap.NewNop()
+	a, err := NewAggregator(testOpts, logger, metricsFactory, mockEP, mockStorage)
+	require.NoError(t, err)
+
+	// Testing non-root span
+	span := &model.Span{References: []model.SpanRef{{SpanID: model.NewSpanID(1), RefType: model.ChildOf}}}
+	a.HandleRootSpan(span, logger)
+	require.Empty(t, a.(*aggregator).currentThroughput)
+
+	// Testing span with service name but no operation
+	span.References = []model.SpanRef{}
+	span.Process = &model.Process{
+		ServiceName: "A",
+	}
+	a.HandleRootSpan(span, logger)
+	require.Empty(t, a.(*aggregator).currentThroughput)
+
+	// Testing span with service name and operation but no probabilistic sampling tags
+	span.OperationName = "GET"
+	a.HandleRootSpan(span, logger)
+	require.Empty(t, a.(*aggregator).currentThroughput)
+
+	// Testing span with service name, operation, and probabilistic sampling tags
+	span.Tags = model.KeyValues{
+		model.String("sampler.type", "probabilistic"),
+		model.String("sampler.param", "0.001"),
+	}
+	a.HandleRootSpan(span, logger)
+	assert.EqualValues(t, 1, a.(*aggregator).currentThroughput["A"]["GET"].Count)
+}