1use std::{borrow::Cow, fmt, str::FromStr};
10
11#[derive(Clone, Copy, Debug)]
15pub struct SamplingDecision {
16 pub priority: Option<SamplingPriority>,
18 pub mechanism: Option<SamplingMechanism>,
20}
21
22#[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 #[inline(always)]
59 pub fn is_keep(&self) -> bool {
60 self.value > 0
61 }
62}
63
64pub mod priority {
69 use super::SamplingPriority;
70
71 pub const USER_REJECT: SamplingPriority = SamplingPriority::from_i8(-1);
73 pub const USER_KEEP: SamplingPriority = SamplingPriority::from_i8(2);
75 pub const AUTO_REJECT: SamplingPriority = SamplingPriority::from_i8(0);
77 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#[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 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
170pub mod mechanism {
174 use super::SamplingMechanism;
175
176 pub const DEFAULT: SamplingMechanism = SamplingMechanism::from_u8(0);
178 pub const AGENT_RATE_BY_SERVICE: SamplingMechanism = SamplingMechanism::from_u8(1);
180 pub const REMOTE_RATE: SamplingMechanism = SamplingMechanism::from_u8(2);
182 pub const LOCAL_USER_TRACE_SAMPLING_RULE: SamplingMechanism = SamplingMechanism::from_u8(3);
184 pub const MANUAL: SamplingMechanism = SamplingMechanism::from_u8(4);
186 pub const APPSEC: SamplingMechanism = SamplingMechanism::from_u8(5);
188 pub const REMOTE_RATE_USER: SamplingMechanism = SamplingMechanism::from_u8(6);
190 pub const REMOTE_RATE_DATADOG: SamplingMechanism = SamplingMechanism::from_u8(7);
192 pub const SPAN_SAMPLING_RULE: SamplingMechanism = SamplingMechanism::from_u8(8);
194 pub const OTLP_INGEST_PROBABILISTIC_SAMPLING: SamplingMechanism = SamplingMechanism::from_u8(9);
196 pub const DATA_JOBS_MONITORING: SamplingMechanism = SamplingMechanism::from_u8(10);
198 pub const REMOTE_USER_TRACE_SAMPLING_RULE: SamplingMechanism = SamplingMechanism::from_u8(11);
200 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 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 #[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()); }
284
285 #[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 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 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 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()); assert!("-99999".parse::<SamplingMechanism>().is_err()); }
408}