Skip to main content

rs_zero/core/logging/
sampling.rs

1use std::{
2    collections::HashMap,
3    sync::{Arc, Mutex},
4};
5
6/// Log sampling configuration.
7#[derive(Debug, Clone, PartialEq, Eq)]
8pub struct SamplingConfig {
9    /// Whether sampling is enabled.
10    pub enabled: bool,
11    /// Always log the first N occurrences for each key.
12    pub first_n: u64,
13    /// After `first_n`, log every Nth occurrence.
14    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/// In-memory deterministic log sampler keyed by low-cardinality fingerprints.
28#[derive(Debug, Clone, Default)]
29pub struct LogSampler {
30    config: SamplingConfig,
31    counts: Arc<Mutex<HashMap<String, u64>>>,
32}
33
34impl LogSampler {
35    /// Creates a sampler.
36    pub fn new(config: SamplingConfig) -> Self {
37        Self {
38            config,
39            counts: Arc::new(Mutex::new(HashMap::new())),
40        }
41    }
42
43    /// Returns whether a log for the key should be emitted.
44    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}