From 1b80fdaa1ec488d0f8caa93e931fc1d45168f129 Mon Sep 17 00:00:00 2001
From: Baptiste Girard-Carrabin <baptiste.girardcarrabin@datadoghq.com>
Date: Fri, 14 Jun 2024 11:58:38 +0200
Subject: [PATCH 1/2] [tracing] Create NameSampler

Simple sampler that samples every span having a certain name.
It's best used in conjunction with the ParentBased sampler so that the child spans are also sampled.
---
 pkg/tracing/plugin/sampler.go | 51 +++++++++++++++++++++++++++++++++++
 1 file changed, 51 insertions(+)
 create mode 100644 pkg/tracing/plugin/sampler.go

diff --git a/pkg/tracing/plugin/sampler.go b/pkg/tracing/plugin/sampler.go
new file mode 100644
index 0000000000000..6f87a0d1ddb2e
--- /dev/null
+++ b/pkg/tracing/plugin/sampler.go
@@ -0,0 +1,51 @@
+package plugin
+
+import (
+	"fmt"
+
+	sdkTrace "go.opentelemetry.io/otel/sdk/trace"
+	"go.opentelemetry.io/otel/trace"
+)
+
+const (
+	samplerNameBased       = "namebased"
+	samplerParentBasedName = "parentbased_name"
+)
+
+type NameSampler struct {
+	// allow is a set of names that should be sampled.
+	// Uses a map of empty structs for O(1) lookups and no memory overhead.
+	allow map[string]struct{}
+}
+
+// NameBased returns a Sampler that samples every span having a certain name.
+// It should be used in conjunction with the ParentBased sampler so that the child spans are also sampled.
+func NameBased(allowedNames []string) NameSampler {
+	allowedNamesMap := make(map[string]struct{}, len(allowedNames))
+	for _, name := range allowedNames {
+		allowedNamesMap[name] = struct{}{}
+	}
+	return NameSampler{
+		allow: allowedNamesMap,
+	}
+}
+
+func (ns NameSampler) ShouldSample(parameters sdkTrace.SamplingParameters) sdkTrace.SamplingResult {
+	psc := trace.SpanContextFromContext(parameters.ParentContext)
+
+	if _, ok := ns.allow[parameters.Name]; ok {
+		return sdkTrace.SamplingResult{
+			Decision:   sdkTrace.RecordAndSample,
+			Tracestate: psc.TraceState(),
+		}
+	}
+
+	return sdkTrace.SamplingResult{
+		Decision:   sdkTrace.Drop,
+		Tracestate: psc.TraceState(),
+	}
+}
+
+func (ns NameSampler) Description() string {
+	return fmt.Sprintf("NameBased:{%v}", ns.allow)
+}

From a65109a854079ad3031d80247f8ebf651fcbc553 Mon Sep 17 00:00:00 2001
From: Baptiste Girard-Carrabin <baptiste.girardcarrabin@datadoghq.com>
Date: Wed, 17 Jul 2024 14:55:16 +0200
Subject: [PATCH 2/2] [tracing] Configure otlp plugin to use NameSampler

Allow the plugin to use the custom NameSampler if specified in the usual env vars used to specify the sampler `OTEL_TRACES_SAMPLER` and `OTEL_TRACES_SAMPLER_ARG`. This sampler will allow to avoid emitting a lot of the noisy traces to only keep a subset of meaningful traces.
---
 pkg/tracing/plugin/otlp.go | 35 +++++++++++++++++++++++++++++++++--
 1 file changed, 33 insertions(+), 2 deletions(-)

diff --git a/pkg/tracing/plugin/otlp.go b/pkg/tracing/plugin/otlp.go
index 52586913eae3c..bf8f574966925 100644
--- a/pkg/tracing/plugin/otlp.go
+++ b/pkg/tracing/plugin/otlp.go
@@ -22,6 +22,7 @@ import (
 	"io"
 	"os"
 	"strconv"
+	"strings"
 	"time"
 
 	"github.com/containerd/containerd/v2/pkg/deprecation"
@@ -52,8 +53,10 @@ const (
 	otlpProtocolEnv       = "OTEL_EXPORTER_OTLP_PROTOCOL"
 	otlpTracesProtocolEnv = "OTEL_EXPORTER_OTLP_TRACES_PROTOCOL"
 
-	otelTracesExporterEnv = "OTEL_TRACES_EXPORTER"
-	otelServiceNameEnv    = "OTEL_SERVICE_NAME"
+	otelTracesExporterEnv   = "OTEL_TRACES_EXPORTER"
+	otelServiceNameEnv      = "OTEL_SERVICE_NAME"
+	otelTracesSamplerEnv    = "OTEL_TRACES_SAMPLER"
+	otelTracesSamplerArgEnv = "OTEL_TRACES_SAMPLER_ARG"
 )
 
 func init() {
@@ -193,6 +196,15 @@ func newTracer(ctx context.Context, procs []trace.SpanProcessor) (io.Closer, err
 	for _, proc := range procs {
 		opts = append(opts, trace.WithSpanProcessor(proc))
 	}
+
+	// Configure custom NameBased sampler if specified in the env
+	if nameSampler := nameSamplerFromEnv(); nameSampler != nil {
+		opts = append(opts, trace.WithSampler(nameSampler))
+		// Unset the env vars so that otel sdk does not attempt to configure the sampler automatically
+		os.Unsetenv(otelTracesSamplerEnv)
+		os.Unsetenv(otelTracesSamplerArgEnv)
+	}
+
 	provider := trace.NewTracerProvider(opts...)
 	otel.SetTracerProvider(provider)
 
@@ -202,6 +214,25 @@ func newTracer(ctx context.Context, procs []trace.SpanProcessor) (io.Closer, err
 
 }
 
+func nameSamplerFromEnv() trace.Sampler {
+	sampler, ok := os.LookupEnv(otelTracesSamplerEnv)
+	if !ok {
+		return nil
+	}
+
+	sampler = strings.ToLower(strings.TrimSpace(sampler))
+	allowedNames := strings.Split(strings.TrimSpace(os.Getenv(otelTracesSamplerArgEnv)), ",")
+
+	switch sampler {
+	case samplerNameBased:
+		return NameBased(allowedNames)
+	case samplerParentBasedName:
+		return trace.ParentBased(NameBased(allowedNames))
+	default:
+		return nil
+	}
+}
+
 func warnTraceConfig(ic *plugin.InitContext) error {
 	ctx := ic.Context
 	cfg := ic.Config.(*TraceConfig)