offline_protocol/telemetry/
sampling.rs

1//! Generic fixed-window rate limiter for event sampling.
2
3use std::collections::HashMap;
4use std::sync::Mutex;
5
6const DEFAULT_WINDOW_MS: i64 = 1_000;
7const DEFAULT_MAX_EVENTS: u32 = 10;
8const DEFAULT_MAX_KEYS: usize = 4096;
9
10#[derive(Debug, Clone)]
11struct EventWindow {
12    window_start_ms: i64,
13    count: u32,
14}
15
16/// Best-effort fixed-window sampler that limits events per category key.
17///
18/// Each unique key gets its own window. When a key's budget is exhausted for
19/// the current window, further events with that key are suppressed until the
20/// window resets. An eviction policy prevents unbounded memory growth when
21/// the number of distinct keys exceeds `max_keys`.
22#[derive(Debug)]
23pub struct CategorySampler {
24    max_events_per_window: u32,
25    window_ms: i64,
26    max_keys: usize,
27    windows: Mutex<HashMap<String, EventWindow>>,
28}
29
30impl Default for CategorySampler {
31    fn default() -> Self {
32        Self {
33            max_events_per_window: DEFAULT_MAX_EVENTS,
34            window_ms: DEFAULT_WINDOW_MS,
35            max_keys: DEFAULT_MAX_KEYS,
36            windows: Mutex::new(HashMap::new()),
37        }
38    }
39}
40
41impl CategorySampler {
42    /// Creates a sampler with custom parameters.
43    pub fn new(max_events_per_window: u32, window_ms: i64, max_keys: usize) -> Self {
44        Self {
45            max_events_per_window,
46            window_ms,
47            max_keys,
48            windows: Mutex::new(HashMap::new()),
49        }
50    }
51
52    /// Returns `true` when an event with the given key should be emitted.
53    ///
54    /// If the mutex is poisoned the call conservatively returns `true` so
55    /// events are never silently dropped due to an unrelated panic.
56    pub fn allow(&self, key: &str, now_ms: i64) -> bool {
57        let Ok(mut windows) = self.windows.lock() else {
58            return true;
59        };
60
61        if windows.len() >= self.max_keys && !windows.contains_key(key) {
62            windows.retain(|_, window| {
63                now_ms.saturating_sub(window.window_start_ms) <= self.window_ms
64            });
65            if windows.len() >= self.max_keys {
66                windows.clear();
67            }
68        }
69
70        let window = windows.entry(key.to_string()).or_insert(EventWindow {
71            window_start_ms: now_ms,
72            count: 0,
73        });
74
75        if now_ms.saturating_sub(window.window_start_ms) >= self.window_ms {
76            window.window_start_ms = now_ms;
77            window.count = 0;
78        }
79
80        if window.count >= self.max_events_per_window {
81            return false;
82        }
83
84        window.count = window.count.saturating_add(1);
85        true
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use super::*;
92
93    #[test]
94    fn blocks_after_budget_exhausted() {
95        let sampler = CategorySampler::default();
96        let now = 1000_i64;
97        let key = "k";
98        for _ in 0..DEFAULT_MAX_EVENTS {
99            assert!(sampler.allow(key, now));
100        }
101        assert!(!sampler.allow(key, now));
102    }
103
104    #[test]
105    fn window_resets_after_elapsed() {
106        let sampler = CategorySampler::default();
107        let now = 1000_i64;
108        for _ in 0..DEFAULT_MAX_EVENTS {
109            sampler.allow("k", now);
110        }
111        assert!(!sampler.allow("k", now));
112        assert!(sampler.allow("k", now + DEFAULT_WINDOW_MS));
113    }
114
115    #[test]
116    fn independent_keys_have_separate_budgets() {
117        let sampler = CategorySampler::default();
118        let now = 1000_i64;
119        for _ in 0..DEFAULT_MAX_EVENTS {
120            sampler.allow("a", now);
121        }
122        assert!(!sampler.allow("a", now));
123        assert!(sampler.allow("b", now));
124    }
125
126    #[test]
127    fn eviction_under_key_pressure() {
128        let sampler = CategorySampler::new(10, 1_000, 2);
129        let now = 1000_i64;
130        assert!(sampler.allow("a", now));
131        assert!(sampler.allow("b", now));
132        // Third key triggers eviction — should still be accepted.
133        assert!(sampler.allow("c", now));
134    }
135}
offline_protocol/telemetry/sampling.rs

offline_protocol/telemetry/
sampling.rs
