Skip to main content

rust_zero_core/
rolling.rs

1use std::{
2    collections::VecDeque,
3    sync::Mutex,
4    time::{Duration, Instant},
5};
6
7/// Time-windowed numeric observations split into fixed-duration buckets.
8pub struct RollingWindow {
9    bucket_count: usize,
10    bucket_width: Duration,
11    state: Mutex<WindowState>,
12}
13
14impl RollingWindow {
15    pub fn new(bucket_count: usize, bucket_width: Duration) -> Self {
16        assert!(bucket_count > 0, "bucket count must be greater than zero");
17        assert!(
18            !bucket_width.is_zero(),
19            "bucket width must be greater than zero"
20        );
21
22        let now = Instant::now();
23        Self {
24            bucket_count,
25            bucket_width,
26            state: Mutex::new(WindowState {
27                current_started: now,
28                buckets: VecDeque::from([Bucket::default()]),
29            }),
30        }
31    }
32
33    pub fn record(&self, value: f64) {
34        assert!(value.is_finite(), "rolling-window values must be finite");
35        let mut state = self.state.lock().expect("rolling window mutex poisoned");
36        self.rotate(&mut state, Instant::now());
37        state
38            .buckets
39            .back_mut()
40            .expect("rolling window always contains its current bucket")
41            .record(value);
42    }
43
44    pub fn snapshot(&self) -> RollingSnapshot {
45        let mut state = self.state.lock().expect("rolling window mutex poisoned");
46        self.rotate(&mut state, Instant::now());
47
48        let mut snapshot = RollingSnapshot::default();
49        for bucket in &state.buckets {
50            snapshot.count += bucket.count;
51            snapshot.sum += bucket.sum;
52            if let Some(minimum) = bucket.minimum {
53                snapshot.minimum =
54                    Some(snapshot.minimum.map_or(minimum, |value| value.min(minimum)));
55            }
56            if let Some(maximum) = bucket.maximum {
57                snapshot.maximum =
58                    Some(snapshot.maximum.map_or(maximum, |value| value.max(maximum)));
59            }
60        }
61        snapshot
62    }
63
64    pub fn reset(&self) {
65        let mut state = self.state.lock().expect("rolling window mutex poisoned");
66        state.current_started = Instant::now();
67        state.buckets.clear();
68        state.buckets.push_back(Bucket::default());
69    }
70
71    fn rotate(&self, state: &mut WindowState, now: Instant) {
72        let elapsed = now.saturating_duration_since(state.current_started);
73        let elapsed_buckets_u128 = elapsed.as_nanos() / self.bucket_width.as_nanos();
74        let elapsed_buckets = usize::try_from(elapsed_buckets_u128).unwrap_or(self.bucket_count);
75        if elapsed_buckets == 0 {
76            return;
77        }
78
79        if elapsed_buckets >= self.bucket_count {
80            state.buckets.clear();
81            state.buckets.push_back(Bucket::default());
82            state.current_started = now;
83            return;
84        }
85
86        for _ in 0..elapsed_buckets {
87            state.buckets.push_back(Bucket::default());
88            if state.buckets.len() > self.bucket_count {
89                state.buckets.pop_front();
90            }
91        }
92        state.current_started += self.bucket_width * elapsed_buckets as u32;
93    }
94}
95
96struct WindowState {
97    current_started: Instant,
98    buckets: VecDeque<Bucket>,
99}
100
101#[derive(Default)]
102struct Bucket {
103    count: u64,
104    sum: f64,
105    minimum: Option<f64>,
106    maximum: Option<f64>,
107}
108
109impl Bucket {
110    fn record(&mut self, value: f64) {
111        self.count += 1;
112        self.sum += value;
113        self.minimum = Some(self.minimum.map_or(value, |minimum| minimum.min(value)));
114        self.maximum = Some(self.maximum.map_or(value, |maximum| maximum.max(value)));
115    }
116}
117
118#[derive(Debug, Clone, Copy, Default, PartialEq)]
119pub struct RollingSnapshot {
120    pub count: u64,
121    pub sum: f64,
122    pub minimum: Option<f64>,
123    pub maximum: Option<f64>,
124}
125
126impl RollingSnapshot {
127    pub fn average(&self) -> Option<f64> {
128        (self.count > 0).then(|| self.sum / self.count as f64)
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135    use std::thread;
136
137    #[test]
138    fn aggregates_observations_in_the_active_window() {
139        let window = RollingWindow::new(3, Duration::from_secs(1));
140        window.record(2.0);
141        window.record(4.0);
142
143        assert_eq!(
144            window.snapshot(),
145            RollingSnapshot {
146                count: 2,
147                sum: 6.0,
148                minimum: Some(2.0),
149                maximum: Some(4.0),
150            }
151        );
152        assert_eq!(window.snapshot().average(), Some(3.0));
153    }
154
155    #[test]
156    fn expires_old_buckets() {
157        let window = RollingWindow::new(2, Duration::from_millis(5));
158        window.record(10.0);
159        thread::sleep(Duration::from_millis(15));
160
161        assert_eq!(window.snapshot(), RollingSnapshot::default());
162    }
163}