rs_zero/core/logging/
sampling.rs1use std::{
2 collections::HashMap,
3 sync::{Arc, Mutex},
4};
5
6#[derive(Debug, Clone, PartialEq, Eq)]
8pub struct SamplingConfig {
9 pub enabled: bool,
11 pub first_n: u64,
13 pub thereafter: u64,
15}
16
17impl Default for SamplingConfig {
18 fn default() -> Self {
19 Self {
20 enabled: false,
21 first_n: 1,
22 thereafter: 100,
23 }
24 }
25}
26
27#[derive(Debug, Clone, Default)]
29pub struct LogSampler {
30 config: SamplingConfig,
31 counts: Arc<Mutex<HashMap<String, u64>>>,
32}
33
34impl LogSampler {
35 pub fn new(config: SamplingConfig) -> Self {
37 Self {
38 config,
39 counts: Arc::new(Mutex::new(HashMap::new())),
40 }
41 }
42
43 pub fn should_log(&self, key: &str) -> bool {
45 if !self.config.enabled {
46 return true;
47 }
48 let mut counts = self.counts.lock().expect("sampler mutex");
49 let count = counts.entry(key.to_string()).or_default();
50 *count += 1;
51 if *count <= self.config.first_n {
52 return true;
53 }
54 let thereafter = self.config.thereafter.max(1);
55 (*count - self.config.first_n).is_multiple_of(thereafter)
56 }
57}