Restrict upgrade of sample decision in composite sampler

signal/trace/composite_sampler.go

@@ -8,6 +8,20 @@ import (
 	"go.opentelemetry.io/otel/trace"
 )
 
+// CompositeSampler decorates samplers.
+//
+// Mixed decision resolving strategy is restrict-only.
+//
+// For example:
+//
+// If any sampler returns sdktrace.Drop, the
+// decision is immediately Drop. If a sampler returns sdktrace.RecordOnly,
+// that decision is recorded, and subsequent samplers cannot upgrade it to
+// sdktrace.RecordAndSample. If a sampler returns sdktrace.RecordAndSample,
+// it is only considered if no prior sampler decided sdktrace.RecordOnly.
+// This ensures that restrictions are respected and the behavior is restrict-only
+// where any sampler can restrict the decision, but not upgrade it beyond
+// what a previous sampler allowed.
 type CompositeSampler struct {
 	samplers []sdktrace.Sampler
 }
@@ -23,6 +37,9 @@ func NewCompositeSampler(samplers ...sdktrace.Sampler) *CompositeSampler {
 	return &CompositeSampler{samplers: copied}
 }
 
+// ShouldSample determines strictest decision in configured samplers.
+//
+// See CompositeSampler
 func (r CompositeSampler) ShouldSample(parameters sdktrace.SamplingParameters) sdktrace.SamplingResult {
 	if len(r.samplers) == 0 {
 		return sdktrace.SamplingResult{
@@ -32,15 +49,24 @@ func (r CompositeSampler) ShouldSample(parameters sdktrace.SamplingParameters) s
 		}
 	}
 
-	var res sdktrace.SamplingResult
-	for _, sampler := range r.samplers {
-		res = sampler.ShouldSample(parameters)
-		if res.Decision == sdktrace.Drop {
-			return res
+	// Relies on OTel SDK ordering: Drop(0) < RecordOnly(1) < RecordAndSample(2).
+	// See TestSamplingDecisionOrdering.
+	const strictest = sdktrace.Drop
+
+	result := r.samplers[0].ShouldSample(parameters)
+
+	for _, sampler := range r.samplers[1:] {
+		if result.Decision == strictest {
+			break
+		}
+
+		current := sampler.ShouldSample(parameters)
+		if current.Decision < result.Decision {
+			result = current
 		}
 	}
 
-	return res
+	return result
 }
 
 func (r CompositeSampler) Description() string {

signal/trace/composite_sampler_test.go

@@ -5,94 +5,102 @@ import (
 
 	"github.com/thumbrise/otelext/signal/trace"
 	sdktrace "go.opentelemetry.io/otel/sdk/trace"
-
-	"github.com/thumbrise/otelext/internal/mock"
 )
 
-func TestCompositeSampler(t *testing.T) {
-	tests := []struct {
-		name     string
-		samplers []sdktrace.Sampler
-		params   sdktrace.SamplingParameters
-		want     sdktrace.SamplingDecision
-	}{
-		{
-			name:     "no samplers",
-			samplers: []sdktrace.Sampler{},
-			params:   sdktrace.SamplingParameters{},
-			want:     sdktrace.Drop,
-		},
-		{
-			name: "first sampler drops",
-			samplers: []sdktrace.Sampler{
-				mock.NewSampler(sdktrace.Drop, "drop sampler"),
-				mock.NewSampler(sdktrace.RecordAndSample, "record sampler"),
-			},
-			params: sdktrace.SamplingParameters{},
-			want:   sdktrace.Drop,
-		},
-		{
-			name: "all samplers record and sample",
-			samplers: []sdktrace.Sampler{
-				mock.NewSampler(sdktrace.RecordAndSample, "record sampler 1"),
-				mock.NewSampler(sdktrace.RecordAndSample, "record sampler 2"),
-			},
-			params: sdktrace.SamplingParameters{},
-			want:   sdktrace.RecordAndSample,
-		},
-		{
-			name: "mixed decisions",
-			samplers: []sdktrace.Sampler{
-				mock.NewSampler(sdktrace.RecordOnly, "record only"),
-				mock.NewSampler(sdktrace.Drop, "drop"),
-				mock.NewSampler(sdktrace.RecordAndSample, "record and sample"),
-			},
-			params: sdktrace.SamplingParameters{},
-			want:   sdktrace.Drop,
-		},
+type staticSampler struct {
+	d    sdktrace.SamplingDecision
+	name string
+}
+
+func (s staticSampler) ShouldSample(parameters sdktrace.SamplingParameters) sdktrace.SamplingResult {
+	return sdktrace.SamplingResult{Decision: s.d}
+}
+
+func (s staticSampler) Description() string {
+	return s.name
+}
+
+func TestCompositeSampler_MixedDecisions_RecordOnlyThenRecordAndSample(t *testing.T) {
+	s1 := staticSampler{d: sdktrace.RecordOnly, name: "RecordOnly"}
+	s2 := staticSampler{d: sdktrace.RecordAndSample, name: "RecordAndSample"}
+	cs := trace.NewCompositeSampler(s1, s2)
+
+	res := cs.ShouldSample(sdktrace.SamplingParameters{})
+	if res.Decision != sdktrace.RecordOnly {
+		t.Fatalf("expected RecordOnly, got %v", res.Decision)
 	}
+}
 
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			sampler := trace.NewCompositeSampler(tt.samplers...)
-			result := sampler.ShouldSample(tt.params)
+func TestCompositeSampler_MixedDecisions_RecordAndSampleThenRecordOnly(t *testing.T) {
+	s1 := staticSampler{d: sdktrace.RecordAndSample, name: "RecordAndSample"}
+	s2 := staticSampler{d: sdktrace.RecordOnly, name: "RecordOnly"}
+	cs := trace.NewCompositeSampler(s1, s2)
 
-			if result.Decision != tt.want {
-				t.Errorf("CompositeSampler.ShouldSample() = %v, want %v", result.Decision, tt.want)
-			}
-		})
+	res := cs.ShouldSample(sdktrace.SamplingParameters{})
+	if res.Decision != sdktrace.RecordOnly {
+		t.Fatalf("expected RecordOnly, got %v", res.Decision)
 	}
 }
 
-func TestCompositeSamplerDescription(t *testing.T) {
-	tests := []struct {
-		name     string
-		samplers []sdktrace.Sampler
-		want     string
-	}{
-		{
-			name:     "no samplers",
-			samplers: []sdktrace.Sampler{},
-			want:     "no samplers passed in composite sampler",
-		},
-		{
-			name: "with samplers",
-			samplers: []sdktrace.Sampler{
-				mock.NewSampler(sdktrace.RecordAndSample, "sampler1"),
-				mock.NewSampler(sdktrace.RecordAndSample, "sampler2"),
-			},
-			want: "Decorates chain of samplers: sampler1\nsampler2",
-		},
+func TestCompositeSampler_MixedDecisions_RecordAndSampleBoth(t *testing.T) {
+	s1 := staticSampler{d: sdktrace.RecordAndSample, name: "R&A1"}
+	s2 := staticSampler{d: sdktrace.RecordAndSample, name: "R&A2"}
+	cs := trace.NewCompositeSampler(s1, s2)
+
+	res := cs.ShouldSample(sdktrace.SamplingParameters{})
+	if res.Decision != sdktrace.RecordAndSample {
+		t.Fatalf("expected RecordAndSample, got %v", res.Decision)
 	}
+}
 
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			sampler := trace.NewCompositeSampler(tt.samplers...)
-			got := sampler.Description()
+func TestCompositeSampler_MixedDecisions_WithDrop(t *testing.T) {
+	s1 := staticSampler{d: sdktrace.RecordAndSample, name: "R&A"}
+	s2 := staticSampler{d: sdktrace.Drop, name: "Drop"}
+	cs := trace.NewCompositeSampler(s1, s2)
+
+	res := cs.ShouldSample(sdktrace.SamplingParameters{})
+	if res.Decision != sdktrace.Drop {
+		t.Fatalf("expected Drop, got %v", res.Decision)
+	}
+}
+
+func TestCompositeSampler_NoSamplers(t *testing.T) {
+	cs := trace.NewCompositeSampler()
+
+	res := cs.ShouldSample(sdktrace.SamplingParameters{})
+	if res.Decision != sdktrace.Drop {
+		t.Fatalf("expected Drop, got %v", res.Decision)
+	}
+}
+
+func TestCompositeSampler_Description(t *testing.T) {
+	t.Run("no samplers", func(t *testing.T) {
+		cs := trace.NewCompositeSampler()
+		got := cs.Description()
+		want := "no samplers passed in composite sampler"
+		if got != want {
+			t.Fatalf("Description() = %q, want %q", got, want)
+		}
+	})
+
+	t.Run("with samplers", func(t *testing.T) {
+		s1 := staticSampler{d: sdktrace.RecordAndSample, name: "sampler1"}
+		s2 := staticSampler{d: sdktrace.RecordAndSample, name: "sampler2"}
+		cs := trace.NewCompositeSampler(s1, s2)
+		got := cs.Description()
+		want := "Decorates chain of samplers: sampler1\nsampler2"
+		if got != want {
+			t.Fatalf("Description() = %q, want %q", got, want)
+		}
+	})
+}
 
-			if got != tt.want {
-				t.Errorf("CompositeSampler.Description() = %q, want %q", got, tt.want)
-			}
-		})
+// TestSamplingDecisionOrdering asserts the OTel SDK enum ordering that
+// CompositeSampler's restrict-only logic depends on:
+// Drop(0) < RecordOnly(1) < RecordAndSample(2).
+func TestSamplingDecisionOrdering(t *testing.T) {
+	if !(sdktrace.Drop < sdktrace.RecordOnly && sdktrace.RecordOnly < sdktrace.RecordAndSample) {
+		t.Fatalf("unexpected SamplingDecision ordering: Drop=%d, RecordOnly=%d, RecordAndSample=%d",
+			sdktrace.Drop, sdktrace.RecordOnly, sdktrace.RecordAndSample)
 	}
 }

signal/trace/filter_based_sampler.go

@@ -32,6 +32,7 @@ func (s FilterBasedSampler) ShouldSample(parameters sdktrace.SamplingParameters)
 
 	// No filters dropped, record and sample
 	psc := trace.SpanContextFromContext(parameters.ParentContext)
+
 	return sdktrace.SamplingResult{
 		Decision:   sdktrace.RecordAndSample,
 		Tracestate: psc.TraceState(),
@@ -40,6 +41,7 @@ func (s FilterBasedSampler) ShouldSample(parameters sdktrace.SamplingParameters)
 
 func (s FilterBasedSampler) Description() string {
 	registered := signal.RegisteredFilters()
+
 	descriptions := make([]string, 0, len(registered))
 	for _, f := range registered {
 		descriptions = append(descriptions, f.Description(context.Background()))