Replace Sampler with ExtendedSampler (#1206)

* Replace Sampler with ExtendedSampler

* Update changelog

CHANGELOG.asciidoc

@@ -34,6 +34,7 @@ https://github.com/elastic/apm-agent-go/compare/v1.15.0...main[View commits]
 - Replace apm.DefaultTracer with an initialization function {pull}1189[#(1189)]
 - Remove transport.Default, construct a new Transport in each new tracer {pull}1195[#(1195)]
 - Add service name and version to User-Agent header {pull}1196[#(1196)]
+- Replace Sampler with ExtendedSampler {pull}1206[#(1206)]
 
 [[release-notes-1.x]]
 === Go Agent version 1.x

config.go

@@ -492,7 +492,6 @@ func (t *Tracer) updateRemoteConfig(logger WarningLogger, old, attrs map[string]
 			} else {
 				updates = append(updates, func(cfg *instrumentationConfig) {
 					cfg.sampler = sampler
-					cfg.extendedSampler, _ = sampler.(ExtendedSampler)
 				})
 			}
 		case apmlog.EnvLogLevel:
@@ -635,7 +634,6 @@ type instrumentationConfigValues struct {
 	recording             bool
 	captureBody           CaptureBodyMode
 	captureHeaders        bool
-	extendedSampler       ExtendedSampler
 	maxSpans              int
 	sampler               Sampler
 	spanFramesMinDuration time.Duration

sampler.go

@@ -28,28 +28,15 @@ import (
 // Sampler provides a means of sampling transactions.
 type Sampler interface {
 	// Sample indicates whether or not a transaction
-	// should be sampled. This method will be invoked
+	// should be sampled, and the sampling rate in
+	// effect at the time.This method will be invoked
 	// by calls to Tracer.StartTransaction for the root
 	// of a trace, so it must be goroutine-safe, and
 	// should avoid synchronization as far as possible.
-	Sample(TraceContext) bool
+	Sample(SampleParams) SampleResult
 }
 
-// ExtendedSampler may be implemented by Samplers, providing
-// a method for sampling and returning an extended SampleResult.
-//
-// TODO(axw) in v2.0.0, replace the Sampler interface with this.
-type ExtendedSampler interface {
-	// SampleExtended indicates whether or not a transaction
-	// should be sampled, and the sampling rate in effect at
-	// the time. This method will be invoked by calls to
-	// Tracer.StartTransaction for the root of a trace, so it
-	// must be goroutine-safe, and should avoid synchronization
-	// as far as possible.
-	SampleExtended(SampleParams) SampleResult
-}
-
-// SampleParams holds parameters for SampleExtended.
+// SampleParams holds parameters for Sampler.Sample.
 type SampleParams struct {
 	// TraceContext holds the newly-generated TraceContext
 	// for the root transaction which is being sampled.
@@ -105,13 +92,7 @@ type ratioSampler struct {
 
 // Sample samples the transaction according to the configured
 // ratio and pseudo-random source.
-func (s ratioSampler) Sample(c TraceContext) bool {
-	return s.SampleExtended(SampleParams{TraceContext: c}).Sampled
-}
-
-// SampleExtended samples the transaction according to the configured
-// ratio and pseudo-random source.
-func (s ratioSampler) SampleExtended(args SampleParams) SampleResult {
+func (s ratioSampler) Sample(args SampleParams) SampleResult {
 	v := binary.BigEndian.Uint64(args.TraceContext.Span[:])
 	result := SampleResult{
 		Sampled:    v > 0 && v-1 < s.ceil,

sampler_test.go

@@ -49,7 +49,9 @@ func TestRatioSampler(t *testing.T) {
 			for j := 0; j < numIterations; j++ {
 				var traceContext apm.TraceContext
 				binary.LittleEndian.PutUint64(traceContext.Span[:], rng.Uint64())
-				if s.Sample(traceContext) {
+				result := s.Sample(apm.SampleParams{TraceContext: traceContext})
+				assert.Equal(t, ratio, result.SampleRate)
+				if result.Sampled {
 					sampled[i]++
 				}
 			}
@@ -67,44 +69,24 @@ func TestRatioSampler(t *testing.T) {
 
 func TestRatioSamplerAlways(t *testing.T) {
 	s := apm.NewRatioSampler(1.0)
-	assert.False(t, s.Sample(apm.TraceContext{})) // invalid span ID
-	assert.True(t, s.Sample(apm.TraceContext{
+	assert.False(t, s.Sample(apm.SampleParams{TraceContext: apm.TraceContext{}}).Sampled) // invalid span ID
+	assert.True(t, s.Sample(apm.SampleParams{TraceContext: apm.TraceContext{
 		Span: apm.SpanID{0, 0, 0, 0, 0, 0, 0, 1},
-	}))
-	assert.True(t, s.Sample(apm.TraceContext{
+	}}).Sampled)
+	assert.True(t, s.Sample(apm.SampleParams{TraceContext: apm.TraceContext{
 		Span: apm.SpanID{255, 255, 255, 255, 255, 255, 255, 255},
-	}))
+	}}).Sampled)
 }
 
 func TestRatioSamplerNever(t *testing.T) {
 	s := apm.NewRatioSampler(0)
-	assert.False(t, s.Sample(apm.TraceContext{})) // invalid span ID
-	assert.False(t, s.Sample(apm.TraceContext{
+	assert.False(t, s.Sample(apm.SampleParams{TraceContext: apm.TraceContext{}}).Sampled) // invalid span ID
+	assert.False(t, s.Sample(apm.SampleParams{TraceContext: apm.TraceContext{
 		Span: apm.SpanID{0, 0, 0, 0, 0, 0, 0, 1},
-	}))
-	assert.False(t, s.Sample(apm.TraceContext{
+	}}).Sampled)
+	assert.False(t, s.Sample(apm.SampleParams{TraceContext: apm.TraceContext{
 		Span: apm.SpanID{255, 255, 255, 255, 255, 255, 255, 255},
-	}))
-}
-
-func TestRatioSamplerExtended(t *testing.T) {
-	s := apm.NewRatioSampler(0.5).(apm.ExtendedSampler)
-
-	result := s.SampleExtended(apm.SampleParams{
-		TraceContext: apm.TraceContext{Span: apm.SpanID{255, 0, 0, 0, 0, 0, 0, 0}},
-	})
-	assert.Equal(t, apm.SampleResult{
-		Sampled:    false,
-		SampleRate: 0.5,
-	}, result)
-
-	result = s.SampleExtended(apm.SampleParams{
-		TraceContext: apm.TraceContext{Span: apm.SpanID{1, 0, 0, 0, 0, 0, 0, 0}},
-	})
-	assert.Equal(t, apm.SampleResult{
-		Sampled:    true,
-		SampleRate: 0.5,
-	}, result)
+	}}).Sampled)
 }
 
 func TestRatioSamplerPrecision(t *testing.T) {
@@ -115,8 +97,8 @@ func TestRatioSamplerPrecision(t *testing.T) {
 		0.55556: 0.5556,
 	}
 	for r, want := range ratios {
-		s := apm.NewRatioSampler(r).(apm.ExtendedSampler)
-		got := s.SampleExtended(apm.SampleParams{}).SampleRate
+		s := apm.NewRatioSampler(r)
+		got := s.Sample(apm.SampleParams{}).SampleRate
 		assert.Equal(t, want, got)
 	}
 }

tracer.go

@@ -473,7 +473,6 @@ func newTracer(opts TracerOptions) *Tracer {
 	})
 	t.setLocalInstrumentationConfig(envTransactionSampleRate, func(cfg *instrumentationConfigValues) {
 		cfg.sampler = opts.sampler
-		cfg.extendedSampler, _ = opts.sampler.(ExtendedSampler)
 	})
 	t.setLocalInstrumentationConfig(envSpanFramesMinDuration, func(cfg *instrumentationConfigValues) {
 		cfg.spanFramesMinDuration = opts.spanFramesMinDuration
@@ -741,7 +740,6 @@ func (t *Tracer) SetRecording(r bool) {
 func (t *Tracer) SetSampler(s Sampler) {
 	t.setLocalInstrumentationConfig(envTransactionSampleRate, func(cfg *instrumentationConfigValues) {
 		cfg.sampler = s
-		cfg.extendedSampler, _ = s.(ExtendedSampler)
 	})
 }
 

transaction.go

@@ -108,8 +108,8 @@ func (t *Tracer) StartTransactionOptions(name, transactionType string, opts Tran
 
 	if root {
 		var result SampleResult
-		if instrumentationConfig.extendedSampler != nil {
-			result = instrumentationConfig.extendedSampler.SampleExtended(SampleParams{
+		if instrumentationConfig.sampler != nil {
+			result = instrumentationConfig.sampler.Sample(SampleParams{
 				TraceContext: tx.traceContext,
 			})
 			if !result.Sampled {
@@ -125,8 +125,6 @@ func (t *Tracer) StartTransactionOptions(name, transactionType string, opts Tran
 				Key:   elasticTracestateVendorKey,
 				Value: formatElasticTracestateValue(sampleRate),
 			})
-		} else if instrumentationConfig.sampler != nil {
-			result.Sampled = instrumentationConfig.sampler.Sample(tx.traceContext)
 		} else {
 			result.Sampled = true
 		}

transaction_test.go

@@ -45,7 +45,7 @@ func TestStartTransactionTraceContextOptions(t *testing.T) {
 func testStartTransactionTraceContextOptions(t *testing.T, recorded bool) {
 	tracer, _ := transporttest.NewRecorderTracer()
 	defer tracer.Close()
-	tracer.SetSampler(samplerFunc(func(apm.TraceContext) bool {
+	tracer.SetSampler(samplerFunc(func(apm.SampleParams) apm.SampleResult {
 		panic("nope")
 	}))
 
@@ -75,9 +75,9 @@ func startTransactionInvalidTraceContext(t *testing.T, traceContext apm.TraceCon
 	defer tracer.Close()
 
 	var samplerCalled bool
-	tracer.SetSampler(samplerFunc(func(apm.TraceContext) bool {
+	tracer.SetSampler(samplerFunc(func(apm.SampleParams) apm.SampleResult {
 		samplerCalled = true
-		return true
+		return apm.SampleResult{Sampled: true}
 	}))
 
 	opts := apm.TransactionOptions{TraceContext: traceContext}
@@ -214,7 +214,9 @@ func TestTransactionParentIDWithEnsureParent(t *testing.T) {
 func TestTransactionContextNotSampled(t *testing.T) {
 	tracer := apmtest.NewRecordingTracer()
 	defer tracer.Close()
-	tracer.SetSampler(samplerFunc(func(apm.TraceContext) bool { return false }))
+	tracer.SetSampler(samplerFunc(func(apm.SampleParams) apm.SampleResult {
+		return apm.SampleResult{Sampled: false}
+	}))
 
 	tx := tracer.StartTransaction("name", "type")
 	tx.Context.SetLabel("foo", "bar")
@@ -230,7 +232,7 @@ func TestTransactionNotRecording(t *testing.T) {
 	tracer := apmtest.NewRecordingTracer()
 	defer tracer.Close()
 	tracer.SetRecording(false)
-	tracer.SetSampler(samplerFunc(func(apm.TraceContext) bool {
+	tracer.SetSampler(samplerFunc(func(apm.SampleParams) apm.SampleResult {
 		panic("should not be called")
 	}))
 
@@ -575,8 +577,8 @@ func BenchmarkTransaction(b *testing.B) {
 	})
 }
 
-type samplerFunc func(apm.TraceContext) bool
+type samplerFunc func(apm.SampleParams) apm.SampleResult
 
-func (f samplerFunc) Sample(t apm.TraceContext) bool {
-	return f(t)
+func (f samplerFunc) Sample(p apm.SampleParams) apm.SampleResult {
+	return f(p)
 }