rate_limiter/
local.rs

1use crate::{
2    AcquireResult, PeekResult, RateLimitError, RateLimitPolicy, RateLimiter, TryAcquireResult,
3};
4use async_trait::async_trait;
5use std::collections::HashMap;
6use std::sync::Arc;
7use std::time::{Duration, Instant};
8use tokio::sync::Mutex;
9
10#[derive(Debug, Clone)]
11pub struct LocalRateLimiter {
12    interval: Duration,
13    state: Arc<Mutex<HashMap<String, Instant>>>,
14}
15
16impl LocalRateLimiter {
17    pub fn new(policy: RateLimitPolicy) -> Result<Self, RateLimitError> {
18        let interval = policy.interval()?;
19        Ok(Self {
20            interval,
21            state: Arc::new(Mutex::new(HashMap::new())),
22        })
23    }
24
25    async fn reserve_slot(&self, key: &str) -> (Duration, Instant) {
26        let now = Instant::now();
27        let mut state = self.state.lock().await;
28        let next_allowed = state.get(key).copied().unwrap_or(now);
29        let wait_duration = next_allowed.saturating_duration_since(now);
30        let acquired_at = now + wait_duration;
31        state.insert(key.to_owned(), acquired_at + self.interval);
32        (wait_duration, acquired_at)
33    }
34
35    async fn try_reserve_slot(&self, key: &str) -> TryAcquireResult {
36        let now = Instant::now();
37        let mut state = self.state.lock().await;
38        let next_allowed = state.get(key).copied().unwrap_or(now);
39        let wait_duration = next_allowed.saturating_duration_since(now);
40        if !wait_duration.is_zero() {
41            return TryAcquireResult::Limited { wait_duration };
42        }
43
44        state.insert(key.to_owned(), now + self.interval);
45        TryAcquireResult::Acquired
46    }
47
48    async fn inspect(&self, key: &str) -> PeekResult {
49        let now = Instant::now();
50        let state = self.state.lock().await;
51        let wait_duration = state
52            .get(key)
53            .copied()
54            .unwrap_or(now)
55            .saturating_duration_since(now);
56        PeekResult {
57            allowed: wait_duration.is_zero(),
58            wait_duration,
59        }
60    }
61}
62
63#[async_trait]
64impl RateLimiter for LocalRateLimiter {
65    async fn acquire(&self, key: &str) -> Result<AcquireResult, RateLimitError> {
66        let (waited, _) = self.reserve_slot(key).await;
67        if !waited.is_zero() {
68            tokio::time::sleep(waited).await;
69        }
70        Ok(AcquireResult {
71            key: key.to_owned(),
72            waited,
73        })
74    }
75
76    async fn try_acquire(&self, key: &str) -> Result<TryAcquireResult, RateLimitError> {
77        Ok(self.try_reserve_slot(key).await)
78    }
79
80    async fn peek(&self, key: &str) -> Result<PeekResult, RateLimitError> {
81        Ok(self.inspect(key).await)
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use super::LocalRateLimiter;
88    use crate::{RateLimitPolicy, RateLimiter, TryAcquireResult};
89    use std::time::Duration;
90
91    fn assert_duration_at_least(actual: Duration, expected: Duration) {
92        assert!(actual >= expected.saturating_sub(Duration::from_millis(1)));
93        assert!(actual <= expected);
94    }
95
96    #[tokio::test(start_paused = true)]
97    async fn same_key_waits_on_second_acquire() {
98        let limiter =
99            LocalRateLimiter::new(RateLimitPolicy::min_interval(Duration::from_secs(3)).unwrap())
100                .unwrap();
101
102        let first = limiter.acquire("same").await.unwrap();
103        assert_eq!(first.waited, Duration::ZERO);
104
105        let limiter_clone = limiter.clone();
106        let task = tokio::spawn(async move { limiter_clone.acquire("same").await.unwrap() });
107
108        tokio::task::yield_now().await;
109        assert!(!task.is_finished());
110
111        tokio::time::advance(Duration::from_secs(3)).await;
112        let second = task.await.unwrap();
113        assert_duration_at_least(second.waited, Duration::from_secs(3));
114    }
115
116    #[tokio::test(start_paused = true)]
117    async fn different_keys_do_not_block_each_other() {
118        let limiter =
119            LocalRateLimiter::new(RateLimitPolicy::min_interval(Duration::from_secs(2)).unwrap())
120                .unwrap();
121
122        limiter.acquire("a").await.unwrap();
123        let second = limiter.acquire("b").await.unwrap();
124        assert_eq!(second.waited, Duration::ZERO);
125    }
126
127    #[tokio::test(start_paused = true)]
128    async fn try_acquire_reports_wait_duration() {
129        let limiter =
130            LocalRateLimiter::new(RateLimitPolicy::min_interval(Duration::from_secs(5)).unwrap())
131                .unwrap();
132
133        assert!(matches!(
134            limiter.try_acquire("same").await.unwrap(),
135            TryAcquireResult::Acquired
136        ));
137
138        match limiter.try_acquire("same").await.unwrap() {
139            TryAcquireResult::Limited { wait_duration } => {
140                assert_duration_at_least(wait_duration, Duration::from_secs(5));
141            }
142            TryAcquireResult::Acquired => panic!("expected limited result"),
143        }
144    }
145
146    #[tokio::test(start_paused = true)]
147    async fn peek_does_not_consume_slot() {
148        let limiter =
149            LocalRateLimiter::new(RateLimitPolicy::min_interval(Duration::from_secs(4)).unwrap())
150                .unwrap();
151
152        let peek = limiter.peek("diag").await.unwrap();
153        assert!(peek.allowed);
154        assert_eq!(peek.wait_duration, Duration::ZERO);
155
156        let acquire = limiter.acquire("diag").await.unwrap();
157        assert_eq!(acquire.waited, Duration::ZERO);
158    }
159}
rate_limiter/local.rs

rate_limiter/
local.rs
