#![allow(
clippy::cast_possible_truncation,
clippy::cast_sign_loss,
clippy::cast_precision_loss,
clippy::cast_possible_wrap
)]
use crate::rate_limit::RateLimitConfig;
use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Mutex;
use tokio::time::Instant;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DistributedAlgorithm {
TokenBucket,
SlidingWindow,
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct RateLimitParams {
pub algorithm: DistributedAlgorithm,
pub rate: f64,
pub burst: f64,
pub window_secs: u64,
}
impl RateLimitParams {
#[must_use]
pub fn from_config(config: &RateLimitConfig) -> Self {
if config.sliding_window {
Self {
algorithm: DistributedAlgorithm::SlidingWindow,
rate: config.rate,
burst: f64::from(config.effective_burst()),
window_secs: config.window_size.max(1),
}
} else {
Self {
algorithm: DistributedAlgorithm::TokenBucket,
rate: config.rate,
burst: f64::from(config.effective_burst()),
window_secs: config.window_size.max(1),
}
}
}
#[inline]
#[must_use]
pub fn max_window_events(&self) -> u64 {
(self.rate * self.window_secs as f64).ceil().max(0.0) as u64
}
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct AcquireOutcome {
pub allowed: bool,
pub remaining: f64,
pub retry_after: Duration,
}
impl AcquireOutcome {
#[inline]
#[must_use]
pub fn allowed(remaining: f64) -> Self {
Self {
allowed: true,
remaining,
retry_after: Duration::ZERO,
}
}
#[inline]
#[must_use]
pub fn denied(remaining: f64, retry_after: Duration) -> Self {
Self {
allowed: false,
remaining,
retry_after,
}
}
}
#[async_trait]
pub trait DistributedRateLimitBackend: Send + Sync {
async fn acquire(
&self,
key: &str,
cost: f64,
params: RateLimitParams,
) -> crate::Result<AcquireOutcome>;
async fn available(&self, key: &str, params: RateLimitParams) -> crate::Result<f64>;
async fn time_until_available(
&self,
key: &str,
cost: f64,
params: RateLimitParams,
) -> crate::Result<Duration>;
async fn reset(&self, key: &str) -> crate::Result<()>;
fn backend_name(&self) -> &str;
}
#[derive(Debug, Clone)]
enum KeyState {
Bucket { tokens: f64, last_refill: Instant },
Window { stamps: Vec<Instant> },
}
#[derive(Debug, Default)]
pub struct InMemoryDistributedBackend {
state: Mutex<HashMap<String, KeyState>>,
}
impl InMemoryDistributedBackend {
#[must_use]
pub fn new() -> Self {
Self {
state: Mutex::new(HashMap::new()),
}
}
pub async fn key_count(&self) -> usize {
self.state.lock().await.len()
}
#[inline]
fn refill_bucket(
tokens: f64,
last_refill: Instant,
now: Instant,
params: RateLimitParams,
) -> f64 {
let elapsed = now.saturating_duration_since(last_refill).as_secs_f64();
(tokens + elapsed * params.rate).min(params.burst)
}
#[inline]
fn prune_window(stamps: &mut Vec<Instant>, now: Instant, params: RateLimitParams) {
let window = Duration::from_secs(params.window_secs);
if let Some(cutoff) = now.checked_sub(window) {
stamps.retain(|&t| t > cutoff);
}
}
}
#[async_trait]
impl DistributedRateLimitBackend for InMemoryDistributedBackend {
async fn acquire(
&self,
key: &str,
cost: f64,
params: RateLimitParams,
) -> crate::Result<AcquireOutcome> {
if cost < 0.0 {
return Err(crate::CelersError::Configuration(
"rate limit acquire cost must be non-negative".to_string(),
));
}
let now = Instant::now();
let mut guard = self.state.lock().await;
match params.algorithm {
DistributedAlgorithm::TokenBucket => {
let entry = guard
.entry(key.to_string())
.or_insert_with(|| KeyState::Bucket {
tokens: params.burst,
last_refill: now,
});
if !matches!(entry, KeyState::Bucket { .. }) {
*entry = KeyState::Bucket {
tokens: params.burst,
last_refill: now,
};
}
let KeyState::Bucket {
tokens,
last_refill,
} = entry
else {
unreachable!("entry coerced to Bucket above")
};
*tokens = Self::refill_bucket(*tokens, *last_refill, now, params);
*last_refill = now;
if *tokens + f64::EPSILON >= cost {
*tokens -= cost;
Ok(AcquireOutcome::allowed(*tokens))
} else {
let deficit = cost - *tokens;
let retry_after = if params.rate > 0.0 {
Duration::from_secs_f64(deficit / params.rate)
} else {
Duration::MAX
};
Ok(AcquireOutcome::denied(*tokens, retry_after))
}
}
DistributedAlgorithm::SlidingWindow => {
let entry = guard
.entry(key.to_string())
.or_insert_with(|| KeyState::Window { stamps: Vec::new() });
if !matches!(entry, KeyState::Window { .. }) {
*entry = KeyState::Window { stamps: Vec::new() };
}
let KeyState::Window { stamps } = entry else {
unreachable!("entry coerced to Window above")
};
Self::prune_window(stamps, now, params);
let max = params.max_window_events();
let needed = cost.ceil() as u64;
let used = stamps.len() as u64;
if used + needed <= max {
for _ in 0..needed {
stamps.push(now);
}
let remaining = max.saturating_sub(used + needed) as f64;
Ok(AcquireOutcome::allowed(remaining))
} else {
let remaining = max.saturating_sub(used) as f64;
let retry_after = stamps.first().map_or(Duration::ZERO, |&oldest| {
let expires = oldest + Duration::from_secs(params.window_secs);
expires.saturating_duration_since(now)
});
Ok(AcquireOutcome::denied(remaining, retry_after))
}
}
}
}
async fn available(&self, key: &str, params: RateLimitParams) -> crate::Result<f64> {
let now = Instant::now();
let mut guard = self.state.lock().await;
match params.algorithm {
DistributedAlgorithm::TokenBucket => match guard.get_mut(key) {
Some(KeyState::Bucket {
tokens,
last_refill,
}) => {
*tokens = Self::refill_bucket(*tokens, *last_refill, now, params);
*last_refill = now;
Ok(*tokens)
}
_ => Ok(params.burst),
},
DistributedAlgorithm::SlidingWindow => {
let max = params.max_window_events() as f64;
match guard.get_mut(key) {
Some(KeyState::Window { stamps }) => {
Self::prune_window(stamps, now, params);
Ok(max - stamps.len() as f64)
}
_ => Ok(max),
}
}
}
}
async fn time_until_available(
&self,
key: &str,
cost: f64,
params: RateLimitParams,
) -> crate::Result<Duration> {
let now = Instant::now();
let mut guard = self.state.lock().await;
match params.algorithm {
DistributedAlgorithm::TokenBucket => {
let tokens = match guard.get_mut(key) {
Some(KeyState::Bucket {
tokens,
last_refill,
}) => {
*tokens = Self::refill_bucket(*tokens, *last_refill, now, params);
*last_refill = now;
*tokens
}
_ => params.burst,
};
if tokens + f64::EPSILON >= cost {
Ok(Duration::ZERO)
} else if params.rate > 0.0 {
Ok(Duration::from_secs_f64((cost - tokens) / params.rate))
} else {
Ok(Duration::MAX)
}
}
DistributedAlgorithm::SlidingWindow => {
let max = params.max_window_events();
let needed = cost.ceil() as u64;
match guard.get_mut(key) {
Some(KeyState::Window { stamps }) => {
Self::prune_window(stamps, now, params);
if stamps.len() as u64 + needed <= max {
Ok(Duration::ZERO)
} else {
Ok(stamps.first().map_or(Duration::ZERO, |&oldest| {
let expires = oldest + Duration::from_secs(params.window_secs);
expires.saturating_duration_since(now)
}))
}
}
_ => Ok(Duration::ZERO),
}
}
}
}
async fn reset(&self, key: &str) -> crate::Result<()> {
self.state.lock().await.remove(key);
Ok(())
}
fn backend_name(&self) -> &str {
"in-memory"
}
}
#[derive(Clone)]
pub struct DistributedRateLimiter {
backend: Arc<dyn DistributedRateLimitBackend>,
key: String,
config: RateLimitConfig,
params: RateLimitParams,
}
impl std::fmt::Debug for DistributedRateLimiter {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("DistributedRateLimiter")
.field("key", &self.key)
.field("backend", &self.backend.backend_name())
.field("config", &self.config)
.finish()
}
}
impl DistributedRateLimiter {
pub fn new(
backend: Arc<dyn DistributedRateLimitBackend>,
key: impl Into<String>,
config: RateLimitConfig,
) -> Self {
let params = RateLimitParams::from_config(&config);
Self {
backend,
key: key.into(),
config,
params,
}
}
#[inline]
#[must_use]
pub fn key(&self) -> &str {
&self.key
}
#[inline]
#[must_use]
pub fn config(&self) -> &RateLimitConfig {
&self.config
}
#[inline]
#[must_use]
pub fn algorithm(&self) -> DistributedAlgorithm {
self.params.algorithm
}
pub async fn try_acquire(&self) -> crate::Result<bool> {
Ok(self
.backend
.acquire(&self.key, 1.0, self.params)
.await?
.allowed)
}
pub async fn try_acquire_n(&self, cost: f64) -> crate::Result<AcquireOutcome> {
self.backend.acquire(&self.key, cost, self.params).await
}
pub async fn acquire(&self, max_sleep: Duration) -> crate::Result<Duration> {
let start = Instant::now();
loop {
let outcome = self.backend.acquire(&self.key, 1.0, self.params).await?;
if outcome.allowed {
return Ok(start.elapsed());
}
let sleep_for = outcome
.retry_after
.min(max_sleep)
.max(Duration::from_millis(1));
tokio::time::sleep(sleep_for).await;
}
}
pub async fn available(&self) -> crate::Result<f64> {
self.backend.available(&self.key, self.params).await
}
pub async fn time_until_available(&self) -> crate::Result<Duration> {
self.backend
.time_until_available(&self.key, 1.0, self.params)
.await
}
pub async fn reset(&self) -> crate::Result<()> {
self.backend.reset(&self.key).await
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_token_bucket_burst_then_deny() {
let backend = Arc::new(InMemoryDistributedBackend::new());
let config = RateLimitConfig::new(10.0).with_burst(5);
let limiter = DistributedRateLimiter::new(backend, "task", config);
for _ in 0..5 {
assert!(limiter.try_acquire().await.unwrap());
}
assert!(!limiter.try_acquire().await.unwrap());
assert_eq!(limiter.algorithm(), DistributedAlgorithm::TokenBucket);
}
#[tokio::test]
async fn test_token_bucket_refill() {
let backend = Arc::new(InMemoryDistributedBackend::new());
let config = RateLimitConfig::new(100.0).with_burst(10);
let limiter = DistributedRateLimiter::new(backend, "task", config);
for _ in 0..10 {
assert!(limiter.try_acquire().await.unwrap());
}
assert!(!limiter.try_acquire().await.unwrap());
tokio::time::sleep(Duration::from_millis(30)).await;
assert!(limiter.try_acquire().await.unwrap());
}
#[tokio::test]
async fn test_sliding_window_basic_and_deny() {
let backend = Arc::new(InMemoryDistributedBackend::new());
let config = RateLimitConfig::new(5.0).with_sliding_window(1);
let limiter = DistributedRateLimiter::new(backend, "task", config);
assert_eq!(limiter.algorithm(), DistributedAlgorithm::SlidingWindow);
for _ in 0..5 {
assert!(limiter.try_acquire().await.unwrap());
}
assert!(!limiter.try_acquire().await.unwrap());
let wait = limiter.time_until_available().await.unwrap();
assert!(wait > Duration::ZERO);
}
#[tokio::test]
async fn test_cost_n_acquire() {
let backend = Arc::new(InMemoryDistributedBackend::new());
let config = RateLimitConfig::new(10.0).with_burst(10);
let limiter = DistributedRateLimiter::new(backend, "task", config);
let outcome = limiter.try_acquire_n(4.0).await.unwrap();
assert!(outcome.allowed);
assert!((outcome.remaining - 6.0).abs() < 1e-6);
let denied = limiter.try_acquire_n(7.0).await.unwrap();
assert!(!denied.allowed);
assert!(denied.retry_after > Duration::ZERO);
}
#[tokio::test]
async fn test_available_and_reset() {
let backend = Arc::new(InMemoryDistributedBackend::new());
let config = RateLimitConfig::new(0.0).with_burst(8);
let limiter = DistributedRateLimiter::new(backend.clone(), "task", config);
assert!((limiter.available().await.unwrap() - 8.0).abs() < 1e-6);
for _ in 0..3 {
assert!(limiter.try_acquire().await.unwrap());
}
let avail = limiter.available().await.unwrap();
assert!((avail - 5.0).abs() < 1e-6, "expected 5 tokens, got {avail}");
limiter.reset().await.unwrap();
assert!((limiter.available().await.unwrap() - 8.0).abs() < 1e-6);
}
#[tokio::test]
async fn test_concurrent_acquire_does_not_exceed_burst() {
let backend = Arc::new(InMemoryDistributedBackend::new());
let config = RateLimitConfig::new(0.0).with_burst(20);
let limiter = DistributedRateLimiter::new(backend, "task", config);
let mut handles = Vec::new();
for _ in 0..8 {
let l = limiter.clone();
handles.push(tokio::spawn(async move {
let mut count = 0u32;
for _ in 0..10 {
if l.try_acquire().await.unwrap_or(false) {
count += 1;
}
}
count
}));
}
let mut total = 0u32;
for h in handles {
total += h.await.unwrap();
}
assert_eq!(total, 20, "exactly the burst capacity should be granted");
}
#[tokio::test]
async fn test_concurrent_sliding_window_cap() {
let backend = Arc::new(InMemoryDistributedBackend::new());
let config = RateLimitConfig::new(3.0).with_sliding_window(100);
let limiter = DistributedRateLimiter::new(backend, "task", config);
let max = limiter.params.max_window_events();
let mut handles = Vec::new();
for _ in 0..6 {
let l = limiter.clone();
handles.push(tokio::spawn(async move {
let mut count = 0u32;
for _ in 0..200 {
if l.try_acquire().await.unwrap_or(false) {
count += 1;
}
}
count
}));
}
let mut total = 0u64;
for h in handles {
total += u64::from(h.await.unwrap());
}
assert_eq!(total, max, "sliding window must not exceed max events");
}
#[tokio::test]
async fn test_acquire_waits_then_succeeds() {
let backend = Arc::new(InMemoryDistributedBackend::new());
let config = RateLimitConfig::new(50.0).with_burst(1);
let limiter = DistributedRateLimiter::new(backend, "task", config);
assert!(limiter.try_acquire().await.unwrap());
let waited = limiter.acquire(Duration::from_millis(100)).await.unwrap();
assert!(waited > Duration::ZERO);
}
#[tokio::test]
async fn test_shared_backend_across_keys() {
let backend = Arc::new(InMemoryDistributedBackend::new());
let a = DistributedRateLimiter::new(
backend.clone(),
"a",
RateLimitConfig::new(0.0).with_burst(2),
);
let b = DistributedRateLimiter::new(
backend.clone(),
"b",
RateLimitConfig::new(0.0).with_burst(2),
);
assert!(a.try_acquire().await.unwrap());
assert!(a.try_acquire().await.unwrap());
assert!(!a.try_acquire().await.unwrap());
assert!(b.try_acquire().await.unwrap());
assert!(b.try_acquire().await.unwrap());
assert!(!b.try_acquire().await.unwrap());
assert_eq!(backend.key_count().await, 2);
}
#[tokio::test]
async fn test_negative_cost_rejected() {
let backend = Arc::new(InMemoryDistributedBackend::new());
let config = RateLimitConfig::new(10.0).with_burst(5);
let limiter = DistributedRateLimiter::new(backend, "task", config);
assert!(limiter.try_acquire_n(-1.0).await.is_err());
}
#[tokio::test]
async fn test_params_from_config() {
let bucket = RateLimitParams::from_config(&RateLimitConfig::new(10.0).with_burst(20));
assert_eq!(bucket.algorithm, DistributedAlgorithm::TokenBucket);
assert!((bucket.burst - 20.0).abs() < 1e-6);
let window =
RateLimitParams::from_config(&RateLimitConfig::new(4.0).with_sliding_window(5));
assert_eq!(window.algorithm, DistributedAlgorithm::SlidingWindow);
assert_eq!(window.window_secs, 5);
assert_eq!(window.max_window_events(), 20);
}
}