1use std::cell::Cell;
18
19use obs_proto::obs::v1::Severity;
20
21use crate::config::SamplingConfig;
22
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25#[non_exhaustive]
26pub enum SamplingDecision {
27 Drop,
29 Keep,
31 ParentSet {
34 sampled: bool,
36 },
37}
38
39#[must_use]
43pub fn decide(
44 cfg: &SamplingConfig,
45 full_name: &str,
46 severity: Severity,
47 inbound_sampled: Option<bool>,
48) -> SamplingDecision {
49 if cfg.honour_traceparent_sampled
50 && let Some(s) = inbound_sampled
51 {
52 return SamplingDecision::ParentSet { sampled: s };
53 }
54 if severity >= cfg.always_log_at_or_above {
55 return SamplingDecision::Keep;
56 }
57 let rate = cfg
58 .per_event
59 .get(full_name)
60 .copied()
61 .unwrap_or(cfg.default_rate);
62 if rate >= 1.0 {
63 return SamplingDecision::Keep;
64 }
65 if rate <= 0.0 {
66 return SamplingDecision::Drop;
67 }
68 if rand_unit_f64() < rate {
69 SamplingDecision::Keep
70 } else {
71 SamplingDecision::Drop
72 }
73}
74
75thread_local! {
76 static SHIFT_STATE: Cell<u64> = Cell::new(seed_state());
77}
78
79fn seed_state() -> u64 {
80 use std::time::{SystemTime, UNIX_EPOCH};
81 let nanos = SystemTime::now()
82 .duration_since(UNIX_EPOCH)
83 .map(|d| d.as_nanos() as u64)
84 .unwrap_or(0x9E37_79B9_7F4A_7C15);
85 nanos | 1
86}
87
88fn rand_unit_f64() -> f64 {
92 SHIFT_STATE.with(|cell| {
93 let mut x = cell.get();
94 if x == 0 {
95 x = seed_state();
96 }
97 x ^= x << 13;
98 x ^= x >> 7;
99 x ^= x << 17;
100 cell.set(x);
101 let top = x >> 11;
103 (top as f64) / ((1u64 << 53) as f64)
104 })
105}
106
107#[cfg(test)]
108mod tests {
109 use super::*;
110 use crate::config::SamplingConfig;
111
112 #[test]
113 fn test_should_keep_at_or_above_floor() {
114 let cfg = SamplingConfig {
115 default_rate: 0.0,
116 ..Default::default()
117 };
118 assert_eq!(
119 decide(&cfg, "x", Severity::Error, None),
120 SamplingDecision::Keep
121 );
122 }
123
124 #[test]
125 fn test_should_drop_below_floor_with_zero_rate() {
126 let cfg = SamplingConfig {
127 default_rate: 0.0,
128 ..Default::default()
129 };
130 assert_eq!(
131 decide(&cfg, "x", Severity::Trace, None),
132 SamplingDecision::Drop
133 );
134 }
135
136 #[test]
137 fn test_should_honour_parent_sampled() {
138 let cfg = SamplingConfig::default();
139 match decide(&cfg, "x", Severity::Trace, Some(true)) {
140 SamplingDecision::ParentSet { sampled } => assert!(sampled),
141 d => panic!("unexpected decision: {d:?}"),
142 }
143 match decide(&cfg, "x", Severity::Error, Some(false)) {
144 SamplingDecision::ParentSet { sampled } => assert!(!sampled),
145 d => panic!("unexpected decision: {d:?}"),
146 }
147 }
148
149 #[test]
150 fn test_should_use_per_event_override() {
151 let mut cfg = SamplingConfig {
152 default_rate: 0.0,
153 ..Default::default()
154 };
155 cfg.per_event.insert("x".to_string(), 1.0);
156 assert_eq!(
157 decide(&cfg, "x", Severity::Trace, None),
158 SamplingDecision::Keep
159 );
160 }
161}