Port name sampler to containerd v2

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)

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