use super::traits::{Limiter, validate_cost};
use crate::clock::{Clock, SystemClock};
use crate::error::LimiteronError;
use async_trait::async_trait;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
pub struct FixedWindowLimiter {
window_size: Duration,
max_requests: u64,
count: AtomicU64,
window_start: AtomicU64,
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,
count: AtomicU64::new(0),
window_start: AtomicU64::new(now),
clock,
}
}
fn check_and_reset_window(&self) {
let now = self.clock.unix_timestamp_nanos();
let window_size_nanos = self.window_size.as_nanos() as u64;
loop {
let current_start = self.window_start.load(Ordering::Acquire);
let window_end = current_start.saturating_add(window_size_nanos);
if now < window_end {
break;
}
let elapsed = now.saturating_sub(current_start);
let windows_passed = elapsed / window_size_nanos;
let new_start = current_start.saturating_add(windows_passed * window_size_nanos);
match self.window_start.compare_exchange(
current_start,
new_start,
Ordering::Release,
Ordering::Relaxed,
) {
Ok(_) => {
self.count.store(0, Ordering::Release);
break;
}
Err(_) => continue,
}
}
}
#[cfg(test)]
fn get_count(&self) -> u64 {
self.check_and_reset_window();
self.count.load(Ordering::Acquire)
}
}
#[async_trait]
impl Limiter for FixedWindowLimiter {
async fn allow(&self, cost: u64) -> Result<bool, LimiteronError> {
let cost = validate_cost(cost)?;
self.check_and_reset_window();
loop {
let current = self.count.load(Ordering::Acquire);
if current + cost > self.max_requests {
return Ok(false);
}
match self.count.compare_exchange(
current,
current + cost,
Ordering::Release,
Ordering::Relaxed,
) {
Ok(_) => return Ok(true),
Err(_) => continue,
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::clock::MockClock;
#[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());
}
}