Skip to main content

arcly_http/resilience/
rate_limit.rs

1//! Lock-free fixed-window rate limiter.
2//!
3//! Lives under `resilience` (not `auth/guards`) because rate limiting is a
4//! load-shedding concern, independent of who the caller is. It still
5//! implements [`Guard`] so it composes with auth guards at handler level:
6//!
7//! ```ignore
8//! static PUBLIC_RATE: RateLimit = RateLimit::new(200, 60);
9//! PUBLIC_RATE.check(&ctx)?;
10//! ```
11
12use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
13use std::time::{SystemTime, UNIX_EPOCH};
14
15use crate::auth::guards::Guard;
16use crate::web::{Error, RequestContext};
17
18/// Per-instance fixed-window limiter. Zero locks: a `(window_start, count)` pair
19/// is packed into a single `AtomicU64` and updated via CAS.
20pub struct RateLimit {
21    state: AtomicU64, // high 32 bits = window-start seconds, low 32 = count
22    window_secs: u32,
23    max_per_window: u32,
24}
25
26impl RateLimit {
27    pub const fn new(max_per_window: u32, window_secs: u32) -> Self {
28        Self {
29            state: AtomicU64::new(0),
30            window_secs,
31            max_per_window,
32        }
33    }
34
35    fn now_secs() -> u32 {
36        SystemTime::now()
37            .duration_since(UNIX_EPOCH)
38            .map(|d| d.as_secs() as u32)
39            .unwrap_or(0)
40    }
41}
42
43impl RateLimit {
44    /// CAS admission core, parameterised by the clock for testability.
45    /// `true` = admitted, `false` = over the window limit.
46    fn admit_at(&self, now: u32) -> bool {
47        loop {
48            let cur = self.state.load(Relaxed);
49            let (start, count) = ((cur >> 32) as u32, cur as u32);
50            let (new_start, new_count) = if now.saturating_sub(start) >= self.window_secs {
51                (now, 1)
52            } else {
53                (start, count.saturating_add(1))
54            };
55            if new_count > self.max_per_window {
56                return false;
57            }
58            let next = ((new_start as u64) << 32) | (new_count as u64);
59            if self
60                .state
61                .compare_exchange(cur, next, Relaxed, Relaxed)
62                .is_ok()
63            {
64                return true;
65            }
66        }
67    }
68}
69
70impl Guard for RateLimit {
71    fn check(&self, _ctx: &RequestContext) -> Result<(), Error> {
72        if self.admit_at(Self::now_secs()) {
73            Ok(())
74        } else {
75            metrics::counter!("rate_limited_total").increment(1);
76            Err(Error::TooManyRequests)
77        }
78    }
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84
85    #[test]
86    fn enforces_limit_and_resets_on_window_roll() {
87        let rl = RateLimit::new(2, 60);
88        assert!(rl.admit_at(1_000));
89        assert!(rl.admit_at(1_000));
90        assert!(!rl.admit_at(1_030), "third hit in the window is rejected");
91        // Window rolls: counter resets.
92        assert!(rl.admit_at(1_060));
93        assert!(rl.admit_at(1_061));
94        assert!(!rl.admit_at(1_062));
95    }
96}