feat: Add trace based sampling for go plugin.

Author: kinyoklion

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
+	})
+}

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
 }
 

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,

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))
+		}),
+	}
+}