api_rate_limiter/
limiter.rs

1use std::sync::{Arc, Mutex};
2use std::time::{Duration, Instant};
3
4pub struct RateLimiter {
5    capacity: u32,
6    tokens: u32,
7    last_refill: Instant,
8    refill_rate: Duration,
9}
10
11impl RateLimiter {
12    pub fn new(capacity: u32, refill_rate: Duration) -> Arc<Mutex<Self>> {
13        Arc::new(Mutex::new(Self {
14            capacity,
15            tokens: capacity,
16            last_refill: Instant::now(),
17            refill_rate,
18        }))
19    }
20
21    pub fn allow(&mut self) -> bool {
22        self.refill();
23
24        if self.tokens > 0 {
25            self.tokens -= 1;
26            true
27        } else {
28            false
29        }
30    }
31
32    fn refill(&mut self) {
33        let now = Instant::now();
34        let time_elapsed = now.duration_since(self.last_refill);
35
36        let new_tokens = (time_elapsed.as_secs_f64() / self.refill_rate.as_secs_f64()) * self.capacity as f64;
37        self.tokens = (self.tokens as f64 + new_tokens).min(self.capacity as f64) as u32;
38        self.last_refill = now;
39    }
40}
41
42#[cfg(test)]
43mod tests {
44    use super::*;
45    use std::thread::sleep;
46
47    #[test]
48    fn test_rate_limiter() {
49        let limiter = RateLimiter::new(5, Duration::from_secs(1));
50        let mut limiter = limiter.lock().unwrap();
51
52        for _ in 0..5 {
53            assert!(limiter.allow());
54        }
55
56        assert!(!limiter.allow());
57
58        sleep(Duration::from_secs(1));
59
60        assert!(limiter.allow());
61    }
62}
api_rate_limiter/limiter.rs

api_rate_limiter/
limiter.rs
