github-ops
270 lines
8.4 KiB
Rust
270 lines
8.4 KiB
Rust
use std::{
|
|
collections::HashMap,
|
|
sync::{Arc, Mutex},
|
|
time::{Duration, Instant, SystemTime, UNIX_EPOCH},
|
|
};
|
|
|
|
use crank_core::{RateLimitBucketState, RateLimitStateStore};
|
|
use thiserror::Error;
|
|
|
|
const STALE_KEY_TTL: Duration = Duration::from_secs(300);
|
|
const TOKEN_SCALE: u64 = 1_000_000;
|
|
|
|
#[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)]
|
|
pub struct RequestRateLimiter {
|
|
config: RequestRateLimitConfig,
|
|
backend: RequestRateLimiterBackend,
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
enum RequestRateLimiterBackend {
|
|
Local {
|
|
states: Arc<Mutex<HashMap<String, LocalBucketState>>>,
|
|
},
|
|
Shared {
|
|
store: Arc<dyn RateLimitStateStore>,
|
|
},
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug)]
|
|
struct LocalBucketState {
|
|
tokens: f64,
|
|
last_refill: Instant,
|
|
last_seen: Instant,
|
|
}
|
|
|
|
impl RequestRateLimiter {
|
|
pub fn new(config: RequestRateLimitConfig) -> Self {
|
|
Self {
|
|
config,
|
|
backend: RequestRateLimiterBackend::Local {
|
|
states: Arc::new(Mutex::new(HashMap::new())),
|
|
},
|
|
}
|
|
}
|
|
|
|
pub fn new_shared(config: RequestRateLimitConfig, store: Arc<dyn RateLimitStateStore>) -> Self {
|
|
Self {
|
|
config,
|
|
backend: RequestRateLimiterBackend::Shared { store },
|
|
}
|
|
}
|
|
|
|
pub async fn check(&self, key: &str) -> Result<(), RateLimitRejection> {
|
|
match &self.backend {
|
|
RequestRateLimiterBackend::Local { .. } => self.check_local_at(key, Instant::now()),
|
|
RequestRateLimiterBackend::Shared { store } => {
|
|
self.check_shared_at(store.as_ref(), key, now_unix_ms())
|
|
.await
|
|
}
|
|
}
|
|
}
|
|
|
|
fn check_local_at(&self, key: &str, now: Instant) -> Result<(), RateLimitRejection> {
|
|
let RequestRateLimiterBackend::Local { states } = &self.backend else {
|
|
panic!("check_local_at called for non-local limiter");
|
|
};
|
|
|
|
let mut states = 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(LocalBucketState {
|
|
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 })
|
|
}
|
|
|
|
async fn check_shared_at(
|
|
&self,
|
|
store: &dyn RateLimitStateStore,
|
|
key: &str,
|
|
now_unix_ms: i64,
|
|
) -> Result<(), RateLimitRejection> {
|
|
let burst_tokens = u64::from(self.config.burst) * TOKEN_SCALE;
|
|
let refill_per_second = u64::from(self.config.requests_per_second) * TOKEN_SCALE;
|
|
let mut state =
|
|
store
|
|
.get_bucket(key)
|
|
.await
|
|
.unwrap_or(None)
|
|
.unwrap_or(RateLimitBucketState {
|
|
tokens_micros: burst_tokens,
|
|
last_refill_unix_ms: now_unix_ms,
|
|
});
|
|
|
|
let elapsed_ms = (now_unix_ms - state.last_refill_unix_ms).max(0) as u64;
|
|
let replenished =
|
|
state.tokens_micros + (elapsed_ms.saturating_mul(refill_per_second) / 1000);
|
|
state.tokens_micros = replenished.min(burst_tokens);
|
|
state.last_refill_unix_ms = now_unix_ms;
|
|
|
|
if state.tokens_micros >= TOKEN_SCALE {
|
|
state.tokens_micros -= TOKEN_SCALE;
|
|
let _ = store.put_bucket(key, state, STALE_KEY_TTL).await;
|
|
return Ok(());
|
|
}
|
|
|
|
let missing_tokens = TOKEN_SCALE.saturating_sub(state.tokens_micros);
|
|
let retry_after_ms = missing_tokens.div_ceil(refill_per_second).max(1);
|
|
let _ = store.put_bucket(key, state, STALE_KEY_TTL).await;
|
|
Err(RateLimitRejection { retry_after_ms })
|
|
}
|
|
}
|
|
|
|
fn now_unix_ms() -> i64 {
|
|
SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.unwrap_or_else(|_| Duration::from_secs(0))
|
|
.as_millis()
|
|
.try_into()
|
|
.unwrap_or(i64::MAX)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use std::{
|
|
sync::Arc,
|
|
time::{Duration, Instant},
|
|
};
|
|
|
|
use crate::InMemoryRateLimitStateStore;
|
|
|
|
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_for_local_limiter() {
|
|
let limiter = RequestRateLimiter::new(RequestRateLimitConfig::new(2, 2).unwrap());
|
|
let start = Instant::now();
|
|
|
|
assert!(limiter.check_local_at("key", start).is_ok());
|
|
assert!(limiter.check_local_at("key", start).is_ok());
|
|
|
|
let rejection = limiter.check_local_at("key", start).unwrap_err();
|
|
assert_eq!(rejection.retry_after_ms, 500);
|
|
|
|
assert!(
|
|
limiter
|
|
.check_local_at("key", start + Duration::from_millis(500))
|
|
.is_ok()
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn allows_burst_then_rejects_until_refilled_for_shared_limiter() {
|
|
let limiter = RequestRateLimiter::new_shared(
|
|
RequestRateLimitConfig::new(2, 2).unwrap(),
|
|
Arc::new(InMemoryRateLimitStateStore::default()),
|
|
);
|
|
|
|
assert!(limiter.check_shared_at_store("key", 0).await.is_ok());
|
|
assert!(limiter.check_shared_at_store("key", 0).await.is_ok());
|
|
|
|
let rejection = limiter.check_shared_at_store("key", 0).await.unwrap_err();
|
|
assert_eq!(rejection.retry_after_ms, 1);
|
|
|
|
assert!(limiter.check_shared_at_store("key", 500).await.is_ok());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn shared_limiter_state_is_visible_across_instances() {
|
|
let store = Arc::new(InMemoryRateLimitStateStore::default());
|
|
let first = RequestRateLimiter::new_shared(
|
|
RequestRateLimitConfig::new(1, 1).unwrap(),
|
|
store.clone(),
|
|
);
|
|
let second =
|
|
RequestRateLimiter::new_shared(RequestRateLimitConfig::new(1, 1).unwrap(), store);
|
|
|
|
assert!(first.check_shared_at_store("shared", 0).await.is_ok());
|
|
assert!(second.check_shared_at_store("shared", 0).await.is_err());
|
|
assert!(second.check_shared_at_store("shared", 1000).await.is_ok());
|
|
}
|
|
|
|
impl RequestRateLimiter {
|
|
async fn check_shared_at_store(
|
|
&self,
|
|
key: &str,
|
|
now_unix_ms: i64,
|
|
) -> Result<(), super::RateLimitRejection> {
|
|
let super::RequestRateLimiterBackend::Shared { store } = &self.backend else {
|
|
panic!("check_shared_at_store called for non-shared limiter");
|
|
};
|
|
self.check_shared_at(store.as_ref(), key, now_unix_ms).await
|
|
}
|
|
}
|
|
}
|