rust_zero_core/
rolling.rs1use std::{
2 collections::VecDeque,
3 sync::Mutex,
4 time::{Duration, Instant},
5};
6
7pub 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}