From 5102235d26976e2b9927f745c16e01cae57b9502 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Mon, 15 Jun 2026 11:49:56 +0800 Subject: [PATCH 1/2] feat(ratelimit): cluster-level rate limiting via shared Redis MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rate-limit counters lived in per-process memory, so an N-replica DP cluster enforced N× every configured limit (a key capped at rpm:1 got one request per replica per minute). Add a Redis-backed shared store so the whole cluster enforces one global window. - Introduce a `RateStore` backend behind `Limiter`: `LocalStore` (unchanged in-memory default) and `RedisStore` (Lua check-and-increment over wall-clock-aligned fixed windows, `redis.call('TIME')` for cross-replica window consistency, hash-tagged keys for Cluster slot co-location). All dimensions are shared — rps/rpm/rph/rpd/tpm/tpd plus concurrency, tracked as a crash-safe ZSET semaphore reclaimed after `concurrency_ttl_secs`. On a Redis outage the store fails open to per-replica counting. - Make the enforcement path async (`pre_commit`/`commit_tokens`/`peek`); concurrency release stays a sync `Drop` (Redis detaches a ZREM). - New `ratelimit` config block (`backend: memory|redis`, `redis`, `concurrency_ttl_secs`), enabled via env on managed deployments. Tests: Rust integration (gated on RATELIMIT_TEST_REDIS_URL) for shared rpm/rps/tpm/concurrency + TTL reclaim; DP e2e spins two real binaries on one Redis (A→200, B→429) plus a memory-backend regression (both 200). Fixes api7/AISIX-Cloud#798 --- .github/workflows/ci.yml | 4 + Cargo.lock | 3 + config.example.yaml | 18 + config.managed.yaml | 6 + crates/aisix-core/src/config.rs | 113 ++ crates/aisix-core/src/lib.rs | 3 +- crates/aisix-proxy/src/chat.rs | 23 +- crates/aisix-proxy/src/embeddings.rs | 8 +- crates/aisix-proxy/src/quota.rs | 25 +- crates/aisix-ratelimit/Cargo.toml | 3 + crates/aisix-ratelimit/src/lib.rs | 4 + crates/aisix-ratelimit/src/limiter.rs | 972 ++++++------------ crates/aisix-ratelimit/src/store/local.rs | 245 +++++ crates/aisix-ratelimit/src/store/mod.rs | 122 +++ crates/aisix-ratelimit/src/store/redis.rs | 407 ++++++++ .../tests/redis_integration.rs | 198 ++++ crates/aisix-server/src/main.rs | 26 +- docs/configuration/rate-limits.md | 26 + .../src/cases/ratelimit-cluster-e2e.test.ts | 214 ++++ 19 files changed, 1744 insertions(+), 676 deletions(-) create mode 100644 crates/aisix-ratelimit/src/store/local.rs create mode 100644 crates/aisix-ratelimit/src/store/mod.rs create mode 100644 crates/aisix-ratelimit/src/store/redis.rs create mode 100644 crates/aisix-ratelimit/tests/redis_integration.rs create mode 100644 tests/e2e/src/cases/ratelimit-cluster-e2e.test.ts diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 71809045..71d7d622 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -79,6 +79,10 @@ jobs: # every Config::load_from_path test as `redis_url` and breaks # the entire `aisix-core::config::tests` module. CACHE_TEST_REDIS_URL: redis://127.0.0.1:6379 + # Picked up by crates/aisix-ratelimit/tests/redis_integration.rs + # (shared cluster-level counters, #798). Same skip-if-unset / + # no-AISIX_-prefix rules as CACHE_TEST_REDIS_URL above. + RATELIMIT_TEST_REDIS_URL: redis://127.0.0.1:6379 # Picked up by crates/aisix-admin/tests/etcd_integration.rs. # Same skip-if-unset pattern as the Redis case above. ADMIN_TEST_ETCD_URL: http://127.0.0.1:2379 diff --git a/Cargo.lock b/Cargo.lock index 3ab27f58..a6b49565 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -362,15 +362,18 @@ name = "aisix-ratelimit" version = "0.1.0" dependencies = [ "aisix-core", + "async-trait", "chrono", "dashmap", "parking_lot", + "redis", "rstest", "serde", "serde_json", "thiserror 1.0.69", "tokio", "tracing", + "uuid", ] [[package]] diff --git a/config.example.yaml b/config.example.yaml index 8b88dbb7..5021991e 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -91,6 +91,24 @@ cache: # url: "redis://127.0.0.1:6379" # mode: "single" # single | cluster | sentinel +# Rate-limit counter backend (api7/AISIX-Cloud#798). +# +# `memory` (default) keeps counters in this process, so a cluster of N +# replicas enforces N× every configured limit. `redis` shares the +# counters across every replica via one Redis, so the whole cluster +# enforces ONE global window — set this on multi-replica deployments. +# May point at the same Redis as `cache` (keys are namespaced +# `aisix:rl:`). On a Redis outage the limiter fails open to per-replica +# in-memory counting (logged) so traffic keeps flowing. +ratelimit: + backend: "memory" # memory | redis + # redis: + # url: "redis://127.0.0.1:6379" + # mode: "single" # single | cluster | sentinel + # Seconds before an unreleased concurrency slot (crashed replica / + # hung upstream) is reclaimed. Redis backend only. + # concurrency_ttl_secs: 300 + # Models, API keys, provider keys, guardrails, cache policies, and # observability exporters are NOT defined in this file. They are stored # in etcd and managed via the Admin API (see docs/api-admin.md). This diff --git a/config.managed.yaml b/config.managed.yaml index 1c6aa381..95f00bad 100644 --- a/config.managed.yaml +++ b/config.managed.yaml @@ -22,6 +22,12 @@ # `AISIX_CACHE__BACKEND`, etc. — every config field is reachable via # `AISIX___` (see crates/aisix-core/src/config.rs). # +# For a multi-replica deployment, enable cluster-level rate limiting so +# the cluster enforces one global window instead of N× per replica +# (api7/AISIX-Cloud#798): +# AISIX_RATELIMIT__BACKEND=redis +# AISIX_RATELIMIT__REDIS__URL=redis://:6379 +# # Subsequent boots re-use the mTLS bundle written under # `managed.mtls_dir`. diff --git a/crates/aisix-core/src/config.rs b/crates/aisix-core/src/config.rs index ddeef3a7..0e934561 100644 --- a/crates/aisix-core/src/config.rs +++ b/crates/aisix-core/src/config.rs @@ -43,6 +43,12 @@ pub struct Config { pub observability: ObservabilityConfig, #[serde(default)] pub cache: CacheConfig, + /// Rate-limit counter backend. Defaults to per-process memory + /// (historical behaviour). Set `backend: redis` with a `redis` block + /// to share counters across every DP replica so a cluster enforces + /// one global window instead of one-per-replica (api7/AISIX-Cloud#798). + #[serde(default)] + pub ratelimit: RateLimitConfig, /// Optional managed-mode configuration. When `managed.enabled = true` /// the admin API and Playground endpoints are **not** bound — the DP /// is a pure etcd reader driven by the aisix.cloud control plane. @@ -559,6 +565,43 @@ impl RedisCacheConfig { } } +/// Rate-limit counter backend (api7/AISIX-Cloud#798). +/// +/// `Memory` is the default: per-process fixed-window counters, so an +/// N-replica cluster enforces N× the configured limit. `Redis` shares +/// the counters across replicas via a single Redis so the whole cluster +/// enforces one global window. The `redis` block is required iff +/// `backend = redis` (validated at boot). Reuses [`RedisCacheConfig`] +/// for the connection shape; may point at the same Redis as `cache` +/// (keys are namespaced `aisix:rl:`). +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields, default)] +pub struct RateLimitConfig { + pub backend: RateLimitBackend, + pub redis: Option, + /// Seconds after which an unreleased concurrency slot is reclaimed + /// (crashed replica / hung upstream). Generous enough for a long + /// streaming response. Redis backend only. + pub concurrency_ttl_secs: u64, +} + +impl Default for RateLimitConfig { + fn default() -> Self { + Self { + backend: RateLimitBackend::Memory, + redis: None, + concurrency_ttl_secs: 300, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum RateLimitBackend { + Memory, + Redis, +} + impl Config { /// Load + merge + validate. /// @@ -658,6 +701,11 @@ impl Config { "observability.metrics.prometheus.addr invalid socket address: {metrics_addr}" ))); } + if self.ratelimit.backend == RateLimitBackend::Redis && self.ratelimit.redis.is_none() { + return Err(BootstrapError::Config( + "ratelimit.backend = redis requires a ratelimit.redis block".into(), + )); + } Ok(()) } } @@ -784,6 +832,71 @@ admin: assert!(err.to_string().contains("admin.admin_keys")); } + #[test] + fn ratelimit_defaults_to_memory_backend() { + let f = write_yaml( + r#" +etcd: + endpoints: ["http://localhost:2379"] +proxy: + addr: "0.0.0.0:3000" +admin: + addr: "127.0.0.1:3001" + admin_keys: ["k1"] +"#, + ); + let cfg = Config::load_from_path(Some(f.path())).unwrap(); + assert_eq!(cfg.ratelimit.backend, RateLimitBackend::Memory); + assert!(cfg.ratelimit.redis.is_none()); + assert_eq!(cfg.ratelimit.concurrency_ttl_secs, 300); + } + + #[test] + fn ratelimit_redis_backend_requires_redis_block() { + let f = write_yaml( + r#" +etcd: + endpoints: ["http://localhost:2379"] +proxy: + addr: "0.0.0.0:3000" +admin: + addr: "127.0.0.1:3001" + admin_keys: ["k1"] +ratelimit: + backend: "redis" +"#, + ); + let err = Config::load_from_path(Some(f.path())).unwrap_err(); + assert!(err.to_string().contains("ratelimit.redis")); + } + + #[test] + fn loads_ratelimit_redis_config() { + let f = write_yaml( + r#" +etcd: + endpoints: ["http://localhost:2379"] +proxy: + addr: "0.0.0.0:3000" +admin: + addr: "127.0.0.1:3001" + admin_keys: ["k1"] +ratelimit: + backend: "redis" + redis: + url: "redis://127.0.0.1:6379" + concurrency_ttl_secs: 120 +"#, + ); + let cfg = Config::load_from_path(Some(f.path())).unwrap(); + assert_eq!(cfg.ratelimit.backend, RateLimitBackend::Redis); + assert_eq!( + cfg.ratelimit.redis.as_ref().unwrap().url, + "redis://127.0.0.1:6379" + ); + assert_eq!(cfg.ratelimit.concurrency_ttl_secs, 120); + } + #[test] fn rejects_invalid_bind_addr() { let f = write_yaml( diff --git a/crates/aisix-core/src/lib.rs b/crates/aisix-core/src/lib.rs index 4a20d9ca..f128fc5e 100644 --- a/crates/aisix-core/src/lib.rs +++ b/crates/aisix-core/src/lib.rs @@ -24,7 +24,8 @@ pub mod snapshot; pub use config::{ AdminConfig, CacheBackend, CacheConfig, Config, EtcdConfig, EtcdTlsConfig, ManagedConfig, - ObservabilityConfig, ProxyConfig, RealIpConfig, TlsConfig, + ObservabilityConfig, ProxyConfig, RateLimitBackend, RateLimitConfig, RealIpConfig, + RedisCacheConfig, TlsConfig, }; pub use error::{ AdminError, AdminErrorEnvelope, BootstrapError, ProxyError, ProxyErrorEnvelope, RateLimitScope, diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index 6f28a230..1d1bfc16 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -245,7 +245,7 @@ pub async fn chat_completions( // current window state. We peek *after* the commit so // remaining-requests reflects the post-dispatch tally. let rl_limits = auth.key().rate_limit.clone().unwrap_or_default(); - if let Some(rl_status) = state.limiter.peek(&api_key_id, &rl_limits) { + if let Some(rl_status) = state.limiter.peek(&api_key_id, &rl_limits).await { crate::render::inject_ratelimit_headers(&mut success.response, &rl_status); state.metrics.set_rate_limit_remaining( &api_key_id, @@ -843,8 +843,9 @@ async fn dispatch( &virtual_entry.id, &virtual_entry.value, ); - let reservation = - crate::quota::enforce_rate_limit(state, auth, Some(&model_rl)).map_err(&with_model)?; + let reservation = crate::quota::enforce_rate_limit(state, auth, Some(&model_rl)) + .await + .map_err(&with_model)?; let now = created_ts(); @@ -1073,7 +1074,7 @@ async fn dispatch( // permit was released here, letting a key capped at N run far more // than N simultaneous streams (#450). let post_stream_keys = reservation.keys(); - let stream_concurrency_hold = reservation.into_stream_hold(Arc::clone(&state.limiter)); + let stream_concurrency_hold = reservation.into_stream_hold(); // Capture everything the stream-completion callback needs so // it can fire `emit_usage_event` once the terminal SSE chunk // has yielded its `usage` block. Telemetry emission has to @@ -1379,7 +1380,7 @@ async fn dispatch( if let (Some(cache), Some(key)) = (policy_cache.as_ref(), cache_key.as_ref()) { match cache.get(key).await { Ok(Some(cached)) => { - reservation.commit_tokens(0); + reservation.commit_tokens(0).await; // #448: a cache hit is client-visible output just like a // fresh upstream response, so it must run output guardrails // before being returned — not bypass them. @@ -1690,7 +1691,7 @@ async fn dispatch( let provider_request_id = upstream.id.clone(); let provider_model_version = upstream.model.clone(); let finish_reason = finish_reason_label(&upstream.finish_reason); - reservation.commit_tokens(total); + reservation.commit_tokens(total).await; // cp-api recomputes cost server-side from its pricing catalog when // ingesting telemetry; the DP just records 0.0 on the wire. @@ -1860,14 +1861,14 @@ async fn dispatch( /// response. `reservation` is the SINGLE entry-level reservation taken /// in `dispatch` — ensemble does not add per-sub-call reservations. #[allow(clippy::too_many_arguments)] -async fn dispatch_ensemble<'a>( - state: &'a ProxyState, +async fn dispatch_ensemble( + state: &ProxyState, snapshot: &aisix_core::AisixSnapshot, virtual_entry: &aisix_core::ResourceEntry, req: &ChatFormat, request_id: &str, created_ts: i64, - reservation: aisix_ratelimit::MultiReservation<'a, aisix_ratelimit::SystemClock>, + reservation: aisix_ratelimit::MultiReservation, resolved_chain: &Arc, applied_guardrails: &[AppliedGuardrail], mut bypass_reason: Option, @@ -2006,7 +2007,7 @@ async fn dispatch_ensemble<'a>( ), }; let survivor_total: u64 = panel.iter().map(|p| u64::from(p.usage.total_tokens)).sum(); - reservation.commit_tokens(survivor_total); + reservation.commit_tokens(survivor_total).await; for (index, member) in panel.iter().enumerate() { emit_panel_member( member, index, /* blocked */ false, /* bypass */ "", @@ -2030,7 +2031,7 @@ async fn dispatch_ensemble<'a>( .sum(); let judge_usage = outcome.response.usage.clone(); let total_tokens = panel_total + u64::from(judge_usage.total_tokens); - reservation.commit_tokens(total_tokens); + reservation.commit_tokens(total_tokens).await; // Emit one usage event per sub-call (each panel member + the judge), // all sharing `request_id`. `attempt_kind` is `"panel"` / `"judge"`; diff --git a/crates/aisix-proxy/src/embeddings.rs b/crates/aisix-proxy/src/embeddings.rs index 7ba098bc..2b6f4636 100644 --- a/crates/aisix-proxy/src/embeddings.rs +++ b/crates/aisix-proxy/src/embeddings.rs @@ -321,7 +321,9 @@ async fn dispatch( // and finalise RPM. Embeddings do report prompt_tokens via // EmbeddingResponse.usage; thread it through so TPM works // here even though other handlers commit 0. - reservation.commit_tokens(embed_resp.usage.total_tokens as u64); + reservation + .commit_tokens(embed_resp.usage.total_tokens as u64) + .await; let provider_label = provider.to_ascii_lowercase(); // Capture the prompt_tokens count BEFORE moving the // embed_resp into the JSON response — the handler needs @@ -342,7 +344,7 @@ async fn dispatch( // (`upstream_called: false` → handler skips emit per the // chat.rs convention that we only attribute usage on a // real upstream completion). - reservation.commit_tokens(0); + reservation.commit_tokens(0).await; let env = ErrorEnvelope::new(msg, "not_implemented"); Ok(EmbedDispatchSuccess { response: (StatusCode::NOT_IMPLEMENTED, Json(env)).into_response(), @@ -357,7 +359,7 @@ async fn dispatch( }) } Err(e) => { - reservation.commit_tokens(0); + reservation.commit_tokens(0).await; Err(ProxyError::Bridge(e)) } } diff --git a/crates/aisix-proxy/src/quota.rs b/crates/aisix-proxy/src/quota.rs index 28b09fdd..4db97421 100644 --- a/crates/aisix-proxy/src/quota.rs +++ b/crates/aisix-proxy/src/quota.rs @@ -121,11 +121,11 @@ fn policy_bucket_key(policy: &RateLimitPolicy, entry_id: &str, auth: &Authentica } /// Reserve across all applicable rate-limit layers (api_key, model, policies). -fn reserve_layers<'a>( - state: &'a ProxyState, +async fn reserve_layers( + state: &ProxyState, auth: &AuthenticatedKey, model_rl: Option<&ModelRateLimit>, -) -> Result, ProxyError> { +) -> Result { let mut reservations = Vec::with_capacity(8); // Layer 1: API key inline rate limit. @@ -134,6 +134,7 @@ fn reserve_layers<'a>( let r = state .limiter .pre_commit(&auth.entry.id, &key_limits) + .await .map_err(ProxyError::from)?; reservations.push(r); } @@ -145,6 +146,7 @@ fn reserve_layers<'a>( let r = state .limiter .pre_commit(&key, limits) + .await .map_err(ProxyError::from)?; reservations.push(r); } @@ -179,6 +181,7 @@ fn reserve_layers<'a>( let r = state .limiter .pre_commit(&bucket_key, &rl) + .await .map_err(ProxyError::from)?; reservations.push(r); } @@ -190,11 +193,11 @@ fn reserve_layers<'a>( /// `model_rl` carries the resolved model identity for policy matching /// and optional inline limits. Pass `None` only for endpoints that /// don't resolve a model (e.g. passthrough). -pub(crate) async fn enforce<'a>( - state: &'a ProxyState, +pub(crate) async fn enforce( + state: &ProxyState, auth: &AuthenticatedKey, model_rl: Option<&ModelRateLimit>, -) -> Result, ProxyError> { +) -> Result { let decision = state.budgets.check(&auth.entry.id).await; let budget_labels = aisix_obs::BudgetLabels { api_key_id: &auth.entry.id, @@ -222,17 +225,17 @@ pub(crate) async fn enforce<'a>( ))); } - reserve_layers(state, auth, model_rl) + reserve_layers(state, auth, model_rl).await } /// Rate-limit-only enforcement (no budget check). Used by `chat.rs` /// which handles budget separately. -pub(crate) fn enforce_rate_limit<'a>( - state: &'a ProxyState, +pub(crate) async fn enforce_rate_limit( + state: &ProxyState, auth: &AuthenticatedKey, model_rl: Option<&ModelRateLimit>, -) -> Result, ProxyError> { - reserve_layers(state, auth, model_rl) +) -> Result { + reserve_layers(state, auth, model_rl).await } #[cfg(test)] diff --git a/crates/aisix-ratelimit/Cargo.toml b/crates/aisix-ratelimit/Cargo.toml index 9bd1cde5..2619ce65 100644 --- a/crates/aisix-ratelimit/Cargo.toml +++ b/crates/aisix-ratelimit/Cargo.toml @@ -18,6 +18,9 @@ serde_json.workspace = true thiserror.workspace = true tracing.workspace = true chrono.workspace = true +async-trait.workspace = true +redis.workspace = true +uuid.workspace = true [dev-dependencies] rstest.workspace = true diff --git a/crates/aisix-ratelimit/src/lib.rs b/crates/aisix-ratelimit/src/lib.rs index c5a342f6..9f37f6b0 100644 --- a/crates/aisix-ratelimit/src/lib.rs +++ b/crates/aisix-ratelimit/src/lib.rs @@ -16,6 +16,7 @@ pub mod clock; mod error; mod limiter; +pub mod store; mod window; pub use clock::{Clock, SystemClock, TestClock}; @@ -23,4 +24,7 @@ pub use error::RateLimitError; pub use limiter::{ Limiter, MultiReservation, RateLimitStatus, Reservation, StreamConcurrencyGuard, }; +pub use store::local::LocalStore; +pub use store::redis::RedisStore; +pub use store::RateStore; pub use window::{FixedWindowCounter, WindowCheck}; diff --git a/crates/aisix-ratelimit/src/limiter.rs b/crates/aisix-ratelimit/src/limiter.rs index f8a23669..d1dd0837 100644 --- a/crates/aisix-ratelimit/src/limiter.rs +++ b/crates/aisix-ratelimit/src/limiter.rs @@ -1,64 +1,35 @@ -//! Two-phase limiter keyed on an opaque `key` (the caller's ApiKey id -//! in production). +//! Two-phase limiter keyed on an opaque `key` (the caller's ApiKey id / +//! policy bucket in production), backed by a pluggable [`RateStore`]. //! //! Phase 1 — **pre-commit**, called before the upstream request fires: -//! - check concurrency (acquire a permit or fail) -//! - check + increment RPM / RPD counters +//! - check concurrency (acquire a slot or fail) +//! - check + increment RPS / RPM / RPH / 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 +//! - release the concurrency slot //! -//! 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. +//! The returned [`Reservation`] handle releases the concurrency slot on +//! drop if `commit_tokens` isn't called, so callers can't forget on the +//! error path. +//! +//! The counters live wherever the [`RateStore`] keeps them: the default +//! [`crate::store::local::LocalStore`] is per-process (historical +//! behaviour), while [`crate::store::redis::RedisStore`] shares them +//! across every DP replica so a cluster enforces one global window +//! (api7/AISIX-Cloud#798). -use aisix_core::{RateLimit, RateLimitScope}; -use dashmap::DashMap; -use parking_lot::Mutex; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; -use crate::clock::{Clock, SystemClock}; -use crate::error::RateLimitError; -use crate::window::{FixedWindowCounter, WindowCheck}; - -const SECOND_SECS: u64 = 1; -const MINUTE_SECS: u64 = 60; -const HOUR_SECS: u64 = 60 * 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). -/// -/// `rps`/`rph` counters added in api7/AISIX-Cloud#426 to fix the -/// `policy_to_rate_limit("second" | "hour")` upscaling workaround -/// that allowed 60× / 24× bursts past the operator-declared cap. -#[derive(Debug)] -struct KeyState { - rps: FixedWindowCounter, - rpm: FixedWindowCounter, - rph: FixedWindowCounter, - rpd: FixedWindowCounter, - tpm: FixedWindowCounter, - tpd: FixedWindowCounter, - in_flight: u32, -} +use aisix_core::RateLimit; -impl KeyState { - fn new() -> Self { - Self { - rps: FixedWindowCounter::new(SECOND_SECS), - rpm: FixedWindowCounter::new(MINUTE_SECS), - rph: FixedWindowCounter::new(HOUR_SECS), - rpd: FixedWindowCounter::new(DAY_SECS), - tpm: FixedWindowCounter::new(MINUTE_SECS), - tpd: FixedWindowCounter::new(DAY_SECS), - in_flight: 0, - } - } -} +use crate::clock::Clock; +use crate::error::RateLimitError; +use crate::store::local::LocalStore; +use crate::store::RateStore; /// Current window state for a single key, returned by [`Limiter::peek`]. /// Used by the proxy handlers to inject the `x-ratelimit-*` response @@ -84,222 +55,96 @@ impl RateLimitStatus { } } -pub struct Limiter { - states: DashMap>>, - clock: C, +/// Two-phase limiter over a shared or local [`RateStore`]. +pub struct Limiter { + store: Arc, + /// Process-unique reservation id prefix (`:`), so concurrency + /// members are globally distinct across replicas in the shared store. + member_prefix: String, + seq: AtomicU64, } -impl Limiter { +impl Limiter { + /// Default per-process limiter (in-memory `LocalStore`). pub fn new() -> Self { - Self::with_clock(SystemClock) + Self::with_store(Arc::new(LocalStore::new())) } -} -impl Default for Limiter { - fn default() -> Self { - Self::new() - } -} - -impl Limiter { - pub fn with_clock(clock: C) -> Self { + /// Build over a specific store — the server bootstrap passes a + /// `RedisStore` when a shared backend is configured. + pub fn with_store(store: Arc) -> Self { Self { - states: DashMap::new(), - clock, + store, + member_prefix: format!("{}:", uuid::Uuid::new_v4().simple()), + seq: AtomicU64::new(0), } } - /// Snapshot of the current rate-limit state for a key, used to inject - /// `x-ratelimit-*` response headers. Returns `None` if the key has - /// never been seen (i.e. no counters yet — headers are meaningless). - /// - /// This is a **read-only** operation; it does not affect any counters. - pub fn peek(&self, key: &str, limits: &aisix_core::RateLimit) -> Option { - let now = self.clock.unix_secs(); - let state = self.states.get(key)?; - let mut s = state.lock(); - - // Roll counters so we're looking at the current window. - let rpm_used = s.rpm.current(now); - let tpm_used = s.tpm.current(now); - let in_flight = s.in_flight; - - // Seconds remaining in the current minute-window. Zero if the - // window just started or has already rolled. - let minute_reset = MINUTE_SECS - (now % MINUTE_SECS); - - Some(RateLimitStatus { - rpm_limit: limits.rpm, - rpm_used, - rpm_reset_secs: minute_reset, - tpm_limit: limits.tpm, - tpm_used, - tpm_reset_secs: minute_reset, - concurrency_limit: limits.concurrency, - in_flight, - }) + /// Test helper: a local store driven by an injectable clock. + pub fn local_with_clock(clock: C) -> Self { + Self::with_store(Arc::new(LocalStore::with_clock(clock))) } - /// Add `tokens` to the post-deduct TPM/TPD counters for `key` - /// without going through a [`Reservation`]. Used by the streaming - /// chat path: at `pre_commit` time we don't yet know how many - /// tokens the upstream will return, so the Reservation is dropped - /// (releasing the concurrency permit + leaving TPM at 0). When the - /// SSE stream finishes, the proxy parses the upstream's terminal - /// usage frame and calls this method to retroactively account for - /// the tokens. Without it, TPM caps are silently bypassed for all - /// streaming traffic — issue #108. - /// - /// No-op when `tokens == 0` (avoids creating an empty per-key - /// counter for keys that never streamed). Otherwise, lazily - /// initialises the per-key state via [`Self::state_for`] so the - /// first streamed-after-restart request still gets accounted for. - pub fn add_tokens_post_stream(&self, key: &str, tokens: u64) { - if tokens == 0 { - return; - } - let now = self.clock.unix_secs(); - let state = self.state_for(key); - let mut s = state.lock(); - s.tpm.add(now, tokens); - s.tpd.add(now, tokens); - } - - 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() + fn next_member(&self) -> String { + let n = self.seq.fetch_add(1, Ordering::Relaxed); + format!("{}{n}", self.member_prefix) } /// 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( + /// via [`Reservation::commit_tokens`] or dropped to release the + /// concurrency slot automatically. + pub async 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. Layered chain - // (rps → rpm → rph → rpd) so a tighter window short-circuits - // a looser one without consuming its slot. If any later - // layer rejects, every earlier-incremented counter is rolled - // back by exactly the delta this call contributed — concurrent - // sibling requests' increments survive. Compensator coverage - // tested in `*_rejection_rolls_back_earlier_increments_*` - // unit tests; the chain expansion was forced by the #426 fix - // adding rps and rph (audit HIGH-2). - let mut rps_incremented = false; - if let Some(max) = limits.rps { - if let WindowCheck::Full { retry_after_secs } = s.rps.check_and_increment(now, 1, max) { - return Err(RateLimitError::Requests { - scope: RateLimitScope::Requests, - retry_after_secs, - }); - } - rps_incremented = true; - } - let mut rpm_incremented = false; - if let Some(max) = limits.rpm { - if let WindowCheck::Full { retry_after_secs } = s.rpm.check_and_increment(now, 1, max) { - if rps_incremented { - s.rps.decrement(now, 1); - } - return Err(RateLimitError::Requests { - scope: RateLimitScope::Requests, - retry_after_secs, - }); - } - rpm_incremented = true; - } - let mut rph_incremented = false; - if let Some(max) = limits.rph { - if let WindowCheck::Full { retry_after_secs } = s.rph.check_and_increment(now, 1, max) { - if rpm_incremented { - s.rpm.decrement(now, 1); - } - if rps_incremented { - s.rps.decrement(now, 1); - } - return Err(RateLimitError::Requests { - scope: RateLimitScope::Requests, - retry_after_secs, - }); - } - rph_incremented = true; - } - if let Some(max) = limits.rpd { - if let WindowCheck::Full { retry_after_secs } = s.rpd.check_and_increment(now, 1, max) { - if rph_incremented { - s.rph.decrement(now, 1); - } - if rpm_incremented { - s.rpm.decrement(now, 1); - } - if rps_incremented { - s.rps.decrement(now, 1); - } - return Err(RateLimitError::Requests { - scope: RateLimitScope::Requests, - retry_after_secs, - }); - } - } - - s.in_flight += 1; - drop(s); - + ) -> Result { + let member = self.next_member(); + self.store.acquire(key, limits, &member).await?; Ok(Reservation { - limiter: self, + store: Arc::clone(&self.store), key: key.to_string(), + member, committed: false, }) } + + /// Add `tokens` to the post-deduct TPM/TPD counters for `key` without + /// going through a [`Reservation`]. Used by the streaming chat path: + /// at pre_commit time the upstream token cost isn't known, so the + /// concurrency slot is held by a [`StreamConcurrencyGuard`] and the + /// tokens are accounted here when the terminal SSE usage frame lands + /// (issue #108). No-op on zero tokens. + pub fn add_tokens_post_stream(&self, key: &str, tokens: u64) { + if tokens == 0 { + return; + } + self.store.add_tokens(key, tokens); + } + + /// Snapshot of the current rate-limit state for a key, used to inject + /// `x-ratelimit-*` response headers. Returns `None` when there is + /// nothing meaningful to report. Read-only — affects no counters. + pub async fn peek(&self, key: &str, limits: &RateLimit) -> Option { + self.store.peek(key, limits).await + } +} + +impl Default for Limiter { + fn default() -> Self { + Self::new() + } } /// 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, +/// safe — the concurrency slot is released, just no tokens are counted. +pub struct Reservation { + store: Arc, key: String, + member: String, committed: bool, } -impl<'a, C: Clock> std::fmt::Debug for Reservation<'a, C> { +impl std::fmt::Debug for Reservation { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("Reservation") .field("key", &self.key) @@ -308,47 +153,40 @@ impl<'a, C: Clock> std::fmt::Debug for Reservation<'a, C> { } } -impl<'a, C: Clock> Reservation<'a, C> { +impl Reservation { /// 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); + /// and releases the concurrency slot. + pub async fn commit_tokens(mut self, tokens: u64) { + self.store.commit(&self.key, tokens, &self.member).await; self.committed = true; } } -impl<'a, C: Clock> Drop for Reservation<'a, C> { +impl Drop for Reservation { 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); + self.store.release(&self.key, &self.member); } } /// Wraps multiple [`Reservation`]s across rate-limit layers (api_key, -/// model, team, member). Commits all with the same token count; -/// dropping releases all concurrency permits. -pub struct MultiReservation<'a, C: Clock> { - reservations: Vec>, +/// model, team, member). Commits all with the same token count; dropping +/// releases all concurrency slots. +pub struct MultiReservation { + reservations: Vec, } -impl<'a, C: Clock> MultiReservation<'a, C> { - pub fn new(reservations: Vec>) -> Self { +impl MultiReservation { + pub fn new(reservations: Vec) -> Self { Self { reservations } } /// Commit the actual token cost to every layer. - pub fn commit_tokens(self, tokens: u64) { + pub async fn commit_tokens(self, tokens: u64) { for r in self.reservations { - r.commit_tokens(tokens); + r.commit_tokens(tokens).await; } } @@ -358,84 +196,80 @@ impl<'a, C: Clock> MultiReservation<'a, C> { } /// Convert into an owned [`StreamConcurrencyGuard`] for the streaming - /// path. The per-layer concurrency permits stay held — they are NOT + /// path. The per-layer concurrency slots stay held — they are NOT /// released here — and are released only when the returned guard drops, /// i.e. at stream completion or cancellation. Token accounting still /// happens via [`Limiter::add_tokens_post_stream`]. /// - /// A borrow-based reservation can't outlive the request handler, so the - /// pre-fix streaming path dropped it at handler return; that released - /// the concurrency permit before the stream finished, letting a key - /// capped at N run many more than N simultaneous streams (#450). - /// - /// `limiter` MUST be the same [`Limiter`] this reservation was acquired - /// against (in the proxy there is exactly one, behind `state.limiter`); - /// the guard decrements `in_flight` on it at drop. A borrow can't be - /// upgraded to an `Arc`, so the caller supplies the owned handle. + /// A borrow-based reservation couldn't outlive the request handler, so + /// the pre-fix streaming path dropped it at handler return; that + /// released the slot before the stream finished, letting a key capped + /// at N run many more than N simultaneous streams (#450). #[must_use = "dropping the returned guard immediately releases the concurrency \ - permit, recreating the early-release bug this fixes"] - pub fn into_stream_hold(mut self, limiter: Arc>) -> StreamConcurrencyGuard { - let keys = self.keys(); - // Defuse each reservation's Drop so it doesn't release the permit - // now; the returned guard owns release from here on. - for r in &mut self.reservations { - r.committed = true; - } + slot, recreating the early-release bug this fixes"] + pub fn into_stream_hold(mut self) -> StreamConcurrencyGuard { + let holds = self + .reservations + .iter_mut() + .map(|r| { + // Defuse each reservation's Drop so it doesn't release the + // slot now; the returned guard owns release from here on. + r.committed = true; + (Arc::clone(&r.store), r.key.clone(), r.member.clone()) + }) + .collect(); StreamConcurrencyGuard { - limiter, - keys, + holds, released: false, } } } -/// Owned concurrency hold for the streaming path. Holds an `Arc` -/// and the reserved keys, and releases the concurrency permit(s) on drop — -/// i.e. when the stream completes or is cancelled — instead of at handler -/// return. See [`MultiReservation::into_stream_hold`]. -pub struct StreamConcurrencyGuard { - limiter: Arc>, - keys: Vec, +impl std::fmt::Debug for MultiReservation { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("MultiReservation") + .field("layers", &self.reservations.len()) + .finish() + } +} + +/// Owned concurrency hold for the streaming path. Releases the +/// concurrency slot(s) on drop — i.e. when the stream completes or is +/// cancelled — instead of at handler return. See +/// [`MultiReservation::into_stream_hold`]. +pub struct StreamConcurrencyGuard { + /// `(store, key, member)` per held layer. + holds: Vec<(Arc, String, String)>, released: bool, } -impl StreamConcurrencyGuard { +impl StreamConcurrencyGuard { fn release_now(&mut self) { if self.released { return; } self.released = true; - for key in &self.keys { - let state = self.limiter.state_for(key); - let mut s = state.lock(); - s.in_flight = s.in_flight.saturating_sub(1); + for (store, key, member) in &self.holds { + store.release(key, member); } } } -impl std::fmt::Debug for StreamConcurrencyGuard { +impl std::fmt::Debug for StreamConcurrencyGuard { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("StreamConcurrencyGuard") - .field("keys", &self.keys) + .field("layers", &self.holds.len()) .field("released", &self.released) .finish() } } -impl Drop for StreamConcurrencyGuard { +impl Drop for StreamConcurrencyGuard { fn drop(&mut self) { self.release_now(); } } -impl<'a, C: Clock> std::fmt::Debug for MultiReservation<'a, C> { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("MultiReservation") - .field("layers", &self.reservations.len()) - .finish() - } -} - #[cfg(test)] mod tests { use super::*; @@ -471,15 +305,15 @@ mod tests { } } - #[test] - fn rpm_caps_request_count_in_window() { + #[tokio::test] + async fn rpm_caps_request_count_in_window() { let clock = TestClock::new(100); - let limiter = Limiter::with_clock(clock.clone()); + let limiter = Limiter::local_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(); + let _r1 = limiter.pre_commit("k1", &l).await.unwrap(); + let _r2 = limiter.pre_commit("k1", &l).await.unwrap(); + let err = limiter.pre_commit("k1", &l).await.unwrap_err(); match err { RateLimitError::Requests { retry_after_secs, .. @@ -490,109 +324,112 @@ mod tests { } } - #[test] - fn rpm_resets_after_window_rollover() { + #[tokio::test] + async fn rpm_resets_after_window_rollover() { let clock = TestClock::new(100); - let limiter = Limiter::with_clock(clock.clone()); + let limiter = Limiter::local_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()); + let _r1 = limiter.pre_commit("k1", &l).await.unwrap(); + assert!(limiter.pre_commit("k1", &l).await.is_err()); // Jump past the minute boundary. clock.advance(61); - let _r2 = limiter.pre_commit("k1", &l).unwrap(); + let _r2 = limiter.pre_commit("k1", &l).await.unwrap(); } - #[test] - fn concurrency_limit_blocks_new_reservations() { + #[tokio::test] + async fn concurrency_limit_blocks_new_reservations() { let clock = TestClock::new(0); - let limiter = Limiter::with_clock(clock.clone()); + let limiter = Limiter::local_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(); + let r1 = limiter.pre_commit("k1", &l).await.unwrap(); + let r2 = limiter.pre_commit("k1", &l).await.unwrap(); assert!(matches!( - limiter.pre_commit("k1", &l).unwrap_err(), + limiter.pre_commit("k1", &l).await.unwrap_err(), RateLimitError::Concurrency, )); // Drop r1 — concurrency should free up. drop(r1); - let _r3 = limiter.pre_commit("k1", &l).unwrap(); + let _r3 = limiter.pre_commit("k1", &l).await.unwrap(); drop(r2); } - #[test] - fn token_commit_updates_post_deduct_counters() { + #[tokio::test] + async fn token_commit_updates_post_deduct_counters() { let clock = TestClock::new(100); - let limiter = Limiter::with_clock(clock.clone()); + let limiter = Limiter::local_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); + let r1 = limiter.pre_commit("k1", &l).await.unwrap(); + r1.commit_tokens(600).await; // 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(); + let _r2 = limiter.pre_commit("k1", &l).await.unwrap(); } - #[test] - fn tpm_blocks_next_request_once_previous_exhausted_the_window() { + #[tokio::test] + async fn tpm_blocks_next_request_once_previous_exhausted_the_window() { let clock = TestClock::new(100); - let limiter = Limiter::with_clock(clock.clone()); + let limiter = Limiter::local_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 + let r1 = limiter.pre_commit("k1", &l).await.unwrap(); + r1.commit_tokens(1_500).await; // overshoot — allowed for the in-flight request // Next pre_commit sees tpm > 1000 and refuses. - let err = limiter.pre_commit("k1", &l).unwrap_err(); + let err = limiter.pre_commit("k1", &l).await.unwrap_err(); assert!(matches!(err, RateLimitError::Tokens { .. })); clock.advance(61); // roll the window - let _r2 = limiter.pre_commit("k1", &l).unwrap(); + let _r2 = limiter.pre_commit("k1", &l).await.unwrap(); } - #[test] - fn reservations_for_different_keys_do_not_collide() { + #[tokio::test] + async fn reservations_for_different_keys_do_not_collide() { let clock = TestClock::new(0); - let limiter = Limiter::with_clock(clock); + let limiter = Limiter::local_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(); + let _r_a = limiter.pre_commit("alpha", &l).await.unwrap(); + let _r_b = limiter.pre_commit("beta", &l).await.unwrap(); } - #[test] - fn drop_without_commit_still_releases_concurrency_permit() { + #[tokio::test] + async fn drop_without_commit_still_releases_concurrency_permit() { let clock = TestClock::new(0); - let limiter = Limiter::with_clock(clock); + let limiter = Limiter::local_with_clock(clock); let l = limits(None, None, Some(1)); { - let _r = limiter.pre_commit("k1", &l).unwrap(); + let _r = limiter.pre_commit("k1", &l).await.unwrap(); } // dropped - let _r2 = limiter.pre_commit("k1", &l).unwrap(); + let _r2 = limiter.pre_commit("k1", &l).await.unwrap(); } - #[test] - fn peek_returns_none_for_unknown_key() { + #[tokio::test] + async fn peek_returns_none_for_unknown_key() { let clock = TestClock::new(100); - let limiter = Limiter::with_clock(clock); - assert!(limiter.peek("unknown", &RateLimit::default()).is_none()); + let limiter = Limiter::local_with_clock(clock); + assert!(limiter + .peek("unknown", &RateLimit::default()) + .await + .is_none()); } - #[test] - fn peek_reports_current_window_counts() { + #[tokio::test] + async fn peek_reports_current_window_counts() { let clock = TestClock::new(100); - let limiter = Limiter::with_clock(clock.clone()); + let limiter = Limiter::local_with_clock(clock.clone()); let l = limits(Some(60), Some(100_000), Some(10)); - let r = limiter.pre_commit("k1", &l).unwrap(); - r.commit_tokens(500); + let r = limiter.pre_commit("k1", &l).await.unwrap(); + r.commit_tokens(500).await; - let status = limiter.peek("k1", &l).unwrap(); + let status = limiter.peek("k1", &l).await.unwrap(); assert_eq!(status.rpm_limit, Some(60)); assert_eq!(status.rpm_used, 1); assert_eq!(status.rpm_remaining(), Some(59)); @@ -602,46 +439,42 @@ mod tests { assert_eq!(status.in_flight, 0); // committed → released } - #[test] - fn peek_reflects_in_flight_count_during_dispatch() { + #[tokio::test] + async fn peek_reflects_in_flight_count_during_dispatch() { let clock = TestClock::new(100); - let limiter = Limiter::with_clock(clock); + let limiter = Limiter::local_with_clock(clock); let l = limits(None, None, Some(5)); - let _r1 = limiter.pre_commit("k1", &l).unwrap(); - let _r2 = limiter.pre_commit("k1", &l).unwrap(); - let status = limiter.peek("k1", &l).unwrap(); + let _r1 = limiter.pre_commit("k1", &l).await.unwrap(); + let _r2 = limiter.pre_commit("k1", &l).await.unwrap(); + let status = limiter.peek("k1", &l).await.unwrap(); assert_eq!(status.in_flight, 2); assert_eq!(status.concurrency_limit, Some(5)); } - #[test] - fn no_limits_means_no_rejections() { + #[tokio::test] + async fn no_limits_means_no_rejections() { let clock = TestClock::new(0); - let limiter = Limiter::with_clock(clock); + let limiter = Limiter::local_with_clock(clock); let l = RateLimit::default(); for _ in 0..100 { - let r = limiter.pre_commit("k1", &l).unwrap(); - r.commit_tokens(1_000); + let r = limiter.pre_commit("k1", &l).await.unwrap(); + r.commit_tokens(1_000).await; } } // ---- regression coverage for issue #109 ------------------------- // The previous compensation path overwrote `s.rpm` with a fresh - // FixedWindowCounter, wiping concurrent siblings' increments. The - // fix replaces the reset with a precise -1 decrement; these tests - // pin both the "siblings are preserved" and the "fresh window is - // not granted" properties at the same level the exploit happens. + // counter, wiping concurrent siblings' increments. The fix replaces + // the reset with a precise -1 decrement; these tests pin both the + // "siblings are preserved" and "fresh window is not granted" + // properties at the level the exploit happens. - #[test] - fn rpd_rejection_does_not_grant_fresh_rpm_window() { + #[tokio::test] + async fn rpd_rejection_does_not_grant_fresh_rpm_window() { let clock = TestClock::new(100); - let limiter = Limiter::with_clock(clock.clone()); - // RPM=10, RPD=20. Drive both close to their caps so the next - // request trips RPD, the buggy reset would have masked the - // RPM cap on the *very next* call, and the test exercises - // that follow-up. + let limiter = Limiter::local_with_clock(clock.clone()); let l = RateLimit { rps: None, rpm: Some(10), @@ -656,40 +489,31 @@ mod tests { if i == 10 { clock.advance(61); // roll RPM, keep RPD } - let _r = limiter.pre_commit("k1", &l).unwrap(); + let _r = limiter.pre_commit("k1", &l).await.unwrap(); } - // Now RPM in current minute = 9 (after the rollover), RPD = 19. - // One more goes through (RPM 10/10, RPD 20/20). - let _r = limiter.pre_commit("k1", &l).unwrap(); - // The 21st request must fail — RPD is full. Crucially, the - // pre-fix bug here resets RPM, so the assertion below would - // have falsely succeeded on a buggy build. - let err = limiter.pre_commit("k1", &l).unwrap_err(); + // RPM in current minute = 9 (after rollover), RPD = 19. One more + // goes through (RPM 10/10, RPD 20/20). + let _r = limiter.pre_commit("k1", &l).await.unwrap(); + // The 21st request must fail — RPD is full. + let err = limiter.pre_commit("k1", &l).await.unwrap_err(); assert!( matches!(err, RateLimitError::Requests { .. }), "expected RPD rejection, got {err:?}" ); - // The next request must STILL fail RPM — proving RPM wasn't - // wiped by the rejected request. With the pre-fix reset, this - // would have succeeded (silent rate-limit bypass). - let err2 = limiter.pre_commit("k1", &l).unwrap_err(); + // The next request must STILL fail RPM — proving RPM wasn't wiped. + let err2 = limiter.pre_commit("k1", &l).await.unwrap_err(); assert!( matches!(err2, RateLimitError::Requests { .. }), "RPM should still be capped after RPD rejection; got {err2:?}" ); - // RPM still reads 10 (the cap), not 0 (a wiped counter). - let status = limiter.peek("k1", &l).unwrap(); + let status = limiter.peek("k1", &l).await.unwrap(); assert_eq!(status.rpm_used, 10, "RPM should not have been reset"); } - #[test] - fn rpd_rejection_preserves_concurrent_rpm_increments() { - // Same shape, but exercises the "sibling increments survive" - // angle directly: drive RPM up to 5 with five accepted - // requests, then trip RPD on the sixth. The accepted five - // must remain counted. + #[tokio::test] + async fn rpd_rejection_preserves_concurrent_rpm_increments() { let clock = TestClock::new(100); - let limiter = Limiter::with_clock(clock.clone()); + let limiter = Limiter::local_with_clock(clock.clone()); let l = RateLimit { rps: None, rpm: Some(100), // very high — RPM never trips here @@ -700,13 +524,12 @@ mod tests { concurrency: None, }; for _ in 0..5 { - let _r = limiter.pre_commit("k1", &l).unwrap(); + let _r = limiter.pre_commit("k1", &l).await.unwrap(); } // RPM=5, RPD=5/5. Sixth request fails RPD. - let err = limiter.pre_commit("k1", &l).unwrap_err(); + let err = limiter.pre_commit("k1", &l).await.unwrap_err(); assert!(matches!(err, RateLimitError::Requests { .. })); - // RPM still reflects the FIVE accepted requests, not zero. - let status = limiter.peek("k1", &l).unwrap(); + let status = limiter.peek("k1", &l).await.unwrap(); assert_eq!( status.rpm_used, 5, "rpd rejection wiped concurrent rpm increments" @@ -714,76 +537,56 @@ mod tests { } // ---- regression coverage for issue #108 ------------------------- - // Streaming chat commits 0 tokens up front because total_tokens - // isn't known until the upstream's terminal usage frame. The fix - // exposes `Limiter::add_tokens_post_stream` so the SSE driver can - // retroactively account for tokens at end-of-stream. The tests - // below pin (1) the post-stream add bumps TPM, (2) zero-token - // calls don't create empty per-key state, (3) once enough tokens - // accumulate the next pre_commit fails on TPM. - - #[test] - fn add_tokens_post_stream_increments_tpm() { + + #[tokio::test] + async fn add_tokens_post_stream_increments_tpm() { let clock = TestClock::new(100); - let limiter = Limiter::with_clock(clock); + let limiter = Limiter::local_with_clock(clock); let l = limits(Some(10), Some(1_000), None); - // Pre-commit + drop (mirrors the streaming chat path: rpm - // counted, concurrency released, tpm = 0 at this point). { - let _r = limiter.pre_commit("k1", &l).unwrap(); + let _r = limiter.pre_commit("k1", &l).await.unwrap(); } assert_eq!( - limiter.peek("k1", &l).unwrap().tpm_used, + limiter.peek("k1", &l).await.unwrap().tpm_used, 0, "TPM should be 0 right after pre_commit + drop", ); - // Streaming reports 750 tokens at end-of-stream. limiter.add_tokens_post_stream("k1", 750); assert_eq!( - limiter.peek("k1", &l).unwrap().tpm_used, + limiter.peek("k1", &l).await.unwrap().tpm_used, 750, "TPM should reflect the post-stream commit", ); } - #[test] - fn add_tokens_post_stream_zero_is_a_noop() { + #[tokio::test] + async fn add_tokens_post_stream_zero_is_a_noop() { let clock = TestClock::new(100); - let limiter = Limiter::with_clock(clock); - // No pre_commit — peek would otherwise return None for an - // unknown key. add_tokens_post_stream(0) must NOT create an - // empty state entry. + let limiter = Limiter::local_with_clock(clock); limiter.add_tokens_post_stream("never-seen", 0); assert!( - limiter.peek("never-seen", &RateLimit::default()).is_none(), + limiter + .peek("never-seen", &RateLimit::default()) + .await + .is_none(), "add_tokens_post_stream(0) should not lazily-create state", ); } - #[test] - fn streaming_path_tpm_cap_blocks_next_request_after_post_stream_commit() { - // Drives the issue #108 exploit shape end-to-end at the - // limiter level: streaming "looks free" pre-fix because - // commit_tokens(0) skipped TPM. With the fix, the post- - // stream add should exhaust TPM and the next pre_commit - // must refuse on TPM. + #[tokio::test] + async fn streaming_path_tpm_cap_blocks_next_request_after_post_stream_commit() { let clock = TestClock::new(100); - let limiter = Limiter::with_clock(clock); + let limiter = Limiter::local_with_clock(clock); let l = limits(Some(100), Some(1_000), None); - // Mimic a successful streaming round: pre_commit + drop, then - // post-stream add that overshoots the cap. The "overshoot is - // allowed for the in-flight request" rule is the same as - // commit_tokens — see tpm_blocks_next_request_once_previous_exhausted_the_window. { - let _r = limiter.pre_commit("k1", &l).unwrap(); + let _r = limiter.pre_commit("k1", &l).await.unwrap(); } limiter.add_tokens_post_stream("k1", 1_500); - // Next request sees tpm > 1000 and refuses. - let err = limiter.pre_commit("k1", &l).unwrap_err(); + let err = limiter.pre_commit("k1", &l).await.unwrap_err(); assert!( matches!(err, RateLimitError::Tokens { .. }), "TPM cap should block the next request after streaming over-shoot; got {err:?}", @@ -792,340 +595,213 @@ mod tests { // --- MultiReservation tests ---------------------------------------- - #[test] - fn multi_reservation_commit_tokens_updates_all_layers() { + #[tokio::test] + async fn multi_reservation_commit_tokens_updates_all_layers() { let clock = TestClock::new(100); - let limiter = Limiter::with_clock(clock.clone()); + let limiter = Limiter::local_with_clock(clock.clone()); let l = limits(None, Some(1000), None); - let r1 = limiter.pre_commit("api_key:k1", &l).unwrap(); - let r2 = limiter.pre_commit("model:gpt-4o", &l).unwrap(); + let r1 = limiter.pre_commit("api_key:k1", &l).await.unwrap(); + let r2 = limiter.pre_commit("model:gpt-4o", &l).await.unwrap(); let multi = MultiReservation::new(vec![r1, r2]); - multi.commit_tokens(500); + multi.commit_tokens(500).await; - let s1 = limiter.peek("api_key:k1", &l).unwrap(); - let s2 = limiter.peek("model:gpt-4o", &l).unwrap(); + let s1 = limiter.peek("api_key:k1", &l).await.unwrap(); + let s2 = limiter.peek("model:gpt-4o", &l).await.unwrap(); assert_eq!(s1.tpm_used, 500); assert_eq!(s2.tpm_used, 500); } - #[test] - fn multi_reservation_drop_releases_all_concurrency() { + #[tokio::test] + async fn multi_reservation_drop_releases_all_concurrency() { let clock = TestClock::new(100); - let limiter = Limiter::with_clock(clock.clone()); + let limiter = Limiter::local_with_clock(clock.clone()); let l = limits(None, None, Some(1)); - let r1 = limiter.pre_commit("k1", &l).unwrap(); - let r2 = limiter.pre_commit("k2", &l).unwrap(); + let r1 = limiter.pre_commit("k1", &l).await.unwrap(); + let r2 = limiter.pre_commit("k2", &l).await.unwrap(); let multi = MultiReservation::new(vec![r1, r2]); - assert!(limiter.pre_commit("k1", &l).is_err()); - assert!(limiter.pre_commit("k2", &l).is_err()); + assert!(limiter.pre_commit("k1", &l).await.is_err()); + assert!(limiter.pre_commit("k2", &l).await.is_err()); drop(multi); - assert!(limiter.pre_commit("k1", &l).is_ok()); - assert!(limiter.pre_commit("k2", &l).is_ok()); + assert!(limiter.pre_commit("k1", &l).await.is_ok()); + assert!(limiter.pre_commit("k2", &l).await.is_ok()); } - #[test] - fn stream_hold_keeps_concurrency_until_guard_drop() { - // #450: a streaming request must keep its concurrency permit for the + #[tokio::test] + async fn stream_hold_keeps_concurrency_until_guard_drop() { + // #450: a streaming request must keep its concurrency slot for the // stream's full lifetime, not release it at handler return. let clock = TestClock::new(100); - let limiter = std::sync::Arc::new(Limiter::with_clock(clock.clone())); + let limiter = Limiter::local_with_clock(clock.clone()); let l = limits(None, None, Some(1)); - let r = limiter.pre_commit("k", &l).unwrap(); - let hold = MultiReservation::new(vec![r]).into_stream_hold(std::sync::Arc::clone(&limiter)); + let r = limiter.pre_commit("k", &l).await.unwrap(); + let hold = MultiReservation::new(vec![r]).into_stream_hold(); - // Permit is still held while the stream runs — a second concurrent + // Slot still held while the stream runs — a second concurrent // request is rejected. assert!(matches!( - limiter.pre_commit("k", &l).unwrap_err(), + limiter.pre_commit("k", &l).await.unwrap_err(), RateLimitError::Concurrency )); - // Stream completes/cancels → guard drops → permit released. + // Stream completes/cancels → guard drops → slot released. drop(hold); - assert!(limiter.pre_commit("k", &l).is_ok()); + assert!(limiter.pre_commit("k", &l).await.is_ok()); } - #[test] - fn multi_reservation_keys_returns_all_held_keys() { + #[tokio::test] + async fn multi_reservation_keys_returns_all_held_keys() { let clock = TestClock::new(100); - let limiter = Limiter::with_clock(clock.clone()); + let limiter = Limiter::local_with_clock(clock.clone()); let l = limits(Some(10), None, None); - let r1 = limiter.pre_commit("api_key:k1", &l).unwrap(); - let r2 = limiter.pre_commit("model:m1", &l).unwrap(); - let r3 = limiter.pre_commit("team:t1", &l).unwrap(); + let r1 = limiter.pre_commit("api_key:k1", &l).await.unwrap(); + let r2 = limiter.pre_commit("model:m1", &l).await.unwrap(); + let r3 = limiter.pre_commit("team:t1", &l).await.unwrap(); let multi = MultiReservation::new(vec![r1, r2, r3]); let keys = multi.keys(); assert_eq!(keys, vec!["api_key:k1", "model:m1", "team:t1"]); } - #[test] - fn multi_reservation_partial_failure_releases_acquired_layers() { + #[tokio::test] + async fn multi_reservation_partial_failure_releases_acquired_layers() { let clock = TestClock::new(100); - let limiter = Limiter::with_clock(clock.clone()); + let limiter = Limiter::local_with_clock(clock.clone()); let l_key = limits(None, None, Some(1)); let l_team = limits(None, None, Some(1)); let l_model = limits(Some(1), None, None); // Exhaust model RPM so the third layer will fail. - let _exhaust = limiter.pre_commit("model:m1", &l_model).unwrap(); + let _exhaust = limiter.pre_commit("model:m1", &l_model).await.unwrap(); - // Simulate multi-layer acquisition: key + team succeed, model fails. - let r_key = limiter.pre_commit("k1", &l_key).unwrap(); - let r_team = limiter.pre_commit("team:t1", &l_team).unwrap(); + let r_key = limiter.pre_commit("k1", &l_key).await.unwrap(); + let r_team = limiter.pre_commit("team:t1", &l_team).await.unwrap(); let acquired = vec![r_key, r_team]; - // Both concurrency slots are now taken. - assert!(limiter.pre_commit("k1", &l_key).is_err()); - assert!(limiter.pre_commit("team:t1", &l_team).is_err()); + assert!(limiter.pre_commit("k1", &l_key).await.is_err()); + assert!(limiter.pre_commit("team:t1", &l_team).await.is_err()); - // Model layer fails — drop acquired reservations (simulates error - // path where partially-built MultiReservation is dropped). - assert!(limiter.pre_commit("model:m1", &l_model).is_err()); + // Model layer fails — drop the partially-built reservations. + assert!(limiter.pre_commit("model:m1", &l_model).await.is_err()); drop(MultiReservation::new(acquired)); - // Both earlier layers' concurrency is released. - assert!(limiter.pre_commit("k1", &l_key).is_ok()); - assert!(limiter.pre_commit("team:t1", &l_team).is_ok()); + assert!(limiter.pre_commit("k1", &l_key).await.is_ok()); + assert!(limiter.pre_commit("team:t1", &l_team).await.is_ok()); } // ───────────────────────── #426 rps / rph coverage ───────────────────────── - // - // The api7/AISIX-Cloud#426 fix added two new request-counter - // layers (rps at 1s, rph at 3600s) to close the - // `policy_to_rate_limit` upscaling exploit. Tests below cover: - // - // 1. rps caps at max within 1s — bug repro from the issue body - // 2. rps window rolls over at the 1s boundary - // 3. rph caps at max within 3600s - // 4. Compensator chain — rejection at later layer rolls back - // earlier increments by exactly 1, never wipes the counter - // (regression of the #109-class bug for the new layers) - // 5. Empty rps with rpm set behaves like before #426 (regression - // guard that the rps wiring is gated by `Some(_)`) - - #[test] - fn rps_caps_request_count_within_one_second() { + + #[tokio::test] + async fn rps_caps_request_count_within_one_second() { let clock = TestClock::new(100); - let limiter = Limiter::with_clock(clock.clone()); + let limiter = Limiter::local_with_clock(clock.clone()); let l = limits_full(Some(5), None, None, None); - // 5 within the first second succeed. for i in 0..5 { limiter .pre_commit("k1", &l) + .await .unwrap_or_else(|e| panic!("request {i}: {e:?}")); } - // 6th in the same second rejected. - let err = limiter.pre_commit("k1", &l).unwrap_err(); + let err = limiter.pre_commit("k1", &l).await.unwrap_err(); assert!( matches!(err, RateLimitError::Requests { .. }), "expected rps rejection, got {err:?}" ); } - #[test] - fn rps_window_rolls_at_one_second_boundary() { + #[tokio::test] + async fn rps_window_rolls_at_one_second_boundary() { let clock = TestClock::new(100); - let limiter = Limiter::with_clock(clock.clone()); + let limiter = Limiter::local_with_clock(clock.clone()); let l = limits_full(Some(3), None, None, None); - // Fill the second-100 bucket. for _ in 0..3 { - limiter.pre_commit("k1", &l).unwrap(); + limiter.pre_commit("k1", &l).await.unwrap(); } - assert!(limiter.pre_commit("k1", &l).is_err()); + assert!(limiter.pre_commit("k1", &l).await.is_err()); - // Cross to second-101. Bucket resets, 3 more pass. clock.advance(1); for _ in 0..3 { - limiter.pre_commit("k1", &l).unwrap(); + limiter.pre_commit("k1", &l).await.unwrap(); } - assert!(limiter.pre_commit("k1", &l).is_err()); + assert!(limiter.pre_commit("k1", &l).await.is_err()); } - #[test] - fn rph_caps_request_count_within_one_hour() { + #[tokio::test] + async fn rph_caps_request_count_within_one_hour() { let clock = TestClock::new(100); - let limiter = Limiter::with_clock(clock.clone()); + let limiter = Limiter::local_with_clock(clock.clone()); let l = limits_full(None, None, Some(10), None); - // 10 within the first hour succeed. for i in 0..10 { limiter .pre_commit("k1", &l) + .await .unwrap_or_else(|e| panic!("request {i}: {e:?}")); } - // 11th in the same hour rejected. - let err = limiter.pre_commit("k1", &l).unwrap_err(); + let err = limiter.pre_commit("k1", &l).await.unwrap_err(); assert!( matches!(err, RateLimitError::Requests { .. }), "expected rph rejection, got {err:?}" ); - // Cross to next hour — bucket resets. clock.advance(3601); - limiter.pre_commit("k1", &l).unwrap(); - } - - #[test] - fn rpm_rejection_rolls_back_rps_increment() { - // Audit HIGH-2: when rps passes but a later request-counter - // (rpm/rph/rpd) rejects, the rps increment from THIS call - // must be rolled back — otherwise a customer could burn an - // rps slot for free on every rpm-rejected request. - let clock = TestClock::new(100); - let limiter = Limiter::with_clock(clock.clone()); - let l = limits_full(Some(10), Some(2), None, None); // rps=10/s, rpm=2/min - - // 2 succeed (both layers fit). - limiter.pre_commit("k1", &l).unwrap(); - limiter.pre_commit("k1", &l).unwrap(); - // 3rd: rps would still admit (2/10), but rpm rejects (2/2 used). - let err = limiter.pre_commit("k1", &l).unwrap_err(); - assert!(matches!(err, RateLimitError::Requests { .. })); - - // Critical: rps counter should still read 2 (the two accepted - // requests), not 3 (which would mean the rpm-rejected attempt - // burned an rps slot). Verify by checking that 8 MORE - // attempts in the same second hit rpm (not rps). - for _ in 0..8 { - let err = limiter.pre_commit("k1", &l).unwrap_err(); - assert!(matches!(err, RateLimitError::Requests { .. })); - } - // The 11th attempt should NOW hit rps (10 total rps used: 2 - // accepted + 8 rejected attempts that DID burn rps — wait no, - // rejected attempts should NOT have burned rps either; that's - // the whole point of the compensator. So all 8 rejections - // hit rpm and never increment rps. After those 8, the next - // attempt also hits rpm — rps still at 2. - // We can't easily distinguish "rejected at rps" vs "rejected - // at rpm" from the public API without internal probe. Use - // a different test angle below. - } - - #[test] - fn rpm_rejection_does_not_burn_rps_capacity() { - // Stronger version of `rpm_rejection_rolls_back_rps_increment`: - // construct a scenario where the COMPENSATOR is the only - // thing preventing rps starvation. rpm=2/min, rps=4/s. After - // the rpm cap is hit, repeated attempts must NOT eventually - // trip rps. The scenario is "every rejected attempt would - // burn an rps slot if the compensator didn't roll back". - let clock = TestClock::new(100); - let limiter = Limiter::with_clock(clock.clone()); - let l = limits_full(Some(4), Some(2), None, None); - - // Burn the rpm cap (also burns 2 rps slots). - limiter.pre_commit("k1", &l).unwrap(); - limiter.pre_commit("k1", &l).unwrap(); - - // Fire 100 more attempts — all should reject at rpm. - // Without the compensator, the FIRST 2 would burn rps to 4/4 - // and the rest would reject at rps instead. We can't - // distinguish by error type, but we CAN cross to the next - // minute and check rps still has headroom (would have been - // exhausted if compensator missed). - for _ in 0..100 { - assert!(limiter.pre_commit("k1", &l).is_err()); - } - - clock.advance(60); // roll rpm window - // If compensator worked correctly, rps still has 2 of 4 - // slots free in the current second (the two original successes). - // We can fire 2 more this second. - limiter.pre_commit("k1", &l).unwrap(); - limiter.pre_commit("k1", &l).unwrap(); - // 3rd in the same second hits rps (since rpm has 2/2 again, - // but rps stays under 4? wait: rps is per-second; after - // advance(60) we're in a new second too, so rps reset). - // Verify rps reset by firing 2 more — should pass rps but - // hit rpm. - let err = limiter.pre_commit("k1", &l).unwrap_err(); - assert!(matches!(err, RateLimitError::Requests { .. })); + limiter.pre_commit("k1", &l).await.unwrap(); } - #[test] - fn rpd_rejection_rolls_back_rps_and_rph_increments() { - // Audit HIGH-2: rpd rejection must roll back ALL earlier - // request-counter increments — rps, rpm, and rph. Mirror of - // the existing `rpd_rejection_does_not_grant_fresh_rpm_window` - // for the two new layers. + #[tokio::test] + async fn rpd_rejection_rolls_back_rps_and_rph_increments() { let clock = TestClock::new(100); - let limiter = Limiter::with_clock(clock.clone()); - // High rps/rpm/rph (so they never trip) + low rpd. + let limiter = Limiter::local_with_clock(clock.clone()); let l = limits_full(Some(1000), Some(1000), Some(1000), Some(2)); - limiter.pre_commit("k1", &l).unwrap(); - limiter.pre_commit("k1", &l).unwrap(); - // 3rd hits rpd. rps/rpm/rph should be at 2 each (the two - // accepted requests), NOT 3 (which would happen if rpd - // didn't roll back the just-incremented earlier counters). - let err = limiter.pre_commit("k1", &l).unwrap_err(); + limiter.pre_commit("k1", &l).await.unwrap(); + limiter.pre_commit("k1", &l).await.unwrap(); + let err = limiter.pre_commit("k1", &l).await.unwrap_err(); assert!(matches!(err, RateLimitError::Requests { .. })); - // Verify the counters didn't burn an extra slot via peek. - // peek() exposes rpm_used; the same logic applies to rps/rph - // but they don't surface via peek today (LOW finding in PR - // audit; out of scope to add header surfaces here). - let status = limiter.peek("k1", &l).unwrap(); + let status = limiter.peek("k1", &l).await.unwrap(); assert_eq!( status.rpm_used, 2, "rpd rejection must roll back rpm by exactly 1, leaving the two earlier accepts" ); } - #[test] - fn rph_rejection_rolls_back_rps_and_rpm_increments() { - // Audit #399 M2: the rpd-rejection test covers the tail of - // the chain (rolls back rps + rpm + rph), and the - // rpm-rejection tests cover the head (rolls back rps). The - // MIDDLE layer — rph rejecting after rps+rpm passed — wasn't - // directly covered. Pin it here. + #[tokio::test] + async fn rph_rejection_rolls_back_rps_and_rpm_increments() { let clock = TestClock::new(100); - let limiter = Limiter::with_clock(clock.clone()); - // High rps + rpm so they never trip; low rph; rpd unset. + let limiter = Limiter::local_with_clock(clock.clone()); let l = limits_full(Some(1000), Some(1000), Some(2), None); - limiter.pre_commit("k1", &l).unwrap(); - limiter.pre_commit("k1", &l).unwrap(); - // 3rd hits rph. rpm must roll back so subsequent reads see - // rpm_used = 2 (the two accepted requests), not 3. - let err = limiter.pre_commit("k1", &l).unwrap_err(); + limiter.pre_commit("k1", &l).await.unwrap(); + limiter.pre_commit("k1", &l).await.unwrap(); + let err = limiter.pre_commit("k1", &l).await.unwrap_err(); assert!(matches!(err, RateLimitError::Requests { .. })); - let status = limiter.peek("k1", &l).unwrap(); + let status = limiter.peek("k1", &l).await.unwrap(); assert_eq!( status.rpm_used, 2, "rph rejection must roll back rpm by exactly 1, leaving the two earlier accepts" ); } - #[test] - fn rps_layer_disabled_when_field_unset() { - // Regression guard: without `rps: Some(_)`, the limiter must - // skip the rps branch entirely — pre-#426 callers (api_key - // inline rate_limit, model inline rate_limit) only set rpm/rpd - // and must still work unchanged. + #[tokio::test] + async fn rps_layer_disabled_when_field_unset() { let clock = TestClock::new(100); - let limiter = Limiter::with_clock(clock.clone()); + let limiter = Limiter::local_with_clock(clock.clone()); let l = limits_full(None, Some(5), None, None); for _ in 0..5 { - limiter.pre_commit("k1", &l).unwrap(); + limiter.pre_commit("k1", &l).await.unwrap(); } - let err = limiter.pre_commit("k1", &l).unwrap_err(); + let err = limiter.pre_commit("k1", &l).await.unwrap_err(); assert!(matches!(err, RateLimitError::Requests { .. })); - // The rps counter must NOT have been touched — there's no - // direct observability of rps_used today, so this is more - // of a "doesn't panic / doesn't deadlock" guard than a deep - // equality check. } } diff --git a/crates/aisix-ratelimit/src/store/local.rs b/crates/aisix-ratelimit/src/store/local.rs new file mode 100644 index 00000000..4f0ce41b --- /dev/null +++ b/crates/aisix-ratelimit/src/store/local.rs @@ -0,0 +1,245 @@ +//! In-process counter store — the historical, default backend. +//! +//! Behaviour-identical to the pre-#798 limiter: a `DashMap` of per-key +//! fixed-window counters guarded by one `parking_lot::Mutex` each. State +//! is per-replica and not shared, so a multi-replica cluster multiplies +//! every limit by the replica count — exactly what [`super::redis`] +//! exists to fix. `member` is ignored here (concurrency is a plain +//! `in_flight` counter). + +use aisix_core::{RateLimit, RateLimitScope}; +use async_trait::async_trait; +use dashmap::DashMap; +use parking_lot::Mutex; +use std::sync::Arc; + +use super::{RateStore, DAY_SECS, HOUR_SECS, MINUTE_SECS, SECOND_SECS}; +use crate::clock::{Clock, SystemClock}; +use crate::error::RateLimitError; +use crate::limiter::RateLimitStatus; +use crate::window::{FixedWindowCounter, WindowCheck}; + +/// Per-key state guarded by a single mutex. Hot path locks once per +/// request; each operation inside is O(1). +#[derive(Debug)] +struct KeyState { + rps: FixedWindowCounter, + rpm: FixedWindowCounter, + rph: FixedWindowCounter, + rpd: FixedWindowCounter, + tpm: FixedWindowCounter, + tpd: FixedWindowCounter, + in_flight: u32, +} + +impl KeyState { + fn new() -> Self { + Self { + rps: FixedWindowCounter::new(SECOND_SECS), + rpm: FixedWindowCounter::new(MINUTE_SECS), + rph: FixedWindowCounter::new(HOUR_SECS), + rpd: FixedWindowCounter::new(DAY_SECS), + tpm: FixedWindowCounter::new(MINUTE_SECS), + tpd: FixedWindowCounter::new(DAY_SECS), + in_flight: 0, + } + } +} + +/// Per-process fixed-window store. +pub struct LocalStore { + states: DashMap>>, + clock: C, +} + +impl LocalStore { + pub fn new() -> Self { + Self::with_clock(SystemClock) + } +} + +impl Default for LocalStore { + fn default() -> Self { + Self::new() + } +} + +impl LocalStore { + 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() + } +} + +#[async_trait] +impl RateStore for LocalStore { + async fn acquire( + &self, + key: &str, + limits: &RateLimit, + _member: &str, + ) -> 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. Layered chain + // (rps → rpm → rph → rpd) so a tighter window short-circuits a + // looser one without consuming its slot. If any later layer + // rejects, every earlier-incremented counter is rolled back by + // exactly the delta this call contributed — concurrent sibling + // requests' increments survive. + let mut rps_incremented = false; + if let Some(max) = limits.rps { + if let WindowCheck::Full { retry_after_secs } = s.rps.check_and_increment(now, 1, max) { + return Err(RateLimitError::Requests { + scope: RateLimitScope::Requests, + retry_after_secs, + }); + } + rps_incremented = true; + } + let mut rpm_incremented = false; + if let Some(max) = limits.rpm { + if let WindowCheck::Full { retry_after_secs } = s.rpm.check_and_increment(now, 1, max) { + if rps_incremented { + s.rps.decrement(now, 1); + } + return Err(RateLimitError::Requests { + scope: RateLimitScope::Requests, + retry_after_secs, + }); + } + rpm_incremented = true; + } + let mut rph_incremented = false; + if let Some(max) = limits.rph { + if let WindowCheck::Full { retry_after_secs } = s.rph.check_and_increment(now, 1, max) { + if rpm_incremented { + s.rpm.decrement(now, 1); + } + if rps_incremented { + s.rps.decrement(now, 1); + } + return Err(RateLimitError::Requests { + scope: RateLimitScope::Requests, + retry_after_secs, + }); + } + rph_incremented = true; + } + if let Some(max) = limits.rpd { + if let WindowCheck::Full { retry_after_secs } = s.rpd.check_and_increment(now, 1, max) { + if rph_incremented { + s.rph.decrement(now, 1); + } + if rpm_incremented { + s.rpm.decrement(now, 1); + } + if rps_incremented { + s.rps.decrement(now, 1); + } + return Err(RateLimitError::Requests { + scope: RateLimitScope::Requests, + retry_after_secs, + }); + } + } + + s.in_flight += 1; + Ok(()) + } + + async fn commit(&self, key: &str, tokens: u64, _member: &str) { + let now = self.clock.unix_secs(); + let state = self.state_for(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); + } + + fn release(&self, key: &str, _member: &str) { + // Non-inserting: a release for a never-acquired bucket is a no-op, + // so the Redis store's belt-and-suspenders local release on the + // happy path doesn't pollute the local map with empty state. + if let Some(state) = self.states.get(key) { + let mut s = state.lock(); + s.in_flight = s.in_flight.saturating_sub(1); + } + } + + fn add_tokens(&self, key: &str, tokens: u64) { + if tokens == 0 { + return; + } + let now = self.clock.unix_secs(); + let state = self.state_for(key); + let mut s = state.lock(); + s.tpm.add(now, tokens); + s.tpd.add(now, tokens); + } + + async fn peek(&self, key: &str, limits: &RateLimit) -> Option { + let now = self.clock.unix_secs(); + let state = self.states.get(key)?; + let mut s = state.lock(); + + let rpm_used = s.rpm.current(now); + let tpm_used = s.tpm.current(now); + let in_flight = s.in_flight; + + // Seconds remaining in the current minute-window. Zero if the + // window just started or has already rolled. + let minute_reset = MINUTE_SECS - (now % MINUTE_SECS); + + Some(RateLimitStatus { + rpm_limit: limits.rpm, + rpm_used, + rpm_reset_secs: minute_reset, + tpm_limit: limits.tpm, + tpm_used, + tpm_reset_secs: minute_reset, + concurrency_limit: limits.concurrency, + in_flight, + }) + } +} diff --git a/crates/aisix-ratelimit/src/store/mod.rs b/crates/aisix-ratelimit/src/store/mod.rs new file mode 100644 index 00000000..820457a6 --- /dev/null +++ b/crates/aisix-ratelimit/src/store/mod.rs @@ -0,0 +1,122 @@ +//! Pluggable counter backend behind the [`crate::Limiter`]. +//! +//! The limiter itself only knows about *buckets* (an opaque key) and a +//! [`RateLimit`]; where the counters actually live is a [`RateStore`]. +//! +//! - [`local::LocalStore`] keeps the historical per-process in-memory +//! counters (a `DashMap` of fixed-window counters). This is the default +//! and is behaviour-identical to the pre-#798 limiter. +//! - [`redis::RedisStore`] keeps the counters in a shared Redis so every +//! DP replica in a cluster enforces ONE global window — the fix for +//! api7/AISIX-Cloud#798, where N replicas multiplied every limit by N. +//! +//! Two phases mirror the limiter's contract: +//! - **acquire** (request path, async): concurrency gate + token +//! check-only + request check-and-increment, all-or-nothing per bucket. +//! - **commit** (request path success, async): post-deduct token add + +//! concurrency release. +//! - **release** / **add_tokens** (after-the-fact, sync): concurrency +//! release on drop, and the streaming post-stream token add. These are +//! sync because they run from `Drop` and from the synchronous SSE +//! completion callback; the Redis impl makes them fire-and-forget. + +use aisix_core::RateLimit; +use async_trait::async_trait; + +use crate::error::RateLimitError; +use crate::limiter::RateLimitStatus; + +pub mod local; +pub mod redis; + +pub(crate) const SECOND_SECS: u64 = 1; +pub(crate) const MINUTE_SECS: u64 = 60; +pub(crate) const HOUR_SECS: u64 = 60 * 60; +pub(crate) const DAY_SECS: u64 = 24 * 60 * 60; + +/// A windowed request/token dimension active on a [`RateLimit`]: +/// `(name, window_secs, limit)`. Shared by both stores so the Redis key +/// layout and the local counter set never drift. +pub(crate) struct Dim { + pub name: &'static str, + pub window_secs: u64, + pub limit: u64, +} + +/// Request-count dimensions (rps/rpm/rph/rpd) that carry a limit. +pub(crate) fn request_dims(limits: &RateLimit) -> Vec { + [ + ("rps", SECOND_SECS, limits.rps), + ("rpm", MINUTE_SECS, limits.rpm), + ("rph", HOUR_SECS, limits.rph), + ("rpd", DAY_SECS, limits.rpd), + ] + .into_iter() + .filter_map(|(name, window_secs, limit)| { + limit.map(|limit| Dim { + name, + window_secs, + limit, + }) + }) + .collect() +} + +/// Token-count dimensions (tpm/tpd) that carry a limit. +pub(crate) fn token_dims(limits: &RateLimit) -> Vec { + [ + ("tpm", MINUTE_SECS, limits.tpm), + ("tpd", DAY_SECS, limits.tpd), + ] + .into_iter() + .filter_map(|(name, window_secs, limit)| { + limit.map(|limit| Dim { + name, + window_secs, + limit, + }) + }) + .collect() +} + +/// Backend that holds the rate-limit counters for a bucket. +/// +/// `member` is a process-unique reservation id (`:`) used +/// by distributed backends to track exactly one in-flight slot in the +/// concurrency set; the local backend ignores it (its `in_flight` is a +/// plain counter). +#[async_trait] +pub trait RateStore: Send + Sync + 'static { + /// Pre-commit acquire for a single bucket. Atomically (per bucket): + /// gate concurrency, check (but do not increment) token windows, then + /// check-and-increment every request window. All-or-nothing: on + /// rejection nothing is incremented and the concurrency slot is not + /// taken. + async fn acquire( + &self, + key: &str, + limits: &RateLimit, + member: &str, + ) -> Result<(), RateLimitError>; + + /// Post-deduct: add `tokens` to the tpm/tpd windows AND release the + /// concurrency slot held by `member`. Like the local backend this + /// always touches both token windows; the tpd counter is harmless + /// when no tpd limit is configured (it simply expires unread). + async fn commit(&self, key: &str, tokens: u64, member: &str); + + /// Release the concurrency slot held by `member` without recording + /// tokens. Sync so it can run from `Drop`; the Redis impl spawns a + /// detached release (the concurrency set self-heals via TTL pruning + /// even if the spawn is lost). + fn release(&self, key: &str, member: &str); + + /// Post-stream token accounting: add `tokens` to tpm/tpd only (no + /// concurrency change). Sync so it can run from the synchronous SSE + /// completion callback; the Redis impl makes it fire-and-forget. + fn add_tokens(&self, key: &str, tokens: u64); + + /// Read-only snapshot for the `x-ratelimit-*` headers. Returns `None` + /// when there is nothing meaningful to report for the bucket. + async fn peek(&self, key: &str, limits: &RateLimit) -> Option; +} diff --git a/crates/aisix-ratelimit/src/store/redis.rs b/crates/aisix-ratelimit/src/store/redis.rs new file mode 100644 index 00000000..041afbe4 --- /dev/null +++ b/crates/aisix-ratelimit/src/store/redis.rs @@ -0,0 +1,407 @@ +//! Redis-backed shared counter store — the fix for api7/AISIX-Cloud#798. +//! +//! Every DP replica points at the same Redis, so one global window is +//! enforced across the cluster instead of one-per-replica. The counter +//! math mirrors [`super::local::LocalStore`] / [`crate::window`] exactly +//! (wall-clock-aligned fixed windows: `window_start = now - now % window`), +//! so swapping `memory ↔ redis` doesn't change observable limits — only +//! whether the count is shared. +//! +//! Key layout (hash-tagged so every key for a bucket shares one Redis +//! Cluster slot, keeping the per-bucket Lua atomic): +//! - `aisix:rl:{}::` — plain +//! `INCR`/`GET` counters, `EXPIRE = window + grace`. +//! - `aisix:rl:{}:conc` — a ZSET (`member → score=now`) acting as a +//! crash-safe distributed semaphore: acquire prunes entries older than +//! `conc_ttl` then counts, so a slot leaked by a crashed/hung replica is +//! reclaimed within `conc_ttl`. (LiteLLM's latest tracks parallel +//! requests as a window-TTL counter; we use a ZSET with a request- +//! lifetime ttl because our streaming requests can outlive a 60s window +//! — the same reason `StreamConcurrencyGuard`/#450 exists.) +//! +//! `now` is read from `redis.call('TIME')` inside every script so window +//! boundaries are identical across replicas regardless of host clock skew. +//! +//! On any Redis error the store fails **open** to a per-process +//! [`LocalStore`] (logged once): traffic keeps flowing with per-replica +//! enforcement during an outage instead of being blocked. Counts may +//! diverge from Redis until it recovers — availability over strict global +//! enforcement, matching the cache's "Redis error → proceed" stance. + +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; + +use aisix_core::RateLimit; +use async_trait::async_trait; +use redis::aio::ConnectionManager; +use redis::Script; + +use super::{local::LocalStore, token_dims, Dim, RateStore}; +use crate::error::RateLimitError; +use crate::limiter::RateLimitStatus; + +/// Default key namespace, kept distinct from `aisix:cache` so the same +/// Redis can back both the response cache and the rate limiter. +pub const DEFAULT_PREFIX: &str = "aisix:rl"; +/// Default concurrency-slot lifetime ceiling. A slot not released within +/// this many seconds (crashed replica, hung upstream) is pruned from the +/// count. Generous enough to cover a long streaming response. +pub const DEFAULT_CONC_TTL_SECS: u64 = 300; +/// Extra seconds added to each window counter's TTL so a counter doesn't +/// expire a hair before its window mathematically closes. +pub const DEFAULT_GRACE_SECS: u64 = 5; + +// Result codes returned by the acquire script's first element. +const CODE_OK: i64 = 0; +const CODE_CONCURRENCY: i64 = 1; +const CODE_TOKENS: i64 = 2; +const CODE_REQUESTS: i64 = 3; + +/// Atomic per-bucket acquire: concurrency gate + token check-only + +/// request check-and-increment, all-or-nothing. See module docs for the +/// key layout. Returns `{code, retry_after}`. +const ACQUIRE_LUA: &str = r#" +local prefix = ARGV[1] +local member = ARGV[2] +local conc_max = tonumber(ARGV[3]) +local conc_ttl = tonumber(ARGV[4]) +local grace = tonumber(ARGV[5]) +local t = redis.call('TIME') +local now = tonumber(t[1]) + +local conc_key = prefix .. ':conc' +if conc_max >= 0 then + redis.call('ZREMRANGEBYSCORE', conc_key, 0, now - conc_ttl) + if redis.call('ZCARD', conc_key) >= conc_max then + return {1, 0} + end +end + +local idx = 6 +local nreq = tonumber(ARGV[idx]); idx = idx + 1 +local req = {} +for i = 1, nreq do + req[i] = {ARGV[idx], tonumber(ARGV[idx+1]), tonumber(ARGV[idx+2])} + idx = idx + 3 +end +local ntok = tonumber(ARGV[idx]); idx = idx + 1 +local tok = {} +for i = 1, ntok do + tok[i] = {ARGV[idx], tonumber(ARGV[idx+1]), tonumber(ARGV[idx+2])} + idx = idx + 3 +end + +for i = 1, ntok do + local name, window, limit = tok[i][1], tok[i][2], tok[i][3] + local ws = now - (now % window) + local cur = tonumber(redis.call('GET', prefix .. ':' .. name .. ':' .. ws) or '0') + if cur > limit then + local retry = window - (now - ws); if retry < 1 then retry = 1 end + return {2, retry} + end +end + +for i = 1, nreq do + local name, window, limit = req[i][1], req[i][2], req[i][3] + local ws = now - (now % window) + local cur = tonumber(redis.call('GET', prefix .. ':' .. name .. ':' .. ws) or '0') + if cur + 1 > limit then + local retry = window - (now - ws); if retry < 1 then retry = 1 end + return {3, retry} + end +end + +for i = 1, nreq do + local name, window = req[i][1], req[i][2] + local ws = now - (now % window) + local k = prefix .. ':' .. name .. ':' .. ws + if redis.call('INCR', k) == 1 then + redis.call('EXPIRE', k, window + grace) + end +end +if conc_max >= 0 then + redis.call('ZADD', conc_key, now, member) + redis.call('EXPIRE', conc_key, conc_ttl) +end +return {0, 0} +"#; + +/// Post-deduct: add `tokens` to the tpm/tpd windows AND release the +/// concurrency slot held by `member`. Both token windows are always +/// touched (matching the local backend); an unread tpd counter just +/// expires. ARGV: prefix, member, tokens, grace. +const COMMIT_LUA: &str = r#" +local prefix = ARGV[1] +local member = ARGV[2] +local tokens = tonumber(ARGV[3]) +local grace = tonumber(ARGV[4]) +local t = redis.call('TIME') +local now = tonumber(t[1]) +if tokens > 0 then + for _, d in ipairs({{'tpm', 60}, {'tpd', 86400}}) do + local ws = now - (now % d[2]) + local k = prefix .. ':' .. d[1] .. ':' .. ws + if redis.call('INCRBY', k, tokens) == tokens then + redis.call('EXPIRE', k, d[2] + grace) + end + end +end +redis.call('ZREM', prefix .. ':conc', member) +return 1 +"#; + +/// Post-stream token add only (no concurrency change). ARGV: prefix, +/// tokens, grace. +const ADD_TOKENS_LUA: &str = r#" +local prefix = ARGV[1] +local tokens = tonumber(ARGV[2]) +local grace = tonumber(ARGV[3]) +local t = redis.call('TIME') +local now = tonumber(t[1]) +if tokens > 0 then + for _, d in ipairs({{'tpm', 60}, {'tpd', 86400}}) do + local ws = now - (now % d[2]) + local k = prefix .. ':' .. d[1] .. ':' .. ws + if redis.call('INCRBY', k, tokens) == tokens then + redis.call('EXPIRE', k, d[2] + grace) + end + end +end +return 1 +"#; + +/// Read-only snapshot for headers: current-minute rpm/tpm counts + +/// pruned concurrency count + seconds-to-minute-reset. ARGV: prefix, +/// conc_ttl. Returns {rpm_used, tpm_used, in_flight, minute_reset}. +const PEEK_LUA: &str = r#" +local prefix = ARGV[1] +local conc_ttl = tonumber(ARGV[2]) +local t = redis.call('TIME') +local now = tonumber(t[1]) +local ws = now - (now % 60) +local rpm = tonumber(redis.call('GET', prefix .. ':rpm:' .. ws) or '0') +local tpm = tonumber(redis.call('GET', prefix .. ':tpm:' .. ws) or '0') +redis.call('ZREMRANGEBYSCORE', prefix .. ':conc', 0, now - conc_ttl) +local inflight = redis.call('ZCARD', prefix .. ':conc') +return {rpm, tpm, inflight, 60 - (now % 60)} +"#; + +pub struct RedisStore { + conn: ConnectionManager, + prefix: String, + conc_ttl: u64, + grace: u64, + /// Per-process fallback used when Redis is unreachable (fail-open). + local: Arc, + /// One-shot guard so the degradation warning is logged once, not per + /// request, while Redis stays down. + degraded_logged: AtomicBool, +} + +impl std::fmt::Debug for RedisStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RedisStore") + .field("prefix", &self.prefix) + .field("conc_ttl", &self.conc_ttl) + .finish_non_exhaustive() + } +} + +impl RedisStore { + /// Connect to Redis via [`ConnectionManager`] (transparent reconnect, + /// no per-request handshake). Single-node URL like `redis://host:port/`. + pub async fn connect(url: &str) -> Result { + let client = redis::Client::open(url)?; + let conn = ConnectionManager::new(client).await?; + Ok(Self { + conn, + prefix: DEFAULT_PREFIX.into(), + conc_ttl: DEFAULT_CONC_TTL_SECS, + grace: DEFAULT_GRACE_SECS, + local: Arc::new(LocalStore::new()), + degraded_logged: AtomicBool::new(false), + }) + } + + pub fn with_conc_ttl(mut self, secs: u64) -> Self { + self.conc_ttl = secs.max(1); + self + } + + /// `aisix:rl:{}` — the hash tag co-locates every key for the + /// bucket on one Redis Cluster slot. + fn bucket_prefix(&self, key: &str) -> String { + format!("{}:{{{}}}", self.prefix, key) + } + + /// Log the fail-open transition once per outage. + fn warn_degraded(&self, op: &str, err: &redis::RedisError) { + if !self.degraded_logged.swap(true, Ordering::Relaxed) { + tracing::warn!( + target: "aisix::ratelimit", + op, + error = %err, + "shared rate-limit Redis unavailable; failing open to per-replica \ + in-memory counting until it recovers (cluster limits not enforced \ + during the outage)" + ); + } + } + + /// Mark Redis healthy again after a successful op (re-arms the warn). + fn mark_ok(&self) { + self.degraded_logged.store(false, Ordering::Relaxed); + } +} + +fn push_dims(args: &mut Vec, dims: &[Dim]) { + args.push(dims.len().to_string()); + for d in dims { + args.push(d.name.to_string()); + args.push(d.window_secs.to_string()); + args.push(d.limit.to_string()); + } +} + +#[async_trait] +impl RateStore for RedisStore { + async fn acquire( + &self, + key: &str, + limits: &RateLimit, + member: &str, + ) -> Result<(), RateLimitError> { + let prefix = self.bucket_prefix(key); + let mut args = vec![ + prefix, + member.to_string(), + limits.concurrency.map(i64::from).unwrap_or(-1).to_string(), + self.conc_ttl.to_string(), + self.grace.to_string(), + ]; + push_dims(&mut args, &super::request_dims(limits)); + push_dims(&mut args, &token_dims(limits)); + + let script = Script::new(ACQUIRE_LUA); + let mut invocation = script.prepare_invoke(); + for a in &args { + invocation.arg(a); + } + let mut conn = self.conn.clone(); + match invocation.invoke_async::>(&mut conn).await { + Ok(reply) => { + self.mark_ok(); + let code = reply.first().copied().unwrap_or(CODE_OK); + let retry = reply.get(1).copied().unwrap_or(0).max(0) as u64; + match code { + CODE_OK => Ok(()), + CODE_CONCURRENCY => Err(RateLimitError::Concurrency), + CODE_TOKENS => Err(RateLimitError::Tokens { + scope: aisix_core::RateLimitScope::Tokens, + retry_after_secs: retry, + }), + CODE_REQUESTS => Err(RateLimitError::Requests { + scope: aisix_core::RateLimitScope::Requests, + retry_after_secs: retry, + }), + _ => Ok(()), + } + } + Err(e) => { + self.warn_degraded("acquire", &e); + self.local.acquire(key, limits, member).await + } + } + } + + async fn commit(&self, key: &str, tokens: u64, member: &str) { + let prefix = self.bucket_prefix(key); + let mut conn = self.conn.clone(); + let res: Result = Script::new(COMMIT_LUA) + .arg(&prefix) + .arg(member) + .arg(tokens) + .arg(self.grace) + .invoke_async(&mut conn) + .await; + match res { + Ok(_) => self.mark_ok(), + Err(e) => { + self.warn_degraded("commit", &e); + self.local.commit(key, tokens, member).await; + } + } + } + + fn release(&self, key: &str, member: &str) { + // Drop the local slot first (a cheap no-op when the bucket was + // never acquired locally); covers the degraded-acquire case. + self.local.release(key, member); + let conc_key = format!("{}:conc", self.bucket_prefix(key)); + let mut conn = self.conn.clone(); + let member = member.to_string(); + if let Ok(handle) = tokio::runtime::Handle::try_current() { + handle.spawn(async move { + let _: Result<(), redis::RedisError> = redis::cmd("ZREM") + .arg(&conc_key) + .arg(&member) + .query_async(&mut conn) + .await; + }); + } + } + + fn add_tokens(&self, key: &str, tokens: u64) { + if tokens == 0 { + return; + } + // Best-effort fire-and-forget to Redis; also record locally so the + // count is right if a later request falls back during an outage. + self.local.add_tokens(key, tokens); + let prefix = self.bucket_prefix(key); + let grace = self.grace; + let mut conn = self.conn.clone(); + if let Ok(handle) = tokio::runtime::Handle::try_current() { + handle.spawn(async move { + let _: Result = Script::new(ADD_TOKENS_LUA) + .arg(&prefix) + .arg(tokens) + .arg(grace) + .invoke_async(&mut conn) + .await; + }); + } + } + + async fn peek(&self, key: &str, limits: &RateLimit) -> Option { + if limits.is_unrestricted() { + return None; + } + let prefix = self.bucket_prefix(key); + let mut conn = self.conn.clone(); + let reply: Result, redis::RedisError> = Script::new(PEEK_LUA) + .arg(&prefix) + .arg(self.conc_ttl) + .invoke_async(&mut conn) + .await; + match reply { + Ok(v) => { + self.mark_ok(); + Some(RateLimitStatus { + rpm_limit: limits.rpm, + rpm_used: v.first().copied().unwrap_or(0).max(0) as u64, + rpm_reset_secs: v.get(3).copied().unwrap_or(0).max(0) as u64, + tpm_limit: limits.tpm, + tpm_used: v.get(1).copied().unwrap_or(0).max(0) as u64, + tpm_reset_secs: v.get(3).copied().unwrap_or(0).max(0) as u64, + concurrency_limit: limits.concurrency, + in_flight: v.get(2).copied().unwrap_or(0).max(0) as u32, + }) + } + Err(e) => { + self.warn_degraded("peek", &e); + self.local.peek(key, limits).await + } + } + } +} diff --git a/crates/aisix-ratelimit/tests/redis_integration.rs b/crates/aisix-ratelimit/tests/redis_integration.rs new file mode 100644 index 00000000..dc8ba8ef --- /dev/null +++ b/crates/aisix-ratelimit/tests/redis_integration.rs @@ -0,0 +1,198 @@ +//! Shared-counter tests for `RedisStore` against a live Redis. +//! +//! Runs only when `RATELIMIT_TEST_REDIS_URL` is set (CI spins +//! `redis:7-alpine` as a service; absence is a no-op so local unit runs +//! stay hermetic). Two `RedisStore` instances stand in for two DP +//! replicas pointed at one Redis — the exact api7/AISIX-Cloud#798 shape: +//! a limit hit on one replica must already be hit on the other. + +use std::time::Duration; + +use aisix_core::{RateLimit, RateLimitScope}; +use aisix_ratelimit::{RateStore, RedisStore}; + +fn redis_url() -> Option { + std::env::var("RATELIMIT_TEST_REDIS_URL").ok() +} + +/// Unique bucket key per test so they don't clobber each other (the store +/// prefixes with a fixed `aisix:rl`; isolation comes from the key). +fn unique_key(tag: &str) -> String { + use std::time::{SystemTime, UNIX_EPOCH}; + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_nanos()) + .unwrap_or_default(); + format!("test:{tag}:{nanos:x}") +} + +fn rl() -> RateLimit { + RateLimit::default() +} + +async fn store(url: &str) -> RedisStore { + RedisStore::connect(url).await.expect("redis connect") +} + +#[tokio::test] +async fn rpm_counter_is_shared_across_replicas() { + let Some(url) = redis_url() else { + eprintln!("skipping: RATELIMIT_TEST_REDIS_URL not set"); + return; + }; + let a = store(&url).await; + let b = store(&url).await; + let key = unique_key("rpm"); + let limits = RateLimit { + rpm: Some(1), + ..rl() + }; + + // Replica A burns the only slot in the minute window. + a.acquire(&key, &limits, "a-1") + .await + .expect("first allowed"); + + // Replica B sees the SAME counter → rejected. Pre-#798 (per-replica + // memory) this would have been allowed, doubling the limit. + let err = b + .acquire(&key, &limits, "b-1") + .await + .expect_err("second replica must be rejected by shared counter"); + assert!( + matches!( + err, + aisix_ratelimit::RateLimitError::Requests { + scope: RateLimitScope::Requests, + .. + } + ), + "got {err:?}" + ); +} + +#[tokio::test] +async fn rps_window_rolls_over_on_the_shared_counter() { + let Some(url) = redis_url() else { + eprintln!("skipping: RATELIMIT_TEST_REDIS_URL not set"); + return; + }; + let a = store(&url).await; + let b = store(&url).await; + let key = unique_key("rps"); + let limits = RateLimit { + rps: Some(1), + ..rl() + }; + + a.acquire(&key, &limits, "a-1") + .await + .expect("first allowed"); + assert!( + b.acquire(&key, &limits, "b-1").await.is_err(), + "same second is shared-rejected" + ); + + // Cross the 1s boundary — the next-second key is fresh. + tokio::time::sleep(Duration::from_millis(1_100)).await; + b.acquire(&key, &limits, "b-2") + .await + .expect("next second has a fresh window"); +} + +#[tokio::test] +async fn token_usage_is_shared_across_replicas() { + let Some(url) = redis_url() else { + eprintln!("skipping: RATELIMIT_TEST_REDIS_URL not set"); + return; + }; + let a = store(&url).await; + let b = store(&url).await; + let key = unique_key("tpm"); + let limits = RateLimit { + tpm: Some(1_000), + ..rl() + }; + + // A admits then over-commits the minute's token budget. + a.acquire(&key, &limits, "a-1") + .await + .expect("first allowed"); + a.commit(&key, 1_500, "a-1").await; + + // B's pre-check sees tpm > 1000 on the shared counter → rejected. + let err = b + .acquire(&key, &limits, "b-1") + .await + .expect_err("token cap is shared"); + assert!( + matches!(err, aisix_ratelimit::RateLimitError::Tokens { .. }), + "got {err:?}" + ); +} + +#[tokio::test] +async fn concurrency_slot_is_shared_and_released_across_replicas() { + let Some(url) = redis_url() else { + eprintln!("skipping: RATELIMIT_TEST_REDIS_URL not set"); + return; + }; + let a = store(&url).await; + let b = store(&url).await; + let key = unique_key("conc"); + let limits = RateLimit { + concurrency: Some(1), + ..rl() + }; + + // A takes the only in-flight slot. + a.acquire(&key, &limits, "a-1") + .await + .expect("first allowed"); + // B is blocked while A holds it. + assert!( + matches!( + b.acquire(&key, &limits, "b-1").await, + Err(aisix_ratelimit::RateLimitError::Concurrency) + ), + "concurrency slot must be shared across replicas" + ); + + // A finishes → releases the slot (sync + detached ZREM). + a.release(&key, "a-1"); + tokio::time::sleep(Duration::from_millis(200)).await; // let the detached ZREM land + + b.acquire(&key, &limits, "b-2") + .await + .expect("slot frees up cluster-wide after release"); +} + +#[tokio::test] +async fn stale_concurrency_slot_is_reclaimed_after_ttl() { + let Some(url) = redis_url() else { + eprintln!("skipping: RATELIMIT_TEST_REDIS_URL not set"); + return; + }; + // 1s slot lifetime: a never-released slot (crashed replica) is pruned. + let a = store(&url).await.with_conc_ttl(1); + let b = store(&url).await.with_conc_ttl(1); + let key = unique_key("conc-ttl"); + let limits = RateLimit { + concurrency: Some(1), + ..rl() + }; + + a.acquire(&key, &limits, "a-leaked") + .await + .expect("first allowed"); + // Never release — simulate a crashed replica holding the slot. + assert!( + b.acquire(&key, &limits, "b-1").await.is_err(), + "slot held while fresh" + ); + + tokio::time::sleep(Duration::from_millis(1_300)).await; + b.acquire(&key, &limits, "b-2") + .await + .expect("stale slot reclaimed after conc_ttl"); +} diff --git a/crates/aisix-server/src/main.rs b/crates/aisix-server/src/main.rs index 428b4678..bb32fe99 100644 --- a/crates/aisix-server/src/main.rs +++ b/crates/aisix-server/src/main.rs @@ -36,7 +36,7 @@ use aisix_provider_vertex::VertexBridge; use aisix_proxy::background::run_background_model_check_once; use aisix_proxy::budget::BudgetClient; use aisix_proxy::{CacheBackends, ProxyState}; -use aisix_ratelimit::Limiter; +use aisix_ratelimit::{Limiter, RedisStore}; use clap::Parser; use etcd_client::{Certificate, ConnectOptions, Identity, TlsOptions}; use tokio::sync::watch; @@ -376,7 +376,29 @@ async fn run(mut cfg: Config) -> anyhow::Result<()> { // Steps 7-8: build Hub, shared components, then routers. let hub = Arc::new(build_hub()); - let limiter = Arc::new(Limiter::new()); + // Rate-limit backend (#798). Default `memory` keeps per-process + // counters; `redis` shares them across every replica so a cluster + // enforces one global window instead of one-per-replica. Fail fast on + // `backend = redis` without a `ratelimit.redis` block (validated in + // Config::validate, re-checked here before connecting). + let limiter = Arc::new(match cfg.ratelimit.redis.as_ref() { + Some(redis_cfg) => { + tracing::info!( + target: "aisix::ratelimit", + backend = "redis", + "connecting shared rate-limit backend" + ); + // No URL in the message: redis URLs carry credentials. + let store = RedisStore::connect(&redis_cfg.url) + .await + .map_err(|e| { + anyhow::anyhow!("redis rate-limit connect failed (ratelimit.redis.url): {e}") + })? + .with_conc_ttl(cfg.ratelimit.concurrency_ttl_secs); + Limiter::with_store(Arc::new(store)) + } + None => Limiter::new(), + }); let metrics = Arc::new(Metrics::new(true)); // Cache backends (#519 B.8). The memory cache is always built // (in-process, cheap); the redis cache is built iff `cache.redis` diff --git a/docs/configuration/rate-limits.md b/docs/configuration/rate-limits.md index 0dc5910b..1d6b2db5 100644 --- a/docs/configuration/rate-limits.md +++ b/docs/configuration/rate-limits.md @@ -136,9 +136,35 @@ When any layer rejects the request, the proxy returns `429`. For rate-limit-styl Successful non-streaming chat responses include `x-ratelimit-*` headers based on the post-dispatch limiter state. Those headers are useful for debugging and for client-side adaptive throttling. +## Counter Storage: Single Node vs Cluster + +Every limit above is enforced against a counter. Where that counter lives is set by the `ratelimit` block in the gateway bootstrap config: + +```yaml title="ratelimit backend" +ratelimit: + backend: "memory" # memory | redis + # redis: + # url: "redis://127.0.0.1:6379" + # mode: "single" + # concurrency_ttl_secs: 300 +``` + +- `memory` (default) — counters live in each gateway process. With a single replica this is exact. With **N replicas behind a load balancer, every limit is effectively multiplied by N**: a key capped at `rpm: 60` gets up to `60 × N` per minute, because each replica counts only the traffic it personally served. +- `redis` — counters are shared across every replica through one Redis, so the whole cluster enforces **one global window** regardless of replica count. Enable this on any multi-replica deployment. The Redis may be the same instance used for the response cache; rate-limit keys are namespaced `aisix:rl:` and hash-tagged per bucket. All dimensions are shared — requests, tokens, and `concurrency` (tracked as a crash-safe distributed semaphore; a slot held by a crashed replica is reclaimed after `concurrency_ttl_secs`, default 300s). + +Enable it via config, or via env on a managed/containerized deployment: + +```bash +AISIX_RATELIMIT__BACKEND=redis +AISIX_RATELIMIT__REDIS__URL=redis://my-redis:6379 +``` + +If Redis becomes unreachable, the limiter **fails open** to per-replica in-memory counting (logged once) so requests keep flowing; cluster-wide limits are not enforced for the duration of the outage and resume automatically when Redis recovers. + ## Operator Guidance - put caller-facing safety limits on `ApiKey.rate_limit` +- on multi-replica deployments, set `ratelimit.backend: redis` so configured limits are enforced cluster-wide instead of per replica - use `Model.rate_limit` to protect a specific upstream model alias - use `RateLimitPolicy` rows when the limit applies to a population that is wider than one key or one model — for example, a whole team - keep token-based caps proportionate to the burst-control caps; a tight `rpm` with an unlimited `tpm` lets a single long completion still saturate upstream diff --git a/tests/e2e/src/cases/ratelimit-cluster-e2e.test.ts b/tests/e2e/src/cases/ratelimit-cluster-e2e.test.ts new file mode 100644 index 00000000..e73c9ebc --- /dev/null +++ b/tests/e2e/src/cases/ratelimit-cluster-e2e.test.ts @@ -0,0 +1,214 @@ +import { createHash, randomUUID } from "node:crypto"; +import { connect } from "node:net"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + AdminClient, + EtcdClient, + ProxyClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +// E2E: cluster-level rate limiting (api7/AISIX-Cloud#798). +// +// Two DP replicas behind one shared etcd (same config → same ApiKey +// entry id → same rate-limit bucket) and one shared Redis. With an +// ApiKey capped at RPM=1, the first request to replica A succeeds and a +// second request to replica B — a DIFFERENT process — is already +// rate-limited (429 + Retry-After). This is the exact repro from the +// issue (curl :3000 then :3001). +// +// The contrast suite below runs the same shape with the default +// `memory` backend and shows BOTH replicas serve the request: per- +// process counters multiply the limit by the replica count, which is +// the bug #798 fixes. + +const CALLER_PLAINTEXT = "sk-rl-cluster-e2e-caller"; +const CALLER_KEY_HASH = createHash("sha256") + .update(CALLER_PLAINTEXT) + .digest("hex"); + +const ETCD_ENDPOINT = process.env.AISIX_E2E_ETCD ?? "http://127.0.0.1:2379"; +const REDIS_URL = process.env.AISIX_E2E_REDIS ?? "redis://127.0.0.1:6379"; + +/** RESP-level PING so the suite skips honestly when no redis is reachable + * (CI provisions redis:7-alpine on :6379). */ +async function redisPing(url: string): Promise { + const m = /^redis:\/\/(?:[^@/]*@)?([^:/]+)(?::(\d+))?/.exec(url); + if (!m) return false; + const host = m[1]; + const port = m[2] ? Number(m[2]) : 6379; + return new Promise((resolve) => { + const sock = connect({ host, port }, () => sock.write("PING\r\n")); + const done = (ok: boolean) => { + sock.destroy(); + resolve(ok); + }; + sock.once("data", (buf) => done(buf.toString().startsWith("+PONG"))); + sock.once("error", () => done(false)); + sock.setTimeout(1000, () => done(false)); + }); +} + +/** A shared etcd block so two replicas read ONE config namespace — the + * ApiKey then has a single entry id across both, which is the rate-limit + * bucket key. (`spawnApp` otherwise gives each app a unique prefix.) */ +function sharedEtcd(prefix: string) { + return { + endpoints: [ETCD_ENDPOINT], + prefix, + dial_timeout_ms: 5000, + request_timeout_ms: 5000, + }; +} + +function chatRequest(proxyUrl: string, model: string): Promise { + return fetch(`${proxyUrl}/v1/chat/completions`, { + method: "POST", + headers: { + authorization: `Bearer ${CALLER_PLAINTEXT}`, + "content-type": "application/json", + }, + body: JSON.stringify({ + model, + messages: [{ role: "user", content: "hello" }], + }), + }); +} + +/** Seed one model + an RPM=1 ApiKey via this app's admin API. The peer + * replica picks the same config up over the shared etcd watch. */ +async function seed(app: SpawnedApp, upstreamBase: string, model: string) { + const admin = new AdminClient(app.adminUrl, app.adminKey); + const pk = await admin.createProviderKey({ + display_name: `${model}-pk`, + secret: "sk-mock", + api_base: `${upstreamBase}/v1`, + }); + await admin.createModel({ + display_name: model, + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: pk.id, + }); + await admin.createApiKey({ + key_hash: CALLER_KEY_HASH, + allowed_models: [model], + rate_limit: { rpm: 1 }, + }); +} + +/** Wait until `model` is visible on `proxyUrl` without spending the RPM=1 + * budget (listModels does not consume a request slot). */ +async function waitModelLive(proxyUrl: string, model: string) { + const probe = new ProxyClient(proxyUrl, CALLER_PLAINTEXT); + await waitConfigPropagation(async () => { + const res = await probe.listModels(); + if (res.status !== 200) return false; + const data = (res.body as { data?: Array<{ id?: string }> }).data ?? []; + return data.some((m) => m.id === model); + }); +} + +describe("rate limit is shared across replicas with backend=redis (#798)", () => { + let appA: SpawnedApp | undefined; + let appB: SpawnedApp | undefined; + let upstream: OpenAiUpstream | undefined; + let infraReady = false; + const prefix = `/aisix-e2e-rl-${randomUUID()}`; + const model = "rl-cluster"; + + beforeAll(async () => { + infraReady = (await new EtcdClient().ping()) && (await redisPing(REDIS_URL)); + if (!infraReady) return; + + upstream = await startOpenAiUpstream(); + const extra = { + etcd: sharedEtcd(prefix), + ratelimit: { backend: "redis", redis: { url: REDIS_URL } }, + }; + appA = await spawnApp({ extra }); + appB = await spawnApp({ extra }); + await seed(appA, upstream.baseUrl, model); + await waitModelLive(appA.proxyUrl, model); + await waitModelLive(appB.proxyUrl, model); + }); + + afterAll(async () => { + await appA?.exit(); + await appB?.exit(); + await upstream?.close(); + // The harness cleans the unique prefixes it generated, not our shared + // override — drop it ourselves. + await new EtcdClient().deletePrefix(prefix); + }); + + test("first call on A succeeds, second call on B is 429", async (ctx) => { + if (!infraReady || !appA || !appB) { + ctx.skip(); + return; + } + + const first = await chatRequest(appA.proxyUrl, model); + expect(first.status).toBe(200); + await first.body?.cancel(); + + // Different process, shared Redis counter → already over the cap. + const second = await chatRequest(appB.proxyUrl, model); + expect(second.status).toBe(429); + // Retry-After is the load-bearing SDK back-off contract. + expect(second.headers.get("retry-after")).toBeTruthy(); + await second.body?.cancel(); + }); +}); + +describe("rate limit is NOT shared with backend=memory (per-replica, the #798 bug)", () => { + let appA: SpawnedApp | undefined; + let appB: SpawnedApp | undefined; + let upstream: OpenAiUpstream | undefined; + let etcdReady = false; + const prefix = `/aisix-e2e-rl-mem-${randomUUID()}`; + const model = "rl-cluster-mem"; + + beforeAll(async () => { + etcdReady = await new EtcdClient().ping(); + if (!etcdReady) return; + + upstream = await startOpenAiUpstream(); + // Shared etcd (same ApiKey entry id) but default memory backend — the + // counters live per-process, so the cap does NOT span replicas. + const extra = { etcd: sharedEtcd(prefix) }; + appA = await spawnApp({ extra }); + appB = await spawnApp({ extra }); + await seed(appA, upstream.baseUrl, model); + await waitModelLive(appA.proxyUrl, model); + await waitModelLive(appB.proxyUrl, model); + }); + + afterAll(async () => { + await appA?.exit(); + await appB?.exit(); + await upstream?.close(); + await new EtcdClient().deletePrefix(prefix); + }); + + test("first call on A and first call on B both succeed", async (ctx) => { + if (!etcdReady || !appA || !appB) { + ctx.skip(); + return; + } + + const first = await chatRequest(appA.proxyUrl, model); + expect(first.status).toBe(200); + await first.body?.cancel(); + + // Default memory backend: B has its own counter → still allowed. With + // N replicas the effective limit is N×, which is what #798 reports. + const second = await chatRequest(appB.proxyUrl, model); + expect(second.status).toBe(200); + await second.body?.cancel(); + }); +}); From 3dc567cf75c11e4afc9333d5c1d2a07d81abdf1e Mon Sep 17 00:00:00 2001 From: Jarvis Date: Mon, 15 Jun 2026 12:01:43 +0800 Subject: [PATCH 2/2] =?UTF-8?q?fix(ratelimit):=20address=20review=20?= =?UTF-8?q?=E2=80=94=20backend=20gating,=20ttl=20validation,=20test=20robu?= =?UTF-8?q?stness?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - main.rs: select the rate-limit store on `ratelimit.backend`, not on `ratelimit.redis` presence, so a stray redis block under `backend: memory` no longer silently activates Redis. - config: reject `concurrency_ttl_secs: 0` for the redis backend (a zero TTL prunes a slot in the same second it is taken, disabling concurrency limiting). + unit test. - redis integration test: poll (bounded) for the detached ZREM instead of a fixed 200ms sleep. - cluster e2e: guard the afterAll deletePrefix behind the readiness flag so teardown doesn't fail when infra is unavailable. --- crates/aisix-core/src/config.rs | 39 +++++++++++++++++-- .../tests/redis_integration.rs | 18 ++++++--- crates/aisix-server/src/main.rs | 20 ++++++---- .../src/cases/ratelimit-cluster-e2e.test.ts | 7 ++-- 4 files changed, 64 insertions(+), 20 deletions(-) diff --git a/crates/aisix-core/src/config.rs b/crates/aisix-core/src/config.rs index 0e934561..310573f2 100644 --- a/crates/aisix-core/src/config.rs +++ b/crates/aisix-core/src/config.rs @@ -701,10 +701,19 @@ impl Config { "observability.metrics.prometheus.addr invalid socket address: {metrics_addr}" ))); } - if self.ratelimit.backend == RateLimitBackend::Redis && self.ratelimit.redis.is_none() { - return Err(BootstrapError::Config( - "ratelimit.backend = redis requires a ratelimit.redis block".into(), - )); + if self.ratelimit.backend == RateLimitBackend::Redis { + if self.ratelimit.redis.is_none() { + return Err(BootstrapError::Config( + "ratelimit.backend = redis requires a ratelimit.redis block".into(), + )); + } + // A zero concurrency TTL would prune a slot in the same second + // it was taken, silently disabling concurrency limiting. + if self.ratelimit.concurrency_ttl_secs == 0 { + return Err(BootstrapError::Config( + "ratelimit.concurrency_ttl_secs must be > 0 for the redis backend".into(), + )); + } } Ok(()) } @@ -870,6 +879,28 @@ ratelimit: assert!(err.to_string().contains("ratelimit.redis")); } + #[test] + fn rejects_zero_concurrency_ttl_for_redis_backend() { + let f = write_yaml( + r#" +etcd: + endpoints: ["http://localhost:2379"] +proxy: + addr: "0.0.0.0:3000" +admin: + addr: "127.0.0.1:3001" + admin_keys: ["k1"] +ratelimit: + backend: "redis" + redis: + url: "redis://127.0.0.1:6379" + concurrency_ttl_secs: 0 +"#, + ); + let err = Config::load_from_path(Some(f.path())).unwrap_err(); + assert!(err.to_string().contains("concurrency_ttl_secs")); + } + #[test] fn loads_ratelimit_redis_config() { let f = write_yaml( diff --git a/crates/aisix-ratelimit/tests/redis_integration.rs b/crates/aisix-ratelimit/tests/redis_integration.rs index dc8ba8ef..77a4f2a3 100644 --- a/crates/aisix-ratelimit/tests/redis_integration.rs +++ b/crates/aisix-ratelimit/tests/redis_integration.rs @@ -158,13 +158,19 @@ async fn concurrency_slot_is_shared_and_released_across_replicas() { "concurrency slot must be shared across replicas" ); - // A finishes → releases the slot (sync + detached ZREM). + // A finishes → releases the slot (sync + detached ZREM). The ZREM is + // fire-and-forget, so poll until the slot frees (bounded) rather than + // assuming a fixed propagation delay that could flake on slow CI. a.release(&key, "a-1"); - tokio::time::sleep(Duration::from_millis(200)).await; // let the detached ZREM land - - b.acquire(&key, &limits, "b-2") - .await - .expect("slot frees up cluster-wide after release"); + let mut acquired = false; + for _ in 0..50 { + if b.acquire(&key, &limits, "b-2").await.is_ok() { + acquired = true; + break; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + assert!(acquired, "slot must free up cluster-wide after release"); } #[tokio::test] diff --git a/crates/aisix-server/src/main.rs b/crates/aisix-server/src/main.rs index bb32fe99..c5350e1f 100644 --- a/crates/aisix-server/src/main.rs +++ b/crates/aisix-server/src/main.rs @@ -24,7 +24,7 @@ mod telemetry; use aisix_admin::{AdminState, ConfigStore, EtcdConfigStore}; use aisix_cache::{Cache, MemoryCache, RedisCache}; use aisix_core::models::Adapter; -use aisix_core::{CacheBackend, Config, EtcdConfig, EtcdTlsConfig}; +use aisix_core::{CacheBackend, Config, EtcdConfig, EtcdTlsConfig, RateLimitBackend}; use aisix_etcd::{EtcdConfigProvider, SnapshotCache, Supervisor}; use aisix_gateway::Hub; use aisix_obs::{init_tracing, install_otlp_tracer, Metrics}; @@ -378,11 +378,17 @@ async fn run(mut cfg: Config) -> anyhow::Result<()> { let hub = Arc::new(build_hub()); // Rate-limit backend (#798). Default `memory` keeps per-process // counters; `redis` shares them across every replica so a cluster - // enforces one global window instead of one-per-replica. Fail fast on - // `backend = redis` without a `ratelimit.redis` block (validated in - // Config::validate, re-checked here before connecting). - let limiter = Arc::new(match cfg.ratelimit.redis.as_ref() { - Some(redis_cfg) => { + // enforces one global window instead of one-per-replica. The + // `ratelimit.backend` field is the selector — a stray `redis` block + // under `backend: memory` is ignored (Config::validate already + // guarantees a `redis` block when `backend = redis`). + let limiter = Arc::new(match cfg.ratelimit.backend { + RateLimitBackend::Redis => { + let redis_cfg = cfg + .ratelimit + .redis + .as_ref() + .expect("validated: ratelimit.redis present when backend = redis"); tracing::info!( target: "aisix::ratelimit", backend = "redis", @@ -397,7 +403,7 @@ async fn run(mut cfg: Config) -> anyhow::Result<()> { .with_conc_ttl(cfg.ratelimit.concurrency_ttl_secs); Limiter::with_store(Arc::new(store)) } - None => Limiter::new(), + RateLimitBackend::Memory => Limiter::new(), }); let metrics = Arc::new(Metrics::new(true)); // Cache backends (#519 B.8). The memory cache is always built diff --git a/tests/e2e/src/cases/ratelimit-cluster-e2e.test.ts b/tests/e2e/src/cases/ratelimit-cluster-e2e.test.ts index e73c9ebc..68225c0b 100644 --- a/tests/e2e/src/cases/ratelimit-cluster-e2e.test.ts +++ b/tests/e2e/src/cases/ratelimit-cluster-e2e.test.ts @@ -142,8 +142,9 @@ describe("rate limit is shared across replicas with backend=redis (#798)", () => await appB?.exit(); await upstream?.close(); // The harness cleans the unique prefixes it generated, not our shared - // override — drop it ourselves. - await new EtcdClient().deletePrefix(prefix); + // override — drop it ourselves. Skip when infra was unavailable (the + // suite skipped) so teardown doesn't fail on an unreachable etcd. + if (infraReady) await new EtcdClient().deletePrefix(prefix); }); test("first call on A succeeds, second call on B is 429", async (ctx) => { @@ -192,7 +193,7 @@ describe("rate limit is NOT shared with backend=memory (per-replica, the #798 bu await appA?.exit(); await appB?.exit(); await upstream?.close(); - await new EtcdClient().deletePrefix(prefix); + if (etcdReady) await new EtcdClient().deletePrefix(prefix); }); test("first call on A and first call on B both succeed", async (ctx) => {