use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Mutex;
use tokio::time::Instant;
#[derive(Debug, Clone)]
pub(crate) struct RateLimiter {
interval: Option<Duration>,
next_allowed: Arc<Mutex<HashMap<String, Instant>>>,
}
impl RateLimiter {
pub(crate) fn per_second(requests_per_second: u32) -> Self {
let interval = (requests_per_second > 0)
.then(|| Duration::from_secs_f64(1.0 / f64::from(requests_per_second)));
Self {
interval,
next_allowed: Arc::new(Mutex::new(HashMap::new())),
}
}
pub(crate) async fn acquire(&self, origin: &str) {
let now = Instant::now();
let wait_until = {
let mut next_allowed = self.next_allowed.lock().await;
let slot = match next_allowed.get(origin) {
Some(&prev) if prev > now => prev,
_ => now,
};
let next = self.interval.map_or(slot, |interval| slot + interval);
next_allowed.insert(origin.to_string(), next);
slot
};
if wait_until > now {
tokio::time::sleep_until(wait_until).await;
}
}
pub(crate) async fn backoff(&self, origin: &str, delay: Duration) {
let blocked_until = Instant::now() + delay;
let mut next_allowed = self.next_allowed.lock().await;
next_allowed
.entry(origin.to_string())
.and_modify(|next| *next = (*next).max(blocked_until))
.or_insert(blocked_until);
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn spaces_same_origin_requests_by_the_interval() {
let limiter = RateLimiter::per_second(100);
let start = Instant::now();
for _ in 0..5 {
limiter.acquire("https://skill.example").await;
}
assert!(
start.elapsed() >= Duration::from_millis(35),
"five paced requests should take at least ~40ms, took {:?}",
start.elapsed()
);
}
#[tokio::test]
async fn tracks_origins_independently() {
let limiter = RateLimiter::per_second(1);
let start = Instant::now();
limiter.acquire("https://skill.example").await;
limiter.acquire("https://stats.example").await;
assert!(start.elapsed() < Duration::from_millis(200));
}
#[tokio::test]
async fn zero_rate_never_delays() {
let limiter = RateLimiter::per_second(0);
let start = Instant::now();
for _ in 0..1000 {
limiter.acquire("https://skill.example").await;
}
assert!(start.elapsed() < Duration::from_millis(200));
}
#[tokio::test]
async fn backoff_delays_every_request_to_the_same_origin() {
let limiter = RateLimiter::per_second(0);
limiter
.backoff("https://skill.example", Duration::from_millis(50))
.await;
let start = Instant::now();
limiter.acquire("https://skill.example").await;
assert!(
start.elapsed() >= Duration::from_millis(45),
"backoff should delay requests even when normal pacing is disabled"
);
}
}