use std::collections::VecDeque;
use std::time::{Duration, Instant};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RateLimitConfig {
pub max_per_minute: u32,
pub max_per_hour: u32,
}
impl Default for RateLimitConfig {
fn default() -> Self {
RateLimitConfig {
max_per_minute: 25,
max_per_hour: 250,
}
}
}
impl RateLimitConfig {
pub fn disabled() -> Self {
RateLimitConfig {
max_per_minute: u32::MAX,
max_per_hour: u32::MAX,
}
}
}
pub(crate) struct RateLimiter {
config: RateLimitConfig,
minute_window: VecDeque<Instant>,
hour_window: VecDeque<Instant>,
}
impl RateLimiter {
pub(crate) fn new(config: RateLimitConfig) -> Self {
RateLimiter {
config,
minute_window: VecDeque::new(),
hour_window: VecDeque::new(),
}
}
pub(crate) fn acquire(&mut self) -> Duration {
let now = Instant::now();
if let Some(minute_cutoff) = now.checked_sub(Duration::from_secs(60)) {
while self
.minute_window
.front()
.is_some_and(|&t| t < minute_cutoff)
{
self.minute_window.pop_front();
}
}
if let Some(hour_cutoff) = now.checked_sub(Duration::from_secs(3600)) {
while self.hour_window.front().is_some_and(|&t| t < hour_cutoff) {
self.hour_window.pop_front();
}
}
let mut wait = Duration::ZERO;
if self.minute_window.len() >= self.config.max_per_minute as usize
&& let Some(&oldest) = self.minute_window.front()
{
let minute_wait = (oldest + Duration::from_secs(60)).saturating_duration_since(now);
if minute_wait > wait {
wait = minute_wait;
}
}
if self.hour_window.len() >= self.config.max_per_hour as usize
&& let Some(&oldest) = self.hour_window.front()
{
let hour_wait = (oldest + Duration::from_secs(3600)).saturating_duration_since(now);
if hour_wait > wait {
wait = hour_wait;
}
}
let request_time = now + wait;
self.minute_window.push_back(request_time);
self.hour_window.push_back(request_time);
wait
}
}
#[cfg(test)]
#[cfg_attr(coverage_nightly, coverage(off))]
mod tests {
use super::*;
#[test]
fn test_rate_limit_config_default() {
let config = RateLimitConfig::default();
assert_eq!(config.max_per_minute, 25);
assert_eq!(config.max_per_hour, 250);
}
#[test]
fn test_rate_limit_config_disabled() {
let config = RateLimitConfig::disabled();
assert_eq!(config.max_per_minute, u32::MAX);
assert_eq!(config.max_per_hour, u32::MAX);
}
#[test]
fn test_rate_limit_config_custom() {
let config = RateLimitConfig {
max_per_minute: 10,
max_per_hour: 100,
};
assert_eq!(config.max_per_minute, 10);
assert_eq!(config.max_per_hour, 100);
}
#[test]
fn test_rate_limit_config_clone() {
let config = RateLimitConfig::default();
let cloned = config.clone();
assert_eq!(config, cloned);
}
#[test]
fn test_rate_limit_config_debug() {
let config = RateLimitConfig::default();
let debug = format!("{:?}", config);
assert!(debug.contains("25"));
assert!(debug.contains("250"));
}
#[test]
fn test_rate_limit_config_equality() {
let a = RateLimitConfig::default();
let b = RateLimitConfig::default();
let c = RateLimitConfig::disabled();
assert_eq!(a, b);
assert_ne!(a, c);
}
#[test]
fn test_rate_limiter_acquire_within_limits() {
let config = RateLimitConfig {
max_per_minute: 5,
max_per_hour: 100,
};
let mut limiter = RateLimiter::new(config);
for _ in 0..5 {
let wait = limiter.acquire();
assert_eq!(wait, Duration::ZERO);
}
}
#[test]
fn test_rate_limiter_acquire_at_minute_limit() {
let config = RateLimitConfig {
max_per_minute: 3,
max_per_hour: 100,
};
let mut limiter = RateLimiter::new(config);
for _ in 0..3 {
let wait = limiter.acquire();
assert_eq!(wait, Duration::ZERO);
}
let wait = limiter.acquire();
assert!(wait > Duration::ZERO);
assert!(wait <= Duration::from_secs(60));
}
#[test]
fn test_rate_limiter_acquire_at_hour_limit() {
let config = RateLimitConfig {
max_per_minute: 100,
max_per_hour: 3,
};
let mut limiter = RateLimiter::new(config);
for _ in 0..3 {
let wait = limiter.acquire();
assert_eq!(wait, Duration::ZERO);
}
let wait = limiter.acquire();
assert!(wait > Duration::ZERO);
assert!(wait <= Duration::from_secs(3600));
}
#[test]
fn test_rate_limiter_disabled_no_wait() {
let config = RateLimitConfig::disabled();
let mut limiter = RateLimiter::new(config);
for _ in 0..1000 {
let wait = limiter.acquire();
assert_eq!(wait, Duration::ZERO);
}
}
#[test]
fn test_rate_limiter_records_future_timestamps() {
let config = RateLimitConfig {
max_per_minute: 2,
max_per_hour: 100,
};
let mut limiter = RateLimiter::new(config);
limiter.acquire();
limiter.acquire();
let wait = limiter.acquire();
assert!(wait > Duration::ZERO);
let wait2 = limiter.acquire();
assert!(wait2 > Duration::ZERO);
}
#[test]
fn test_rate_limiter_minute_window_takes_precedence() {
let config = RateLimitConfig {
max_per_minute: 2,
max_per_hour: 10,
};
let mut limiter = RateLimiter::new(config);
limiter.acquire();
limiter.acquire();
let wait = limiter.acquire();
assert!(wait > Duration::ZERO);
assert!(wait <= Duration::from_secs(60));
}
}