use std::{
collections::HashMap,
hash::Hash,
num::{NonZeroU32, NonZeroUsize},
str::FromStr,
sync::{Arc, Mutex, PoisonError, Weak},
time::{Duration, Instant},
};
use governor::{
DefaultDirectRateLimiter, Quota, RateLimiter,
clock::{Clock as _, DefaultClock},
};
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
pub enum BoundedLimiterError {
#[error("rate limit exceeded for key")]
RateLimited,
}
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
pub enum BoundedLimiterDeny {
#[error("rate limit exceeded; retry after {0:?}")]
RateLimited(Duration),
#[error("tracked-key capacity is full")]
CapacityFull,
}
#[non_exhaustive]
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum KeyEvictionPolicy {
#[default]
EvictLru,
RejectNew,
}
impl FromStr for KeyEvictionPolicy {
type Err = ();
fn from_str(value: &str) -> Result<Self, Self::Err> {
match value {
"evict_lru" => Ok(Self::EvictLru),
"reject_new" => Ok(Self::RejectNew),
_ => Err(()),
}
}
}
struct Entry {
limiter: DefaultDirectRateLimiter,
last_seen: Instant,
}
struct Inner<K: Eq + Hash + Clone> {
map: Mutex<HashMap<K, Entry>>,
quota: Quota,
max_tracked_keys: usize,
idle_eviction: Duration,
key_eviction_policy: KeyEvictionPolicy,
}
#[allow(
missing_debug_implementations,
reason = "wraps governor RateLimiter which has no Debug impl"
)]
pub struct BoundedKeyedLimiter<K: Eq + Hash + Clone> {
inner: Arc<Inner<K>>,
}
impl<K: Eq + Hash + Clone> Clone for BoundedKeyedLimiter<K> {
fn clone(&self) -> Self {
Self {
inner: Arc::clone(&self.inner),
}
}
}
impl<K: Eq + Hash + Clone + Send + Sync + 'static> BoundedKeyedLimiter<K> {
#[must_use]
pub(crate) fn new(
quota: Quota,
max_tracked_keys: NonZeroUsize,
idle_eviction: Duration,
) -> Self {
Self::new_with_policy(
quota,
max_tracked_keys,
idle_eviction,
KeyEvictionPolicy::default(),
)
}
#[must_use]
pub(crate) fn new_with_policy(
quota: Quota,
max_tracked_keys: NonZeroUsize,
idle_eviction: Duration,
key_eviction_policy: KeyEvictionPolicy,
) -> Self {
let inner = Arc::new(Inner {
map: Mutex::new(HashMap::new()),
quota,
max_tracked_keys: max_tracked_keys.get(),
idle_eviction,
key_eviction_policy,
});
Self::spawn_prune_task(&inner);
Self { inner }
}
#[must_use]
pub fn with_per_minute(
requests_per_minute: u32,
max_tracked_keys: usize,
idle_eviction: Duration,
) -> Self {
let rate = NonZeroU32::new(requests_per_minute.max(1)).unwrap_or(NonZeroU32::MIN);
let max_tracked_keys = NonZeroUsize::new(max_tracked_keys).unwrap_or(NonZeroUsize::MIN);
Self::new(Quota::per_minute(rate), max_tracked_keys, idle_eviction)
}
#[must_use]
pub fn with_per_minute_and_policy(
requests_per_minute: u32,
max_tracked_keys: usize,
idle_eviction: Duration,
key_eviction_policy: KeyEvictionPolicy,
) -> Self {
let rate = NonZeroU32::new(requests_per_minute.max(1)).unwrap_or(NonZeroU32::MIN);
let max_tracked_keys = NonZeroUsize::new(max_tracked_keys).unwrap_or(NonZeroUsize::MIN);
Self::new_with_policy(
Quota::per_minute(rate),
max_tracked_keys,
idle_eviction,
key_eviction_policy,
)
}
#[must_use]
pub fn with_per_second(
requests_per_second: u32,
max_tracked_keys: usize,
idle_eviction: Duration,
) -> Self {
let rate = NonZeroU32::new(requests_per_second.max(1)).unwrap_or(NonZeroU32::MIN);
let max_tracked_keys = NonZeroUsize::new(max_tracked_keys).unwrap_or(NonZeroUsize::MIN);
Self::new(Quota::per_second(rate), max_tracked_keys, idle_eviction)
}
#[must_use]
pub fn with_per_second_and_policy(
requests_per_second: u32,
max_tracked_keys: usize,
idle_eviction: Duration,
key_eviction_policy: KeyEvictionPolicy,
) -> Self {
let rate = NonZeroU32::new(requests_per_second.max(1)).unwrap_or(NonZeroU32::MIN);
let max_tracked_keys = NonZeroUsize::new(max_tracked_keys).unwrap_or(NonZeroUsize::MIN);
Self::new_with_policy(
Quota::per_second(rate),
max_tracked_keys,
idle_eviction,
key_eviction_policy,
)
}
fn spawn_prune_task(inner: &Arc<Inner<K>>) {
let Ok(handle) = tokio::runtime::Handle::try_current() else {
return;
};
let weak: Weak<Inner<K>> = Arc::downgrade(inner);
let interval = (inner.idle_eviction / 4).max(Duration::from_mins(1));
handle.spawn(async move {
let mut ticker = tokio::time::interval(interval);
ticker.tick().await;
loop {
ticker.tick().await;
let Some(inner) = weak.upgrade() else {
return;
};
Self::prune_idle(&inner);
}
});
}
fn prune_idle(inner: &Inner<K>) {
let mut guard = inner.map.lock().unwrap_or_else(PoisonError::into_inner);
let cutoff = Instant::now()
.checked_sub(inner.idle_eviction)
.unwrap_or_else(Instant::now);
guard.retain(|_, entry| entry.last_seen >= cutoff);
}
fn evict_lru(map: &mut HashMap<K, Entry>) {
let oldest_key = map
.iter()
.min_by_key(|(_, entry)| entry.last_seen)
.map(|(k, _)| k.clone());
if let Some(key) = oldest_key {
map.remove(&key);
}
}
pub fn check_key(&self, key: &K) -> Result<(), BoundedLimiterError> {
self.check_key_wait(key)
.map_err(|_| BoundedLimiterError::RateLimited)
}
pub fn check_key_wait(&self, key: &K) -> Result<(), Duration> {
let mut guard = self
.inner
.map
.lock()
.unwrap_or_else(PoisonError::into_inner);
let now = Instant::now();
if let Some(entry) = guard.get_mut(key) {
entry.last_seen = now;
return entry
.limiter
.check()
.map_err(|not_until| not_until.wait_time_from(DefaultClock::default().now()));
}
if guard.len() >= self.inner.max_tracked_keys {
let cutoff = now
.checked_sub(self.inner.idle_eviction)
.unwrap_or_else(Instant::now);
guard.retain(|_, entry| entry.last_seen >= cutoff);
if guard.len() >= self.inner.max_tracked_keys {
Self::evict_lru(&mut guard);
}
}
let limiter = RateLimiter::direct(self.inner.quota);
let result = limiter
.check()
.map_err(|not_until| not_until.wait_time_from(DefaultClock::default().now()));
guard.insert(
key.clone(),
Entry {
limiter,
last_seen: now,
},
);
result
}
pub fn check_key_detailed(&self, key: &K) -> Result<(), BoundedLimiterDeny> {
let mut guard = self
.inner
.map
.lock()
.unwrap_or_else(PoisonError::into_inner);
let now = Instant::now();
if let Some(entry) = guard.get_mut(key) {
entry.last_seen = now;
return entry.limiter.check().map_err(|not_until| {
BoundedLimiterDeny::RateLimited(
not_until.wait_time_from(DefaultClock::default().now()),
)
});
}
if guard.len() >= self.inner.max_tracked_keys {
let cutoff = now
.checked_sub(self.inner.idle_eviction)
.unwrap_or_else(Instant::now);
guard.retain(|_, entry| entry.last_seen >= cutoff);
if guard.len() >= self.inner.max_tracked_keys {
match self.inner.key_eviction_policy {
KeyEvictionPolicy::EvictLru => Self::evict_lru(&mut guard),
KeyEvictionPolicy::RejectNew => return Err(BoundedLimiterDeny::CapacityFull),
}
}
}
let limiter = RateLimiter::direct(self.inner.quota);
let result = limiter.check().map_err(|not_until| {
BoundedLimiterDeny::RateLimited(not_until.wait_time_from(DefaultClock::default().now()))
});
guard.insert(
key.clone(),
Entry {
limiter,
last_seen: now,
},
);
result
}
#[must_use]
pub fn len(&self) -> usize {
self.inner
.map
.lock()
.unwrap_or_else(PoisonError::into_inner)
.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
#[cfg(test)]
mod tests {
use std::{
net::IpAddr,
num::{NonZeroU32, NonZeroUsize},
time::Duration,
};
use governor::Quota;
use super::{BoundedKeyedLimiter, BoundedLimiterDeny, BoundedLimiterError, KeyEvictionPolicy};
fn ip(n: u32) -> IpAddr {
IpAddr::from(n.to_be_bytes())
}
fn cap(n: usize) -> NonZeroUsize {
NonZeroUsize::new(n).unwrap_or(NonZeroUsize::MIN)
}
#[test]
fn check_key_wait_existing_key_deny_returns_bounded_wait() {
let quota = Quota::per_minute(NonZeroU32::new(1).unwrap());
let limiter: BoundedKeyedLimiter<IpAddr> =
BoundedKeyedLimiter::new(quota, cap(10), Duration::from_hours(1));
assert!(limiter.check_key_wait(&ip(1)).is_ok(), "burst admits first");
let wait = limiter
.check_key_wait(&ip(1))
.expect_err("second call within the window must deny");
assert!(wait > Duration::ZERO, "wait must be positive, got {wait:?}");
assert!(
wait <= Duration::from_secs(60),
"per-minute quota wait must be <= 60s, got {wait:?}"
);
}
#[test]
fn check_key_wait_new_key_first_check_admits() {
let quota = Quota::per_minute(NonZeroU32::new(1).unwrap());
let limiter: BoundedKeyedLimiter<IpAddr> =
BoundedKeyedLimiter::new(quota, cap(10), Duration::from_hours(1));
for i in 0..5_u32 {
assert!(
limiter.check_key_wait(&ip(i)).is_ok(),
"first check for new key {i} must admit"
);
}
}
#[test]
fn check_key_delegates_to_wait_path() {
let quota = Quota::per_minute(NonZeroU32::new(1).unwrap());
let limiter: BoundedKeyedLimiter<IpAddr> =
BoundedKeyedLimiter::new(quota, cap(10), Duration::from_hours(1));
assert!(limiter.check_key(&ip(7)).is_ok());
assert_eq!(
limiter.check_key(&ip(7)),
Err(BoundedLimiterError::RateLimited)
);
}
#[test]
fn check_key_detailed_reports_rate_limit_wait_under_default_policy() {
let quota = Quota::per_minute(NonZeroU32::new(1).unwrap());
let limiter: BoundedKeyedLimiter<IpAddr> =
BoundedKeyedLimiter::new(quota, cap(10), Duration::from_hours(1));
assert!(limiter.check_key_detailed(&ip(7)).is_ok());
let deny = limiter
.check_key_detailed(&ip(7))
.expect_err("second call within the window must deny");
match deny {
BoundedLimiterDeny::RateLimited(wait) => {
assert!(wait > Duration::ZERO, "wait must be positive, got {wait:?}");
assert!(
wait <= Duration::from_secs(60),
"per-minute quota wait must be <= 60s, got {wait:?}"
);
}
BoundedLimiterDeny::CapacityFull => panic!("default policy must not reject capacity"),
}
}
#[test]
fn never_exceeds_max_tracked_keys() {
let quota = Quota::per_minute(NonZeroU32::new(10).unwrap());
let limiter: BoundedKeyedLimiter<IpAddr> =
BoundedKeyedLimiter::new(quota, cap(100), Duration::from_hours(1));
for i in 0..10_000_u32 {
let _ = limiter.check_key(&ip(i));
assert!(
limiter.len() <= 100,
"tracked keys exceeded cap at iteration {i}: {} > 100",
limiter.len()
);
}
assert_eq!(limiter.len(), 100, "table should be full at the cap");
}
#[test]
fn reject_new_at_cap_denies_unseen_key_but_keeps_established_key() {
let quota = Quota::per_minute(NonZeroU32::new(2).unwrap());
let limiter: BoundedKeyedLimiter<IpAddr> = BoundedKeyedLimiter::new_with_policy(
quota,
cap(1),
Duration::from_hours(1),
KeyEvictionPolicy::RejectNew,
);
let established = ip(10);
assert!(limiter.check_key_detailed(&established).is_ok());
assert_eq!(limiter.len(), 1);
let unseen = ip(11);
assert_eq!(
limiter.check_key_detailed(&unseen),
Err(BoundedLimiterDeny::CapacityFull)
);
assert_eq!(limiter.len(), 1);
assert!(
limiter.check_key_detailed(&established).is_ok(),
"established key keeps its existing bucket and remaining quota"
);
}
#[test]
fn evict_lru_policy_at_cap_admits_new_key_and_evicts_lru() {
let quota = Quota::per_minute(NonZeroU32::new(2).unwrap());
let limiter: BoundedKeyedLimiter<IpAddr> = BoundedKeyedLimiter::new_with_policy(
quota,
cap(1),
Duration::from_hours(1),
KeyEvictionPolicy::EvictLru,
);
let first = ip(20);
assert!(limiter.check_key_detailed(&first).is_ok());
assert!(limiter.check_key_detailed(&first).is_ok());
assert!(limiter.check_key_detailed(&first).is_err());
std::thread::sleep(Duration::from_millis(5));
assert!(limiter.check_key_detailed(&ip(21)).is_ok());
assert_eq!(limiter.len(), 1);
assert!(
limiter.check_key_detailed(&first).is_ok(),
"LRU-evicted key returns with fresh quota under EvictLru"
);
}
#[test]
fn evicted_keys_get_fresh_quota() {
let quota = Quota::per_minute(NonZeroU32::new(2).unwrap());
let limiter: BoundedKeyedLimiter<IpAddr> =
BoundedKeyedLimiter::new(quota, cap(2), Duration::from_hours(1));
let target = ip(1);
assert!(limiter.check_key(&target).is_ok(), "first ok");
assert!(limiter.check_key(&target).is_ok(), "second ok");
assert!(limiter.check_key(&target).is_err(), "third blocked");
std::thread::sleep(Duration::from_millis(5));
let _ = limiter.check_key(&ip(2));
std::thread::sleep(Duration::from_millis(5));
let _ = limiter.check_key(&ip(3));
std::thread::sleep(Duration::from_millis(5));
let _ = limiter.check_key(&ip(4));
std::thread::sleep(Duration::from_millis(5));
let _ = limiter.check_key(&ip(5));
assert!(
limiter.check_key(&target).is_ok(),
"evicted key gets a fresh quota on reappearance"
);
}
#[test]
fn active_over_quota_key_not_evicted() {
let quota = Quota::per_minute(NonZeroU32::new(2).unwrap());
let limiter: BoundedKeyedLimiter<IpAddr> =
BoundedKeyedLimiter::new(quota, cap(3), Duration::from_hours(1));
for i in 100..103_u32 {
let _ = limiter.check_key(&ip(i));
}
assert_eq!(limiter.len(), 3);
std::thread::sleep(Duration::from_millis(5));
let attacker = ip(200);
let _ = limiter.check_key(&attacker);
let _ = limiter.check_key(&attacker);
for new_key in 300..310_u32 {
std::thread::sleep(Duration::from_millis(2));
let _ = limiter.check_key(&attacker); std::thread::sleep(Duration::from_millis(2));
let _ = limiter.check_key(&ip(new_key)); }
let _ = limiter.check_key(&attacker);
assert!(
limiter.check_key(&attacker).is_err(),
"actively over-quota attacker must not be evicted into a fresh quota"
);
}
}