Refactor Adaptive Sampling Aggregator & Strategy Store

plugin/sampling/strategystore/adaptive/aggregator.go

@@ -15,13 +15,18 @@
 package adaptive
 
 import (
+	"errors"
 	"sync"
 	"time"
 
+	"go.uber.org/zap"
+
 	"github.com/jaegertracing/jaeger/cmd/collector/app/sampling/model"
 	"github.com/jaegertracing/jaeger/cmd/collector/app/sampling/strategystore"
 	span_model "github.com/jaegertracing/jaeger/model"
+	"github.com/jaegertracing/jaeger/pkg/hostname"
 	"github.com/jaegertracing/jaeger/pkg/metrics"
+	"github.com/jaegertracing/jaeger/plugin/sampling/leaderelection"
 	"github.com/jaegertracing/jaeger/storage/samplingstore"
 )
 
@@ -35,32 +40,51 @@ type aggregator struct {
 	operationsCounter   metrics.Counter
 	servicesCounter     metrics.Counter
 	currentThroughput   serviceOperationThroughput
+	postAggregator      *PostAggregator
 	aggregationInterval time.Duration
 	storage             samplingstore.Store
 	stop                chan struct{}
+	bgFinished          sync.WaitGroup
 }
 
 // NewAggregator creates a throughput aggregator that simply emits metrics
 // about the number of operations seen over the aggregationInterval.
-func NewAggregator(metricsFactory metrics.Factory, interval time.Duration, storage samplingstore.Store) strategystore.Aggregator {
+func NewAggregator(options Options, logger *zap.Logger, metricsFactory metrics.Factory, participant leaderelection.ElectionParticipant, store samplingstore.Store) (strategystore.Aggregator, error) {
+	hostname, err := hostname.AsIdentifier()
+	if err != nil {
+		return nil, err
+	}
+	logger.Info("Using unique participantName in adaptive sampling", zap.String("participantName", hostname))
+
+	postAgg, err := newPostAggregator(options, hostname, store, participant, metricsFactory, logger)
+	if err != nil {
+		return nil, err
+	}
+
 	return &aggregator{
 		operationsCounter:   metricsFactory.Counter(metrics.Options{Name: "sampling_operations"}),
 		servicesCounter:     metricsFactory.Counter(metrics.Options{Name: "sampling_services"}),
 		currentThroughput:   make(serviceOperationThroughput),
-		aggregationInterval: interval,
-		storage:             storage,
+		aggregationInterval: options.CalculationInterval,
+		postAggregator:      postAgg,
+		storage:             store,
 		stop:                make(chan struct{}),
-	}
+	}, nil
 }
 
 func (a *aggregator) runAggregationLoop() {
 	ticker := time.NewTicker(a.aggregationInterval)
+
+	// NB: the first tick will be slightly delayed by the initializeThroughput call.
+	a.postAggregator.lastCheckedTime = time.Now().Add(a.postAggregator.Delay * -1)
+	a.postAggregator.initializeThroughput(a.postAggregator.lastCheckedTime)
 	for {
 		select {
 		case <-ticker.C:
 			a.Lock()
 			a.saveThroughput()
 			a.currentThroughput = make(serviceOperationThroughput)
+			a.postAggregator.runCalculation()
 			a.Unlock()
 		case <-a.stop:
 			ticker.Stop()
@@ -111,10 +135,26 @@ func (a *aggregator) RecordThroughput(service, operation string, samplerType spa
 }
 
 func (a *aggregator) Start() {
-	go a.runAggregationLoop()
+	a.postAggregator.Start()
+	a.runBackground(a.runAggregationLoop)
 }
 
 func (a *aggregator) Close() error {
-	close(a.stop)
-	return nil
+	var errs []error
+	if err := a.postAggregator.Close(); err != nil {
+		errs = append(errs, err)
+	}
+	if a.stop != nil {
+		close(a.stop)
+	}
+	a.bgFinished.Wait()
+	return errors.Join(errs...)
+}
+
+func (a *aggregator) runBackground(f func()) {
+	a.bgFinished.Add(1)
+	go func() {
+		f()
+		a.bgFinished.Done()
+	}()
 }

plugin/sampling/strategystore/adaptive/aggregator_test.go

@@ -20,9 +20,12 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/zap"
 
 	"github.com/jaegertracing/jaeger/internal/metricstest"
 	"github.com/jaegertracing/jaeger/model"
+	epmocks "github.com/jaegertracing/jaeger/plugin/sampling/leaderelection/mocks"
 	"github.com/jaegertracing/jaeger/storage/samplingstore/mocks"
 )
 
@@ -32,8 +35,19 @@ func TestAggregator(t *testing.T) {
 
 	mockStorage := &mocks.Store{}
 	mockStorage.On("InsertThroughput", mock.AnythingOfType("[]*model.Throughput")).Return(nil)
+	mockEP := &epmocks.ElectionParticipant{}
+	mockEP.On("Start").Return(nil)
+	mockEP.On("Close").Return(nil)
+	mockEP.On("IsLeader").Return(true)
+	testOpts := Options{
+		CalculationInterval:   1 * time.Second,
+		AggregationBuckets:    1,
+		BucketsForCalculation: 1,
+	}
+	logger := zap.NewNop()
 
-	a := NewAggregator(metricsFactory, 5*time.Millisecond, mockStorage)
+	a, err := NewAggregator(testOpts, logger, metricsFactory, mockEP, mockStorage)
+	require.NoError(t, err)
 	a.RecordThroughput("A", "GET", model.SamplerTypeProbabilistic, 0.001)
 	a.RecordThroughput("B", "POST", model.SamplerTypeProbabilistic, 0.001)
 	a.RecordThroughput("C", "GET", model.SamplerTypeProbabilistic, 0.001)
@@ -60,15 +74,23 @@ func TestAggregator(t *testing.T) {
 func TestIncrementThroughput(t *testing.T) {
 	metricsFactory := metricstest.NewFactory(0)
 	mockStorage := &mocks.Store{}
-
-	a := NewAggregator(metricsFactory, 5*time.Millisecond, mockStorage)
+	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)
 	// 20 different probabilities
 	for i := 0; i < 20; i++ {
 		a.RecordThroughput("A", "GET", model.SamplerTypeProbabilistic, 0.001*float64(i))
 	}
 	assert.Len(t, a.(*aggregator).currentThroughput["A"]["GET"].Probabilities, 10)
 
-	a = NewAggregator(metricsFactory, 5*time.Millisecond, mockStorage)
+	a, err = NewAggregator(testOpts, logger, metricsFactory, mockEP, mockStorage)
+	require.NoError(t, err)
 	// 20 of the same probabilities
 	for i := 0; i < 20; i++ {
 		a.RecordThroughput("A", "GET", model.SamplerTypeProbabilistic, 0.001)
@@ -79,8 +101,16 @@ func TestIncrementThroughput(t *testing.T) {
 func TestLowerboundThroughput(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 := NewAggregator(metricsFactory, 5*time.Millisecond, mockStorage)
+	a, err := NewAggregator(testOpts, logger, metricsFactory, mockEP, mockStorage)
+	require.NoError(t, err)
 	a.RecordThroughput("A", "GET", model.SamplerTypeLowerBound, 0.001)
 	assert.EqualValues(t, 0, a.(*aggregator).currentThroughput["A"]["GET"].Count)
 	assert.Empty(t, a.(*aggregator).currentThroughput["A"]["GET"].Probabilities["0.001000"])

plugin/sampling/strategystore/adaptive/factory.go

@@ -25,6 +25,7 @@ import (
 	"github.com/jaegertracing/jaeger/pkg/distributedlock"
 	"github.com/jaegertracing/jaeger/pkg/metrics"
 	"github.com/jaegertracing/jaeger/plugin"
+	"github.com/jaegertracing/jaeger/plugin/sampling/leaderelection"
 	"github.com/jaegertracing/jaeger/storage"
 	"github.com/jaegertracing/jaeger/storage/samplingstore"
 )
@@ -83,12 +84,23 @@ func (f *Factory) Initialize(metricsFactory metrics.Factory, ssFactory storage.S
 
 // CreateStrategyStore implements strategystore.Factory
 func (f *Factory) CreateStrategyStore() (strategystore.StrategyStore, strategystore.Aggregator, error) {
-	p, err := NewStrategyStore(*f.options, f.metricsFactory, f.logger, f.lock, f.store)
+	participant := leaderelection.NewElectionParticipant(f.lock, defaultResourceName, leaderelection.ElectionParticipantOptions{
+		FollowerLeaseRefreshInterval: f.options.FollowerLeaseRefreshInterval,
+		LeaderLeaseRefreshInterval:   f.options.LeaderLeaseRefreshInterval,
+		Logger:                       f.logger,
+	})
+
+	ss, err := NewStrategyStore(*f.options, f.logger, participant, f.store)
+	if err != nil {
+		return nil, nil, err
+	}
+	ss.Start()
+
+	a, err := NewAggregator(*f.options, f.logger, f.metricsFactory, participant, f.store)
 	if err != nil {
 		return nil, nil, err
 	}
-	p.Start()
-	a := NewAggregator(f.metricsFactory, f.options.CalculationInterval, f.store)
 	a.Start()
-	return p, a, nil
+
+	return ss, a, nil
 }