arcly_http/resilience/
rate_limit.rs1use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
13use std::time::{SystemTime, UNIX_EPOCH};
14
15use crate::auth::guards::Guard;
16use crate::web::{Error, RequestContext};
17
18pub struct RateLimit {
21 state: AtomicU64, 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 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 assert!(rl.admit_at(1_060));
93 assert!(rl.admit_at(1_061));
94 assert!(!rl.admit_at(1_062));
95 }
96}