[processor/tailsampling] Support extensions (#42573)

<!--Ex. Fixing a bug - Describe the bug and how this fixes the issue.
Ex. Adding a feature - Explain what this achieves.-->
#### Description
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 or updated NewEvaluator will
be called.

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,
configuration can be defined in the policy config under a key matching
the type of the extension for anything related to one instance,
following the pattern for all other policy types. This is useful when
reusing the same extension e.g. different params for different services.
For example:
```yaml
type: custom_policy
custom_policy:
  foo: bar
```

What this ends up with is that extensions look exactly like all other
policy types with no need for users to know about some sort of extension
container.

I have tested this out with an extension in a custom build and it is
working well, I am sure we will iterate on the interfaces a bit as we
write more. Specifically a way to stop the extension when policies are
updated may be necessary to clean up any background processes. @yvrhdn
is working on an extension we plan to add to contrib as well.

<!-- Issue number (e.g. #1234) or full URL to issue, if applicable. -->
#### Link to tracking issue
Part of
#31582

Author: csmarchbanks

.chloggen/tailsampling-extension-support.yaml

@@ -0,0 +1,28 @@
+# Use this changelog template to create an entry for release notes.
+
+# One of 'breaking', 'deprecation', 'new_component', 'enhancement', 'bug_fix'
+change_type: enhancement
+
+# The name of the component, or a single word describing the area of concern, (e.g. filelogreceiver)
+component: tailsamplingprocessor
+
+# A brief description of the change.  Surround your text with quotes ("") if it needs to start with a backtick (`).
+note: Add support for extensions that implement sampling policies.
+
+# Mandatory: One or more tracking issues related to the change. You can use the PR number here if no issue exists.
+issues: [31582]
+
+# (Optional) One or more lines of additional information to render under the primary note.
+# These lines will be padded with 2 spaces and then inserted directly into the document.
+# Use pipe (|) for multiline entries.
+subtext: |
+  Extension support for tailsamplingprocessor is still in development and the interfaces may change at any time.
+
+# If your change doesn't affect end users or the exported elements of any package,
+# you should instead start your pull request title with [chore] or use the "Skip Changelog" label.
+# Optional: The change log or logs in which this entry should be included.
+# e.g. '[user]' or '[user, api]'
+# Include 'user' if the change is relevant to end users.
+# Include 'api' if there is a change to a library API.
+# Default: '[user]'
+change_logs: []

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]map[string]any `mapstructure:",remain"`
 }
 
 // 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(policyName string, 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,20 +216,20 @@ 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))
 
 	switch cfg.Type {
@@ -273,8 +272,17 @@ func getSharedPolicyEvaluator(settings component.TelemetrySettings, cfg *sharedP
 	case OTTLCondition:
 		ottlfCfg := cfg.OTTLConditionCfg
 		return sampling.NewOTTLConditionFilter(settings, ottlfCfg.SpanConditions, ottlfCfg.SpanEventConditions, ottlfCfg.ErrorMode)
-
 	default:
+		t := string(cfg.Type)
+		extension, ok := policyExtensions[t]
+		if ok {
+			evaluator, err := extension.NewEvaluator(cfg.Name, cfg.ExtensionCfg[t])
+			if err != nil {
+				return nil, fmt.Errorf("unable to load extension %s: %w", t, err)
+			}
+			return evaluator, nil
+		}
+
 		return nil, fmt.Errorf("unknown sampling policy type %s", cfg.Type)
 	}
 }
@@ -326,7 +334,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)
 		}
@@ -365,14 +373,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))
@@ -382,17 +390,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))
@@ -648,11 +656,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(_ 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()