diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index 4df92b82..b6c9e5aa 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -60,6 +60,12 @@ pub async fn chat_completions( .get(provider) .ok_or(ProxyError::ProviderUnavailable)?; + // Rate-limit pre-commit. Key on ApiKey id so two different keys get + // independent buckets even if they alias the same upstream credential. + let rl_key = auth.entry.id.clone(); + let rl_limits = auth.key().rate_limit.clone().unwrap_or_default(); + let reservation = state.limiter.pre_commit(&rl_key, &rl_limits)?; + let request_id = format!("req-{}", Uuid::new_v4()); let model_arc = std::sync::Arc::new(model_entry.value.clone()); let ctx = BridgeContext::new(&request_id, model_arc); @@ -67,7 +73,12 @@ pub async fn chat_completions( let now = created_ts(); if req.is_streaming() { + // Streaming: we can't measure tokens before the stream ends, so + // commit zero up front to keep the reservation's drop-guard from + // silently counting nothing. A later PR will tally tokens as the + // stream runs; for now release the permit when the handler returns. let upstream = bridge.chat_stream(&req, &ctx).await?; + reservation.commit_tokens(0); let model_name = req.model.clone(); let sse_stream = build_sse_stream(upstream, model_name, now); let response = @@ -76,6 +87,8 @@ pub async fn chat_completions( } let upstream = bridge.chat(&req, &ctx).await?; + let tokens = upstream.usage.total_tokens as u64; + reservation.commit_tokens(tokens); let rendered = render_response(now, upstream); Ok(Json(rendered).into_response()) } diff --git a/crates/aisix-proxy/src/error.rs b/crates/aisix-proxy/src/error.rs index 36cd887c..6cade265 100644 --- a/crates/aisix-proxy/src/error.rs +++ b/crates/aisix-proxy/src/error.rs @@ -18,7 +18,8 @@ //! JSON shape boilerplate. use aisix_gateway::BridgeError; -use axum::http::StatusCode; +use aisix_ratelimit::RateLimitError; +use axum::http::{HeaderValue, StatusCode}; use axum::response::{IntoResponse, Response}; use axum::Json; use serde::Serialize; @@ -72,6 +73,8 @@ pub enum ProxyError { #[error("no bridge registered for provider")] ProviderUnavailable, #[error(transparent)] + RateLimit(#[from] RateLimitError), + #[error(transparent)] Bridge(#[from] BridgeError), } @@ -83,6 +86,7 @@ impl ProxyError { ProxyError::ModelNotFound(_) => StatusCode::NOT_FOUND, ProxyError::InvalidRequest(_) => StatusCode::BAD_REQUEST, ProxyError::ProviderUnavailable => StatusCode::SERVICE_UNAVAILABLE, + ProxyError::RateLimit(_) => StatusCode::TOO_MANY_REQUESTS, ProxyError::Bridge(b) => { StatusCode::from_u16(b.http_status()).unwrap_or(StatusCode::BAD_GATEWAY) } @@ -96,10 +100,21 @@ impl ProxyError { ProxyError::ModelNotFound(_) => "model_not_found", ProxyError::InvalidRequest(_) => "invalid_request_error", ProxyError::ProviderUnavailable => "provider_unavailable", + ProxyError::RateLimit(_) => "rate_limit_exceeded", ProxyError::Bridge(b) => b.error_type(), } } + /// Seconds the client should wait before retrying. Only present for + /// rate-limit-style rejections so the proxy can emit a `Retry-After` + /// header. + pub fn retry_after_secs(&self) -> Option { + match self { + ProxyError::RateLimit(e) => e.retry_after_secs(), + _ => None, + } + } + pub fn envelope(&self) -> ErrorEnvelope { ErrorEnvelope::new(self.to_string(), self.kind()) } @@ -108,8 +123,15 @@ impl ProxyError { impl IntoResponse for ProxyError { fn into_response(self) -> Response { let status = self.status(); + let retry_after = self.retry_after_secs(); let body = self.envelope(); - (status, Json(body)).into_response() + let mut response = (status, Json(body)).into_response(); + if let Some(secs) = retry_after { + if let Ok(value) = HeaderValue::from_str(&secs.to_string()) { + response.headers_mut().insert("retry-after", value); + } + } + response } } diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index e81354f5..d055023c 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -105,8 +105,20 @@ mod tests { } fn apikey_entry(key: &str, allowed: &[&str]) -> ResourceEntry { + apikey_entry_with_limits(key, allowed, None) + } + + fn apikey_entry_with_limits( + key: &str, + allowed: &[&str], + rate_limit: Option, + ) -> ResourceEntry { let allowed_json = serde_json::to_string(&allowed).unwrap(); - let cfg = format!(r#"{{"key": "{key}", "allowed_models": {allowed_json}}}"#); + let rl_tail = match rate_limit { + Some(v) => format!(", \"rate_limit\": {v}"), + None => String::new(), + }; + let cfg = format!(r#"{{"key": "{key}", "allowed_models": {allowed_json}{rl_tail}}}"#); let apikey: ApiKey = serde_json::from_str(&cfg).unwrap(); ResourceEntry::new("key-id-1", apikey, 1) } @@ -118,6 +130,22 @@ mod tests { snap } + fn seed_snapshot_with_limits( + model: &str, + allowed: &[&str], + api_base: &str, + rate_limit: serde_json::Value, + ) -> AisixSnapshot { + let snap = AisixSnapshot::new(); + snap.models.insert(model_entry(model, api_base)); + snap.apikeys.insert(apikey_entry_with_limits( + "sk-caller", + allowed, + Some(rate_limit), + )); + snap + } + async fn run(app: Router, req: Request) -> axum::http::Response { app.oneshot(req).await.unwrap() } @@ -379,4 +407,120 @@ data: [DONE]\n\n"; "expected at least two chat chunks, got {data_count}" ); } + + #[tokio::test] + async fn rate_limit_rpm_returns_429_with_retry_after_header() { + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "cmpl-up", + "model": "gpt-4o", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop" + }], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} + }))) + .mount(&upstream) + .await; + + let hub = Arc::new(Hub::new()); + hub.register(Provider::Openai, Arc::new(OpenAiBridge::new())); + let snap = seed_snapshot_with_limits( + "my-gpt4", + &["my-gpt4"], + &upstream.uri(), + serde_json::json!({"rpm": 1}), + ); + let state = build_state(snap, hub); + let body = serde_json::json!({ + "model": "my-gpt4", + "messages": [{"role": "user", "content": "hi"}] + }); + let make_req = || { + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", "Bearer sk-caller") + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .unwrap() + }; + + // First request succeeds. + let resp = run(build_router(state.clone()), make_req()).await; + assert_eq!(resp.status(), StatusCode::OK); + + // Second request within the same minute trips rpm=1 → 429. + let resp = run(build_router(state.clone()), make_req()).await; + assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS); + let retry = resp + .headers() + .get("retry-after") + .and_then(|v| v.to_str().ok()) + .and_then(|s| s.parse::().ok()) + .expect("missing or malformed retry-after header"); + assert!(retry >= 1); + let body_bytes = to_bytes(resp.into_body(), 1024).await.unwrap(); + let v: serde_json::Value = serde_json::from_slice(&body_bytes).unwrap(); + assert_eq!(v["error"]["type"], "rate_limit_exceeded"); + } + + #[tokio::test] + async fn rate_limit_tpm_blocks_after_token_commit_exhausts_window() { + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "cmpl-up", + "model": "gpt-4o", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop" + }], + // Deliberately overshoot the TPM cap so the next + // pre_commit observes an exhausted window. + "usage": {"prompt_tokens": 10_000, "completion_tokens": 10_000, "total_tokens": 20_000} + }))) + .mount(&upstream) + .await; + + let hub = Arc::new(Hub::new()); + hub.register(Provider::Openai, Arc::new(OpenAiBridge::new())); + let snap = seed_snapshot_with_limits( + "my-gpt4", + &["my-gpt4"], + &upstream.uri(), + serde_json::json!({"tpm": 1_000}), + ); + let state = build_state(snap, hub); + let body = serde_json::json!({ + "model": "my-gpt4", + "messages": [{"role": "user", "content": "hi"}] + }); + let make_req = || { + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", "Bearer sk-caller") + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .unwrap() + }; + + // First request goes through (pre-commit TPM is unchecked for an + // empty bucket); usage counted at post-deduct overshoots the cap. + let resp = run(build_router(state.clone()), make_req()).await; + assert_eq!(resp.status(), StatusCode::OK); + + // Second request sees TPM > 1000 and rejects. + let resp = run(build_router(state), make_req()).await; + assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS); + let body_bytes = to_bytes(resp.into_body(), 1024).await.unwrap(); + let v: serde_json::Value = serde_json::from_slice(&body_bytes).unwrap(); + assert_eq!(v["error"]["type"], "rate_limit_exceeded"); + } } diff --git a/crates/aisix-proxy/src/state.rs b/crates/aisix-proxy/src/state.rs index e79e60de..ea2585eb 100644 --- a/crates/aisix-proxy/src/state.rs +++ b/crates/aisix-proxy/src/state.rs @@ -4,6 +4,8 @@ //! - the lock-free `SnapshotHandle` for looking up //! Models and ApiKeys on every request //! - the `Hub` for resolving a `Provider` to the Bridge that serves it +//! - the per-key [`Limiter`] — queried before each upstream call and +//! finalised after the response completes //! - the configured request-body size limit //! //! Cheap to clone: every field is either an `Arc` or a small Copy scalar. @@ -11,12 +13,14 @@ use aisix_core::snapshot::SnapshotHandle; use aisix_core::{AisixSnapshot, ProxyConfig}; use aisix_gateway::Hub; +use aisix_ratelimit::Limiter; use std::sync::Arc; #[derive(Clone)] pub struct ProxyState { pub snapshot: SnapshotHandle, pub hub: Arc, + pub limiter: Arc, pub request_body_limit_bytes: usize, } @@ -25,6 +29,23 @@ impl ProxyState { Self { snapshot, hub, + limiter: Arc::new(Limiter::new()), + request_body_limit_bytes: cfg.request_body_limit_bytes, + } + } + + /// Alternative constructor for callers that want to share a preexisting + /// limiter (e.g. tests with a deterministic clock). + pub fn with_limiter( + snapshot: SnapshotHandle, + hub: Arc, + limiter: Arc, + cfg: &ProxyConfig, + ) -> Self { + Self { + snapshot, + hub, + limiter, request_body_limit_bytes: cfg.request_body_limit_bytes, } } diff --git a/crates/aisix-ratelimit/src/clock.rs b/crates/aisix-ratelimit/src/clock.rs new file mode 100644 index 00000000..a99a9172 --- /dev/null +++ b/crates/aisix-ratelimit/src/clock.rs @@ -0,0 +1,75 @@ +//! Injected clock so the fixed-window counters can be tested +//! deterministically. The production [`SystemClock`] delegates to +//! `SystemTime::now()`; [`TestClock`] is a thread-safe stepper a test +//! can advance by hand. + +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::{SystemTime, UNIX_EPOCH}; + +/// Second-resolution wall-clock. That's all the fixed-window counters +/// need — they bucket by minute and by day boundaries. +pub trait Clock: Send + Sync + 'static { + fn unix_secs(&self) -> u64; +} + +#[derive(Debug, Default, Clone, Copy)] +pub struct SystemClock; + +impl Clock for SystemClock { + fn unix_secs(&self) -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0) + } +} + +/// Simple test double. Call [`TestClock::advance`] between operations +/// to jump to the next window without spinning wall-clock time. +#[derive(Debug, Clone, Default)] +pub struct TestClock { + now: Arc, +} + +impl TestClock { + pub fn new(initial_secs: u64) -> Self { + Self { + now: Arc::new(AtomicU64::new(initial_secs)), + } + } + + pub fn advance(&self, secs: u64) { + self.now.fetch_add(secs, Ordering::SeqCst); + } + + pub fn set(&self, secs: u64) { + self.now.store(secs, Ordering::SeqCst); + } +} + +impl Clock for TestClock { + fn unix_secs(&self) -> u64 { + self.now.load(Ordering::SeqCst) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn system_clock_returns_positive_now() { + assert!(SystemClock.unix_secs() > 0); + } + + #[test] + fn test_clock_advances_and_sets() { + let c = TestClock::new(100); + assert_eq!(c.unix_secs(), 100); + c.advance(30); + assert_eq!(c.unix_secs(), 130); + c.set(500); + assert_eq!(c.unix_secs(), 500); + } +} diff --git a/crates/aisix-ratelimit/src/error.rs b/crates/aisix-ratelimit/src/error.rs new file mode 100644 index 00000000..a31289a7 --- /dev/null +++ b/crates/aisix-ratelimit/src/error.rs @@ -0,0 +1,66 @@ +//! Limiter error taxonomy. +//! +//! Uses [`aisix_core::RateLimitScope`] so the proxy layer can plug the +//! error straight into its OpenAI-style envelope without translation. + +use aisix_core::RateLimitScope; + +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum RateLimitError { + #[error("request limit exceeded ({scope})")] + Requests { + scope: RateLimitScope, + retry_after_secs: u64, + }, + #[error("token limit exceeded ({scope})")] + Tokens { + scope: RateLimitScope, + retry_after_secs: u64, + }, + #[error("concurrency limit exceeded")] + Concurrency, +} + +impl RateLimitError { + pub fn scope(&self) -> RateLimitScope { + match self { + RateLimitError::Requests { scope, .. } => *scope, + RateLimitError::Tokens { scope, .. } => *scope, + RateLimitError::Concurrency => RateLimitScope::Concurrency, + } + } + + pub fn retry_after_secs(&self) -> Option { + match self { + RateLimitError::Requests { + retry_after_secs, .. + } + | RateLimitError::Tokens { + retry_after_secs, .. + } => Some(*retry_after_secs), + RateLimitError::Concurrency => None, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn requests_scope_preserved_on_access() { + let e = RateLimitError::Requests { + scope: RateLimitScope::Requests, + retry_after_secs: 42, + }; + assert_eq!(e.scope(), RateLimitScope::Requests); + assert_eq!(e.retry_after_secs(), Some(42)); + } + + #[test] + fn concurrency_has_no_retry_after_hint() { + let e = RateLimitError::Concurrency; + assert_eq!(e.scope(), RateLimitScope::Concurrency); + assert!(e.retry_after_secs().is_none()); + } +} diff --git a/crates/aisix-ratelimit/src/lib.rs b/crates/aisix-ratelimit/src/lib.rs index 3ffe2923..204a38d0 100644 --- a/crates/aisix-ratelimit/src/lib.rs +++ b/crates/aisix-ratelimit/src/lib.rs @@ -1,4 +1,24 @@ -//! aisix-ratelimit — fixed-window rate limiter with commit/read-only modes, plus concurrency permits. +//! aisix-ratelimit — two-phase RPM/TPM/concurrency limiter. +//! +//! The proxy middleware calls [`Limiter::pre_commit`] before dispatching +//! a chat request; the returned [`Reservation`] is finalised with +//! [`Reservation::commit_tokens`] after the upstream response completes. +//! +//! Limits come from [`aisix_core::RateLimit`] — currently attached to +//! `ApiKey` entries. RPM/RPD are checked-and-incremented up front so +//! burst traffic fails fast; TPM/TPD are checked-only up front and +//! incremented on commit because the token cost is only known after +//! the upstream response lands. #![forbid(unsafe_code)] #![deny(rust_2018_idioms)] + +pub mod clock; +mod error; +mod limiter; +mod window; + +pub use clock::{Clock, SystemClock, TestClock}; +pub use error::RateLimitError; +pub use limiter::{Limiter, Reservation}; +pub use window::{FixedWindowCounter, WindowCheck}; diff --git a/crates/aisix-ratelimit/src/limiter.rs b/crates/aisix-ratelimit/src/limiter.rs new file mode 100644 index 00000000..42345e20 --- /dev/null +++ b/crates/aisix-ratelimit/src/limiter.rs @@ -0,0 +1,340 @@ +//! Two-phase limiter keyed on an opaque `key` (the caller's ApiKey id +//! in production). +//! +//! Phase 1 — **pre-commit**, called before the upstream request fires: +//! - check concurrency (acquire a permit or fail) +//! - check + increment RPM / RPD counters +//! - *check-only* TPM / TPD (we don't know the token cost yet) +//! +//! Phase 2 — **post-deduct**, called after the upstream response +//! completes: +//! - add actual `prompt_tokens + completion_tokens` to TPM / TPD +//! - release the concurrency permit +//! +//! The returned [`Reservation`] handle wraps the concurrency permit so +//! callers cannot forget to release on the error path — the permit is +//! released on drop if `commit_tokens` / `abort` isn't called. + +use aisix_core::{RateLimit, RateLimitScope}; +use dashmap::DashMap; +use parking_lot::Mutex; +use std::sync::Arc; + +use crate::clock::{Clock, SystemClock}; +use crate::error::RateLimitError; +use crate::window::{FixedWindowCounter, WindowCheck}; + +const MINUTE_SECS: u64 = 60; +const DAY_SECS: u64 = 24 * 60 * 60; + +/// Per-key state guarded by a single mutex. Hot path locks once per +/// request; each operation inside is O(1). +#[derive(Debug)] +struct KeyState { + rpm: FixedWindowCounter, + rpd: FixedWindowCounter, + tpm: FixedWindowCounter, + tpd: FixedWindowCounter, + in_flight: u32, +} + +impl KeyState { + fn new() -> Self { + Self { + rpm: FixedWindowCounter::new(MINUTE_SECS), + rpd: FixedWindowCounter::new(DAY_SECS), + tpm: FixedWindowCounter::new(MINUTE_SECS), + tpd: FixedWindowCounter::new(DAY_SECS), + in_flight: 0, + } + } +} + +pub struct Limiter { + states: DashMap>>, + clock: C, +} + +impl Limiter { + pub fn new() -> Self { + Self::with_clock(SystemClock) + } +} + +impl Default for Limiter { + fn default() -> Self { + Self::new() + } +} + +impl Limiter { + pub fn with_clock(clock: C) -> Self { + Self { + states: DashMap::new(), + clock, + } + } + + fn state_for(&self, key: &str) -> Arc> { + if let Some(entry) = self.states.get(key) { + return entry.clone(); + } + self.states + .entry(key.to_string()) + .or_insert_with(|| Arc::new(Mutex::new(KeyState::new()))) + .clone() + } + + /// Pre-commit phase. Returns a [`Reservation`] that must be finalised + /// via [`Limiter::commit_tokens`] or dropped to release the + /// concurrency permit automatically. + pub fn pre_commit( + &self, + key: &str, + limits: &RateLimit, + ) -> Result, RateLimitError> { + let now = self.clock.unix_secs(); + let state = self.state_for(key); + let mut s = state.lock(); + + // Concurrency first — cheapest and never consumes a window slot. + if let Some(max) = limits.concurrency { + if s.in_flight >= max { + return Err(RateLimitError::Concurrency); + } + } + + // Token limits — checked but not incremented. We refuse new + // requests if the previous minute/day already overran the cap. + if let Some(max) = limits.tpm { + if let Some(retry) = s.tpm.is_exceeded(now, max) { + return Err(RateLimitError::Tokens { + scope: RateLimitScope::Tokens, + retry_after_secs: retry, + }); + } + } + if let Some(max) = limits.tpd { + if let Some(retry) = s.tpd.is_exceeded(now, max) { + return Err(RateLimitError::Tokens { + scope: RateLimitScope::Tokens, + retry_after_secs: retry, + }); + } + } + + // Request limits — checked AND incremented. + if let Some(max) = limits.rpm { + if let WindowCheck::Full { retry_after_secs } = s.rpm.check_and_increment(now, 1, max) { + return Err(RateLimitError::Requests { + scope: RateLimitScope::Requests, + retry_after_secs, + }); + } + } + if let Some(max) = limits.rpd { + if let WindowCheck::Full { retry_after_secs } = s.rpd.check_and_increment(now, 1, max) { + // Compensate: we already incremented RPM above. Decrement + // it so the caller's retry on a different day still + // counts correctly. RPM would have rolled by then, so + // this is primarily defensive. + if s.rpm.current(now) > 0 { + // Roll back the increment we just made. + s.rpm = FixedWindowCounter::new(MINUTE_SECS); + if let Some(max) = limits.rpm { + let _ = s.rpm.check_and_increment(now, 0, max); + } + } + return Err(RateLimitError::Requests { + scope: RateLimitScope::Requests, + retry_after_secs, + }); + } + } + + s.in_flight += 1; + drop(s); + + Ok(Reservation { + limiter: self, + key: key.to_string(), + committed: false, + }) + } +} + +/// Reservation guard. Dropping without a `commit_tokens` call is still +/// safe — the concurrency permit is released, just no tokens are +/// counted. +pub struct Reservation<'a, C: Clock> { + limiter: &'a Limiter, + key: String, + committed: bool, +} + +impl<'a, C: Clock> std::fmt::Debug for Reservation<'a, C> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Reservation") + .field("key", &self.key) + .field("committed", &self.committed) + .finish() + } +} + +impl<'a, C: Clock> Reservation<'a, C> { + /// Post-deduct phase. Records the actual token cost against TPM/TPD + /// and releases the concurrency permit. + pub fn commit_tokens(mut self, tokens: u64) { + let now = self.limiter.clock.unix_secs(); + let state = self.limiter.state_for(&self.key); + let mut s = state.lock(); + s.tpm.add(now, tokens); + s.tpd.add(now, tokens); + s.in_flight = s.in_flight.saturating_sub(1); + self.committed = true; + } +} + +impl<'a, C: Clock> Drop for Reservation<'a, C> { + fn drop(&mut self) { + if self.committed { + return; + } + let state = self.limiter.state_for(&self.key); + let mut s = state.lock(); + s.in_flight = s.in_flight.saturating_sub(1); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::clock::TestClock; + + fn limits(rpm: Option, tpm: Option, concurrency: Option) -> RateLimit { + RateLimit { + rpm, + rpd: None, + tpm, + tpd: None, + concurrency, + } + } + + #[test] + fn rpm_caps_request_count_in_window() { + let clock = TestClock::new(100); + let limiter = Limiter::with_clock(clock.clone()); + let l = limits(Some(2), None, None); + + let _r1 = limiter.pre_commit("k1", &l).unwrap(); + let _r2 = limiter.pre_commit("k1", &l).unwrap(); + let err = limiter.pre_commit("k1", &l).unwrap_err(); + match err { + RateLimitError::Requests { + retry_after_secs, .. + } => { + assert!(retry_after_secs > 0); + } + other => panic!("expected Requests, got {other:?}"), + } + } + + #[test] + fn rpm_resets_after_window_rollover() { + let clock = TestClock::new(100); + let limiter = Limiter::with_clock(clock.clone()); + let l = limits(Some(1), None, None); + + let _r1 = limiter.pre_commit("k1", &l).unwrap(); + assert!(limiter.pre_commit("k1", &l).is_err()); + + // Jump past the minute boundary. + clock.advance(61); + let _r2 = limiter.pre_commit("k1", &l).unwrap(); + } + + #[test] + fn concurrency_limit_blocks_new_reservations() { + let clock = TestClock::new(0); + let limiter = Limiter::with_clock(clock.clone()); + let l = limits(None, None, Some(2)); + + let r1 = limiter.pre_commit("k1", &l).unwrap(); + let r2 = limiter.pre_commit("k1", &l).unwrap(); + assert!(matches!( + limiter.pre_commit("k1", &l).unwrap_err(), + RateLimitError::Concurrency, + )); + + // Drop r1 — concurrency should free up. + drop(r1); + let _r3 = limiter.pre_commit("k1", &l).unwrap(); + drop(r2); + } + + #[test] + fn token_commit_updates_post_deduct_counters() { + let clock = TestClock::new(100); + let limiter = Limiter::with_clock(clock.clone()); + let l = limits(Some(10), Some(1_000), None); + + let r1 = limiter.pre_commit("k1", &l).unwrap(); + r1.commit_tokens(600); + + // TPM now at 600. Next pre_commit with a strict TPM should still + // succeed because 600 <= 1000. + let _r2 = limiter.pre_commit("k1", &l).unwrap(); + } + + #[test] + fn tpm_blocks_next_request_once_previous_exhausted_the_window() { + let clock = TestClock::new(100); + let limiter = Limiter::with_clock(clock.clone()); + let l = limits(Some(10), Some(1_000), None); + + let r1 = limiter.pre_commit("k1", &l).unwrap(); + r1.commit_tokens(1_500); // overshoot — allowed for the in-flight request + + // Next pre_commit sees tpm > 1000 and refuses. + let err = limiter.pre_commit("k1", &l).unwrap_err(); + assert!(matches!(err, RateLimitError::Tokens { .. })); + + clock.advance(61); // roll the window + let _r2 = limiter.pre_commit("k1", &l).unwrap(); + } + + #[test] + fn reservations_for_different_keys_do_not_collide() { + let clock = TestClock::new(0); + let limiter = Limiter::with_clock(clock); + let l = limits(Some(1), None, None); + + let _r_a = limiter.pre_commit("alpha", &l).unwrap(); + let _r_b = limiter.pre_commit("beta", &l).unwrap(); + } + + #[test] + fn drop_without_commit_still_releases_concurrency_permit() { + let clock = TestClock::new(0); + let limiter = Limiter::with_clock(clock); + let l = limits(None, None, Some(1)); + + { + let _r = limiter.pre_commit("k1", &l).unwrap(); + } // dropped + let _r2 = limiter.pre_commit("k1", &l).unwrap(); + } + + #[test] + fn no_limits_means_no_rejections() { + let clock = TestClock::new(0); + let limiter = Limiter::with_clock(clock); + let l = RateLimit::default(); + + for _ in 0..100 { + let r = limiter.pre_commit("k1", &l).unwrap(); + r.commit_tokens(1_000); + } + } +} diff --git a/crates/aisix-ratelimit/src/window.rs b/crates/aisix-ratelimit/src/window.rs new file mode 100644 index 00000000..8f3b8d90 --- /dev/null +++ b/crates/aisix-ratelimit/src/window.rs @@ -0,0 +1,169 @@ +//! Fixed-window counter. +//! +//! A single counter that resets at the start of every new window. Used +//! for both the per-minute and per-day dimensions; callers instantiate +//! one per dimension with the matching `window_secs`. +//! +//! Not thread-safe on its own — the caller is expected to hold the +//! `KeyState` guard before touching it. That keeps the hot path lock- +//! cheap: one `DashMap` shard + one `parking_lot::Mutex` per key. + +/// Result of attempting to reserve capacity. +#[derive(Debug, PartialEq, Eq)] +pub enum WindowCheck { + Ok, + Full { retry_after_secs: u64 }, +} + +#[derive(Debug)] +pub struct FixedWindowCounter { + window_secs: u64, + window_start: u64, + count: u64, +} + +impl FixedWindowCounter { + pub fn new(window_secs: u64) -> Self { + assert!(window_secs > 0, "window_secs must be positive"); + Self { + window_secs, + window_start: 0, + count: 0, + } + } + + pub fn window_secs(&self) -> u64 { + self.window_secs + } + + fn roll_if_stale(&mut self, now_secs: u64) { + let bucket_start = (now_secs / self.window_secs) * self.window_secs; + if bucket_start != self.window_start { + self.window_start = bucket_start; + self.count = 0; + } + } + + /// Check whether `delta` more units would fit under `limit`. If yes, + /// commit them (increment counter) and return `Ok`. If no, return + /// `Full` with seconds until the window rolls over. + pub fn check_and_increment(&mut self, now_secs: u64, delta: u64, limit: u64) -> WindowCheck { + self.roll_if_stale(now_secs); + let would_be = self.count.saturating_add(delta); + if would_be > limit { + let remainder = self + .window_secs + .saturating_sub(now_secs.saturating_sub(self.window_start)); + return WindowCheck::Full { + retry_after_secs: remainder.max(1), + }; + } + self.count = would_be; + WindowCheck::Ok + } + + /// Add to the counter without a check. Used on the post-deduct side + /// for TPM/TPD — the token count isn't known until the upstream + /// response has completed, so we record after the fact. + pub fn add(&mut self, now_secs: u64, delta: u64) { + self.roll_if_stale(now_secs); + self.count = self.count.saturating_add(delta); + } + + pub fn current(&mut self, now_secs: u64) -> u64 { + self.roll_if_stale(now_secs); + self.count + } + + /// Peek at whether the current count already exceeds the limit. Used + /// on the *next* request's pre-commit to short-circuit before + /// increment: TPM is checked-but-not-incremented at pre-commit, then + /// incremented on post-deduct by the actual token usage. + pub fn is_exceeded(&mut self, now_secs: u64, limit: u64) -> Option { + self.roll_if_stale(now_secs); + if self.count > limit { + let remainder = self + .window_secs + .saturating_sub(now_secs.saturating_sub(self.window_start)); + Some(remainder.max(1)) + } else { + None + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn first_increments_fit_then_subsequent_block() { + let mut w = FixedWindowCounter::new(60); + assert_eq!(w.check_and_increment(100, 1, 3), WindowCheck::Ok); + assert_eq!(w.check_and_increment(100, 1, 3), WindowCheck::Ok); + assert_eq!(w.check_and_increment(100, 1, 3), WindowCheck::Ok); + // 4th attempt overflows the limit. + match w.check_and_increment(100, 1, 3) { + WindowCheck::Full { retry_after_secs } => { + assert!(retry_after_secs > 0); + assert!(retry_after_secs <= 60); + } + _ => panic!("expected Full"), + } + } + + #[test] + fn counter_rolls_over_at_window_boundary() { + let mut w = FixedWindowCounter::new(60); + for _ in 0..3 { + w.check_and_increment(100, 1, 3); + } + // Cross into the next minute. + assert_eq!(w.check_and_increment(161, 1, 3), WindowCheck::Ok); + assert_eq!(w.current(161), 1); + } + + #[test] + fn retry_after_reflects_time_remaining_in_window() { + let mut w = FixedWindowCounter::new(60); + // Fill the bucket at second 100 (bucket starts at 60, ends at 120). + for _ in 0..3 { + w.check_and_increment(100, 1, 3); + } + match w.check_and_increment(110, 1, 3) { + WindowCheck::Full { retry_after_secs } => { + assert_eq!(retry_after_secs, 10); // 60+60 - 110 = 10 + } + _ => panic!("expected Full"), + } + } + + #[test] + fn add_records_post_deduct_usage_and_is_checkable() { + let mut w = FixedWindowCounter::new(60); + w.add(100, 1_000); + w.add(101, 500); + assert_eq!(w.current(101), 1_500); + + assert!(w.is_exceeded(101, 2_000).is_none()); // 1500 <= 2000 + assert!(w.is_exceeded(101, 1_000).is_some()); // 1500 > 1000 + } + + #[test] + fn check_with_zero_delta_is_a_read_only_peek_that_succeeds() { + let mut w = FixedWindowCounter::new(60); + assert_eq!(w.check_and_increment(100, 0, 5), WindowCheck::Ok); + assert_eq!(w.current(100), 0); + } + + #[test] + fn retry_after_is_at_least_one_second() { + let mut w = FixedWindowCounter::new(60); + // Window is [60, 120). Fill it past the cap. + w.add(100, 1_000); + // Query right at the last second of the same window — remainder + // would be 0, but we clamp to 1 so clients don't spin. + let hint = w.is_exceeded(119, 100).unwrap(); + assert!(hint >= 1); + } +}