Respecting -tracing.otel.sample-ratio configuration when enabling OpenTelemetry tracing with X-ray

CHANGELOG.md

@@ -57,6 +57,7 @@
 * [FEATURE] Storage/Bucket: Added `-*.s3.bucket-lookup-type` allowing to configure the s3 bucket lookup type. #4794
 * [BUGFIX] Memberlist: Add join with no retrying when starting service. #4804
 * [BUGFIX] Ruler: Fix /ruler/rule_groups returns YAML with extra fields. #4767
+* [BUGFIX] Respecting `-tracing.otel.sample-ratio` configuration when enabling OpenTelemetry tracing with X-ray. #4862
 
 ## 1.13.0 2022-07-14
 

pkg/tracing/sampler/sampling.go

@@ -0,0 +1,51 @@
+package sampler
+
+import (
+	"fmt"
+
+	sdktrace "go.opentelemetry.io/otel/sdk/trace"
+	"go.opentelemetry.io/otel/trace"
+)
+
+type randGenerator interface {
+	Float64() float64
+}
+
+type RandomRatioBased struct {
+	sdktrace.Sampler
+	rnd      randGenerator
+	fraction float64
+}
+
+// NewRandomRatioBased crea
+// fraction parameter should be between 0 and 1 where:
+// fraction >= 1 it will always sample
+// fraction <= 0 it will never sample
+func NewRandomRatioBased(fraction float64, rnd randGenerator) sdktrace.Sampler {
+	if fraction >= 1 {
+		return sdktrace.AlwaysSample()
+	} else if fraction <= 0 {
+		return sdktrace.NeverSample()
+	}
+
+	return &RandomRatioBased{rnd: rnd, fraction: fraction}
+}
+
+func (s *RandomRatioBased) ShouldSample(p sdktrace.SamplingParameters) sdktrace.SamplingResult {
+	psc := trace.SpanContextFromContext(p.ParentContext)
+	shouldSample := s.rnd.Float64() < s.fraction
+	if shouldSample {
+		return sdktrace.SamplingResult{
+			Decision:   sdktrace.RecordAndSample,
+			Tracestate: psc.TraceState(),
+		}
+	}
+	return sdktrace.SamplingResult{
+		Decision:   sdktrace.Drop,
+		Tracestate: psc.TraceState(),
+	}
+}
+
+func (s *RandomRatioBased) Description() string {
+	return fmt.Sprintf("RandomRatioBased{%g}", s.fraction)
+}

pkg/tracing/sampler/sampling_test.go

@@ -0,0 +1,79 @@
+package sampler
+
+import (
+	"context"
+	"math/rand"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	sdktrace "go.opentelemetry.io/otel/sdk/trace"
+	"go.opentelemetry.io/otel/trace"
+)
+
+type mockGenerator struct {
+	mockedValue float64
+}
+
+func (g *mockGenerator) Float64() float64 {
+	return g.mockedValue
+}
+
+func Test_ShouldSample(t *testing.T) {
+	traceID, _ := trace.TraceIDFromHex("4bf92f3577b34da6a3ce929d0e0e4736")
+	parentCtx := trace.ContextWithSpanContext(
+		context.Background(),
+		trace.NewSpanContext(trace.SpanContextConfig{
+			TraceState: trace.TraceState{},
+		}),
+	)
+
+	tests := []struct {
+		name             string
+		samplingDecision sdktrace.SamplingDecision
+		fraction         float64
+		generator        randGenerator
+	}{
+		{
+			name:             "should always sample",
+			samplingDecision: sdktrace.RecordAndSample,
+			fraction:         1,
+			generator:        rand.New(rand.NewSource(rand.Int63())),
+		},
+		{
+			name:             "should nerver sample",
+			samplingDecision: sdktrace.Drop,
+			fraction:         0,
+			generator:        rand.New(rand.NewSource(rand.Int63())),
+		},
+		{
+			name:             "should sample when fraction is above generated",
+			samplingDecision: sdktrace.RecordAndSample,
+			fraction:         0.5,
+			generator:        &mockGenerator{0.2},
+		},
+		{
+			name:             "should not sample when fraction is not above generated",
+			samplingDecision: sdktrace.Drop,
+			fraction:         0.5,
+			generator:        &mockGenerator{0.8},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			s := NewRandomRatioBased(tt.fraction, tt.generator)
+			for i := 0; i < 100; i++ {
+
+				r := s.ShouldSample(
+					sdktrace.SamplingParameters{
+						ParentContext: parentCtx,
+						TraceID:       traceID,
+						Name:          "test",
+						Kind:          trace.SpanKindServer,
+					})
+
+				require.Equal(t, tt.samplingDecision, r.Decision)
+			}
+		})
+	}
+}

pkg/tracing/tracing.go

@@ -4,7 +4,9 @@ import (
 	"context"
 	"flag"
 	"fmt"
+	"math/rand"
 	"strings"
+	"time"
 
 	"github.com/go-kit/log/level"
 	"github.com/pkg/errors"
@@ -22,6 +24,7 @@ import (
 	semconv "go.opentelemetry.io/otel/semconv/v1.10.0"
 
 	"github.com/cortexproject/cortex/pkg/tracing/migration"
+	"github.com/cortexproject/cortex/pkg/tracing/sampler"
 	util_log "github.com/cortexproject/cortex/pkg/util/log"
 	"github.com/cortexproject/cortex/pkg/util/tls"
 )
@@ -118,13 +121,15 @@ func SetupTracing(ctx context.Context, name string, c Config) (func(context.Cont
 func newTraceProvider(name string, c Config, exporter *otlptrace.Exporter) *sdktrace.TracerProvider {
 	options := []sdktrace.TracerProviderOption{
 		sdktrace.WithBatcher(exporter),
-		sdktrace.WithSampler(sdktrace.ParentBased(sdktrace.TraceIDRatioBased(c.Otel.SampleRatio))),
 		sdktrace.WithResource(newResource(name)),
 	}
 
 	switch strings.ToLower(c.Otel.ExporterType) {
 	case "awsxray":
 		options = append(options, sdktrace.WithIDGenerator(xray.NewIDGenerator()))
+		options = append(options, sdktrace.WithSampler(sdktrace.ParentBased(sampler.NewRandomRatioBased(c.Otel.SampleRatio, rand.New(rand.NewSource(time.Now().Unix()))))))
+	default:
+		options = append(options, sdktrace.WithSampler(sdktrace.ParentBased(sdktrace.TraceIDRatioBased(c.Otel.SampleRatio))))
 	}
 
 	return sdktrace.NewTracerProvider(options...)