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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged

tobz merged 59 commits into main from andrewq/add-trace-sampling

Jan 29, 2026

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

-Original file line number
+Diff line change
@@ Expand Up / @@ -25,7 +25,7 @@ use saluki_components::{ @@
         transforms::{
             AggregateConfiguration, ApmStatsTransformConfiguration, ChainedConfiguration, DogstatsDMapperConfiguration,
             DogstatsDPrefixFilterConfiguration, HostEnrichmentConfiguration, HostTagsConfiguration,
-            TraceObfuscationConfiguration,
+            TraceObfuscationConfiguration, TraceSamplerConfiguration,
         },
     };
     use saluki_config::{ConfigurationLoader, GenericConfiguration};
@@ Expand Down Expand Up / @@ -330,6 +330,8 @@ async fn add_baseline_traces_pipeline_to_blueprint( @@
         let dd_traces_enrich_config = ChainedConfiguration::default()
             .with_transform_builder("apm_onboarding", ApmOnboardingConfiguration)
             .with_transform_builder("trace_obfuscation", trace_obfuscation_config);
+        let trace_sampler_config = TraceSamplerConfiguration::from_configuration(config)
+            .error_context("Failed to configure Trace Sampler transform.")?;
         let apm_stats_transform_config = ApmStatsTransformConfiguration::from_configuration(config)
             .error_context("Failed to configure APM Stats transform.")?
             .with_environment_provider(env_provider.clone())
@@ Expand All / @@ -341,11 +343,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"])?;
@@ Expand Down @@

lib/saluki-components/src/common/datadog/apm.rs

-Original file line number
+Diff line change
@@ Expand Up / @@ -12,6 +12,21 @@ const fn default_target_traces_per_second() -> f64 { @@
     const fn default_errors_per_second() -> f64 {
 .0
     }
+    const fn default_sampling_percentage() -> f64 {
+.0
+    }
+    const fn default_error_sampling_enabled() -> bool {
+        true
+    }
+    const fn default_error_tracking_standalone_enabled() -> bool {
+        false
+    }
+    const fn default_probabilistic_sampling_enabled() -> bool {
+        false
+    }
     const fn default_peer_tags_aggregation() -> bool {
         true
     }
@@ Expand All / @@ -33,6 +48,56 @@ struct ApmConfiguration { @@
         apm_config: ApmConfig,
     }
+    #[derive(Clone, Debug, Deserialize)]
+    struct ProbabilisticSamplerConfig {
+        /// Enables probabilistic sampling.
+        ///
+        /// When enabled, the trace sampler keeps approximately `sampling_percentage` of traces using a
+        /// deterministic hash of the trace ID.
+        ///
+        /// Defaults to `false`.
+        #[serde(default = "default_probabilistic_sampling_enabled")]
+        enabled: bool,
+        /// Sampling percentage (0-100).
+        ///
+        /// Determines the percentage of traces to keep. A value of 100 keeps all traces,
+        /// while 50 keeps approximately half. Values outside 0-100 are treated as 100.
+        ///
+        /// Defaults to 100.0 (keep all traces).
+        #[serde(default = "default_sampling_percentage")]
+        sampling_percentage: f64,
+    }
+    impl Default for ProbabilisticSamplerConfig {
+        fn default() -> Self {
+            Self {
+                enabled: default_probabilistic_sampling_enabled(),
+                sampling_percentage: default_sampling_percentage(),
+            }
+        }
+    }
+    #[derive(Clone, Debug, Deserialize)]
+    struct ErrorTrackingStandaloneConfig {
+        /// Enables Error Tracking Standalone mode.
+        ///
+        /// When enabled, error tracking standalone mode suppresses single-span sampling and analytics
+        /// events for dropped traces.
+        ///
+        /// Defaults to `false`.
+        #[serde(default = "default_error_tracking_standalone_enabled")]
+        enabled: bool,
+    }
+    impl Default for ErrorTrackingStandaloneConfig {
+        fn default() -> Self {
+            Self {
+                enabled: default_error_tracking_standalone_enabled(),
+            }
+        }
+    }
     #[derive(Clone, Debug, Deserialize)]
     pub struct ApmConfig {
         /// Target traces per second for priority sampling.
@@ Expand All / @@ -47,6 +112,27 @@ pub struct ApmConfig { @@
         #[serde(default = "default_errors_per_second")]
         errors_per_second: f64,
+        /// Probabilistic sampler configuration.
+        ///
+        /// Defaults to enabled with `sampling_percentage` set to 100.0 (keep all traces).
+        #[serde(default)]
+        probabilistic_sampler: ProbabilisticSamplerConfig,
+        /// Enable error sampling in the trace sampler.
+        ///
+        /// When enabled, traces containing errors will be kept even if they would be dropped by
+        /// probabilistic sampling. This ensures error visibility at low sampling rates.
+        ///
+        /// Defaults to `true`.
+        #[serde(default = "default_error_sampling_enabled")]
+        error_sampling_enabled: bool,
+        /// Error Tracking Standalone configuration.
+        ///
+        /// Defaults to disabled.
+        #[serde(default)]
+        error_tracking_standalone: ErrorTrackingStandaloneConfig,
         /// Enables an additional stats computation check on spans to see if they have an eligible `span.kind` (server, consumer, client, producer).
         /// If enabled, a span with an eligible `span.kind` will have stats computed. If disabled, only top-level and measured spans will have stats computed.
         ///
@@ Expand Down Expand Up / @@ -100,6 +186,26 @@ impl ApmConfig { @@
             self.errors_per_second
         }
+        /// Returns if probabilistic sampling is enabled.
+        pub const fn probabilistic_sampler_enabled(&self) -> bool {
+            self.probabilistic_sampler.enabled
+        }
+        /// Returns the probabilistic sampler sampling percentage.
+        pub const fn probabilistic_sampler_sampling_percentage(&self) -> f64 {
+            self.probabilistic_sampler.sampling_percentage
+        }
+        /// Returns if error sampling is enabled.
+        pub const fn error_sampling_enabled(&self) -> bool {
+            self.error_sampling_enabled
+        }
+        /// Returns if error tracking standalone mode is enabled.
+        pub const fn error_tracking_standalone_enabled(&self) -> bool {
+            self.error_tracking_standalone.enabled
+        }
         /// Returns if stats computation by span kind is enabled.
         pub const fn compute_stats_by_span_kind(&self) -> bool {
             self.compute_stats_by_span_kind
@@ Expand Down Expand Up / @@ -143,6 +249,9 @@ impl Default for ApmConfig { @@
             Self {
                 target_traces_per_second: default_target_traces_per_second(),
                 errors_per_second: default_errors_per_second(),
+                probabilistic_sampler: ProbabilisticSamplerConfig::default(),
+                error_sampling_enabled: default_error_sampling_enabled(),
+                error_tracking_standalone: ErrorTrackingStandaloneConfig::default(),
                 compute_stats_by_span_kind: default_compute_stats_by_span_kind(),
                 peer_tags_aggregation: default_peer_tags_aggregation(),
                 peer_tags: Vec::new(),
@@ Expand Down @@

lib/saluki-components/src/common/datadog/mod.rs

-Original file line number
+Diff line change
@@ Expand Up / @@ -10,8 +10,40 @@ mod retry; @@
     pub mod telemetry;
     pub mod transaction;
+    /// Metric key used to store Datadog sampling priority (`_sampling_priority_v1`).
+    pub const SAMPLING_PRIORITY_METRIC_KEY: &str = "_sampling_priority_v1";
     /// Default compressed size limit for intake requests.
     pub const DEFAULT_INTAKE_COMPRESSED_SIZE_LIMIT: usize = 3_200_000; // 3 MiB
     /// Default uncompressed size limit for intake requests.
     pub const DEFAULT_INTAKE_UNCOMPRESSED_SIZE_LIMIT: usize = 62_914_560; // 60 MiB
+    /// Metadata tag used to store the sampling decision maker (`_dd.p.dm`).
+    pub const TAG_DECISION_MAKER: &str = "_dd.p.dm";
+    /// Decision maker value for probabilistic sampling (matches Datadog Agent).
+    pub const DECISION_MAKER_PROBABILISTIC: &str = "-9";
+    /// Metadata key used to store the OTEL trace id.
+    pub const OTEL_TRACE_ID_META_KEY: &str = "otel.trace_id";
+    /// Maximum trace id used for deterministic sampling.
+    pub const MAX_TRACE_ID: u64 = u64::MAX;
+    /// Precomputed float form of `MAX_TRACE_ID`.
+    pub const MAX_TRACE_ID_FLOAT: f64 = MAX_TRACE_ID as f64;
+    /// Hasher used for deterministic sampling.
+    pub const SAMPLER_HASHER: u64 = 1111111111111111111;
+    /// Returns whether to keep a trace, based on its ID and a sampling rate.
+    ///
+    /// This assumes trace IDs are nearly uniformly distributed.
+    pub fn sample_by_rate(trace_id: u64, rate: f64) -> bool {
+        if rate < 1.0 {
+            trace_id.wrapping_mul(SAMPLER_HASHER) < (rate * MAX_TRACE_ID_FLOAT) as u64
+        } else {
+            true
+        }
+    }

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

-Original file line number
+Diff line change
@@ Expand Up / @@ -20,6 +20,7 @@ use serde_json::{Map as JsonMap, Value as JsonValue}; @@
     use stringtheory::MetaString;
     use tracing::error;
+    use crate::common::datadog::{OTEL_TRACE_ID_META_KEY, SAMPLING_PRIORITY_METRIC_KEY};
     use crate::common::otlp::attributes::{get_int_attribute, HTTP_MAPPINGS};
     use crate::common::otlp::traces::normalize::{normalize_service, normalize_tag_value};
     use crate::common::otlp::traces::normalize::{truncate_utf8, MAX_RESOURCE_LEN};
@@ Expand All / @@ -29,7 +30,6 @@ use crate::common::otlp::util::{ @@
         DEPLOYMENT_ENVIRONMENT_KEY, KEY_DATADOG_CONTAINER_ID, KEY_DATADOG_ENVIRONMENT, KEY_DATADOG_VERSION,
     };
-    pub(crate) const SAMPLING_PRIORITY_METRIC_KEY: &str = "_sampling_priority_v1";
     const EVENT_EXTRACTION_METRIC_KEY: &str = "_dd1.sr.eausr";
     const ANALYTICS_EVENT_KEY: &str = "analytics.event";
     const HTTP_REQUEST_HEADER_PREFIX: &str = "http.request.header.";
@@ Expand Down Expand Up @@
     const HTTP_STATUS_CODE_KEY: &str = "http.status_code";
     const HTTP_RESPONSE_STATUS_CODE_KEY: &str = "http.response.status_code";
     const SPAN_KIND_META_KEY: &str = "span.kind";
-    const OTEL_TRACE_ID_META_KEY: &str = "otel.trace_id";
     const W3C_TRACESTATE_META_KEY: &str = "w3c.tracestate";
     const OTEL_LIBRARY_NAME_META_KEY: &str = "otel.library.name";
     const OTEL_LIBRARY_VERSION_META_KEY: &str = "otel.library.version";
@@ Expand Down @@

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

-Original file line number
+Diff line change
@@ Expand Up @@
     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::datadog::SAMPLING_PRIORITY_METRIC_KEY;
     use crate::common::otlp::config::TracesConfig;
     use crate::common::otlp::traces::transform::otel_span_to_dd_span;
     use crate::common::otlp::traces::transform::otlp_value_to_string;
@@ Expand Down Expand Up / @@ -52,6 +53,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 {
@@ Expand All / @@ -67,17 +69,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_METRIC_KEY) {
+                        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()
@@ Expand Down @@

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

-Original file line number
+Diff line change
@@ Expand Up / @@ -49,7 +49,7 @@ use crate::common::datadog::{ @@
         io::RB_BUFFER_CHUNK_SIZE,
         request_builder::{EndpointEncoder, RequestBuilder},
         telemetry::ComponentTelemetry,
-        DEFAULT_INTAKE_COMPRESSED_SIZE_LIMIT, DEFAULT_INTAKE_UNCOMPRESSED_SIZE_LIMIT,
+        DEFAULT_INTAKE_COMPRESSED_SIZE_LIMIT, DEFAULT_INTAKE_UNCOMPRESSED_SIZE_LIMIT, TAG_DECISION_MAKER,
     };
     use crate::common::otlp::config::TracesConfig;
     use crate::common::otlp::util::{
@@ Expand All / @@ -62,6 +62,10 @@ 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 DEFAULT_CHUNK_PRIORITY: i32 = 1; // PRIORITY_AUTO_KEEP
     fn default_serializer_compressor_kind() -> String {
         "zstd".to_string()
     }
@@ Expand Down Expand Up / @@ -547,28 +551,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);
@@ Expand Down @@

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

-Original file line number
+Diff line change
@@ Expand Up / @@ -21,6 +21,9 @@ 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;
@@ Expand Down @@

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

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

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!