YAAAAS - Yet Another Attempt At Adaptive Sampling

CHANGELOG.md

@@ -43,6 +43,10 @@ Unreleased
 -------------------
 ### Backend Changes
 
+#### New Features
+
+* Add support for adaptive sampling with a Cassandra backend. ([#2966](https://github.com/jaegertracing/jaeger/pull/2966), [@joe-elliott](https://github.com/joe-elliott)):
+
 #### Breaking Changes
 
 * Remove unused `--es-archive.max-span-age` flag ([#2865](https://github.com/jaegertracing/jaeger/pull/2865), [@albertteoh](https://github.com/albertteoh)):

cmd/all-in-one/main.go

@@ -68,7 +68,11 @@ func main() {
 	if err != nil {
 		log.Fatalf("Cannot initialize storage factory: %v", err)
 	}
-	strategyStoreFactory, err := ss.NewFactory(ss.FactoryConfigFromEnv())
+	strategyStoreFactoryConfig, err := ss.FactoryConfigFromEnv()
+	if err != nil {
+		log.Fatalf("Cannot initialize sampling strategy store factory config: %v", err)
+	}
+	strategyStoreFactory, err := ss.NewFactory(*strategyStoreFactoryConfig)
 	if err != nil {
 		log.Fatalf("Cannot initialize sampling strategy store factory: %v", err)
 	}
@@ -121,11 +125,16 @@ by default uses only in-memory database.`,
 				logger.Fatal("Failed to create metrics reader", zap.Error(err))
 			}
 
+			lock, samplingStore, err := storageFactory.CreateLockAndSamplingStore()
+			if err != nil {
+				logger.Fatal("Failed to create lock and sampling store for adaptive sampling", zap.Error(err))
+			}
+
 			strategyStoreFactory.InitFromViper(v, logger)
-			if err := strategyStoreFactory.Initialize(metricsFactory, logger); err != nil {
+			if err := strategyStoreFactory.Initialize(metricsFactory, lock, samplingStore, logger); err != nil {
 				logger.Fatal("Failed to init sampling strategy store factory", zap.Error(err))
 			}
-			strategyStore, err := strategyStoreFactory.CreateStrategyStore()
+			strategyStore, aggregator, err := strategyStoreFactory.CreateStrategyStore()
 			if err != nil {
 				logger.Fatal("Failed to create sampling strategy store", zap.Error(err))
 			}
@@ -143,6 +152,7 @@ by default uses only in-memory database.`,
 				MetricsFactory: metricsFactory,
 				SpanWriter:     spanWriter,
 				StrategyStore:  strategyStore,
+				Aggregator:     aggregator,
 				HealthCheck:    svc.HC(),
 			})
 			if err := c.Start(cOpts); err != nil {

cmd/collector/app/collector.go

@@ -40,6 +40,7 @@ type Collector struct {
 	metricsFactory metrics.Factory
 	spanWriter     spanstore.Writer
 	strategyStore  strategystore.StrategyStore
+	aggregator     strategystore.Aggregator
 	hCheck         *healthcheck.HealthCheck
 	spanProcessor  processor.SpanProcessor
 	spanHandlers   *SpanHandlers
@@ -59,6 +60,7 @@ type CollectorParams struct {
 	MetricsFactory metrics.Factory
 	SpanWriter     spanstore.Writer
 	StrategyStore  strategystore.StrategyStore
+	Aggregator     strategystore.Aggregator
 	HealthCheck    *healthcheck.HealthCheck
 }
 
@@ -70,6 +72,7 @@ func New(params *CollectorParams) *Collector {
 		metricsFactory: params.MetricsFactory,
 		spanWriter:     params.SpanWriter,
 		strategyStore:  params.StrategyStore,
+		aggregator:     params.Aggregator,
 		hCheck:         params.HealthCheck,
 	}
 }
@@ -83,7 +86,12 @@ func (c *Collector) Start(builderOpts *CollectorOptions) error {
 		MetricsFactory: c.metricsFactory,
 	}
 
-	c.spanProcessor = handlerBuilder.BuildSpanProcessor()
+	var additionalProcessors []ProcessSpan
+	if c.aggregator != nil {
+		additionalProcessors = append(additionalProcessors, handleRootSpan(c.aggregator, c.logger))
+	}
+
+	c.spanProcessor = handlerBuilder.BuildSpanProcessor(additionalProcessors...)
 	c.spanHandlers = handlerBuilder.BuildHandlers(c.spanProcessor)
 
 	grpcServer, err := server.StartGRPCServer(&server.GRPCServerParams{
@@ -168,6 +176,13 @@ func (c *Collector) Close() error {
 		c.logger.Error("failed to close span processor.", zap.Error(err))
 	}
 
+	// aggregator does not exist for all strategy stores. only Close() if exists.
+	if c.aggregator != nil {
+		if err := c.aggregator.Close(); err != nil {
+			c.logger.Error("failed to close aggregator.", zap.Error(err))
+		}
+	}
+
 	// watchers actually never return errors from Close
 	_ = c.tlsGRPCCertWatcherCloser.Close()
 	_ = c.tlsHTTPCertWatcherCloser.Close()

cmd/collector/app/collector_test.go

@@ -25,6 +25,8 @@ import (
 	"github.com/uber/jaeger-lib/metrics/metricstest"
 	"go.uber.org/zap"
 
+	"github.com/jaegertracing/jaeger/cmd/collector/app/processor"
+	"github.com/jaegertracing/jaeger/model"
 	"github.com/jaegertracing/jaeger/pkg/healthcheck"
 	"github.com/jaegertracing/jaeger/thrift-gen/sampling"
 )
@@ -98,3 +100,60 @@ func TestCollector_PublishOpts(t *testing.T) {
 		Value: 42,
 	})
 }
+
+func TestAggregator(t *testing.T) {
+	// prepare
+	hc := healthcheck.New()
+	logger := zap.NewNop()
+	baseMetrics := metricstest.NewFactory(time.Hour)
+	spanWriter := &fakeSpanWriter{}
+	strategyStore := &mockStrategyStore{}
+	agg := &mockAggregator{}
+
+	c := New(&CollectorParams{
+		ServiceName:    "collector",
+		Logger:         logger,
+		MetricsFactory: baseMetrics,
+		SpanWriter:     spanWriter,
+		StrategyStore:  strategyStore,
+		HealthCheck:    hc,
+		Aggregator:     agg,
+	})
+	collectorOpts := &CollectorOptions{
+		QueueSize:  10,
+		NumWorkers: 10,
+	}
+
+	// test
+	c.Start(collectorOpts)
+
+	// assert that aggregator was added to the collector
+	_, err := c.spanProcessor.ProcessSpans([]*model.Span{
+		{
+			OperationName: "y",
+			Process: &model.Process{
+				ServiceName: "x",
+			},
+			Tags: []model.KeyValue{
+				{
+					Key:  "sampler.type",
+					VStr: "probabilistic",
+				},
+				{
+					Key:  "sampler.param",
+					VStr: "1",
+				},
+			},
+		},
+	}, processor.SpansOptions{SpanFormat: processor.JaegerSpanFormat})
+	assert.NoError(t, err)
+
+	// verify
+	assert.NoError(t, c.Close())
+
+	// assert that aggregator was used
+	assert.Equal(t, 1, agg.callCount)
+
+	// assert that aggregator close was called
+	assert.Equal(t, 1, agg.closeCount)
+}

plugin/sampling/strategystore/adaptive/root_span_handler.go → cmd/collector/app/root_span_handler.go

@@ -12,17 +12,17 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package adaptive
+package app
 
 import (
 	"go.uber.org/zap"
 
-	"github.com/jaegertracing/jaeger/cmd/collector/app"
+	"github.com/jaegertracing/jaeger/cmd/collector/app/sampling/strategystore"
 	"github.com/jaegertracing/jaeger/model"
 )
 
-// HandleRootSpan returns a function that records throughput for root spans
-func HandleRootSpan(aggregator Aggregator, logger *zap.Logger) app.ProcessSpan {
+// handleRootSpan returns a function that records throughput for root spans
+func handleRootSpan(aggregator strategystore.Aggregator, logger *zap.Logger) ProcessSpan {
 	return func(span *model.Span) {
 		// TODO 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.
@@ -33,7 +33,7 @@ func HandleRootSpan(aggregator Aggregator, logger *zap.Logger) app.ProcessSpan {
 		if service == "" || span.OperationName == "" {
 			return
 		}
-		samplerType, samplerParam := GetSamplerParams(span, logger)
+		samplerType, samplerParam := span.GetSamplerParams(logger)
 		if samplerType == "" {
 			return
 		}

plugin/sampling/strategystore/adaptive/root_span_handler_test.go → cmd/collector/app/root_span_handler_test.go

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package adaptive
+package app
 
 import (
 	"testing"
@@ -24,18 +24,22 @@ import (
 )
 
 type mockAggregator struct {
-	callCount int
+	callCount  int
+	closeCount int
 }
 
 func (t *mockAggregator) RecordThroughput(service, operation, samplerType string, probability float64) {
 	t.callCount++
 }
 func (t *mockAggregator) Start() {}
-func (t *mockAggregator) Stop()  {}
+func (t *mockAggregator) Close() error {
+	t.closeCount++
+	return nil
+}
 
 func TestHandleRootSpan(t *testing.T) {
 	aggregator := &mockAggregator{}
-	processor := HandleRootSpan(aggregator, zap.NewNop())
+	processor := handleRootSpan(aggregator, zap.NewNop())
 
 	// Testing non-root span
 	span := &model.Span{References: []model.SpanRef{{SpanID: model.NewSpanID(1), RefType: model.ChildOf}}}

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

@@ -17,6 +17,9 @@ package strategystore
 import (
 	"github.com/uber/jaeger-lib/metrics"
 	"go.uber.org/zap"
+
+	"github.com/jaegertracing/jaeger/pkg/distributedlock"
+	"github.com/jaegertracing/jaeger/storage/samplingstore"
 )
 
 // Factory defines an interface for a factory that can create implementations of different strategy storage components.
@@ -27,8 +30,8 @@ import (
 // plugin.Configurable
 type Factory interface {
 	// Initialize performs internal initialization of the factory.
-	Initialize(metricsFactory metrics.Factory, logger *zap.Logger) error
+	Initialize(metricsFactory metrics.Factory, lock distributedlock.Lock, store samplingstore.Store, logger *zap.Logger) error
 
 	// CreateStrategyStore initializes the StrategyStore and returns it.
-	CreateStrategyStore() (StrategyStore, error)
+	CreateStrategyStore() (StrategyStore, Aggregator, error)
 }

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

@@ -16,6 +16,7 @@ package strategystore
 
 import (
 	"context"
+	"io"
 
 	"github.com/jaegertracing/jaeger/thrift-gen/sampling"
 )
@@ -25,3 +26,15 @@ type StrategyStore interface {
 	// GetSamplingStrategy retrieves the sampling strategy for the specified service.
 	GetSamplingStrategy(ctx context.Context, serviceName string) (*sampling.SamplingStrategyResponse, error)
 }
+
+// Aggregator defines an interface used to aggregate operation throughput.
+type Aggregator interface {
+	// Close() from io.Closer stops the aggregator from aggregating throughput.
+	io.Closer
+
+	// RecordThroughput records throughput for an operation for aggregation.
+	RecordThroughput(service, operation, samplerType string, probability float64)
+
+	// Start starts aggregating operation throughput.
+	Start()
+}

cmd/collector/app/span_handler_builder.go

@@ -44,13 +44,14 @@ type SpanHandlers struct {
 }
 
 // BuildSpanProcessor builds the span processor to be used with the handlers
-func (b *SpanHandlerBuilder) BuildSpanProcessor() processor.SpanProcessor {
+func (b *SpanHandlerBuilder) BuildSpanProcessor(additional ...ProcessSpan) processor.SpanProcessor {
 	hostname, _ := os.Hostname()
 	svcMetrics := b.metricsFactory()
 	hostMetrics := svcMetrics.Namespace(metrics.NSOptions{Tags: map[string]string{"host": hostname}})
 
 	return NewSpanProcessor(
 		b.SpanWriter,
+		additional,
 		Options.ServiceMetrics(svcMetrics),
 		Options.HostMetrics(hostMetrics),
 		Options.Logger(b.logger()),
@@ -61,7 +62,6 @@ func (b *SpanHandlerBuilder) BuildSpanProcessor() processor.SpanProcessor {
 		Options.DynQueueSizeWarmup(uint(b.CollectorOpts.QueueSize)), // same as queue size for now
 		Options.DynQueueSizeMemory(b.CollectorOpts.DynQueueSizeMemory),
 	)
-
 }
 
 // BuildHandlers builds span handlers (Zipkin, Jaeger)

cmd/collector/app/span_processor.go

@@ -66,9 +66,10 @@ type queueItem struct {
 // NewSpanProcessor returns a SpanProcessor that preProcesses, filters, queues, sanitizes, and processes spans
 func NewSpanProcessor(
 	spanWriter spanstore.Writer,
+	additional []ProcessSpan,
 	opts ...Option,
 ) processor.SpanProcessor {
-	sp := newSpanProcessor(spanWriter, opts...)
+	sp := newSpanProcessor(spanWriter, additional, opts...)
 
 	sp.queue.StartConsumers(sp.numWorkers, func(item interface{}) {
 		value := item.(*queueItem)
@@ -84,7 +85,7 @@ func NewSpanProcessor(
 	return sp
 }
 
-func newSpanProcessor(spanWriter spanstore.Writer, opts ...Option) *spanProcessor {
+func newSpanProcessor(spanWriter spanstore.Writer, additional []ProcessSpan, opts ...Option) *spanProcessor {
 	options := Options.apply(opts...)
 	handlerMetrics := NewSpanProcessorMetrics(
 		options.serviceMetrics,
@@ -122,6 +123,8 @@ func newSpanProcessor(spanWriter spanstore.Writer, opts ...Option) *spanProcesso
 		processSpanFuncs = append(processSpanFuncs, sp.countSpan)
 	}
 
+	processSpanFuncs = append(processSpanFuncs, additional...)
+
 	sp.processSpan = ChainedProcessSpan(processSpanFuncs...)
 	return &sp
 }