saito_core/core/consensus/peers/
rate_limiter.rs

1use crate::core::defs::Timestamp;
2use std::time::Duration;
3
4#[derive(Debug, Clone)]
5pub struct RateLimiter {
6    /// Max allowed requests in the window
7    limit: u64,
8    window: Timestamp,            // Window duration in milliseconds
9    request_count: u64,           // Number of requests made in the current window
10    last_request_time: Timestamp, // Timestamp of the last request in milliseconds
11}
12
13impl RateLimiter {
14    pub fn set_limit(&mut self, limit: u64) {
15        self.limit = limit;
16    }
17
18    pub fn set_window(&mut self, window: Timestamp) {
19        self.window = window;
20    }
21
22    pub fn has_limit_exceeded(&mut self, current_time: Timestamp) -> bool {
23        // TODO : current implementation allows twice the peers from spikes. a sliding window implementation would be better I think.
24        if current_time.saturating_sub(self.last_request_time) > self.window {
25            self.request_count = 0;
26            self.last_request_time = current_time;
27        }
28        self.request_count >= self.limit
29    }
30    pub fn increase(&mut self) {
31        self.request_count += 1;
32    }
33
34    pub fn builder(count: u64, duration: Duration) -> Self {
35        Self {
36            limit: count,
37            window: duration.as_millis() as Timestamp,
38            request_count: 0,
39            last_request_time: 0,
40        }
41    }
42}
43
44#[cfg(test)]
45mod tests {
46    use super::*;
47
48    #[test]
49    fn default_rate_limiter_test() {
50        let rate_limiter = RateLimiter::builder(10, Duration::from_secs(1));
51
52        assert_eq!(rate_limiter.limit, 10);
53        assert_eq!(rate_limiter.window, 1_000);
54        assert_eq!(rate_limiter.request_count, 0);
55        assert_eq!(rate_limiter.last_request_time, 0);
56    }
57
58    #[test]
59    fn set_limit_test() {
60        let mut rate_limiter = RateLimiter::builder(10, Duration::from_secs(1));
61        rate_limiter.set_limit(5);
62
63        assert_eq!(rate_limiter.limit, 5);
64    }
65
66    #[test]
67    fn set_window_test() {
68        let mut rate_limiter = RateLimiter::builder(10, Duration::from_secs(1));
69        rate_limiter.set_window(120_000); // 120 seconds
70
71        assert_eq!(rate_limiter.window, 120_000);
72    }
73
74    #[test]
75    fn can_make_request_within_limit_test() {
76        let mut rate_limiter = RateLimiter::builder(10, Duration::from_secs(1));
77        let current_time = 10_000;
78        rate_limiter.increase();
79        assert_eq!(rate_limiter.request_count, 1);
80        assert!(!rate_limiter.has_limit_exceeded(current_time));
81        assert_eq!(rate_limiter.request_count, 0);
82    }
83
84    #[test]
85    fn can_make_request_exceeding_limit_test() {
86        let mut rate_limiter = RateLimiter::builder(10, Duration::from_secs(1));
87        let current_time = 10_000;
88        for _ in 0..rate_limiter.limit {
89            rate_limiter.increase();
90            assert!(!rate_limiter.has_limit_exceeded(current_time));
91        }
92        rate_limiter.increase();
93        assert!(rate_limiter.has_limit_exceeded(current_time));
94        assert_eq!(rate_limiter.request_count, rate_limiter.limit);
95    }
96
97    #[test]
98    fn can_make_request_after_window_reset_test() {
99        let mut rate_limiter = RateLimiter::builder(10, Duration::from_secs(1));
100        let current_time = 10_000;
101
102        for _ in 0..rate_limiter.limit {
103            rate_limiter.increase();
104            assert!(!rate_limiter.has_limit_exceeded(current_time));
105        }
106        let new_time = current_time + rate_limiter.window + 1;
107        rate_limiter.increase();
108        assert!(!rate_limiter.has_limit_exceeded(new_time));
109        assert_eq!(rate_limiter.request_count, 0);
110        assert_eq!(rate_limiter.last_request_time, new_time);
111    }
112
113    #[test]
114    fn can_make_request_reset_on_new_time_test() {
115        let mut rate_limiter = RateLimiter::builder(10, Duration::from_secs(1));
116        let initial_time = 10_000;
117
118        for _ in 0..rate_limiter.limit {
119            rate_limiter.increase();
120            assert!(!rate_limiter.has_limit_exceeded(initial_time));
121        }
122
123        let later_time = initial_time + rate_limiter.window + 500;
124        rate_limiter.increase();
125        assert!(!rate_limiter.has_limit_exceeded(later_time));
126        dbg!("{}", &rate_limiter);
127        assert_eq!(rate_limiter.request_count, 0);
128        assert_eq!(rate_limiter.last_request_time, later_time);
129    }
130}
saito_core/core/consensus/peers/rate_limiter.rs

saito_core/core/consensus/peers/
rate_limiter.rs
