Skip to main content

obs_core/
sampling.rs

1//! Head sampler — per `(full_name, sev)` rate decision.
2//!
3//! Spec 13 § 6, spec 11 § 4.1 step 4.
4//!
5//! Decision order, mirroring spec 13 § 6:
6//!
7//! 1. Inbound `traceparent.sampled` — when set on the active scope frame, the upstream caller
8//!    already decided. Returns `SamplingDecision::ParentSet { sampled }` so the caller stamps
9//!    `sampling_reason = OVERRIDE` on emit.
10//! 2. Severity floor (`always_log_at_or_above`) — bypasses sampling.
11//! 3. Per-event rate from config; otherwise the global default rate.
12//!
13//! The implementation uses `fastrand::f64()`-equivalent behaviour
14//! via `rand`-free path — we cheat with a per-thread XorShift to keep
15//! the runtime out of the `rand` dependency tree on the hot path.
16
17use std::cell::Cell;
18
19use obs_proto::obs::v1::Severity;
20
21use crate::config::SamplingConfig;
22
23/// Outcome of the head-sampler decision.
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25#[non_exhaustive]
26pub enum SamplingDecision {
27    /// The envelope is dropped before reaching the worker.
28    Drop,
29    /// The envelope is kept under the local rate (or severity floor).
30    Keep,
31    /// The decision was forced by an upstream `traceparent.sampled`;
32    /// the emitter should stamp `SAMPLING_REASON_OVERRIDE`.
33    ParentSet {
34        /// Whether the parent decided to sample.
35        sampled: bool,
36    },
37}
38
39/// Run the head sampler. `inbound_sampled` is the `traceparent.sampled`
40/// bit lifted off the active scope frame; pass `None` when no scope is
41/// active or the caller did not propagate W3C trace context.
42#[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
88/// XorShift64-based uniform `[0.0, 1.0)` sample. Cheap and per-thread;
89/// suitable for sampling decisions where statistical independence
90/// across threads matters more than cryptographic strength.
91fn 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        // Use the top 53 bits to construct an f64 in [0,1). 53 = mantissa.
102        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}