feat(sampling): datadog sampler

datadog-opentelemetry/Cargo.toml

@@ -11,6 +11,7 @@ description.workspace = true
 [dependencies]
 # Workspace dependencies
 dd-trace = { path = "../dd-trace" }
+dd-trace-sampling = { path = "../dd-trace-sampling" }
 dd-trace-propagation = { path = "../dd-trace-propagation", features = ["opentelemetry"] }
 
 # External dependencies
@@ -31,6 +32,7 @@ assert_unordered = "0.3"
 
 # Libdatadog dependencies - change to a stable version once we release
 datadog-trace-utils = { workspace = true, features = ["test-utils"] }
+dd-trace-sampling = { path = "../dd-trace-sampling" }
 
 # Depend on ourselves to have APIs exposed only during tests
 datadog-opentelemetry = { path = ".", features = ["test-utils"] }

datadog-opentelemetry/src/lib.rs

@@ -2,15 +2,17 @@
 // SPDX-License-Identifier: Apache-2.0
 
 mod ddtrace_transform;
+mod sampler;
 mod span_exporter;
 mod span_processor;
 mod text_map_propagator;
 mod trace_id;
 mod transform;
 
-use std::sync::Arc;
+use std::sync::{Arc, RwLock};
 
-use opentelemetry_sdk::trace::SdkTracerProvider;
+use opentelemetry_sdk::{trace::SdkTracerProvider, Resource};
+use sampler::Sampler;
 use span_processor::{DatadogSpanProcessor, TraceRegistry};
 use text_map_propagator::DatadogPropagator;
 
@@ -41,41 +43,49 @@ pub fn init_datadog(
     // all parameters and has an install method?
     tracer_provider_builder: opentelemetry_sdk::trace::TracerProviderBuilder,
 ) -> SdkTracerProvider {
+    let (tracer_provider, propagator) = make_tracer(config, tracer_provider_builder, None);
+
+    opentelemetry::global::set_text_map_propagator(propagator);
+    opentelemetry::global::set_tracer_provider(tracer_provider.clone());
+    tracer_provider
+}
+
+/// Create an instance of the tracer provider
+fn make_tracer(
+    config: dd_trace::Config,
+    mut tracer_provider_builder: opentelemetry_sdk::trace::TracerProviderBuilder,
+    resource: Option<Resource>,
+) -> (SdkTracerProvider, DatadogPropagator) {
     let registry = Arc::new(TraceRegistry::new());
+    let resource_slot = Arc::new(RwLock::new(Resource::builder_empty().build()));
+    let sampler = Sampler::new(&config, resource_slot.clone(), registry.clone());
+
+    if let Some(resource) = resource {
+        tracer_provider_builder = tracer_provider_builder.with_resource(resource)
+    }
 
     let propagator = DatadogPropagator::new(&config, registry.clone());
-    opentelemetry::global::set_text_map_propagator(propagator);
 
+    let span_processor = DatadogSpanProcessor::new(config, registry.clone(), resource_slot.clone());
     let tracer_provider = tracer_provider_builder
-        .with_span_processor(DatadogSpanProcessor::new(config, registry))
+        .with_span_processor(span_processor)
+        .with_sampler(sampler) // Use the sampler created above
         .with_id_generator(trace_id::TraceidGenerator)
-        // TODO: hookup additional components
-        // .with_sampler(sampler)
         .build();
-    opentelemetry::global::set_tracer_provider(tracer_provider.clone());
-    tracer_provider
+
+    (tracer_provider, propagator)
 }
 
 #[cfg(feature = "test-utils")]
-/// Create a local instance of the tracer provider
-pub fn make_tracer(
+pub fn make_test_tracer(
     config: dd_trace::Config,
     tracer_provider_builder: opentelemetry_sdk::trace::TracerProviderBuilder,
-) -> SdkTracerProvider {
-    use opentelemetry::KeyValue;
-    use opentelemetry_sdk::Resource;
-
-    let registry = Arc::new(TraceRegistry::new());
-
-    tracer_provider_builder
-        .with_resource(
-            Resource::builder()
-                .with_attribute(KeyValue::new("service.name", config.service().to_string()))
-                .build(),
-        )
-        .with_span_processor(DatadogSpanProcessor::new(config, registry))
-        .with_id_generator(trace_id::TraceidGenerator)
-        // TODO: hookup additional components
-        // .with_sampler(sampler)
-        .build()
+) -> (SdkTracerProvider, DatadogPropagator) {
+    let resource = Resource::builder()
+        .with_attribute(opentelemetry::KeyValue::new(
+            "service.name",
+            config.service().to_string(),
+        ))
+        .build();
+    make_tracer(config, tracer_provider_builder, Some(resource))
 }

datadog-opentelemetry/src/sampler.rs

@@ -0,0 +1,229 @@
+// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/
+// SPDX-License-Identifier: Apache-2.0
+
+use dd_trace::Config;
+use dd_trace_sampling::DatadogSampler;
+use opentelemetry::trace::TraceContextExt;
+use opentelemetry_sdk::{trace::ShouldSample, Resource};
+use std::{
+    collections::HashMap,
+    sync::{Arc, RwLock},
+};
+
+use crate::{
+    span_processor::{RegisterTracePropagationResult, SamplingDecision},
+    TraceRegistry,
+};
+
+#[derive(Debug, Clone)]
+pub struct Sampler {
+    sampler: DatadogSampler,
+    trace_registry: Arc<TraceRegistry>,
+}
+
+impl Sampler {
+    pub fn new(
+        cfg: &Config,
+        resource: Arc<RwLock<Resource>>,
+        trace_registry: Arc<TraceRegistry>,
+    ) -> Self {
+        let rules = cfg
+            .trace_sampling_rules()
+            .iter()
+            .map(|r| {
+                dd_trace_sampling::SamplingRule::new(
+                    r.sample_rate,
+                    r.service.clone(),
+                    r.name.clone(),
+                    r.resource.clone(),
+                    Some(r.tags.clone()),
+                    Some(r.provenance.clone()),
+                )
+            })
+            .collect::<Vec<_>>();
+        let sampler =
+            dd_trace_sampling::DatadogSampler::new(rules, cfg.trace_rate_limit(), resource);
+        Self {
+            sampler,
+            trace_registry,
+        }
+    }
+}
+
+impl ShouldSample for Sampler {
+    fn should_sample(
+        &self,
+        parent_context: Option<&opentelemetry::Context>,
+        trace_id: opentelemetry::trace::TraceId,
+        name: &str,
+        span_kind: &opentelemetry::trace::SpanKind,
+        attributes: &[opentelemetry::KeyValue],
+        _links: &[opentelemetry::trace::Link],
+    ) -> opentelemetry::trace::SamplingResult {
+        let result = self.sampler.sample(
+            parent_context
+                .filter(|c| c.has_active_span())
+                .map(|c| c.span().span_context().is_sampled()),
+            trace_id,
+            name,
+            span_kind,
+            attributes,
+        );
+        if let Some(trace_root_info) = &result.trace_root_info {
+            match self.trace_registry.register_trace_propagation_data(
+                trace_id.to_bytes(),
+                SamplingDecision {
+                    decision: trace_root_info.sampling_priority(result.is_sampled).value(),
+                    // TODO: unify these types with decision maker with the one in the span
+                    // processor
+                    decision_maker: trace_root_info.mechanism.value() as i8,
+                },
+                None,
+                // TODO(paullgdc): This is here so the injector adds the t.dm tag to
+                // tracecontext. The injector should probably inject it from
+                // the trace propagation data instead of tags.
+                Some(HashMap::from_iter([(
+                    "_dd.p.dm".to_string(),
+                    format!("{}", -(trace_root_info.mechanism.value() as i32)),
+                )])),
+            ) {
+                RegisterTracePropagationResult::Existing(sampling_decision) => {
+                    return opentelemetry::trace::SamplingResult {
+                        decision: if sampling_decision.decision > 0 {
+                            opentelemetry::trace::SamplingDecision::RecordAndSample
+                        } else {
+                            opentelemetry::trace::SamplingDecision::RecordOnly
+                        },
+                        attributes: Vec::new(),
+                        trace_state: parent_context
+                            .map(|c| c.span().span_context().trace_state().clone())
+                            .unwrap_or_default(),
+                    }
+                }
+                RegisterTracePropagationResult::New => {}
+            }
+        }
+
+        opentelemetry::trace::SamplingResult {
+            decision: result.to_otel_decision(),
+            attributes: result.to_dd_sampling_tags(),
+            trace_state: parent_context
+                .map(|c| c.span().span_context().trace_state().clone())
+                .unwrap_or_default(),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use dd_trace::configuration::SamplingRuleConfig;
+    use opentelemetry::{
+        trace::{SamplingDecision, SpanContext, SpanKind, TraceId, TraceState},
+        Context, SpanId, TraceFlags,
+    };
+    use opentelemetry_sdk::trace::ShouldSample;
+    use std::env;
+
+    #[test]
+    fn test_create_sampler_with_sampling_rules() {
+        // Build a fresh config to pick up the env var
+        let mut config = Config::builder();
+        config.set_trace_sampling_rules(vec![SamplingRuleConfig {
+            sample_rate: 0.5,
+            service: Some("test-service".to_string()),
+            name: None,
+            resource: None,
+            tags: HashMap::new(),
+            provenance: "customer".to_string(),
+        }]);
+        let config = config.build();
+
+        let test_resource = Arc::new(RwLock::new(Resource::builder().build()));
+        let sampler = Sampler::new(&config, test_resource, Arc::new(TraceRegistry::new()));
+
+        let trace_id_bytes = [1; 16];
+        let trace_id = TraceId::from_bytes(trace_id_bytes);
+
+        // Basic assertion: Check if the attributes added by the sampler are not empty,
+        // implying some sampling logic (like adding priority tags) ran.
+        assert!(
+            !sampler
+                .should_sample(None, trace_id, "test", &SpanKind::Client, &[], &[])
+                .attributes
+                .is_empty(),
+            "Sampler should add attributes even if decision is complex"
+        );
+
+        // Clean up environment
+        env::remove_var("DD_TRACE_SAMPLING_RULES");
+    }
+
+    #[test]
+    fn test_create_default_sampler() {
+        // Create a default config (no rules, no specific rate limit)
+        let config = Config::builder().build();
+
+        let test_resource = Arc::new(RwLock::new(Resource::builder_empty().build()));
+        let sampler = Sampler::new(&config, test_resource, Arc::new(TraceRegistry::new()));
+
+        let trace_id_bytes = [2; 16];
+        let trace_id = TraceId::from_bytes(trace_id_bytes);
+
+        // Verify the default sampler behavior
+        let result = sampler.should_sample(None, trace_id, "test", &SpanKind::Client, &[], &[]);
+        assert_eq!(
+            result.decision,
+            SamplingDecision::RecordAndSample,
+            "Default sampler should record and sample by default"
+        );
+    }
+
+    #[test]
+    fn test_trace_state_propagation() {
+        let config = Config::builder().build();
+
+        let test_resource = Arc::new(RwLock::new(Resource::builder_empty().build()));
+        let sampler = Sampler::new(&config, test_resource, Arc::new(TraceRegistry::new()));
+
+        let trace_id = TraceId::from_bytes([2; 16]);
+        let span_id = SpanId::from_bytes([3; 8]);
+
+        for is_sampled in [true, false] {
+            let trace_state = TraceState::from_key_value([("test_key", "test_value")]).unwrap();
+            let span_context = SpanContext::new(
+                trace_id,
+                span_id,
+                is_sampled
+                    .then_some(TraceFlags::SAMPLED)
+                    .unwrap_or_default(),
+                true,
+                trace_state.clone(),
+            );
+
+            // Verify the sampler with a parent context
+            let result = sampler.should_sample(
+                Some(&Context::new().with_remote_span_context(span_context)),
+                trace_id,
+                "test",
+                &SpanKind::Client,
+                &[],
+                &[],
+            );
+            assert_eq!(
+                result.decision,
+                if is_sampled {
+                    SamplingDecision::RecordAndSample
+                } else {
+                    SamplingDecision::RecordOnly
+                },
+                "Sampler should respect parent context sampling decision"
+            );
+            assert_eq!(
+                result.trace_state.header(),
+                "test_key=test_value",
+                "Sampler should propagate trace state from parent context"
+            );
+        }
+    }
+}