use std::{
collections::HashMap,
sync::{Arc, Mutex},
};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SamplingConfig {
pub enabled: bool,
pub first_n: u64,
pub thereafter: u64,
}
impl Default for SamplingConfig {
fn default() -> Self {
Self {
enabled: false,
first_n: 1,
thereafter: 100,
}
}
}
#[derive(Debug, Clone, Default)]
pub struct LogSampler {
config: SamplingConfig,
counts: Arc<Mutex<HashMap<String, u64>>>,
}
impl LogSampler {
pub fn new(config: SamplingConfig) -> Self {
Self {
config,
counts: Arc::new(Mutex::new(HashMap::new())),
}
}
pub fn should_log(&self, key: &str) -> bool {
if !self.config.enabled {
return true;
}
let mut counts = self.counts.lock().expect("sampler mutex");
let count = counts.entry(key.to_string()).or_default();
*count += 1;
if *count <= self.config.first_n {
return true;
}
let thereafter = self.config.thereafter.max(1);
(*count - self.config.first_n).is_multiple_of(thereafter)
}
}