diff --git a/go/config.go b/go/config.go index bc10b470a..99e033a39 100644 --- a/go/config.go +++ b/go/config.go @@ -4,17 +4,19 @@ import ( "context" "github.com/launchdarkly/observability-sdk/go/internal/defaults" + "go.opentelemetry.io/otel/trace" ) type observabilityConfig struct { - serviceName string - serviceVersion string - environment string - backendURL string - otlpEndpoint string - manualStart bool - context context.Context - debug bool + serviceName string + serviceVersion string + environment string + backendURL string + otlpEndpoint string + manualStart bool + context context.Context + debug bool + samplingRateMap map[trace.SpanKind]float64 } func defaultConfig() observabilityConfig { @@ -86,3 +88,12 @@ func WithDebug() Option { c.debug = true } } + +// WithSamplingRateMap sets the sampling rate for each span kind. +// This setting can influence the quality of metrics used for experiments and guarded +// releases and should only be adjusted with consultation. +func WithSamplingRateMap(rates map[trace.SpanKind]float64) Option { + return Option(func(conf *observabilityConfig) { + conf.samplingRateMap = rates + }) +} diff --git a/go/internal/otel/otel.go b/go/internal/otel/otel.go index db08f6d0e..e91324d25 100644 --- a/go/internal/otel/otel.go +++ b/go/internal/otel/otel.go @@ -210,7 +210,6 @@ func createTracerProvider( return nil, fmt.Errorf("creating OTLP trace exporter: %w", err) } opts = append([]sdktrace.TracerProviderOption{ - sdktrace.WithSampler(sampler), sdktrace.WithBatcher(newTraceExporter(exporter, customSampler), sdktrace.WithBatchTimeout(time.Second), sdktrace.WithExportTimeout(30*time.Second), @@ -219,6 +218,10 @@ func createTracerProvider( ), sdktrace.WithResource(resources), }, opts...) + // Only configure a sampler when there is a sampling configuration. + if sampler != nil { + opts = append(opts, sdktrace.WithSampler(sampler)) + } return sdktrace.NewTracerProvider(opts...), nil } diff --git a/go/plugin.go b/go/plugin.go index aef2576fb..66b933eac 100644 --- a/go/plugin.go +++ b/go/plugin.go @@ -56,20 +56,6 @@ func (p ObservabilityPlugin) Metadata() ldplugins.Metadata { return ldplugins.NewMetadata("launchdarkly-observability") } -type allSampler struct{} - -// Description provide a description of the sampler. -func (a *allSampler) Description() string { - return "samples all traces" -} - -// ShouldSample determines if the trace should be sampled. -func (a *allSampler) ShouldSample(parameters trace.SamplingParameters) trace.SamplingResult { - return trace.SamplingResult{ - Decision: trace.RecordAndSample, - } -} - func (p ObservabilityPlugin) getSamplingConfig(projectId string) (*gql.GetSamplingConfigResponse, error) { var ctx context.Context if p.config.context != nil { @@ -101,7 +87,12 @@ func (p ObservabilityPlugin) Register(client interfaces.LDClientInterface, ldmd logging.SetLogger(logging.ConsoleLogger{}) } - var s trace.Sampler = &allSampler{} + var s trace.Sampler + if len(p.config.samplingRateMap) > 0 { + s = getSampler(p.config.samplingRateMap) + } else { + s = nil + } otel.SetConfig(otel.Config{ OtlpEndpoint: p.config.otlpEndpoint, ResourceAttributes: attributes, diff --git a/go/trace_sampler.go b/go/trace_sampler.go new file mode 100644 index 000000000..2d523ba74 --- /dev/null +++ b/go/trace_sampler.go @@ -0,0 +1,64 @@ +package ldobserve + +import ( + "encoding/binary" + "fmt" + + "github.com/samber/lo" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/trace" +) + +type traceSampler struct { + traceIDUpperBounds map[trace.SpanKind]uint64 + description string +} + +func (ts traceSampler) ShouldSample(p sdktrace.SamplingParameters) sdktrace.SamplingResult { + psc := trace.SpanContextFromContext(p.ParentContext) + if psc.IsSampled() { + return sdktrace.SamplingResult{ + Decision: sdktrace.RecordAndSample, + Tracestate: psc.TraceState(), + } + } + bound, ok := ts.traceIDUpperBounds[p.Kind] + if !ok { + bound, ok = ts.traceIDUpperBounds[trace.SpanKindUnspecified] + // If there are no bounds specified, then we sample all + // Avoiding doing work here versus having default bounds which would + // would require additional work per span. + if !ok { + return sdktrace.SamplingResult{ + Decision: sdktrace.RecordAndSample, + Tracestate: psc.TraceState(), + } + } + } + + x := binary.BigEndian.Uint64(p.TraceID[8:16]) >> 1 + if x < bound { + return sdktrace.SamplingResult{ + Decision: sdktrace.RecordAndSample, + Tracestate: psc.TraceState(), + } + } + return sdktrace.SamplingResult{ + Decision: sdktrace.Drop, + Tracestate: psc.TraceState(), + } +} + +func (ts traceSampler) Description() string { + return ts.description +} + +// creates a per-span-kind sampler that samples each kind at a provided fraction. +func getSampler(rates map[trace.SpanKind]float64) traceSampler { + return traceSampler{ + description: fmt.Sprintf("TraceIDRatioBased{%+v}", rates), + traceIDUpperBounds: lo.MapEntries(rates, func(key trace.SpanKind, value float64) (trace.SpanKind, uint64) { + return key, uint64(value * (1 << 63)) + }), + } +} diff --git a/go/trace_sampler_test.go b/go/trace_sampler_test.go new file mode 100644 index 000000000..1a94b172c --- /dev/null +++ b/go/trace_sampler_test.go @@ -0,0 +1,662 @@ +package ldobserve + +import ( + "context" + "encoding/binary" + "fmt" + "math" + "math/rand" + "testing" + + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/trace" +) + +func TestTraceSampler_ShouldSample_WithSampledParent(t *testing.T) { + // Create a sampler with some rates + rates := map[trace.SpanKind]float64{ + trace.SpanKindServer: 0.5, + } + sampler := getSampler(rates) + + // Create a parent context with sampled trace + parentTraceID := trace.TraceID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16} + parentSpanID := trace.SpanID{1, 2, 3, 4, 5, 6, 7, 8} + parentContext := trace.NewSpanContext(trace.SpanContextConfig{ + TraceID: parentTraceID, + SpanID: parentSpanID, + TraceFlags: trace.FlagsSampled, + }) + ctx := trace.ContextWithSpanContext(context.Background(), parentContext) + + // Test sampling parameters + params := sdktrace.SamplingParameters{ + ParentContext: ctx, + TraceID: trace.TraceID{17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32}, + Name: "test-span", + Kind: trace.SpanKindServer, + } + + result := sampler.ShouldSample(params) + + // Should always sample when parent is sampled + if result.Decision != sdktrace.RecordAndSample { + t.Errorf("Expected decision %v, got %v", sdktrace.RecordAndSample, result.Decision) + } +} + +func TestTraceSampler_ShouldSample_WithUnsampledParent(t *testing.T) { + // Create a sampler with specific rates + rates := map[trace.SpanKind]float64{ + trace.SpanKindServer: 0.5, + } + sampler := getSampler(rates) + + // Create a parent context without sampled trace + parentTraceID := trace.TraceID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16} + parentSpanID := trace.SpanID{1, 2, 3, 4, 5, 6, 7, 8} + parentContext := trace.NewSpanContext(trace.SpanContextConfig{ + TraceID: parentTraceID, + SpanID: parentSpanID, + TraceFlags: 0, // Not sampled + }) + ctx := trace.ContextWithSpanContext(context.Background(), parentContext) + + // Test with trace ID that should be sampled (lower than threshold) + // For 0.5 rate, threshold is 0.5 * (1 << 63) = 0x4000000000000000 + // We'll use a trace ID with upper 8 bytes that when shifted right by 1 gives a value < threshold + traceID := trace.TraceID{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0x20} // Small value + + params := sdktrace.SamplingParameters{ + ParentContext: ctx, + TraceID: traceID, + Name: "test-span", + Kind: trace.SpanKindServer, + } + + result := sampler.ShouldSample(params) + + // Should sample based on trace ID ratio + if result.Decision != sdktrace.RecordAndSample { + t.Errorf("Expected decision %v, got %v", sdktrace.RecordAndSample, result.Decision) + } +} + +func TestTraceSampler_ShouldSample_WithUnsampledParent_AboveThreshold(t *testing.T) { + // Create a sampler with specific rates + rates := map[trace.SpanKind]float64{ + trace.SpanKindServer: 0.5, + } + sampler := getSampler(rates) + + // Create a parent context without sampled trace + parentTraceID := trace.TraceID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16} + parentSpanID := trace.SpanID{1, 2, 3, 4, 5, 6, 7, 8} + parentContext := trace.NewSpanContext(trace.SpanContextConfig{ + TraceID: parentTraceID, + SpanID: parentSpanID, + TraceFlags: 0, // Not sampled + }) + ctx := trace.ContextWithSpanContext(context.Background(), parentContext) + + // Test with trace ID that should NOT be sampled (above threshold) + // For 0.5 rate, threshold is 0.5 * (1 << 63) = 0x4000000000000000 + // The condition is: if x < bound then sample, so we need x >= bound to NOT sample + // We need (x >> 1) >= threshold, so x >= threshold * 2 + threshold := uint64(0.5 * (1 << 63)) + x := threshold*2 + 1 // This ensures (x >> 1) = threshold + 0.5 >= threshold + + // Create trace ID with the calculated value in the upper 8 bytes + traceIDBytes := make([]byte, 16) + binary.BigEndian.PutUint64(traceIDBytes[8:16], x) + traceID := trace.TraceID{} + copy(traceID[:], traceIDBytes) + + params := sdktrace.SamplingParameters{ + ParentContext: ctx, + TraceID: traceID, + Name: "test-span", + Kind: trace.SpanKindServer, + } + + result := sampler.ShouldSample(params) + + // Should NOT sample based on trace ID ratio + if result.Decision != sdktrace.Drop { + t.Errorf("Expected decision %v, got %v", sdktrace.Drop, result.Decision) + } +} + +func TestTraceSampler_ShouldSample_WithUnsampledParent_NoParentContext(t *testing.T) { + // Create a sampler with specific rates + rates := map[trace.SpanKind]float64{ + trace.SpanKindServer: 0.5, + } + sampler := getSampler(rates) + + // Test with no parent context + params := sdktrace.SamplingParameters{ + ParentContext: context.Background(), + TraceID: trace.TraceID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, + Name: "test-span", + Kind: trace.SpanKindServer, + } + + result := sampler.ShouldSample(params) + + // Should sample based on trace ID ratio since no parent context + if result.Decision != sdktrace.RecordAndSample { + t.Errorf("Expected decision %v, got %v", sdktrace.RecordAndSample, result.Decision) + } +} + +func TestTraceSampler_ShouldSample_WithUnsampledParent_UnspecifiedKind(t *testing.T) { + // Create a sampler with specific rates but no Unspecified kind + rates := map[trace.SpanKind]float64{ + trace.SpanKindServer: 0.5, + } + sampler := getSampler(rates) + + // Create a parent context without sampled trace + parentTraceID := trace.TraceID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16} + parentSpanID := trace.SpanID{1, 2, 3, 4, 5, 6, 7, 8} + parentContext := trace.NewSpanContext(trace.SpanContextConfig{ + TraceID: parentTraceID, + SpanID: parentSpanID, + TraceFlags: 0, // Not sampled + }) + ctx := trace.ContextWithSpanContext(context.Background(), parentContext) + + // Test with Unspecified kind + params := sdktrace.SamplingParameters{ + ParentContext: ctx, + TraceID: trace.TraceID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, + Name: "test-span", + Kind: trace.SpanKindUnspecified, + } + + result := sampler.ShouldSample(params) + + // Should sample since no bounds specified for Unspecified kind + if result.Decision != sdktrace.RecordAndSample { + t.Errorf("Expected decision %v, got %v", sdktrace.RecordAndSample, result.Decision) + } +} + +func TestTraceSampler_ShouldSample_WithUnsampledParent_UnspecifiedKindWithDefault(t *testing.T) { + // Create a sampler with specific rates including Unspecified kind + rates := map[trace.SpanKind]float64{ + trace.SpanKindServer: 0.5, + trace.SpanKindUnspecified: 0.25, + } + sampler := getSampler(rates) + + // Create a parent context without sampled trace + parentTraceID := trace.TraceID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16} + parentSpanID := trace.SpanID{1, 2, 3, 4, 5, 6, 7, 8} + parentContext := trace.NewSpanContext(trace.SpanContextConfig{ + TraceID: parentTraceID, + SpanID: parentSpanID, + TraceFlags: 0, // Not sampled + }) + ctx := trace.ContextWithSpanContext(context.Background(), parentContext) + + // Test with Unspecified kind + params := sdktrace.SamplingParameters{ + ParentContext: ctx, + TraceID: trace.TraceID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, + Name: "test-span", + Kind: trace.SpanKindUnspecified, + } + + result := sampler.ShouldSample(params) + + // Should sample based on Unspecified kind rate + if result.Decision != sdktrace.RecordAndSample { + t.Errorf("Expected decision %v, got %v", sdktrace.RecordAndSample, result.Decision) + } +} + +func TestTraceSampler_ShouldSample_WithUnsampledParent_UnknownKind(t *testing.T) { + // Create a sampler with specific rates but no Client kind + rates := map[trace.SpanKind]float64{ + trace.SpanKindServer: 0.5, + } + sampler := getSampler(rates) + + // Create a parent context without sampled trace + parentTraceID := trace.TraceID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16} + parentSpanID := trace.SpanID{1, 2, 3, 4, 5, 6, 7, 8} + parentContext := trace.NewSpanContext(trace.SpanContextConfig{ + TraceID: parentTraceID, + SpanID: parentSpanID, + TraceFlags: 0, // Not sampled + }) + ctx := trace.ContextWithSpanContext(context.Background(), parentContext) + + // Test with Client kind (not in rates) + params := sdktrace.SamplingParameters{ + ParentContext: ctx, + TraceID: trace.TraceID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, + Name: "test-span", + Kind: trace.SpanKindClient, + } + + result := sampler.ShouldSample(params) + + // Should sample since no bounds specified for Client kind + if result.Decision != sdktrace.RecordAndSample { + t.Errorf("Expected decision %v, got %v", sdktrace.RecordAndSample, result.Decision) + } +} + +func TestTraceSampler_ShouldSample_WithUnsampledParent_UnknownKindWithUnspecifiedFallback(t *testing.T) { + // Create a sampler with specific rates including Unspecified kind as fallback + rates := map[trace.SpanKind]float64{ + trace.SpanKindServer: 0.5, + trace.SpanKindUnspecified: 0.25, // This should be used as fallback for unknown kinds + } + sampler := getSampler(rates) + + // Create a parent context without sampled trace + parentTraceID := trace.TraceID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16} + parentSpanID := trace.SpanID{1, 2, 3, 4, 5, 6, 7, 8} + parentContext := trace.NewSpanContext(trace.SpanContextConfig{ + TraceID: parentTraceID, + SpanID: parentSpanID, + TraceFlags: 0, // Not sampled + }) + ctx := trace.ContextWithSpanContext(context.Background(), parentContext) + + // Test with Client kind (not in rates, but should fall back to Unspecified rate) + params := sdktrace.SamplingParameters{ + ParentContext: ctx, + TraceID: trace.TraceID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, + Name: "test-span", + Kind: trace.SpanKindClient, + } + + result := sampler.ShouldSample(params) + + // Should sample based on Unspecified kind rate (0.25) since Client kind has no rate configured + if result.Decision != sdktrace.RecordAndSample { + t.Errorf("Expected decision %v, got %v", sdktrace.RecordAndSample, result.Decision) + } +} + +func TestTraceSampler_ShouldSample_WithUnsampledParent_UnknownKindWithUnspecifiedFallback_AboveThreshold(t *testing.T) { + // Create a sampler with specific rates including Unspecified kind as fallback + rates := map[trace.SpanKind]float64{ + trace.SpanKindServer: 0.5, + trace.SpanKindUnspecified: 0.25, // This should be used as fallback for unknown kinds + } + sampler := getSampler(rates) + + // Create a parent context without sampled trace + parentTraceID := trace.TraceID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16} + parentSpanID := trace.SpanID{1, 2, 3, 4, 5, 6, 7, 8} + parentContext := trace.NewSpanContext(trace.SpanContextConfig{ + TraceID: parentTraceID, + SpanID: parentSpanID, + TraceFlags: 0, // Not sampled + }) + ctx := trace.ContextWithSpanContext(context.Background(), parentContext) + + // Test with Client kind (not in rates, but should fall back to Unspecified rate) + // For 0.25 rate, threshold is 0.25 * (1 << 63) = 0x2000000000000000 + // We need (x >> 1) >= threshold, so x >= threshold * 2 + threshold := uint64(0.25 * (1 << 63)) + x := threshold*2 + 1 // This ensures (x >> 1) = threshold + 0.5 >= threshold + + // Create trace ID with the calculated value in the upper 8 bytes + traceIDBytes := make([]byte, 16) + binary.BigEndian.PutUint64(traceIDBytes[8:16], x) + traceID := trace.TraceID{} + copy(traceID[:], traceIDBytes) + + params := sdktrace.SamplingParameters{ + ParentContext: ctx, + TraceID: traceID, + Name: "test-span", + Kind: trace.SpanKindClient, + } + + result := sampler.ShouldSample(params) + + // Should NOT sample based on Unspecified kind rate (0.25) since trace ID is above threshold + if result.Decision != sdktrace.Drop { + t.Errorf("Expected decision %v, got %v", sdktrace.Drop, result.Decision) + } +} + +func TestTraceSampler_Description(t *testing.T) { + rates := map[trace.SpanKind]float64{ + trace.SpanKindServer: 0.5, + trace.SpanKindClient: 0.25, + } + sampler := getSampler(rates) + + // The description format is "TraceIDRatioBased{map[key:value]}" + // Since map order is not guaranteed, we'll check that it contains the expected parts + description := sampler.Description() + if !contains(description, "TraceIDRatioBased{map[") { + t.Errorf("Expected description to contain 'TraceIDRatioBased{map[', got %s", description) + } + if !contains(description, "}") { + t.Errorf("Expected description to end with '}', got %s", description) + } + + // Check that it contains the expected key-value pairs + // The actual format uses lowercase keys like "server:0.5" + if !contains(description, "server:0.5") { + t.Errorf("Expected description to contain server rate, got %s", description) + } + if !contains(description, "client:0.25") { + t.Errorf("Expected description to contain client rate, got %s", description) + } +} + +// Helper function to check if a string contains a substring +func contains(s, substr string) bool { + return len(s) >= len(substr) && (s == substr || + (len(s) > len(substr) && (s[:len(substr)] == substr || + s[len(s)-len(substr):] == substr || + contains(s[1:], substr)))) +} + +func TestGetSampler_EmptyRates(t *testing.T) { + rates := map[trace.SpanKind]float64{} + sampler := getSampler(rates) + + // Should create sampler with empty bounds + if len(sampler.traceIDUpperBounds) != 0 { + t.Errorf("Expected empty bounds, got %d", len(sampler.traceIDUpperBounds)) + } + if sampler.description != "TraceIDRatioBased{map[]}" { + t.Errorf("Expected description %s, got %s", "TraceIDRatioBased{map[]}", sampler.description) + } +} + +func TestGetSampler_WithRates(t *testing.T) { + rates := map[trace.SpanKind]float64{ + trace.SpanKindServer: 0.5, + trace.SpanKindClient: 0.25, + } + sampler := getSampler(rates) + + // Check that bounds are calculated correctly + // 0.5 * (1 << 63) = 0x4000000000000000 + // 0.25 * (1 << 63) = 0x2000000000000000 + expectedServerBound := uint64(0.5 * (1 << 63)) + expectedClientBound := uint64(0.25 * (1 << 63)) + + if sampler.traceIDUpperBounds[trace.SpanKindServer] != expectedServerBound { + t.Errorf("Expected server bound %d, got %d", expectedServerBound, sampler.traceIDUpperBounds[trace.SpanKindServer]) + } + if sampler.traceIDUpperBounds[trace.SpanKindClient] != expectedClientBound { + t.Errorf("Expected client bound %d, got %d", expectedClientBound, sampler.traceIDUpperBounds[trace.SpanKindClient]) + } +} + +func TestTraceSampler_ShouldSample_EdgeCase_ZeroRate(t *testing.T) { + // Create a sampler with zero rate + rates := map[trace.SpanKind]float64{ + trace.SpanKindServer: 0.0, + } + sampler := getSampler(rates) + + // Create a parent context without sampled trace + parentTraceID := trace.TraceID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16} + parentSpanID := trace.SpanID{1, 2, 3, 4, 5, 6, 7, 8} + parentContext := trace.NewSpanContext(trace.SpanContextConfig{ + TraceID: parentTraceID, + SpanID: parentSpanID, + TraceFlags: 0, // Not sampled + }) + ctx := trace.ContextWithSpanContext(context.Background(), parentContext) + + params := sdktrace.SamplingParameters{ + ParentContext: ctx, + TraceID: trace.TraceID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, + Name: "test-span", + Kind: trace.SpanKindServer, + } + + result := sampler.ShouldSample(params) + + // With 0.0 rate, threshold is 0, so no trace ID should be sampled + if result.Decision != sdktrace.Drop { + t.Errorf("Expected decision %v, got %v", sdktrace.Drop, result.Decision) + } +} + +func TestTraceSampler_ShouldSample_EdgeCase_OneRate(t *testing.T) { + // Create a sampler with 100% rate + rates := map[trace.SpanKind]float64{ + trace.SpanKindServer: 1.0, + } + sampler := getSampler(rates) + + // Create a parent context without sampled trace + parentTraceID := trace.TraceID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16} + parentSpanID := trace.SpanID{1, 2, 3, 4, 5, 6, 7, 8} + parentContext := trace.NewSpanContext(trace.SpanContextConfig{ + TraceID: parentTraceID, + SpanID: parentSpanID, + TraceFlags: 0, // Not sampled + }) + ctx := trace.ContextWithSpanContext(context.Background(), parentContext) + + params := sdktrace.SamplingParameters{ + ParentContext: ctx, + TraceID: trace.TraceID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, + Name: "test-span", + Kind: trace.SpanKindServer, + } + + result := sampler.ShouldSample(params) + + // With 1.0 rate, threshold is max uint64, so all trace IDs should be sampled + if result.Decision != sdktrace.RecordAndSample { + t.Errorf("Expected decision %v, got %v", sdktrace.RecordAndSample, result.Decision) + } +} + +func TestTraceSampler_ShouldSample_TraceIDCalculation(t *testing.T) { + // Create a sampler with 0.5 rate + rates := map[trace.SpanKind]float64{ + trace.SpanKindServer: 0.5, + } + sampler := getSampler(rates) + + // Create a parent context without sampled trace + parentTraceID := trace.TraceID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16} + parentSpanID := trace.SpanID{1, 2, 3, 4, 5, 6, 7, 8} + parentContext := trace.NewSpanContext(trace.SpanContextConfig{ + TraceID: parentTraceID, + SpanID: parentSpanID, + TraceFlags: 0, // Not sampled + }) + ctx := trace.ContextWithSpanContext(context.Background(), parentContext) + + // Test the exact calculation from the code + // The code does: binary.BigEndian.Uint64(p.TraceID[8:16]) >> 1 + // For 0.5 rate, threshold is 0.5 * (1 << 63) = 0x4000000000000000 + // The condition is: if x < bound then sample + + // Create a trace ID where the upper 8 bytes when shifted right by 1 equals exactly the threshold + threshold := uint64(0.5 * (1 << 63)) + // We need (x >> 1) = threshold, so x = threshold * 2 + x := threshold * 2 + + // Create trace ID with the calculated value in the upper 8 bytes + traceIDBytes := make([]byte, 16) + binary.BigEndian.PutUint64(traceIDBytes[8:16], x) + traceID := trace.TraceID{} + copy(traceID[:], traceIDBytes) + + params := sdktrace.SamplingParameters{ + ParentContext: ctx, + TraceID: traceID, + Name: "test-span", + Kind: trace.SpanKindServer, + } + + result := sampler.ShouldSample(params) + + // This should be exactly at the threshold, so (x >> 1) = threshold + // Since the condition is x < bound, and x = threshold, this should NOT sample + if result.Decision != sdktrace.Drop { + t.Errorf("Expected decision %v, got %v", sdktrace.Drop, result.Decision) + } +} + +func TestTraceSampler_ShouldSample_TraceIDCalculation_JustBelowThreshold(t *testing.T) { + // Create a sampler with 0.5 rate + rates := map[trace.SpanKind]float64{ + trace.SpanKindServer: 0.5, + } + sampler := getSampler(rates) + + // Create a parent context without sampled trace + parentTraceID := trace.TraceID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16} + parentSpanID := trace.SpanID{1, 2, 3, 4, 5, 6, 7, 8} + parentContext := trace.NewSpanContext(trace.SpanContextConfig{ + TraceID: parentTraceID, + SpanID: parentSpanID, + TraceFlags: 0, // Not sampled + }) + ctx := trace.ContextWithSpanContext(context.Background(), parentContext) + + // Test with a value just below the threshold + threshold := uint64(0.5 * (1 << 63)) + // We need (x >> 1) < threshold, so x < threshold * 2 + x := threshold*2 - 2 // This ensures (x >> 1) = threshold - 1 < threshold + + // Create trace ID with the calculated value in the upper 8 bytes + traceIDBytes := make([]byte, 16) + binary.BigEndian.PutUint64(traceIDBytes[8:16], x) + traceID := trace.TraceID{} + copy(traceID[:], traceIDBytes) + + params := sdktrace.SamplingParameters{ + ParentContext: ctx, + TraceID: traceID, + Name: "test-span", + Kind: trace.SpanKindServer, + } + + result := sampler.ShouldSample(params) + + // This should be just below the threshold, so it should sample + if result.Decision != sdktrace.RecordAndSample { + t.Errorf("Expected decision %v, got %v", sdktrace.RecordAndSample, result.Decision) + } +} + +func TestTraceSampler_ShouldSample_TraceIDCalculation_JustAboveThreshold(t *testing.T) { + // Create a sampler with 0.5 rate + rates := map[trace.SpanKind]float64{ + trace.SpanKindServer: 0.5, + } + sampler := getSampler(rates) + + // Create a parent context without sampled trace + parentTraceID := trace.TraceID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16} + parentSpanID := trace.SpanID{1, 2, 3, 4, 5, 6, 7, 8} + parentContext := trace.NewSpanContext(trace.SpanContextConfig{ + TraceID: parentTraceID, + SpanID: parentSpanID, + TraceFlags: 0, // Not sampled + }) + ctx := trace.ContextWithSpanContext(context.Background(), parentContext) + + // Test with a value just above the threshold + threshold := uint64(0.5 * (1 << 63)) + // We need (x >> 1) >= threshold, so x >= threshold * 2 + x := threshold*2 + 2 // This ensures (x >> 1) = threshold + 1 >= threshold + + // Create trace ID with the calculated value in the upper 8 bytes + traceIDBytes := make([]byte, 16) + binary.BigEndian.PutUint64(traceIDBytes[8:16], x) + traceID := trace.TraceID{} + copy(traceID[:], traceIDBytes) + + params := sdktrace.SamplingParameters{ + ParentContext: ctx, + TraceID: traceID, + Name: "test-span", + Kind: trace.SpanKindServer, + } + + result := sampler.ShouldSample(params) + + // This should be just above the threshold, so it should NOT sample + if result.Decision != sdktrace.Drop { + t.Errorf("Expected decision %v, got %v", sdktrace.Drop, result.Decision) + } +} + +func TestTraceSampler_StatisticalSampling_50PercentRate(t *testing.T) { + // Create a sampler with 50% rate + rates := map[trace.SpanKind]float64{ + trace.SpanKindServer: 0.5, + } + sampler := getSampler(rates) + + // The test is non-deterministic, so this tries to reach a good balance of + // trials versus speed. When running tests with the race detector performance + // is very slow, so we run fewer trials and make the deviation tolerance larger. + // These number can be adjusted if the test proves flaky. + const numSpans = 100_000 + const expectedRate = 0.5 + const tolerance = 0.005 // Allow 0.5% deviation from expected rate + + sampledCount := 0 + + for i := range numSpans { + // Generate a random trace ID for each span + traceID := generateRandomTraceID() + + params := sdktrace.SamplingParameters{ + ParentContext: context.Background(), + TraceID: traceID, + Name: fmt.Sprintf("test-span-%d", i), + Kind: trace.SpanKindServer, + } + + result := sampler.ShouldSample(params) + if result.Decision == sdktrace.RecordAndSample { + sampledCount++ + } + } + + // Calculate observed sampling rate + observedRate := float64(sampledCount) / float64(numSpans) + + // Check if observed rate is within acceptable tolerance + if math.Abs(observedRate-expectedRate) > tolerance { + t.Errorf("Observed sampling rate %.6f does not match expected rate %.6f within tolerance %.6f", + observedRate, expectedRate, tolerance) + t.Errorf("Sampled %d out of %d spans (expected approximately %d)", + sampledCount, numSpans, int(float64(numSpans)*expectedRate)) + } + + // Log the results for verification + t.Logf("Statistical sampling test results:") + t.Logf(" Total spans tested: %d", numSpans) + t.Logf(" Spans sampled: %d", sampledCount) + t.Logf(" Observed rate: %.6f", observedRate) + t.Logf(" Expected rate: %.6f", expectedRate) + t.Logf(" Difference: %.6f", math.Abs(observedRate-expectedRate)) + t.Logf(" Tolerance: %.6f", tolerance) +} + +// generateRandomTraceID creates a random trace ID for testing +func generateRandomTraceID() trace.TraceID { + var traceID trace.TraceID + for i := 0; i < 16; i++ { + traceID[i] = byte(rand.Intn(256)) + } + return traceID +}