Align jaeger probabilistic sampler with trace.TraceIDRatioBased (#6892)

Co-authored-by: Damien Mathieu <[email protected]>

CHANGELOG.md

@@ -8,6 +8,10 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ## [Unreleased]
 
+### Changed
+
+- Jaeger remote sampler's probabilistic strategy now uses the same sampling algorithm as `trace.TraceIDRatioBased` in `go.opentelemetry.io/contrib/samplers/jaegerremote`. (#6892)
+
 ### Removed
 
 - Drop support for [Go 1.22]. (#6853)

samplers/jaegerremote/sampler.go

@@ -19,7 +19,6 @@
 package jaegerremote // import "go.opentelemetry.io/contrib/samplers/jaegerremote"
 
 import (
-	"encoding/binary"
 	"fmt"
 	"math"
 	"sync"
@@ -39,25 +38,20 @@ const (
 // probabilisticSampler is a sampler that randomly samples a certain percentage
 // of traces.
 type probabilisticSampler struct {
-	samplingRate     float64
-	samplingBoundary uint64
+	samplingRate float64
+	sampler      trace.Sampler
 }
 
-const maxRandomNumber = ^(uint64(1) << 63) // i.e. 0x7fffffffffffffff
-
 // newProbabilisticSampler creates a sampler that randomly samples a certain percentage of traces specified by the
-// samplingRate, in the range between 0.0 and 1.0.
-//
-// It relies on the fact that new trace IDs are 63bit random numbers themselves, thus making the sampling decision
-// without generating a new random number, but simply calculating if traceID < (samplingRate * 2^63).
+// samplingRate, in the range between 0.0 and 1.0. it utilizes the SDK `trace.TraceIDRatioBased` sampler.
 func newProbabilisticSampler(samplingRate float64) *probabilisticSampler {
 	s := new(probabilisticSampler)
 	return s.init(samplingRate)
 }
 
 func (s *probabilisticSampler) init(samplingRate float64) *probabilisticSampler {
 	s.samplingRate = math.Max(0.0, math.Min(samplingRate, 1.0))
-	s.samplingBoundary = uint64(float64(maxRandomNumber) * s.samplingRate)
+	s.sampler = trace.TraceIDRatioBased(s.samplingRate)
 	return s
 }
 
@@ -67,24 +61,13 @@ func (s *probabilisticSampler) SamplingRate() float64 {
 }
 
 func (s *probabilisticSampler) ShouldSample(p trace.SamplingParameters) trace.SamplingResult {
-	psc := oteltrace.SpanContextFromContext(p.ParentContext)
-	traceID := binary.BigEndian.Uint64(p.TraceID[0:8])
-	if s.samplingBoundary >= traceID&maxRandomNumber {
-		return trace.SamplingResult{
-			Decision:   trace.RecordAndSample,
-			Tracestate: psc.TraceState(),
-		}
-	}
-	return trace.SamplingResult{
-		Decision:   trace.Drop,
-		Tracestate: psc.TraceState(),
-	}
+	return s.sampler.ShouldSample(p)
 }
 
 // Equal compares with another sampler.
 func (s *probabilisticSampler) Equal(other trace.Sampler) bool {
 	if o, ok := other.(*probabilisticSampler); ok {
-		return s.samplingBoundary == o.samplingBoundary
+		return math.Abs(s.samplingRate-o.samplingRate) < 1e-9 // consider equal if within 0.000001%
 	}
 	return false
 }
@@ -99,7 +82,7 @@ func (s *probabilisticSampler) Update(samplingRate float64) error {
 }
 
 func (s *probabilisticSampler) Description() string {
-	return "probabilisticSampler{}"
+	return s.sampler.Description()
 }
 
 // -----------------------

samplers/jaegerremote/sampler_remote_test.go

@@ -165,7 +165,7 @@ func initAgent(t *testing.T) (*testutils.MockAgent, *Sampler) {
 
 func makeSamplingParameters(id uint64, operationName string) trace.SamplingParameters {
 	var traceID oteltrace.TraceID
-	binary.BigEndian.PutUint64(traceID[:], id)
+	binary.BigEndian.PutUint64(traceID[8:], id)
 
 	return trace.SamplingParameters{
 		TraceID: traceID,

samplers/jaegerremote/sampler_test.go

@@ -19,7 +19,10 @@
 package jaegerremote
 
 import (
+	crand "crypto/rand"
 	"encoding/binary"
+	"math"
+	"math/rand"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -34,29 +37,78 @@ const (
 	testFirstTimeOperationName = "firstTimeOp"
 
 	testDefaultSamplingProbability = 0.5
-	testMaxID                      = uint64(1) << 62
+	testMaxID                      = uint64(1) << 63
 	testDefaultMaxOperations       = 10
 )
 
+type randomIDGenerator struct {
+	randSource *rand.Rand
+}
+
+// NewTraceID returns a non-zero trace ID from a randomly-chosen sequence.
+func (gen *randomIDGenerator) NewTraceID() oteltrace.TraceID {
+	tid := oteltrace.TraceID{}
+	for {
+		_, _ = gen.randSource.Read(tid[:])
+		if tid.IsValid() {
+			break
+		}
+	}
+	return tid
+}
+
+func defaultIDGenerator() *randomIDGenerator {
+	gen := &randomIDGenerator{}
+	var rngSeed int64
+	_ = binary.Read(crand.Reader, binary.LittleEndian, &rngSeed)
+	gen.randSource = rand.New(rand.NewSource(rngSeed))
+	return gen
+}
+
 func TestProbabilisticSampler(t *testing.T) {
 	var traceID oteltrace.TraceID
 
 	sampler := newProbabilisticSampler(0.5)
-	binary.BigEndian.PutUint64(traceID[:], testMaxID+10)
+	binary.BigEndian.PutUint64(traceID[8:], testMaxID+10)
 	result := sampler.ShouldSample(trace.SamplingParameters{TraceID: traceID})
 	assert.Equal(t, trace.Drop, result.Decision)
-	binary.BigEndian.PutUint64(traceID[:], testMaxID-20)
+	binary.BigEndian.PutUint64(traceID[8:], testMaxID-20)
 	result = sampler.ShouldSample(trace.SamplingParameters{TraceID: traceID})
 	assert.Equal(t, trace.RecordAndSample, result.Decision)
 
 	t.Run("test_64bit_id", func(t *testing.T) {
-		binary.BigEndian.PutUint64(traceID[:], (testMaxID+10)|1<<63)
+		binary.BigEndian.PutUint64(traceID[:8], math.MaxUint64)
+		binary.BigEndian.PutUint64(traceID[8:], testMaxID+10)
 		result = sampler.ShouldSample(trace.SamplingParameters{TraceID: traceID})
 		assert.Equal(t, trace.Drop, result.Decision)
-		binary.BigEndian.PutUint64(traceID[:], (testMaxID-20)|1<<63)
+		binary.BigEndian.PutUint64(traceID[8:], testMaxID-20)
 		result = sampler.ShouldSample(trace.SamplingParameters{TraceID: traceID})
 		assert.Equal(t, trace.RecordAndSample, result.Decision)
 	})
+
+	t.Run("test_parity", func(t *testing.T) {
+		numTests := 1000
+
+		sampler := newProbabilisticSampler(0.5)
+		oracle := trace.TraceIDRatioBased(0.5)
+		idGenerator := defaultIDGenerator()
+
+		for range numTests {
+			traceID := idGenerator.NewTraceID()
+			assert.Equal(t,
+				oracle.ShouldSample(trace.SamplingParameters{TraceID: traceID}),
+				sampler.ShouldSample(trace.SamplingParameters{TraceID: traceID}),
+			)
+		}
+	})
+
+	t.Run("Equals", func(t *testing.T) {
+		sampler := newProbabilisticSampler(0.5)
+		assert.True(t, sampler.Equal(newProbabilisticSampler(0.5)))
+		assert.False(t, sampler.Equal(newProbabilisticSampler(0.0)))
+		assert.False(t, sampler.Equal(newProbabilisticSampler(0.75)))
+		assert.False(t, sampler.Equal(newProbabilisticSampler(1.0)))
+	})
 }
 
 func TestRateLimitingSampler(t *testing.T) {