[processor/tailsampling] Support extensions

Allow the Tail Sampling Processor to use extensions that implement an
interface simply containing NewEvaluator. To use an extension all that
needs to be done is specify the extension's type/name such as
my_extension or my_extension/my_name in the type field of the policy
configuration. When configuration is loaded/updated the extension will
be called in order to create a new evaluator.

Extensions can be configured in two places. First, in the extensions
section like all other extensions. This is useful for configuration such
as if an extension needs to connect to an external service. In addition,
an object can defined under the "extension" key the policy config for
anything related to one instance. This is useful when reusing the same
extension e.g. different params for different services.

Author: csmarchbanks

processor/tailsamplingprocessor/and_helper.go

@@ -10,11 +10,11 @@ import (
 	"github.com/open-telemetry/opentelemetry-collector-contrib/processor/tailsamplingprocessor/pkg/samplingpolicy"
 )
 
-func getNewAndPolicy(settings component.TelemetrySettings, config *AndCfg) (samplingpolicy.Evaluator, error) {
+func getNewAndPolicy(settings component.TelemetrySettings, config *AndCfg, policyExtensions map[string]samplingpolicy.Extension) (samplingpolicy.Evaluator, error) {
 	subPolicyEvaluators := make([]samplingpolicy.Evaluator, len(config.SubPolicyCfg))
 	for i := range config.SubPolicyCfg {
 		policyCfg := &config.SubPolicyCfg[i]
-		policy, err := getAndSubPolicyEvaluator(settings, policyCfg)
+		policy, err := getAndSubPolicyEvaluator(settings, policyCfg, policyExtensions)
 		if err != nil {
 			return nil, err
 		}
@@ -24,6 +24,6 @@ func getNewAndPolicy(settings component.TelemetrySettings, config *AndCfg) (samp
 }
 
 // Return instance of and sub-policy
-func getAndSubPolicyEvaluator(settings component.TelemetrySettings, cfg *AndSubPolicyCfg) (samplingpolicy.Evaluator, error) {
-	return getSharedPolicyEvaluator(settings, &cfg.sharedPolicyCfg)
+func getAndSubPolicyEvaluator(settings component.TelemetrySettings, cfg *AndSubPolicyCfg, policyExtensions map[string]samplingpolicy.Extension) (samplingpolicy.Evaluator, error) {
+	return getSharedPolicyEvaluator(settings, &cfg.sharedPolicyCfg, policyExtensions)
 }

processor/tailsamplingprocessor/and_helper_test.go

@@ -27,7 +27,7 @@ func TestAndHelper(t *testing.T) {
 					},
 				},
 			},
-		})
+		}, nil)
 		require.NoError(t, err)
 
 		expected := sampling.NewAnd(zap.NewNop(), []samplingpolicy.Evaluator{
@@ -46,7 +46,7 @@ func TestAndHelper(t *testing.T) {
 					},
 				},
 			},
-		})
+		}, nil)
 		require.EqualError(t, err, "unknown sampling policy type and")
 	})
 }

processor/tailsamplingprocessor/composite_helper.go

@@ -11,12 +11,12 @@ import (
 	"github.com/open-telemetry/opentelemetry-collector-contrib/processor/tailsamplingprocessor/pkg/samplingpolicy"
 )
 
-func getNewCompositePolicy(settings component.TelemetrySettings, config *CompositeCfg) (samplingpolicy.Evaluator, error) {
+func getNewCompositePolicy(settings component.TelemetrySettings, config *CompositeCfg, policyExtensions map[string]samplingpolicy.Extension) (samplingpolicy.Evaluator, error) {
 	subPolicyEvalParams := make([]sampling.SubPolicyEvalParams, len(config.SubPolicyCfg))
 	rateAllocationsMap := getRateAllocationMap(config)
 	for i := range config.SubPolicyCfg {
 		policyCfg := &config.SubPolicyCfg[i]
-		policy, err := getCompositeSubPolicyEvaluator(settings, policyCfg)
+		policy, err := getCompositeSubPolicyEvaluator(settings, policyCfg, policyExtensions)
 		if err != nil {
 			return nil, err
 		}
@@ -48,11 +48,11 @@ func getRateAllocationMap(config *CompositeCfg) map[string]float64 {
 }
 
 // Return instance of composite sub-policy
-func getCompositeSubPolicyEvaluator(settings component.TelemetrySettings, cfg *CompositeSubPolicyCfg) (samplingpolicy.Evaluator, error) {
+func getCompositeSubPolicyEvaluator(settings component.TelemetrySettings, cfg *CompositeSubPolicyCfg, policyExtensions map[string]samplingpolicy.Extension) (samplingpolicy.Evaluator, error) {
 	switch cfg.Type {
 	case And:
-		return getNewAndPolicy(settings, &cfg.AndCfg)
+		return getNewAndPolicy(settings, &cfg.AndCfg, policyExtensions)
 	default:
-		return getSharedPolicyEvaluator(settings, &cfg.sharedPolicyCfg)
+		return getSharedPolicyEvaluator(settings, &cfg.sharedPolicyCfg, policyExtensions)
 	}
 }

processor/tailsamplingprocessor/composite_helper_test.go

@@ -45,7 +45,7 @@ func TestCompositeHelper(t *testing.T) {
 					Percent: 0, // will be populated with default
 				},
 			},
-		})
+		}, nil)
 		require.NoError(t, err)
 
 		expected := sampling.NewComposite(zap.NewNop(), 1000, []sampling.SubPolicyEvalParams{
@@ -73,7 +73,7 @@ func TestCompositeHelper(t *testing.T) {
 					},
 				},
 			},
-		})
+		}, nil)
 		require.EqualError(t, err, "unknown sampling policy type composite")
 	})
 }

processor/tailsamplingprocessor/config.go

@@ -74,6 +74,8 @@ type sharedPolicyCfg struct {
 	BooleanAttributeCfg BooleanAttributeCfg `mapstructure:"boolean_attribute"`
 	// Configs for OTTL condition filter sampling policy evaluator
 	OTTLConditionCfg OTTLConditionCfg `mapstructure:"ottl_condition"`
+	// Configs for any extensions that are used.
+	ExtensionCfg map[string]any `mapstructure:"extension"`
 }
 
 // CompositeSubPolicyCfg holds the common configuration to all policies under composite policy.

processor/tailsamplingprocessor/drop_helper.go

@@ -10,11 +10,11 @@ import (
 	"github.com/open-telemetry/opentelemetry-collector-contrib/processor/tailsamplingprocessor/pkg/samplingpolicy"
 )
 
-func getNewDropPolicy(settings component.TelemetrySettings, config *DropCfg) (samplingpolicy.Evaluator, error) {
+func getNewDropPolicy(settings component.TelemetrySettings, config *DropCfg, policyExtensions map[string]samplingpolicy.Extension) (samplingpolicy.Evaluator, error) {
 	subPolicyEvaluators := make([]samplingpolicy.Evaluator, len(config.SubPolicyCfg))
 	for i := range config.SubPolicyCfg {
 		policyCfg := &config.SubPolicyCfg[i]
-		policy, err := getDropSubPolicyEvaluator(settings, policyCfg)
+		policy, err := getDropSubPolicyEvaluator(settings, policyCfg, policyExtensions)
 		if err != nil {
 			return nil, err
 		}
@@ -24,6 +24,6 @@ func getNewDropPolicy(settings component.TelemetrySettings, config *DropCfg) (sa
 }
 
 // Return instance of and sub-policy
-func getDropSubPolicyEvaluator(settings component.TelemetrySettings, cfg *AndSubPolicyCfg) (samplingpolicy.Evaluator, error) {
-	return getSharedPolicyEvaluator(settings, &cfg.sharedPolicyCfg)
+func getDropSubPolicyEvaluator(settings component.TelemetrySettings, cfg *AndSubPolicyCfg, policyExtensions map[string]samplingpolicy.Extension) (samplingpolicy.Evaluator, error) {
+	return getSharedPolicyEvaluator(settings, &cfg.sharedPolicyCfg, policyExtensions)
 }

processor/tailsamplingprocessor/drop_helper_test.go

@@ -27,7 +27,7 @@ func TestDropHelper(t *testing.T) {
 					},
 				},
 			},
-		})
+		}, nil)
 		require.NoError(t, err)
 
 		expected := sampling.NewDrop(zap.NewNop(), []samplingpolicy.Evaluator{
@@ -46,7 +46,7 @@ func TestDropHelper(t *testing.T) {
 					},
 				},
 			},
-		})
+		}, nil)
 		require.EqualError(t, err, "unknown sampling policy type drop")
 	})
 }

processor/tailsamplingprocessor/pkg/samplingpolicy/samplingpolicy.go

@@ -61,3 +61,7 @@ type Evaluator interface {
 	// Evaluate looks at the trace data and returns a corresponding SamplingDecision.
 	Evaluate(ctx context.Context, traceID pcommon.TraceID, trace *TraceData) (Decision, error)
 }
+
+type Extension interface {
+	NewEvaluator(cfg map[string]any) (Evaluator, error)
+}

processor/tailsamplingprocessor/processor.go

@@ -68,6 +68,8 @@ type tailSamplingSpanProcessor struct {
 	setPolicyMux       sync.Mutex
 	pendingPolicy      []PolicyCfg
 	sampleOnFirstMatch bool
+
+	host component.Host
 }
 
 type traceLimiter interface {
@@ -156,10 +158,7 @@ func newTracesProcessor(ctx context.Context, set processor.Settings, nextConsume
 	}
 
 	if tsp.policies == nil {
-		err := tsp.loadSamplingPolicy(cfg.PolicyCfgs)
-		if err != nil {
-			return nil, err
-		}
+		tsp.pendingPolicy = cfg.PolicyCfgs
 	}
 
 	if tsp.decisionBatcher == nil {
@@ -217,22 +216,31 @@ func withRecordPolicy() Option {
 	}
 }
 
-func getPolicyEvaluator(settings component.TelemetrySettings, cfg *PolicyCfg) (samplingpolicy.Evaluator, error) {
+func getPolicyEvaluator(settings component.TelemetrySettings, cfg *PolicyCfg, policyExtensions map[string]samplingpolicy.Extension) (samplingpolicy.Evaluator, error) {
 	switch cfg.Type {
 	case Composite:
-		return getNewCompositePolicy(settings, &cfg.CompositeCfg)
+		return getNewCompositePolicy(settings, &cfg.CompositeCfg, policyExtensions)
 	case And:
-		return getNewAndPolicy(settings, &cfg.AndCfg)
+		return getNewAndPolicy(settings, &cfg.AndCfg, policyExtensions)
 	case Drop:
-		return getNewDropPolicy(settings, &cfg.DropCfg)
+		return getNewDropPolicy(settings, &cfg.DropCfg, policyExtensions)
 	default:
-		return getSharedPolicyEvaluator(settings, &cfg.sharedPolicyCfg)
+		return getSharedPolicyEvaluator(settings, &cfg.sharedPolicyCfg, policyExtensions)
 	}
 }
 
-func getSharedPolicyEvaluator(settings component.TelemetrySettings, cfg *sharedPolicyCfg) (samplingpolicy.Evaluator, error) {
+func getSharedPolicyEvaluator(settings component.TelemetrySettings, cfg *sharedPolicyCfg, policyExtensions map[string]samplingpolicy.Extension) (samplingpolicy.Evaluator, error) {
 	settings.Logger = settings.Logger.With(zap.Any("policy", cfg.Type))
 
+	extension, ok := policyExtensions[string(cfg.Type)]
+	if ok {
+		evaluator, err := extension.NewEvaluator(cfg.ExtensionCfg)
+		if err != nil {
+			return nil, fmt.Errorf("unable to load extension %s: %w", string(cfg.Type), err)
+		}
+		return evaluator, nil
+	}
+
 	switch cfg.Type {
 	case AlwaysSample:
 		return sampling.NewAlwaysSample(settings), nil
@@ -325,7 +333,7 @@ func (tsp *tailSamplingSpanProcessor) loadSamplingPolicy(cfgs []PolicyCfg) error
 		}
 		policyNames[cfg.Name] = struct{}{}
 
-		eval, err := getPolicyEvaluator(telemetrySettings, &cfg)
+		eval, err := getPolicyEvaluator(telemetrySettings, &cfg, tsp.extensions())
 		if err != nil {
 			return fmt.Errorf("failed to create policy evaluator for %q: %w", cfg.Name, err)
 		}
@@ -364,14 +372,14 @@ func (tsp *tailSamplingSpanProcessor) SetSamplingPolicy(cfgs []PolicyCfg) {
 	tsp.pendingPolicy = cfgs
 }
 
-func (tsp *tailSamplingSpanProcessor) loadPendingSamplingPolicy() {
+func (tsp *tailSamplingSpanProcessor) loadPendingSamplingPolicy() error {
 	tsp.setPolicyMux.Lock()
 	defer tsp.setPolicyMux.Unlock()
 
 	// Nothing pending, do nothing.
 	pLen := len(tsp.pendingPolicy)
 	if pLen == 0 {
-		return
+		return nil
 	}
 
 	tsp.logger.Debug("Loading pending sampling policy", zap.Int("pending.len", pLen))
@@ -381,17 +389,17 @@ func (tsp *tailSamplingSpanProcessor) loadPendingSamplingPolicy() {
 	// Empty pending regardless of error. If policy is invalid, it will fail on
 	// every tick, no need to do extra work and flood the log with errors.
 	tsp.pendingPolicy = nil
-
-	if err != nil {
-		tsp.logger.Error("Failed to load pending sampling policy", zap.Error(err))
-		tsp.logger.Debug("Continuing to use the previously loaded sampling policy")
-	}
+	return err
 }
 
 func (tsp *tailSamplingSpanProcessor) samplingPolicyOnTick() {
 	tsp.logger.Debug("Sampling Policy Evaluation ticked")
 
-	tsp.loadPendingSamplingPolicy()
+	err := tsp.loadPendingSamplingPolicy()
+	if err != nil {
+		tsp.logger.Error("Failed to load pending sampling policy", zap.Error(err))
+		tsp.logger.Debug("Continuing to use the previously loaded sampling policy")
+	}
 
 	ctx := context.Background()
 	metrics := newPolicyMetrics(len(tsp.policies))
@@ -644,11 +652,35 @@ func (*tailSamplingSpanProcessor) Capabilities() consumer.Capabilities {
 }
 
 // Start is invoked during service startup.
-func (tsp *tailSamplingSpanProcessor) Start(context.Context, component.Host) error {
+func (tsp *tailSamplingSpanProcessor) Start(ctx context.Context, host component.Host) error {
+	// We need to store the host before loading sampling policies in order to load any extensions.
+	tsp.host = host
+	if tsp.policies == nil {
+		err := tsp.loadPendingSamplingPolicy()
+		if err != nil {
+			tsp.logger.Error("Failed to load initial sampling policy", zap.Error(err))
+			return err
+		}
+	}
 	tsp.policyTicker.Start(tsp.tickerFrequency)
 	return nil
 }
 
+func (tsp *tailSamplingSpanProcessor) extensions() map[string]samplingpolicy.Extension {
+	if tsp.host == nil {
+		return nil
+	}
+	extensions := tsp.host.GetExtensions()
+	scoped := map[string]samplingpolicy.Extension{}
+	for id, ext := range extensions {
+		evaluator, ok := ext.(samplingpolicy.Extension)
+		if ok {
+			scoped[id.String()] = evaluator
+		}
+	}
+	return scoped
+}
+
 // Shutdown is invoked during service shutdown.
 func (tsp *tailSamplingSpanProcessor) Shutdown(context.Context) error {
 	tsp.decisionBatcher.Stop()