103 lines
2.8 KiB
Rust
103 lines
2.8 KiB
Rust
use std::{sync::Mutex, time::Instant};
|
|
|
|
// This challenge uses a simple in-memory global limiter. In production,
|
|
// prefer API gateway rate limiting so limits are shared across instances and
|
|
// identity-aware policies can be keyed by API key, user, or client IP.
|
|
#[derive(Debug)]
|
|
pub struct RateLimiter {
|
|
bucket: Mutex<TokenBucket>,
|
|
}
|
|
|
|
impl RateLimiter {
|
|
pub fn new(refill_per_second: f64, burst: u32) -> Result<Self, RateLimitError> {
|
|
if !refill_per_second.is_finite() || refill_per_second <= 0.0 || burst == 0 {
|
|
return Err(RateLimitError::InvalidConfiguration);
|
|
}
|
|
|
|
Ok(Self {
|
|
bucket: Mutex::new(TokenBucket {
|
|
capacity: f64::from(burst),
|
|
tokens: f64::from(burst),
|
|
refill_per_second,
|
|
last_refill: Instant::now(),
|
|
}),
|
|
})
|
|
}
|
|
|
|
#[must_use]
|
|
pub fn try_acquire(&self) -> bool {
|
|
self.bucket
|
|
.lock()
|
|
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
|
.try_acquire()
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, thiserror::Error)]
|
|
pub enum RateLimitError {
|
|
#[error("rate limit configuration must be positive")]
|
|
InvalidConfiguration,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
struct TokenBucket {
|
|
capacity: f64,
|
|
tokens: f64,
|
|
refill_per_second: f64,
|
|
last_refill: Instant,
|
|
}
|
|
|
|
impl TokenBucket {
|
|
fn try_acquire(&mut self) -> bool {
|
|
self.refill();
|
|
if self.tokens >= 1.0 {
|
|
self.tokens -= 1.0;
|
|
true
|
|
} else {
|
|
false
|
|
}
|
|
}
|
|
|
|
fn refill(&mut self) {
|
|
let elapsed = self.last_refill.elapsed().as_secs_f64();
|
|
self.tokens = self
|
|
.capacity
|
|
.min(self.tokens + elapsed.mul_add(self.refill_per_second, 0.0));
|
|
self.last_refill = Instant::now();
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::RateLimiter;
|
|
use std::time::Duration;
|
|
|
|
#[test]
|
|
fn rejects_non_positive_or_non_finite_configuration() {
|
|
assert!(RateLimiter::new(0.0, 1).is_err());
|
|
assert!(RateLimiter::new(f64::INFINITY, 1).is_err());
|
|
assert!(RateLimiter::new(1.0, 0).is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn allows_requests_up_to_the_burst_capacity() {
|
|
let limiter = RateLimiter::new(100.0, 2).expect("rate limiter should build");
|
|
|
|
assert!(limiter.try_acquire());
|
|
assert!(limiter.try_acquire());
|
|
assert!(!limiter.try_acquire());
|
|
}
|
|
|
|
#[test]
|
|
fn refills_tokens_over_time() {
|
|
let limiter = RateLimiter::new(100.0, 1).expect("rate limiter should build");
|
|
|
|
assert!(limiter.try_acquire());
|
|
assert!(!limiter.try_acquire());
|
|
// A production-grade limiter would inject a clock; a short sleep keeps
|
|
// this challenge implementation simple without coupling tests to internals.
|
|
std::thread::sleep(Duration::from_millis(20));
|
|
assert!(limiter.try_acquire());
|
|
}
|
|
}
|