From 8f442c95db09bc80c03065f02e2d37a14a527610 Mon Sep 17 00:00:00 2001
From: Baptiste Girard-Carrabin <baptiste.girardcarrabin@datadoghq.com>
Date: Fri, 14 Jun 2024 11:22:04 +0200
Subject: [PATCH] [tracing] POC for a simple attribute based sampler

Can be used to filter out traces based on any attribute present on the span
---
 pkg/tracing/plugin/otlp.go    |  6 +++++
 pkg/tracing/plugin/sampler.go | 46 +++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+)
 create mode 100644 pkg/tracing/plugin/sampler.go

diff --git a/pkg/tracing/plugin/otlp.go b/pkg/tracing/plugin/otlp.go
index 52586913eae3c..c21cc57360314 100644
--- a/pkg/tracing/plugin/otlp.go
+++ b/pkg/tracing/plugin/otlp.go
@@ -38,6 +38,7 @@ import (
 	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp"
 	"go.opentelemetry.io/otel/propagation"
 	"go.opentelemetry.io/otel/sdk/trace"
+	semconv "go.opentelemetry.io/otel/semconv/v1.21.0"
 )
 
 const exporterPlugin = "otlp"
@@ -193,6 +194,11 @@ func newTracer(ctx context.Context, procs []trace.SpanProcessor) (io.Closer, err
 	for _, proc := range procs {
 		opts = append(opts, trace.WithSpanProcessor(proc))
 	}
+
+	attributes := map[string][]string{
+		string(semconv.RPCMethodKey): {"PullImage", "Create"},
+	}
+	opts = append(opts, trace.WithSampler(trace.ParentBased(AttributeBased(attributes))))
 	provider := trace.NewTracerProvider(opts...)
 	otel.SetTracerProvider(provider)
 
diff --git a/pkg/tracing/plugin/sampler.go b/pkg/tracing/plugin/sampler.go
new file mode 100644
index 0000000000000..6135c9fb1480d
--- /dev/null
+++ b/pkg/tracing/plugin/sampler.go
@@ -0,0 +1,46 @@
+package plugin
+
+import (
+	"fmt"
+	"slices"
+
+	sdkTrace "go.opentelemetry.io/otel/sdk/trace"
+	"go.opentelemetry.io/otel/trace"
+)
+
+type attributeSampler struct {
+	selectedAttrs map[string][]string
+}
+
+// AttributeBased returns a Sampler that samples every trace matching a set of attributes.
+// It only need to find one matching key-value pairs present in the attributes of the span.
+func AttributeBased(selectedAttrs map[string][]string) attributeSampler {
+	return attributeSampler{
+		selectedAttrs: selectedAttrs,
+	}
+}
+
+func (ss attributeSampler) ShouldSample(parameters sdkTrace.SamplingParameters) sdkTrace.SamplingResult {
+	psc := trace.SpanContextFromContext(parameters.ParentContext)
+	println(parameters.Name)
+	for _, attr := range parameters.Attributes {
+		println(attr.Key, attr.Value.Emit())
+		if selectedValues, ok := ss.selectedAttrs[string(attr.Key)]; ok {
+			if slices.Contains(selectedValues, attr.Value.Emit()) {
+				println("Sampled")
+				return sdkTrace.SamplingResult{
+					Decision:   sdkTrace.RecordAndSample,
+					Tracestate: psc.TraceState(),
+				}
+			}
+		}
+	}
+	return sdkTrace.SamplingResult{
+		Decision:   sdkTrace.Drop,
+		Tracestate: psc.TraceState(),
+	}
+}
+
+func (ss attributeSampler) Description() string {
+	return fmt.Sprintf("AttributeBased:{%v}", ss.selectedAttrs)
+}