use super::traits::{Limiter, RateLimitSnapshot, validate_cost};
use crate::clock::{Clock, SystemClock};
use crate::error::LimiteronError;
use async_trait::async_trait;
use parking_lot::Mutex;
use std::sync::Arc;
use std::time::Duration;
#[derive(Debug)]
struct WindowState {
count: u64,
window_start: u64,
}
pub struct FixedWindowLimiter {
window_size: Duration,
max_requests: u64,
state: Mutex<WindowState>,
clock: Arc<dyn Clock>,
}
impl FixedWindowLimiter {
pub fn new(window_size: Duration, max_requests: u64) -> Self {
Self::with_clock(window_size, max_requests, Arc::new(SystemClock))
}
pub fn with_clock(window_size: Duration, max_requests: u64, clock: Arc<dyn Clock>) -> Self {
let now = clock.unix_timestamp_nanos();
Self {
window_size,
max_requests,
state: Mutex::new(WindowState {
count: 0,
window_start: now,
}),
clock,
}
}
fn advance_expired_window(&self, state: &mut WindowState) {
let now = self.clock.unix_timestamp_nanos();
let window_size_nanos = self.window_size.as_nanos().max(1) as u64;
let window_end = state.window_start.saturating_add(window_size_nanos);
if now >= window_end {
let elapsed = now.saturating_sub(state.window_start);
let windows_passed = elapsed / window_size_nanos;
state.window_start = state
.window_start
.saturating_add(windows_passed.saturating_mul(window_size_nanos));
state.count = 0;
}
}
}
#[async_trait]
impl Limiter for FixedWindowLimiter {
async fn allow(&self, cost: u64) -> Result<bool, LimiteronError> {
let cost = validate_cost(cost)?;
let mut state = self.state.lock();
self.advance_expired_window(&mut state);
if cost > self.max_requests.saturating_sub(state.count) {
return Ok(false);
}
state.count += cost;
Ok(true)
}
async fn peek(&self, cost: u64) -> Result<RateLimitSnapshot, LimiteronError> {
let cost = validate_cost(cost)?;
let mut state = self.state.lock();
self.advance_expired_window(&mut state);
let remaining = self.max_requests.saturating_sub(state.count);
let reset_secs = self.window_reset_secs_locked(&state);
let _ = cost;
Ok(RateLimitSnapshot {
limit: self.max_requests,
remaining,
reset_secs,
})
}
async fn remaining(&self) -> Result<RateLimitSnapshot, LimiteronError> {
let mut state = self.state.lock();
self.advance_expired_window(&mut state);
let reset_secs = self.window_reset_secs_locked(&state);
Ok(RateLimitSnapshot {
limit: self.max_requests,
remaining: self.max_requests.saturating_sub(state.count),
reset_secs,
})
}
}
impl FixedWindowLimiter {
fn window_reset_secs_locked(&self, state: &WindowState) -> u64 {
let now_ns = self.clock.unix_timestamp_nanos();
let window_ns = self.window_size.as_nanos() as u64;
let window_end_ns = state.window_start.saturating_add(window_ns);
if window_end_ns <= now_ns {
return 0;
}
(window_end_ns - now_ns).div_ceil(1_000_000_000)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::clock::MockClock;
use std::sync::atomic::{AtomicU64, Ordering};
impl FixedWindowLimiter {
fn get_count(&self) -> u64 {
let mut state = self.state.lock();
self.advance_expired_window(&mut state);
state.count
}
}
#[tokio::test]
async fn test_fixed_window_basic() {
let limiter = FixedWindowLimiter::new(Duration::from_secs(60), 10);
for _ in 0..5 {
assert!(limiter.allow(1).await.unwrap());
}
assert_eq!(limiter.get_count(), 5);
}
#[tokio::test]
async fn test_fixed_window_exceed() {
let limiter = FixedWindowLimiter::new(Duration::from_secs(60), 3);
assert!(limiter.allow(1).await.unwrap());
assert!(limiter.allow(1).await.unwrap());
assert!(limiter.allow(1).await.unwrap());
assert!(!limiter.allow(1).await.unwrap());
}
#[tokio::test]
async fn test_fixed_window_with_mock_clock() {
let mock_clock = Arc::new(MockClock::new());
let clock: Arc<dyn Clock> = mock_clock.clone();
let limiter = FixedWindowLimiter::with_clock(Duration::from_secs(10), 5, clock);
for _ in 0..5 {
assert!(limiter.allow(1).await.unwrap());
}
assert!(!limiter.allow(1).await.unwrap());
mock_clock.advance(Duration::from_secs(11));
assert!(limiter.allow(1).await.unwrap());
}
#[tokio::test]
async fn test_fixed_window_no_lost_increments_across_rollover() {
const MAX: u64 = 64;
const THREADS: usize = 64;
const ROUNDS: usize = 25;
for round in 0..ROUNDS {
let mock = Arc::new(MockClock::new());
let clock: Arc<dyn Clock> = mock.clone();
let limiter = Arc::new(FixedWindowLimiter::with_clock(
Duration::from_secs(10),
MAX,
clock,
));
for _ in 0..MAX {
assert!(
limiter.allow(1).await.unwrap(),
"round {round} fill phase false rejection"
);
}
assert!(!limiter.allow(1).await.unwrap());
mock.advance(Duration::from_secs(11));
let allowed = Arc::new(AtomicU64::new(0));
let barrier = Arc::new(tokio::sync::Barrier::new(THREADS));
let mut handles = Vec::with_capacity(THREADS);
for _ in 0..THREADS {
let limiter = Arc::clone(&limiter);
let barrier = Arc::clone(&barrier);
let allowed = Arc::clone(&allowed);
handles.push(tokio::spawn(async move {
barrier.wait().await;
if limiter.allow(1).await.unwrap() {
allowed.fetch_add(1, Ordering::SeqCst);
}
}));
}
for h in handles {
h.await.unwrap();
}
assert_eq!(
allowed.load(Ordering::SeqCst),
MAX,
"round {round} new-window allowance deviates"
);
assert_eq!(
limiter.get_count(),
MAX,
"round {round} ledger mismatch: the increment was swallowed by the flip reset"
);
assert!(
!limiter.allow(1).await.unwrap(),
"round {round} flip race causes over-issuance"
);
}
}
}