use std::collections::HashMap;
use std::fmt;
use std::sync::Arc;
use std::time::{Duration, Instant};
use crate::clock::{Clock, SystemClock};
use parking_lot::Mutex;
const DEFAULT_MAX_TRACKED: usize = 65_536;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RateLimitRejection {
pub limit: u64,
pub window: Duration,
}
impl fmt::Display for RateLimitRejection {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
f,
"rate limit exceeded: {} requests per {}s window",
self.limit,
self.window.as_secs()
)
}
}
impl std::error::Error for RateLimitRejection {}
#[derive(Clone, Copy, Debug)]
struct WindowEntry {
window_start: Instant,
count: u64,
}
pub struct SyncFixedWindowLimiter {
limit: u64,
window: Duration,
max_tracked: usize,
clock: Arc<dyn Clock>,
state: Mutex<HashMap<String, WindowEntry>>,
}
impl fmt::Debug for SyncFixedWindowLimiter {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("SyncFixedWindowLimiter")
.field("limit", &self.limit)
.field("window", &self.window)
.field("max_tracked", &self.max_tracked)
.finish_non_exhaustive()
}
}
impl SyncFixedWindowLimiter {
pub fn new(limit: u64, window: Duration) -> Self {
Self::with_clock(Arc::new(SystemClock), limit, window)
}
pub fn with_clock(clock: Arc<dyn Clock>, limit: u64, window: Duration) -> Self {
Self {
limit: limit.max(1),
window: if window.is_zero() {
Duration::from_nanos(1)
} else {
window
},
max_tracked: DEFAULT_MAX_TRACKED,
clock,
state: Mutex::new(HashMap::new()),
}
}
pub fn with_max_tracked(mut self, max_tracked: usize) -> Self {
self.max_tracked = max_tracked.max(1);
self
}
pub fn check(&self, identifier: &str) -> Result<(), RateLimitRejection> {
self.admit_at(identifier, self.clock.now())
}
pub fn peek(&self, identifier: &str) -> Option<(u64, u64)> {
self.peek_at(identifier, self.clock.now())
}
pub fn peek_at(&self, identifier: &str, now: Instant) -> Option<(u64, u64)> {
let state = self.state.lock();
state.get(identifier).map(|entry| {
if now.duration_since(entry.window_start) >= self.window {
(0, self.limit)
} else {
(entry.count, self.limit.saturating_sub(entry.count))
}
})
}
pub fn admit_at(&self, identifier: &str, now: Instant) -> Result<(), RateLimitRejection> {
let rejection = RateLimitRejection {
limit: self.limit,
window: self.window,
};
let mut state = self.state.lock();
let entry = match state.get_mut(identifier) {
Some(entry) => entry,
None => {
Self::make_room(&mut state, now, self.window, self.max_tracked);
state.entry(identifier.to_string()).or_insert(WindowEntry {
window_start: now,
count: 0,
})
}
};
if now.duration_since(entry.window_start) >= self.window {
*entry = WindowEntry {
window_start: now,
count: 0,
};
}
if entry.count >= self.limit {
return Err(rejection);
}
entry.count += 1;
Ok(())
}
fn make_room(
state: &mut HashMap<String, WindowEntry>,
now: Instant,
window: Duration,
max_tracked: usize,
) {
if state.len() < max_tracked {
return;
}
state.retain(|_, entry| now.duration_since(entry.window_start) < window);
if state.len() >= max_tracked {
let oldest = state
.iter()
.min_by_key(|(_, entry)| entry.window_start)
.map(|(key, _)| key.clone());
if let Some(key) = oldest {
state.remove(&key);
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::clock::MockClock;
fn t0() -> Instant {
Instant::now()
}
#[test]
fn test_admit_within_limit_and_reject_beyond() {
let limiter = SyncFixedWindowLimiter::new(2, Duration::from_secs(60));
assert!(limiter.check("ip1").is_ok());
assert!(limiter.check("ip1").is_ok());
let err = limiter.check("ip1").unwrap_err();
assert_eq!(
err,
RateLimitRejection {
limit: 2,
window: Duration::from_secs(60)
}
);
assert!(limiter.check("ip1").is_err());
assert!(limiter.check("ip2").is_ok());
}
#[test]
fn test_peek_is_non_consuming() {
let limiter = SyncFixedWindowLimiter::new(3, Duration::from_secs(60));
assert_eq!(limiter.peek("nope"), None);
assert!(limiter.check("a").is_ok());
assert!(limiter.check("a").is_ok());
assert!(limiter.check("a").is_ok());
let err = limiter.check("a").unwrap_err();
assert_eq!(err.limit, 3);
assert_eq!(limiter.peek("a"), Some((3, 0)));
assert_eq!(limiter.peek("a"), Some((3, 0)));
assert!(limiter.check("a").is_err());
assert_eq!(limiter.peek("b"), None);
}
#[test]
fn test_window_roll_resets_count() {
let clock = Arc::new(MockClock::with_instant(t0(), 0));
let limiter = SyncFixedWindowLimiter::with_clock(clock.clone(), 1, Duration::from_secs(60));
assert!(limiter.check("ip1").is_ok());
assert!(limiter.check("ip1").is_err());
clock.advance(Duration::from_secs(61));
assert!(
limiter.check("ip1").is_ok(),
"allow should resume after the window rolls"
);
}
#[test]
fn test_expired_entries_purged_when_full() {
let clock = Arc::new(MockClock::with_instant(t0(), 0));
let limiter = SyncFixedWindowLimiter::with_clock(clock.clone(), 2, Duration::from_secs(60))
.with_max_tracked(2);
assert!(limiter.check("a").is_ok());
assert!(limiter.check("b").is_ok());
clock.advance(Duration::from_secs(61));
assert!(limiter.check("c").is_ok());
assert!(
limiter.check("d").is_ok(),
"there is still room after expired entries are cleared"
);
assert!(limiter.check("c").is_ok());
assert!(limiter.check("d").is_ok());
assert!(
limiter.check("c").is_err(),
"c's budget has already been consumed twice"
);
}
#[test]
fn test_oldest_entry_evicted_when_full_in_window() {
let clock = Arc::new(MockClock::with_instant(t0(), 0));
let limiter =
SyncFixedWindowLimiter::with_clock(clock.clone(), 10, Duration::from_secs(60))
.with_max_tracked(2);
assert!(limiter.check("a").is_ok());
clock.advance(Duration::from_secs(1));
assert!(limiter.check("b").is_ok());
clock.advance(Duration::from_secs(1));
assert!(limiter.check("c").is_ok());
assert_eq!(
limiter.state.lock().len(),
2,
"the identifier table stays bounded"
);
assert!(limiter.check("a").is_ok());
}
#[test]
fn test_limit_below_one_raised_to_one() {
let limiter = SyncFixedWindowLimiter::new(0, Duration::from_secs(60));
assert!(limiter.check("ip1").is_ok());
assert!(limiter.check("ip1").is_err());
}
#[test]
fn test_rejection_display() {
let err = RateLimitRejection {
limit: 120,
window: Duration::from_secs(60),
};
assert_eq!(
err.to_string(),
"rate limit exceeded: 120 requests per 60s window"
);
}
}