mysten_common/
moving_window.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::{collections::VecDeque, fmt::Debug, time::Duration};
5
6/// A moving window that maintains the last N values of type `T` and calculates their arithmetic
7/// mean. All values in the window have equal weight and the oldest value is dropped when the
8/// window exceeds its configured capacity.
9#[derive(Debug, Clone)]
10pub struct MovingWindow<T: MovingWindowValue> {
11    values: VecDeque<T>,
12    max_size: usize,
13    sum: T,
14}
15
16impl<T: MovingWindowValue> MovingWindow<T> {
17    /// Creates a new `MovingWindow` with the specified maximum size and an `init_value`.
18    /// The provided `max_size` must be greater than 0.
19    pub fn new(init_value: T, max_size: usize) -> Self {
20        assert!(max_size > 0, "Window size must be greater than 0");
21        let mut window = Self {
22            values: VecDeque::with_capacity(max_size),
23            max_size,
24            sum: T::zero(),
25        };
26        window.add_value(init_value);
27        window
28    }
29
30    /// Adds a new value to the window. If the window is at capacity, the oldest value is removed
31    /// before adding the new value.
32    pub fn add_value(&mut self, value: T) {
33        if self.values.len() == self.max_size
34            && let Some(old_value) = self.values.pop_front()
35        {
36            T::sub_assign(&mut self.sum, old_value);
37        }
38
39        self.values.push_back(value);
40        T::add_assign(&mut self.sum, value);
41    }
42
43    /// Get the current average of all values in the window. Returns the value's zero if the
44    /// window is empty.
45    pub fn get(&self) -> T {
46        if self.values.is_empty() {
47            T::zero()
48        } else {
49            T::average(self.sum, self.values.len())
50        }
51    }
52
53    /// Get the number of values currently in the window.
54    pub fn len(&self) -> usize {
55        self.values.len()
56    }
57
58    pub fn is_empty(&self) -> bool {
59        self.values.is_empty()
60    }
61}
62
63pub trait MovingWindowValue: Copy + Debug {
64    fn zero() -> Self;
65    fn add_assign(target: &mut Self, value: Self);
66    fn sub_assign(target: &mut Self, value: Self);
67    fn average(total: Self, divisor: usize) -> Self;
68}
69
70impl MovingWindowValue for Duration {
71    fn zero() -> Self {
72        Duration::ZERO
73    }
74
75    fn add_assign(target: &mut Self, value: Self) {
76        *target += value;
77    }
78
79    fn sub_assign(target: &mut Self, value: Self) {
80        *target -= value;
81    }
82
83    fn average(total: Self, divisor: usize) -> Self {
84        let divisor = u32::try_from(divisor).expect("window size too large for Duration average");
85        total / divisor
86    }
87}
88
89impl MovingWindowValue for f64 {
90    fn zero() -> Self {
91        0.0
92    }
93
94    fn add_assign(target: &mut Self, value: Self) {
95        *target += value;
96    }
97
98    fn sub_assign(target: &mut Self, value: Self) {
99        *target -= value;
100    }
101
102    fn average(total: Self, divisor: usize) -> Self {
103        total / divisor as f64
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110
111    #[test]
112    fn test_with_initial_value() {
113        let window = MovingWindow::new(10.0, 3);
114        assert_eq!(window.len(), 1);
115        assert_eq!(window.get(), 10.0);
116    }
117
118    #[test]
119    fn test_duration_window() {
120        let mut window = MovingWindow::new(Duration::ZERO, 3);
121        assert_eq!(window.get(), Duration::ZERO);
122
123        // Adding value within the window size.
124        window.add_value(Duration::from_millis(100));
125        assert_eq!(window.get(), Duration::from_millis(50));
126        assert_eq!(window.len(), 2);
127
128        // Adding value within the window size.
129        window.add_value(Duration::from_millis(200));
130        assert_eq!(window.get(), Duration::from_millis(100));
131        assert_eq!(window.len(), 3);
132
133        // Adding values exceeding the window size.
134        window.add_value(Duration::from_millis(300));
135        assert_eq!(window.get(), Duration::from_millis(200));
136        assert_eq!(window.len(), 3);
137
138        // Adding values exceeding the window size.
139        window.add_value(Duration::from_millis(400));
140        assert_eq!(window.get(), Duration::from_millis(300));
141        assert_eq!(window.len(), 3);
142    }
143
144    #[test]
145    fn test_float_window() {
146        let mut window = MovingWindow::new(0.0_f64, 3);
147        assert_eq!(window.get(), 0.0);
148
149        // Adding value within the window size.
150        window.add_value(1.0);
151        assert_eq!(window.get(), 0.5);
152        assert_eq!(window.len(), 2);
153
154        // Adding value within the window size.
155        window.add_value(2.0);
156        assert_eq!(window.get(), 1.0);
157        assert_eq!(window.len(), 3);
158
159        // Adding value exceeding the window size.
160        window.add_value(3.0);
161        assert_eq!(window.get(), 2.0);
162        assert_eq!(window.len(), 3);
163
164        // Adding value exceeding the window size.
165        window.add_value(4.0);
166        assert_eq!(window.get(), 3.0);
167        assert_eq!(window.len(), 3);
168    }
169
170    #[test]
171    #[should_panic(expected = "Window size must be greater than 0")]
172    fn test_zero_size_panics() {
173        let _window = MovingWindow::new(0.0_f64, 0);
174    }
175}
mysten_common/moving_window.rs

mysten_common/
moving_window.rs
