[AGNTLOG-468] chore(datadog): add trace sampling

bin/agent-data-plane/src/cli/run.rs

@@ -25,6 +25,7 @@ use saluki_components::{
     transforms::{
         AggregateConfiguration, ApmStatsTransformConfiguration, ChainedConfiguration, DogstatsDMapperConfiguration,
         DogstatsDPrefixFilterConfiguration, HostEnrichmentConfiguration, HostTagsConfiguration,
+        TraceSamplerConfiguration,
     },
 };
 use saluki_config::{ConfigurationLoader, GenericConfiguration};
@@ -327,6 +328,7 @@ async fn add_baseline_traces_pipeline_to_blueprint(
         .await?;
     let dd_traces_enrich_config =
         ChainedConfiguration::default().with_transform_builder("apm_onboarding", ApmOnboardingConfiguration);
+    let trace_sampler_config = TraceSamplerConfiguration::default();
     let apm_stats_transform_config = ApmStatsTransformConfiguration::from_configuration(config)
         .error_context("Failed to configure APM Stats transform.")?
         .with_environment_provider(env_provider.clone())
@@ -338,11 +340,13 @@ async fn add_baseline_traces_pipeline_to_blueprint(
 
     blueprint
         .add_transform("traces_enrich", dd_traces_enrich_config)?
+        .add_transform("trace_sampler", trace_sampler_config)?
         .add_transform("dd_apm_stats", apm_stats_transform_config)?
         .add_encoder("dd_stats_encode", dd_apm_stats_encoder)?
         .add_encoder("dd_traces_encode", dd_traces_config)?
-        .connect_component("dd_apm_stats", ["traces_enrich"])?
-        .connect_component("dd_traces_encode", ["traces_enrich"])?
+        .connect_component("trace_sampler", ["traces_enrich"])?
+        .connect_component("dd_apm_stats", ["trace_sampler"])?
+        .connect_component("dd_traces_encode", ["trace_sampler"])?
         .connect_component("dd_stats_encode", ["dd_apm_stats"])?
         .connect_component("dd_out", ["dd_traces_encode", "dd_stats_encode"])?;
 

lib/saluki-components/src/common/otlp/traces/translator.rs

@@ -3,7 +3,7 @@ use otlp_protos::opentelemetry::proto::resource::v1::Resource as OtlpResource;
 use otlp_protos::opentelemetry::proto::trace::v1::ResourceSpans;
 use saluki_common::collections::FastHashMap;
 use saluki_context::tags::TagSet;
-use saluki_core::data_model::event::trace::{Span as DdSpan, Trace};
+use saluki_core::data_model::event::trace::{Span as DdSpan, Trace, TraceSampling};
 use saluki_core::data_model::event::Event;
 
 use crate::common::otlp::config::TracesConfig;
@@ -52,6 +52,7 @@ impl OtlpTracesTranslator {
         let resource: OtlpResource = resource_spans.resource.unwrap_or_default();
         let resource_tags: TagSet = resource_attributes_to_tagset(&resource.attributes);
         let mut traces_by_id: FastHashMap<u64, Vec<DdSpan>> = FastHashMap::default();
+        let mut priorities_by_id: FastHashMap<u64, i32> = FastHashMap::default();
         let ignore_missing_fields = self.config.ignore_missing_datadog_fields;
 
         for scope_spans in resource_spans.scope_spans {
@@ -67,17 +68,30 @@ impl OtlpTracesTranslator {
                     ignore_missing_fields,
                     self.config.enable_otlp_compute_top_level_by_span_kind,
                 );
+
+                // Track last-seen priority for this trace (overwrites previous values)
+                if let Some(&priority) = dd_span.metrics().get("_sampling_priority_v1") {
+                    priorities_by_id.insert(trace_id, priority as i32);
+                }
+
                 traces_by_id.entry(trace_id).or_default().push(dd_span);
             }
         }
 
         traces_by_id
             .into_iter()
-            .filter_map(|(_, spans)| {
+            .filter_map(|(trace_id, spans)| {
                 if spans.is_empty() {
                     None
                 } else {
-                    Some(Event::Trace(Trace::new(spans, resource_tags.clone())))
+                    let mut trace = Trace::new(spans, resource_tags.clone());
+
+                    // Set the trace-level sampling priority if one was found
+                    if let Some(&priority) = priorities_by_id.get(&trace_id) {
+                        trace.set_sampling(Some(TraceSampling::new(false, Some(priority), None, None)));
+                    }
+
+                    Some(Event::Trace(trace))
                 }
             })
             .collect()

lib/saluki-components/src/encoders/datadog/traces/mod.rs

@@ -62,6 +62,12 @@ const CONTAINER_TAGS_META_KEY: &str = "_dd.tags.container";
 const MAX_TRACES_PER_PAYLOAD: usize = 10000;
 static CONTENT_TYPE_PROTOBUF: HeaderValue = HeaderValue::from_static("application/x-protobuf");
 
+// Sampling metadata keys / values.
+const TAG_OTLP_SAMPLING_RATE: &str = "_dd.otlp_sr";
+const TAG_DECISION_MAKER: &str = "_dd.p.dm";
+const SAMPLING_PRIORITY_METRIC_KEY: &str = "_sampling_priority_v1";
+const DEFAULT_CHUNK_PRIORITY: i32 = 1; // PRIORITY_AUTO_KEEP
+
 fn default_serializer_compressor_kind() -> String {
     "zstd".to_string()
 }
@@ -547,28 +553,40 @@ impl TraceEndpointEncoder {
     }
 
     fn build_trace_chunk(&self, trace: &Trace) -> TraceChunk {
-        let mut spans: Vec<ProtoSpan> = trace.spans().iter().map(convert_span).collect();
+        let spans: Vec<ProtoSpan> = trace.spans().iter().map(convert_span).collect();
         let mut chunk = TraceChunk::new();
 
-        let rate = self.sampling_rate();
         let mut tags = std::collections::HashMap::new();
-        tags.insert("_dd.otlp_sr".to_string(), format!("{:.2}", rate));
-
-        // TODO: Remove this once we have sampling. We have to hardcode the priority to 1 for now so that intake does not drop the trace.
-        const PRIORITY_AUTO_KEEP: i32 = 1;
-        chunk.set_priority(PRIORITY_AUTO_KEEP);
-
-        // Set _dd.p.dm (decision maker)
-        // Only set if sampling priority is "keep" (which it is, since we set PRIORITY_AUTO_KEEP)
-        // Decision maker "-9" indicates probabilistic sampler made the decision
-        const DECISION_MAKER: &str = "-9";
-        if let Some(first_span) = spans.first_mut() {
-            let mut meta = first_span.take_meta();
-            meta.insert("_dd.p.dm".to_string(), DECISION_MAKER.to_string());
-            first_span.set_meta(meta);
+
+        // Use trace-level sampling metadata if available (set by the trace sampler transform).
+        // This provides explicit trace-level sampling information without needing to scan spans.
+        if let Some(sampling) = trace.sampling() {
+            // Set priority from trace metadata
+            chunk.set_priority(sampling.priority.unwrap_or(DEFAULT_CHUNK_PRIORITY));
+            chunk.set_droppedTrace(sampling.dropped_trace);
+
+            // Set decision maker tag if present
+            if let Some(dm) = &sampling.decision_maker {
+                tags.insert(TAG_DECISION_MAKER.to_string(), dm.to_string());
+            }
+
+            // Set OTLP sampling rate tag if present (from sampler)
+            if let Some(otlp_sr) = &sampling.otlp_sampling_rate {
+                tags.insert(TAG_OTLP_SAMPLING_RATE.to_string(), otlp_sr.to_string());
+            } else {
+                // Fallback to encoder's computed rate
+                let rate = self.sampling_rate();
+                tags.insert(TAG_OTLP_SAMPLING_RATE.to_string(), format!("{:.2}", rate));
+            }
+        } else {
+            // Fallback: if trace.sampling is None, use defaults
+            // (No span scanning per the plan's "no fallback scan" requirement)
+            chunk.set_priority(DEFAULT_CHUNK_PRIORITY);
+            chunk.set_droppedTrace(false);
+            let rate = self.sampling_rate();
+            tags.insert(TAG_OTLP_SAMPLING_RATE.to_string(), format!("{:.2}", rate));
         }
 
-        tags.insert("_dd.p.dm".to_string(), DECISION_MAKER.to_string());
         chunk.set_tags(tags);
 
         chunk.set_spans(spans);

lib/saluki-components/src/transforms/mod.rs

@@ -21,5 +21,8 @@ pub use self::dogstatsd_mapper::DogstatsDMapperConfiguration;
 mod metric_router;
 pub use self::metric_router::MetricRouterConfiguration;
 
+mod trace_sampler;
+pub use self::trace_sampler::TraceSamplerConfiguration;
+
 mod apm_stats;
 pub use self::apm_stats::ApmStatsTransformConfiguration;