use std::cell::Cell;
use obs_types::Severity;
use crate::config::SamplingConfig;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum SamplingDecision {
Drop,
Keep,
ParentSet {
sampled: bool,
},
}
#[must_use]
pub fn decide(
cfg: &SamplingConfig,
full_name: &str,
severity: Severity,
inbound_sampled: Option<bool>,
) -> SamplingDecision {
if cfg.honour_traceparent_sampled
&& let Some(s) = inbound_sampled
{
return SamplingDecision::ParentSet { sampled: s };
}
if severity >= cfg.always_log_at_or_above {
return SamplingDecision::Keep;
}
let rate = cfg
.per_event
.get(full_name)
.copied()
.unwrap_or(cfg.default_rate);
if rate >= 1.0 {
return SamplingDecision::Keep;
}
if rate <= 0.0 {
return SamplingDecision::Drop;
}
if rand_unit_f64() < rate {
SamplingDecision::Keep
} else {
SamplingDecision::Drop
}
}
thread_local! {
static SHIFT_STATE: Cell<u64> = Cell::new(seed_state());
}
fn seed_state() -> u64 {
use std::time::{SystemTime, UNIX_EPOCH};
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_nanos() as u64)
.unwrap_or(0x9E37_79B9_7F4A_7C15);
nanos | 1
}
fn rand_unit_f64() -> f64 {
SHIFT_STATE.with(|cell| {
let mut x = cell.get();
if x == 0 {
x = seed_state();
}
x ^= x << 13;
x ^= x >> 7;
x ^= x << 17;
cell.set(x);
let top = x >> 11;
(top as f64) / ((1u64 << 53) as f64)
})
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::SamplingConfig;
#[test]
fn test_should_keep_at_or_above_floor() {
let cfg = SamplingConfig {
default_rate: 0.0,
..Default::default()
};
assert_eq!(
decide(&cfg, "x", Severity::Error, None),
SamplingDecision::Keep
);
}
#[test]
fn test_should_drop_below_floor_with_zero_rate() {
let cfg = SamplingConfig {
default_rate: 0.0,
..Default::default()
};
assert_eq!(
decide(&cfg, "x", Severity::Trace, None),
SamplingDecision::Drop
);
}
#[test]
fn test_should_honour_parent_sampled() {
let cfg = SamplingConfig::default();
match decide(&cfg, "x", Severity::Trace, Some(true)) {
SamplingDecision::ParentSet { sampled } => assert!(sampled),
d => panic!("unexpected decision: {d:?}"),
}
match decide(&cfg, "x", Severity::Error, Some(false)) {
SamplingDecision::ParentSet { sampled } => assert!(!sampled),
d => panic!("unexpected decision: {d:?}"),
}
}
#[test]
fn test_should_use_per_event_override() {
let mut cfg = SamplingConfig {
default_rate: 0.0,
..Default::default()
};
cfg.per_event.insert("x".to_string(), 1.0);
assert_eq!(
decide(&cfg, "x", Severity::Trace, None),
SamplingDecision::Keep
);
}
}