a.tolmachev
146 lines
4.0 KiB
Rust
146 lines
4.0 KiB
Rust
use std::{
|
|
collections::HashMap,
|
|
sync::{Arc, Mutex},
|
|
time::{Duration, Instant},
|
|
};
|
|
|
|
use thiserror::Error;
|
|
|
|
const STALE_KEY_TTL: Duration = Duration::from_secs(300);
|
|
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
|
pub struct RequestRateLimitConfig {
|
|
pub requests_per_second: u32,
|
|
pub burst: u32,
|
|
}
|
|
|
|
impl RequestRateLimitConfig {
|
|
pub fn new(requests_per_second: u32, burst: u32) -> Result<Self, RequestRateLimitConfigError> {
|
|
if requests_per_second == 0 {
|
|
return Err(RequestRateLimitConfigError::InvalidRequestsPerSecond(
|
|
requests_per_second,
|
|
));
|
|
}
|
|
if burst == 0 {
|
|
return Err(RequestRateLimitConfigError::InvalidBurst(burst));
|
|
}
|
|
|
|
Ok(Self {
|
|
requests_per_second,
|
|
burst,
|
|
})
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Error, PartialEq, Eq)]
|
|
pub enum RequestRateLimitConfigError {
|
|
#[error("requests_per_second must be greater than zero, got {0}")]
|
|
InvalidRequestsPerSecond(u32),
|
|
#[error("burst must be greater than zero, got {0}")]
|
|
InvalidBurst(u32),
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
|
pub struct RateLimitRejection {
|
|
pub retry_after_ms: u64,
|
|
}
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub struct RequestRateLimiter {
|
|
config: RequestRateLimitConfig,
|
|
states: Arc<Mutex<HashMap<String, BucketState>>>,
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug)]
|
|
struct BucketState {
|
|
tokens: f64,
|
|
last_refill: Instant,
|
|
last_seen: Instant,
|
|
}
|
|
|
|
impl RequestRateLimiter {
|
|
pub fn new(config: RequestRateLimitConfig) -> Self {
|
|
Self {
|
|
config,
|
|
states: Arc::new(Mutex::new(HashMap::new())),
|
|
}
|
|
}
|
|
|
|
pub fn check(&self, key: &str) -> Result<(), RateLimitRejection> {
|
|
self.check_at(key, Instant::now())
|
|
}
|
|
|
|
pub fn check_at(&self, key: &str, now: Instant) -> Result<(), RateLimitRejection> {
|
|
let mut states = self
|
|
.states
|
|
.lock()
|
|
.unwrap_or_else(|poisoned| poisoned.into_inner());
|
|
states.retain(|_, state| now.duration_since(state.last_seen) <= STALE_KEY_TTL);
|
|
|
|
let burst = self.config.burst as f64;
|
|
let requests_per_second = self.config.requests_per_second as f64;
|
|
let state = states.entry(key.to_owned()).or_insert(BucketState {
|
|
tokens: burst,
|
|
last_refill: now,
|
|
last_seen: now,
|
|
});
|
|
|
|
let elapsed = now.duration_since(state.last_refill).as_secs_f64();
|
|
state.tokens = (state.tokens + elapsed * requests_per_second).min(burst);
|
|
state.last_refill = now;
|
|
state.last_seen = now;
|
|
|
|
if state.tokens >= 1.0 {
|
|
state.tokens -= 1.0;
|
|
return Ok(());
|
|
}
|
|
|
|
let missing_tokens = (1.0 - state.tokens).max(0.0);
|
|
let retry_after_ms = ((missing_tokens / requests_per_second) * 1000.0)
|
|
.ceil()
|
|
.max(1.0) as u64;
|
|
Err(RateLimitRejection { retry_after_ms })
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use std::time::{Duration, Instant};
|
|
|
|
use super::{RequestRateLimitConfig, RequestRateLimitConfigError, RequestRateLimiter};
|
|
|
|
#[test]
|
|
fn rejects_zero_requests_per_second() {
|
|
assert_eq!(
|
|
RequestRateLimitConfig::new(0, 1),
|
|
Err(RequestRateLimitConfigError::InvalidRequestsPerSecond(0))
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn rejects_zero_burst() {
|
|
assert_eq!(
|
|
RequestRateLimitConfig::new(1, 0),
|
|
Err(RequestRateLimitConfigError::InvalidBurst(0))
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn allows_burst_then_rejects_until_refilled() {
|
|
let limiter = RequestRateLimiter::new(RequestRateLimitConfig::new(2, 2).unwrap());
|
|
let start = Instant::now();
|
|
|
|
assert!(limiter.check_at("key", start).is_ok());
|
|
assert!(limiter.check_at("key", start).is_ok());
|
|
|
|
let rejection = limiter.check_at("key", start).unwrap_err();
|
|
assert_eq!(rejection.retry_after_ms, 500);
|
|
|
|
assert!(
|
|
limiter
|
|
.check_at("key", start + Duration::from_millis(500))
|
|
.is_ok()
|
|
);
|
|
}
|
|
}
|