Skip to main content

libdd_sampling/
dd_sampling.rs

1// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/
2// SPDX-License-Identifier: Apache-2.0
3
4//! Sampling types and mechanisms for Datadog distributed tracing.
5//!
6//! This module provides types for representing sampling decisions, priorities,
7//! and mechanisms used in Datadog's trace sampling system.
8
9use std::{borrow::Cow, fmt, str::FromStr};
10
11/// Represents a sampling decision for a trace.
12///
13/// Contains the priority level and the mechanism that made the decision.
14#[derive(Clone, Copy, Debug)]
15pub struct SamplingDecision {
16    /// The sampling priority indicating whether the trace should be kept or rejected.
17    pub priority: Option<SamplingPriority>,
18    /// The mechanism that made the sampling decision.
19    pub mechanism: Option<SamplingMechanism>,
20}
21
22/// Represents the sampling priority of a trace.
23///
24/// Positive values indicate the trace should be kept, while zero or negative
25/// values indicate rejection. Use the constants in the [`priority`] module
26/// for standard priority values.
27#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
28pub struct SamplingPriority {
29    value: i8,
30}
31
32impl SamplingPriority {
33    pub const fn from_i8(value: i8) -> Self {
34        Self { value }
35    }
36
37    pub fn into_i8(self) -> i8 {
38        self.value
39    }
40
41    /// Returns whether this sampling priority indicates the trace should be kept.
42    ///
43    /// # Returns
44    ///
45    /// `true` if the priority value is positive (indicating the trace should be kept),
46    /// `false` otherwise (indicating the trace should be dropped).
47    ///
48    /// # Examples
49    ///
50    /// ```
51    /// use libdd_sampling::priority;
52    ///
53    /// assert!(priority::AUTO_KEEP.is_keep());
54    /// assert!(priority::USER_KEEP.is_keep());
55    /// assert!(!priority::AUTO_REJECT.is_keep());
56    /// assert!(!priority::USER_REJECT.is_keep());
57    /// ```
58    #[inline(always)]
59    pub fn is_keep(&self) -> bool {
60        self.value > 0
61    }
62}
63
64/// Sampling priority constants.
65///
66/// These values indicate whether a trace should be kept or rejected,
67/// and whether the decision was made automatically or by the user.
68pub mod priority {
69    use super::SamplingPriority;
70
71    /// User explicitly rejected this trace (priority -1).
72    pub const USER_REJECT: SamplingPriority = SamplingPriority::from_i8(-1);
73    /// User explicitly requested to keep this trace (priority 2).
74    pub const USER_KEEP: SamplingPriority = SamplingPriority::from_i8(2);
75    /// Automatically rejected by the sampler (priority 0).
76    pub const AUTO_REJECT: SamplingPriority = SamplingPriority::from_i8(0);
77    /// Automatically kept by the sampler (priority 1).
78    pub const AUTO_KEEP: SamplingPriority = SamplingPriority::from_i8(1);
79}
80
81impl fmt::Display for SamplingPriority {
82    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
83        write!(f, "{}", self.value)
84    }
85}
86
87impl FromStr for SamplingPriority {
88    type Err = ();
89
90    fn from_str(s: &str) -> Result<Self, Self::Err> {
91        match s.parse::<i8>() {
92            Ok(value) => Ok(SamplingPriority::from_i8(value)),
93            Err(_) => Err(()),
94        }
95    }
96}
97
98/// Represents the mechanism that made a sampling decision.
99///
100/// The sampling mechanism identifies which component or rule determined
101/// whether a trace should be sampled (kept or rejected).
102#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
103pub struct SamplingMechanism {
104    value: u8,
105}
106
107impl SamplingMechanism {
108    pub const fn from_u8(value: u8) -> Self {
109        Self { value }
110    }
111
112    pub fn into_u8(self) -> u8 {
113        self.value
114    }
115
116    pub fn to_priority(self, is_keep: bool) -> SamplingPriority {
117        const AUTO_PAIR: PriorityPair = PriorityPair {
118            keep: priority::AUTO_KEEP,
119            reject: priority::AUTO_REJECT,
120        };
121        const USER_PAIR: PriorityPair = PriorityPair {
122            keep: priority::USER_KEEP,
123            reject: priority::USER_REJECT,
124        };
125        let pair = match self {
126            mechanism::AGENT_RATE_BY_SERVICE | mechanism::DEFAULT => AUTO_PAIR,
127            mechanism::MANUAL
128            | mechanism::LOCAL_USER_TRACE_SAMPLING_RULE
129            | mechanism::REMOTE_USER_TRACE_SAMPLING_RULE
130            | mechanism::REMOTE_DYNAMIC_TRACE_SAMPLING_RULE
131            | mechanism::SPAN_SAMPLING_RULE
132            | mechanism::DATA_JOBS_MONITORING => USER_PAIR,
133            mechanism::APPSEC => AUTO_PAIR,
134
135            _ => AUTO_PAIR,
136        };
137        if is_keep {
138            pair.keep
139        } else {
140            pair.reject
141        }
142    }
143
144    /// Returns the string representation of the sampling mechanism.
145    ///
146    /// The format is `"-N"` (e.g. `"-4"` for manual sampling). The leading `-` comes from the
147    /// propagation tags RFC, which initially had a prefix component before the `-`; that prefix
148    /// was dropped, but the `-` was retained as-is for backwards compatibility with existing
149    /// existing tracers.
150    pub fn to_cow(self) -> Cow<'static, str> {
151        match self {
152            mechanism::DEFAULT => Cow::Borrowed("-0"),
153            mechanism::AGENT_RATE_BY_SERVICE => Cow::Borrowed("-1"),
154            mechanism::REMOTE_RATE => Cow::Borrowed("-2"),
155            mechanism::LOCAL_USER_TRACE_SAMPLING_RULE => Cow::Borrowed("-3"),
156            mechanism::MANUAL => Cow::Borrowed("-4"),
157            mechanism::APPSEC => Cow::Borrowed("-5"),
158            mechanism::REMOTE_RATE_USER => Cow::Borrowed("-6"),
159            mechanism::REMOTE_RATE_DATADOG => Cow::Borrowed("-7"),
160            mechanism::SPAN_SAMPLING_RULE => Cow::Borrowed("-8"),
161            mechanism::OTLP_INGEST_PROBABILISTIC_SAMPLING => Cow::Borrowed("-9"),
162            mechanism::DATA_JOBS_MONITORING => Cow::Borrowed("-10"),
163            mechanism::REMOTE_USER_TRACE_SAMPLING_RULE => Cow::Borrowed("-11"),
164            mechanism::REMOTE_DYNAMIC_TRACE_SAMPLING_RULE => Cow::Borrowed("-12"),
165            _ => Cow::Owned(self.to_string()),
166        }
167    }
168}
169
170/// Sampling mechanism constants.
171///
172/// These constants identify which component or rule made a sampling decision.
173pub mod mechanism {
174    use super::SamplingMechanism;
175
176    /// Default sampling mechanism.
177    pub const DEFAULT: SamplingMechanism = SamplingMechanism::from_u8(0);
178    /// Agent-based rate sampling by service.
179    pub const AGENT_RATE_BY_SERVICE: SamplingMechanism = SamplingMechanism::from_u8(1);
180    /// Remote configuration rate sampling.
181    pub const REMOTE_RATE: SamplingMechanism = SamplingMechanism::from_u8(2);
182    /// Local user-defined trace sampling rule.
183    pub const LOCAL_USER_TRACE_SAMPLING_RULE: SamplingMechanism = SamplingMechanism::from_u8(3);
184    /// Manual sampling decision via API.
185    pub const MANUAL: SamplingMechanism = SamplingMechanism::from_u8(4);
186    /// Application Security (AppSec) sampling.
187    pub const APPSEC: SamplingMechanism = SamplingMechanism::from_u8(5);
188    /// Remote user rate sampling.
189    pub const REMOTE_RATE_USER: SamplingMechanism = SamplingMechanism::from_u8(6);
190    /// Remote Datadog rate sampling.
191    pub const REMOTE_RATE_DATADOG: SamplingMechanism = SamplingMechanism::from_u8(7);
192    /// Span-level sampling rule.
193    pub const SPAN_SAMPLING_RULE: SamplingMechanism = SamplingMechanism::from_u8(8);
194    /// OTLP ingest probabilistic sampling.
195    pub const OTLP_INGEST_PROBABILISTIC_SAMPLING: SamplingMechanism = SamplingMechanism::from_u8(9);
196    /// Data Jobs Monitoring sampling.
197    pub const DATA_JOBS_MONITORING: SamplingMechanism = SamplingMechanism::from_u8(10);
198    /// Remote user-defined trace sampling rule.
199    pub const REMOTE_USER_TRACE_SAMPLING_RULE: SamplingMechanism = SamplingMechanism::from_u8(11);
200    /// Remote dynamic trace sampling rule.
201    pub const REMOTE_DYNAMIC_TRACE_SAMPLING_RULE: SamplingMechanism =
202        SamplingMechanism::from_u8(12);
203}
204
205impl fmt::Display for SamplingMechanism {
206    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
207        write!(f, "-{}", self.into_u8())
208    }
209}
210
211impl FromStr for SamplingMechanism {
212    type Err = ();
213
214    /// Gets the sampling mechanism from it's string representation.
215    fn from_str(s: &str) -> Result<Self, ()> {
216        let val: i16 = s.parse().map_err(drop)?;
217        if val > 0 {
218            return Err(());
219        }
220        let val = -val;
221        if val > u8::MAX as i16 {
222            return Err(());
223        }
224        Ok(SamplingMechanism::from_u8(val as u8))
225    }
226}
227
228struct PriorityPair {
229    keep: SamplingPriority,
230    reject: SamplingPriority,
231}
232
233#[cfg(test)]
234mod tests {
235    use super::*;
236
237    // --- SamplingPriority ---
238
239    #[test]
240    fn test_priority_into_i8() {
241        assert_eq!(priority::AUTO_KEEP.into_i8(), 1);
242        assert_eq!(priority::AUTO_REJECT.into_i8(), 0);
243        assert_eq!(priority::USER_KEEP.into_i8(), 2);
244        assert_eq!(priority::USER_REJECT.into_i8(), -1);
245    }
246
247    #[test]
248    fn test_priority_is_keep() {
249        assert!(priority::AUTO_KEEP.is_keep());
250        assert!(priority::USER_KEEP.is_keep());
251        assert!(!priority::AUTO_REJECT.is_keep());
252        assert!(!priority::USER_REJECT.is_keep());
253    }
254
255    #[test]
256    fn test_priority_display() {
257        assert_eq!(priority::AUTO_KEEP.to_string(), "1");
258        assert_eq!(priority::AUTO_REJECT.to_string(), "0");
259        assert_eq!(priority::USER_KEEP.to_string(), "2");
260        assert_eq!(priority::USER_REJECT.to_string(), "-1");
261    }
262
263    #[test]
264    fn test_priority_from_str() {
265        assert_eq!(
266            "1".parse::<SamplingPriority>().unwrap(),
267            priority::AUTO_KEEP
268        );
269        assert_eq!(
270            "0".parse::<SamplingPriority>().unwrap(),
271            priority::AUTO_REJECT
272        );
273        assert_eq!(
274            "2".parse::<SamplingPriority>().unwrap(),
275            priority::USER_KEEP
276        );
277        assert_eq!(
278            "-1".parse::<SamplingPriority>().unwrap(),
279            priority::USER_REJECT
280        );
281        assert!("not_a_number".parse::<SamplingPriority>().is_err());
282        assert!("999".parse::<SamplingPriority>().is_err()); // overflows i8
283    }
284
285    // --- SamplingMechanism ---
286
287    #[test]
288    fn test_mechanism_into_u8() {
289        assert_eq!(mechanism::DEFAULT.into_u8(), 0);
290        assert_eq!(mechanism::AGENT_RATE_BY_SERVICE.into_u8(), 1);
291        assert_eq!(mechanism::LOCAL_USER_TRACE_SAMPLING_RULE.into_u8(), 3);
292        assert_eq!(mechanism::REMOTE_USER_TRACE_SAMPLING_RULE.into_u8(), 11);
293        assert_eq!(mechanism::REMOTE_DYNAMIC_TRACE_SAMPLING_RULE.into_u8(), 12);
294    }
295
296    #[test]
297    fn test_mechanism_to_priority_auto_pair() {
298        // DEFAULT and AGENT_RATE_BY_SERVICE use AUTO priority
299        assert_eq!(mechanism::DEFAULT.to_priority(true), priority::AUTO_KEEP);
300        assert_eq!(mechanism::DEFAULT.to_priority(false), priority::AUTO_REJECT);
301        assert_eq!(
302            mechanism::AGENT_RATE_BY_SERVICE.to_priority(true),
303            priority::AUTO_KEEP
304        );
305        assert_eq!(
306            mechanism::AGENT_RATE_BY_SERVICE.to_priority(false),
307            priority::AUTO_REJECT
308        );
309        assert_eq!(mechanism::APPSEC.to_priority(true), priority::AUTO_KEEP);
310        assert_eq!(mechanism::APPSEC.to_priority(false), priority::AUTO_REJECT);
311    }
312
313    #[test]
314    fn test_mechanism_to_priority_user_pair() {
315        // Rule-based mechanisms use USER priority
316        assert_eq!(
317            mechanism::LOCAL_USER_TRACE_SAMPLING_RULE.to_priority(true),
318            priority::USER_KEEP
319        );
320        assert_eq!(
321            mechanism::LOCAL_USER_TRACE_SAMPLING_RULE.to_priority(false),
322            priority::USER_REJECT
323        );
324        assert_eq!(
325            mechanism::REMOTE_USER_TRACE_SAMPLING_RULE.to_priority(true),
326            priority::USER_KEEP
327        );
328        assert_eq!(
329            mechanism::REMOTE_DYNAMIC_TRACE_SAMPLING_RULE.to_priority(true),
330            priority::USER_KEEP
331        );
332        assert_eq!(mechanism::MANUAL.to_priority(true), priority::USER_KEEP);
333        assert_eq!(
334            mechanism::SPAN_SAMPLING_RULE.to_priority(false),
335            priority::USER_REJECT
336        );
337        assert_eq!(
338            mechanism::DATA_JOBS_MONITORING.to_priority(true),
339            priority::USER_KEEP
340        );
341    }
342
343    #[test]
344    fn test_mechanism_to_priority_unknown_falls_back_to_auto() {
345        let unknown = SamplingMechanism::from_u8(99);
346        assert_eq!(unknown.to_priority(true), priority::AUTO_KEEP);
347        assert_eq!(unknown.to_priority(false), priority::AUTO_REJECT);
348    }
349
350    #[test]
351    fn test_mechanism_to_cow() {
352        assert_eq!(mechanism::DEFAULT.to_cow(), "-0");
353        assert_eq!(mechanism::AGENT_RATE_BY_SERVICE.to_cow(), "-1");
354        assert_eq!(mechanism::REMOTE_RATE.to_cow(), "-2");
355        assert_eq!(mechanism::LOCAL_USER_TRACE_SAMPLING_RULE.to_cow(), "-3");
356        assert_eq!(mechanism::MANUAL.to_cow(), "-4");
357        assert_eq!(mechanism::APPSEC.to_cow(), "-5");
358        assert_eq!(mechanism::REMOTE_RATE_USER.to_cow(), "-6");
359        assert_eq!(mechanism::REMOTE_RATE_DATADOG.to_cow(), "-7");
360        assert_eq!(mechanism::SPAN_SAMPLING_RULE.to_cow(), "-8");
361        assert_eq!(mechanism::OTLP_INGEST_PROBABILISTIC_SAMPLING.to_cow(), "-9");
362        assert_eq!(mechanism::DATA_JOBS_MONITORING.to_cow(), "-10");
363        assert_eq!(mechanism::REMOTE_USER_TRACE_SAMPLING_RULE.to_cow(), "-11");
364        assert_eq!(
365            mechanism::REMOTE_DYNAMIC_TRACE_SAMPLING_RULE.to_cow(),
366            "-12"
367        );
368        // Unknown mechanism falls back to Display
369        assert_eq!(SamplingMechanism::from_u8(99).to_cow(), "-99");
370    }
371
372    #[test]
373    fn test_mechanism_display() {
374        assert_eq!(mechanism::DEFAULT.to_string(), "-0");
375        assert_eq!(mechanism::LOCAL_USER_TRACE_SAMPLING_RULE.to_string(), "-3");
376        assert_eq!(
377            mechanism::REMOTE_DYNAMIC_TRACE_SAMPLING_RULE.to_string(),
378            "-12"
379        );
380    }
381
382    #[test]
383    fn test_mechanism_from_str() {
384        assert_eq!(
385            "-0".parse::<SamplingMechanism>().unwrap(),
386            mechanism::DEFAULT
387        );
388        assert_eq!(
389            "-1".parse::<SamplingMechanism>().unwrap(),
390            mechanism::AGENT_RATE_BY_SERVICE
391        );
392        assert_eq!(
393            "-3".parse::<SamplingMechanism>().unwrap(),
394            mechanism::LOCAL_USER_TRACE_SAMPLING_RULE
395        );
396        assert_eq!(
397            "-12".parse::<SamplingMechanism>().unwrap(),
398            mechanism::REMOTE_DYNAMIC_TRACE_SAMPLING_RULE
399        );
400    }
401
402    #[test]
403    fn test_mechanism_from_str_errors() {
404        assert!("not_a_number".parse::<SamplingMechanism>().is_err());
405        assert!("1".parse::<SamplingMechanism>().is_err()); // positive not allowed
406        assert!("-99999".parse::<SamplingMechanism>().is_err()); // overflows u8
407    }
408}