diff --git a/crates/mesh/src/crdt_kv/crdt.rs b/crates/mesh/src/crdt_kv/crdt.rs index d1e8f7d011..e8efdca048 100644 --- a/crates/mesh/src/crdt_kv/crdt.rs +++ b/crates/mesh/src/crdt_kv/crdt.rs @@ -169,6 +169,39 @@ impl CrdtOrMap { all } + /// Live keys under `prefix`, consulting only the engine the prefix + /// routes to. Readers on request paths use this instead of + /// [`Self::keys`], which scans and clones every key of every engine. + /// A `prefix` shorter than a registered prefix could span engines and + /// falls back to the full scan. + pub fn keys_with_prefix(&self, prefix: &str) -> Vec { + let engines = self.engines_snapshot(); + for (registered, engine) in engines.iter() { + if prefix.starts_with(registered.as_str()) { + return engine + .keys() + .into_iter() + .filter(|key| key.starts_with(prefix)) + .collect(); + } + } + if engines + .iter() + .any(|(registered, _)| registered.starts_with(prefix)) + { + return self + .keys() + .into_iter() + .filter(|key| key.starts_with(prefix)) + .collect(); + } + self.default_engine + .keys() + .into_iter() + .filter(|key| key.starts_with(prefix)) + .collect() + } + pub fn all(&self) -> BTreeMap> { let mut all = BTreeMap::new(); for engine in self.all_engines() { diff --git a/crates/mesh/src/kv.rs b/crates/mesh/src/kv.rs index a9574a35ab..1977676748 100644 --- a/crates/mesh/src/kv.rs +++ b/crates/mesh/src/kv.rs @@ -291,13 +291,11 @@ impl CrdtNamespace { } /// List all live keys matching a sub-prefix within this namespace. + /// Routed to this namespace's engine, so the cost scales with this + /// namespace's keys, not the whole store's. pub fn keys(&self, sub_prefix: &str) -> Vec { let full_prefix = format!("{}{}", self.prefix, sub_prefix); - self.store - .keys() - .into_iter() - .filter(|k| k.starts_with(&full_prefix)) - .collect() + self.store.keys_with_prefix(&full_prefix) } /// Subscribe to changes for keys matching a sub-prefix within this namespace. diff --git a/model_gateway/src/mesh/adapters/rate_limit_sync.rs b/model_gateway/src/mesh/adapters/rate_limit_sync.rs index 160617783d..4de6562bf1 100644 --- a/model_gateway/src/mesh/adapters/rate_limit_sync.rs +++ b/model_gateway/src/mesh/adapters/rate_limit_sync.rs @@ -196,6 +196,29 @@ impl RateLimitSyncAdapter { .try_fold(0i64, |acc, s| acc.checked_add(s.count)) .unwrap_or(i64::MAX) } + + /// Cluster-wide aggregate for a counter at exactly `epoch`. Shards from + /// other epochs contribute nothing: the caller pins the window instead + /// of trusting the highest observed epoch, which after a window roll is + /// a stale window nothing has published past yet. + pub fn get_aggregate_at(&self, counter_name: &str, epoch: u64) -> i64 { + debug_assert!( + !counter_name.contains(':'), + "counter_name must not contain ':' (got {counter_name:?})", + ); + let sub_prefix = format!("{counter_name}:"); + let mut total = 0i64; + for key in self.rate_limits.keys(&sub_prefix) { + if let Some(bytes) = self.rate_limits.get(&key) { + if let Some(value) = decode_epoch_count(&bytes) { + if value.epoch == epoch { + total = total.checked_add(value.count).unwrap_or(i64::MAX); + } + } + } + } + total + } } #[cfg(test)] diff --git a/model_gateway/src/mesh/global_rate_limit.rs b/model_gateway/src/mesh/global_rate_limit.rs new file mode 100644 index 0000000000..9a8d4605a2 --- /dev/null +++ b/model_gateway/src/mesh/global_rate_limit.rs @@ -0,0 +1,435 @@ +//! Cluster-wide rate limiting. +//! +//! Each node counts the requests it admits into a wall-clock window and +//! publishes the monotone cumulative count as its `rl:` shard; enforcement +//! compares the epoch-pinned cluster aggregate (plus this node's unflushed +//! delta) against the limit stored under `config:rate_limit`. +//! +//! The EpochMaxWins merge pins the per-shard maximum within an epoch, so a +//! published count must never regress inside a window — the counter only +//! ever increments, and a window roll advances the epoch so the fresh zero +//! replaces (rather than fights) the old peak. +//! +//! Coordination model, stated honestly: windows are wall-clock seconds, so +//! alignment assumes NTP-disciplined clocks. Sub-second boundary jitter +//! costs a brief under-count; a node skewed by a full second or more +//! enforces its own window island (worst case: islands × limit admitted +//! cluster-wide). A peer's spend reaches this node only after that peer's +//! publish (≤ flush interval) AND the next gossip round that ships its +//! op-log (≤ gossip period, ~1s) — the gossip term dominates, and with a +//! 1s window equals it. So the effective cluster guarantee is between +//! `limit` and roughly `limit + nodes × (flush-interval + gossip-period) +//! admits` per window — coordination-free by design, not an exact global +//! counter, and under sustained load it trends toward the per-node-island +//! worst case above. A wall-clock step is absorbed by the rotation's +//! monotone grace: it costs at most ~one grace period, not the step's +//! magnitude. + +use std::{ + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, + }, + time::{Duration, Instant}, +}; + +use parking_lot::Mutex; +use serde::{Deserialize, Serialize}; +use smg_mesh::CrdtNamespace; +use tracing::warn; + +use super::adapters::RateLimitSyncAdapter; + +/// Cluster-wide rate-limit configuration, stored under +/// [`RATE_LIMIT_CONFIG_KEY`] (LWW) and read per request — a DashMap lookup, +/// not a network call. +#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +pub struct RateLimitConfig { + /// Maximum admitted requests per second across the whole cluster. + pub limit_per_second: u64, +} + +/// Config key carrying a bincode-encoded [`RateLimitConfig`]. Absent or +/// undecodable config disables cluster-wide limiting (fail-open; decode +/// failures warn once per failure streak). +pub const RATE_LIMIT_CONFIG_KEY: &str = "config:rate_limit"; + +/// The cluster-wide counter name every node shards into. +const GLOBAL_COUNTER: &str = "global"; + +/// Real time a window must survive a wall clock running behind it before +/// the epoch is re-anchored to the wall. Absorbs NTP steps and spurious +/// readings: a backward step or a forward glitch wedges admission for at +/// most this long instead of the step's magnitude. +const REGRESSION_GRACE: Duration = Duration::from_secs(2); + +/// Pace of inline shard publishing from the admission path. Faster than +/// the 1s window so peers see this node's spend mid-window (tightening the +/// cross-node overshoot), slow enough that the per-publish CRDT put and +/// op-log append stay negligible. Publishing is traffic-driven, so no +/// background task is needed: an admit advances the count and publishes; +/// once admits stop, the last published value is frozen — it lags the +/// window's true final count by at most one flush interval of admits, and +/// a window roll discards that unpublished tail. Harmless for enforcement +/// (this node's own check adds the unflushed delta; peers under-count by a +/// bounded amount in the over-admit direction), but a future consumer of +/// the `rl:` shard history (observability, billing) would under-read every +/// window's tail — wire a flush-on-roll then. +const FLUSH_INTERVAL: Duration = Duration::from_millis(250); + +/// One node's view of the current window. +struct Window { + /// Wall-clock second this window counts (the EpochMaxWins epoch). + epoch: u64, + /// Requests admitted this window on this node. Monotone within the + /// window — EpochMaxWins pins the per-shard maximum, so a published + /// count must never regress inside an epoch. + count: i64, + /// The count last published via `sync_counter`; the delta above it is + /// what the cluster aggregate does not see yet. + flushed: i64, + /// Monotone moment of the last rotation, the regression detector. + rotated_at: Instant, + /// Monotone moment of the last publish, pacing the inline flush. + flushed_at: Instant, +} + +impl Window { + /// Roll forward when the wall clock entered a new epoch — or re-anchor + /// to the wall when it has run *behind* the epoch for a full grace of + /// real time (a backward NTP step, or a prior spurious forward + /// reading). Re-anchoring may republish a lower count at an epoch this + /// node already used; EpochMaxWins masks it by keeping the old peak, + /// which over-counts this node for that one window — the conservative + /// direction. + fn rotate(&mut self, now_secs: u64, grace: Duration) { + let regressed = now_secs < self.epoch && self.rotated_at.elapsed() >= grace; + if now_secs > self.epoch || regressed { + self.epoch = now_secs; + self.count = 0; + self.flushed = 0; + self.rotated_at = Instant::now(); + } + } +} + +/// Cluster-wide request limiter. One per gateway, owned by `MeshAdapters`; +/// the concurrency middleware checks it before the local token bucket and +/// records only requests the node actually admits. +pub struct GlobalRateLimiter { + rate_limit: Arc, + configs: Arc, + window: Mutex, + regression_grace: Duration, + /// Latches decode failures so the warn fires once per failure streak. + config_decode_failed: AtomicBool, +} + +impl std::fmt::Debug for GlobalRateLimiter { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("GlobalRateLimiter").finish_non_exhaustive() + } +} + +impl GlobalRateLimiter { + pub(crate) fn new( + rate_limit: Arc, + configs: Arc, + ) -> Arc { + Self::with_regression_grace(rate_limit, configs, REGRESSION_GRACE) + } + + fn with_regression_grace( + rate_limit: Arc, + configs: Arc, + regression_grace: Duration, + ) -> Arc { + Arc::new(Self { + rate_limit, + configs, + window: Mutex::new(Window { + epoch: 0, + count: 0, + flushed: 0, + rotated_at: Instant::now(), + flushed_at: Instant::now(), + }), + regression_grace, + config_decode_failed: AtomicBool::new(false), + }) + } + + /// The configured cluster-wide limit, if any. Fail-open on a missing + /// key; a present-but-undecodable value also fails open but warns once + /// per failure streak — silently disabled limiting must not pass for + /// unconfigured. + fn limit(&self) -> Option { + let bytes = self.configs.get(RATE_LIMIT_CONFIG_KEY)?; + match bincode::deserialize::(&bytes) { + Ok(config) => { + self.config_decode_failed.store(false, Ordering::Release); + Some(config.limit_per_second) + } + Err(err) => { + if !self.config_decode_failed.swap(true, Ordering::AcqRel) { + warn!( + %err, + "config:rate_limit present but undecodable; cluster-wide \ + rate limiting is DISABLED until it decodes" + ); + } + None + } + } + } + + /// Whether a request arriving at `now_secs` fits the cluster budget. + /// Does NOT consume budget — the caller counts via + /// [`Self::record_admit`] only after the request clears local admission + /// too, so requests the local limiter then rejects never burn the + /// cluster's budget (a saturated node must not starve its peers). + /// + /// `check` and `record_admit` are separate lock acquisitions, so up to + /// the node's in-flight concurrency can pass `check` before any records; + /// per-node overshoot is bounded by that concurrency. Deliberate — the + /// split is what keeps locally-rejected requests from burning budget — + /// and dwarfed by the gossip-cadence slack documented on the module. + pub fn check(&self, now_secs: u64) -> bool { + let Some(limit) = self.limit() else { + return true; + }; + let mut window = self.window.lock(); + window.rotate(now_secs, self.regression_grace); + // Pinned to this window's epoch: after a roll, the stale window's + // shards must not gate admission (nothing would ever publish the + // new epoch and the cluster would reject forever). + let aggregate = self + .rate_limit + .get_aggregate_at(GLOBAL_COUNTER, window.epoch); + let unflushed = (window.count - window.flushed).max(0); + aggregate.saturating_add(unflushed) < limit as i64 + } + + /// Count one locally-admitted request into the current window, + /// publishing the grown shard at most once per [`FLUSH_INTERVAL`]. + pub fn record_admit(&self, now_secs: u64) { + if self.limit().is_none() { + return; + } + let mut window = self.window.lock(); + window.rotate(now_secs, self.regression_grace); + window.count += 1; + if window.flushed_at.elapsed() >= FLUSH_INTERVAL { + self.publish(&mut window); + } + } + + /// Force publication regardless of pacing. Production publishes inline + /// from `record_admit`; only tests need to flush on demand. + #[cfg(test)] + pub(crate) fn flush(&self, now_secs: u64) { + let mut window = self.window.lock(); + window.rotate(now_secs, self.regression_grace); + self.publish(&mut window); + } + + // Runs under the window Mutex and calls into the CRDT store + // (window-lock -> store-lock order). Safe because the store's subscriber + // notify is `try_send` with no synchronous callback back into the + // limiter; if a future store path called back in, that order would need + // revisiting. EpochMaxWins tolerates same-epoch out-of-order publishes, + // so the put could move outside the lock should this ever contend. + fn publish(&self, window: &mut Window) { + if window.count > window.flushed { + self.rate_limit + .sync_counter(GLOBAL_COUNTER, window.epoch, window.count); + window.flushed = window.count; + window.flushed_at = Instant::now(); + } + } +} + +/// Wall-clock seconds since the unix epoch; the cross-node window identity. +pub(crate) fn unix_now_secs() -> u64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|elapsed| elapsed.as_secs()) + .unwrap_or(0) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use smg_mesh::{MergeStrategy, MeshKV}; + + use super::*; + + fn limiter_with(limit: Option) -> Arc { + limiter_with_grace(limit, REGRESSION_GRACE) + } + + fn limiter_with_grace(limit: Option, grace: Duration) -> Arc { + let mesh = MeshKV::new("node-a".into()); + let rl_ns = mesh.configure_crdt_prefix("rl:", MergeStrategy::EpochMaxWins); + let rate_limit = RateLimitSyncAdapter::new(rl_ns, "node-a".into()); + let configs = mesh.configs(); + if let Some(limit_per_second) = limit { + let config = RateLimitConfig { limit_per_second }; + configs.put( + RATE_LIMIT_CONFIG_KEY, + bincode::serialize(&config).expect("config encodes"), + ); + } + GlobalRateLimiter::with_regression_grace(rate_limit, configs, grace) + } + + /// The middleware's two-step admission as one call. + fn admit(limiter: &GlobalRateLimiter, now_secs: u64) -> bool { + if limiter.check(now_secs) { + limiter.record_admit(now_secs); + true + } else { + false + } + } + + #[test] + fn no_config_means_unlimited() { + let limiter = limiter_with(None); + for _ in 0..1000 { + assert!(admit(&limiter, 100)); + } + } + + #[test] + fn undecodable_config_fails_open() { + let limiter = limiter_with(None); + limiter + .configs + .put(RATE_LIMIT_CONFIG_KEY, b"not-bincode!".to_vec()); + assert!(admit(&limiter, 100), "corrupt config fails open"); + } + + #[test] + fn admits_up_to_the_limit_then_rejects() { + let limiter = limiter_with(Some(5)); + for _ in 0..5 { + assert!(admit(&limiter, 100), "under the limit admits"); + } + assert!(!admit(&limiter, 100), "at the limit rejects"); + assert!(!admit(&limiter, 100), "rejections consume no budget"); + } + + #[test] + fn check_alone_consumes_no_budget() { + // The middleware checks before local admission and records only + // after it: requests the local limiter rejects must not burn the + // cluster budget. + let limiter = limiter_with(Some(2)); + for _ in 0..100 { + assert!(limiter.check(100)); + } + assert!(admit(&limiter, 100)); + assert!(admit(&limiter, 100)); + assert!(!admit(&limiter, 100)); + } + + #[test] + fn window_roll_resets_the_budget() { + let limiter = limiter_with(Some(2)); + assert!(admit(&limiter, 100)); + assert!(admit(&limiter, 100)); + assert!(!admit(&limiter, 100)); + + assert!(admit(&limiter, 101), "a new window admits again"); + } + + #[test] + fn flush_publishes_monotone_cumulative_counts() { + let limiter = limiter_with(Some(100)); + assert!(admit(&limiter, 100)); + assert!(admit(&limiter, 100)); + limiter.flush(100); + assert_eq!( + limiter.rate_limit.get_aggregate("global"), + 2, + "the flushed shard carries the cumulative window count" + ); + + assert!(admit(&limiter, 100)); + limiter.flush(100); + assert_eq!( + limiter.rate_limit.get_aggregate("global"), + 3, + "the count grows monotonically within the window" + ); + } + + #[test] + fn enforcement_combines_aggregate_and_unflushed_delta() { + let limiter = limiter_with(Some(10)); + for _ in 0..4 { + assert!(admit(&limiter, 100)); + } + limiter.flush(100); + + // The aggregate carries the flushed 4; six unflushed admits exhaust + // the cluster budget of 10 without double-counting the shard. + for _ in 0..6 { + assert!(admit(&limiter, 100)); + } + assert!(!admit(&limiter, 100), "cluster budget exhausted"); + } + + #[test] + fn stale_epoch_shards_do_not_inflate_the_budget_check() { + let limiter = limiter_with(Some(2)); + assert!(admit(&limiter, 100)); + assert!(admit(&limiter, 100)); + limiter.flush(100); + assert!(!admit(&limiter, 100), "window 100 exhausted"); + + // The flushed shard (epoch 100, count 2) still sits in the store, + // but the budget check is pinned to the rolled window's epoch — a + // stale window must never wedge the cluster into rejecting forever. + assert!(admit(&limiter, 101)); + limiter.flush(101); + assert_eq!(limiter.rate_limit.get_aggregate("global"), 1); + } + + #[test] + fn backward_clock_step_recovers_after_the_grace() { + // A backward NTP step must wedge admission for at most the + // regression grace, not the step's magnitude. + let limiter = limiter_with_grace(Some(2), Duration::from_millis(50)); + assert!(admit(&limiter, 100)); + assert!(admit(&limiter, 100)); + assert!(!admit(&limiter, 100), "window 100 exhausted"); + + // The wall clock steps back 50 seconds; within the grace the window + // stays pinned (and exhausted). + assert!(!admit(&limiter, 50), "regression inside the grace holds"); + + std::thread::sleep(Duration::from_millis(60)); + assert!( + admit(&limiter, 50), + "past the grace the window re-anchors to the wall clock" + ); + } + + #[test] + fn forward_clock_glitch_recovers_after_the_grace() { + // One spurious far-future reading must not wedge the node until + // real time catches up. + let limiter = limiter_with_grace(Some(2), Duration::from_millis(50)); + assert!(admit(&limiter, 100)); + assert!(admit(&limiter, 4000), "glitch rolls the window forward"); + + std::thread::sleep(Duration::from_millis(60)); + assert!( + admit(&limiter, 101), + "after the grace the window re-anchors to real wall time" + ); + } +} diff --git a/model_gateway/src/mesh/mod.rs b/model_gateway/src/mesh/mod.rs index 3eed824db8..9d305e2c6b 100644 --- a/model_gateway/src/mesh/mod.rs +++ b/model_gateway/src/mesh/mod.rs @@ -3,7 +3,9 @@ //! and shutdown wiring added in later steps. pub mod adapters; +pub mod global_rate_limit; pub mod wiring; pub use adapters::{RateLimitSyncAdapter, TreeDelta, TreeSyncAdapter, WorkerSyncAdapter}; +pub use global_rate_limit::{GlobalRateLimiter, RateLimitConfig, RATE_LIMIT_CONFIG_KEY}; pub use wiring::MeshAdapters; diff --git a/model_gateway/src/mesh/wiring.rs b/model_gateway/src/mesh/wiring.rs index de37e19633..625e226862 100644 --- a/model_gateway/src/mesh/wiring.rs +++ b/model_gateway/src/mesh/wiring.rs @@ -11,7 +11,10 @@ use std::sync::Arc; use smg_mesh::{MergeStrategy, MeshKV}; -use super::adapters::{RateLimitSyncAdapter, WorkerSyncAdapter}; +use super::{ + adapters::{RateLimitSyncAdapter, WorkerSyncAdapter}, + global_rate_limit::GlobalRateLimiter, +}; use crate::worker::WorkerRegistry; /// Owns the started mesh sync adapters. Mesh on means every adapter here is @@ -21,6 +24,7 @@ use crate::worker::WorkerRegistry; pub struct MeshAdapters { worker: Arc, rate_limit: Arc, + global_rate_limit: Arc, } impl MeshAdapters { @@ -49,7 +53,16 @@ impl MeshAdapters { let rate_limit = RateLimitSyncAdapter::new(rl_ns, node_name); worker.start(); rate_limit.start(); - Arc::new(Self { worker, rate_limit }) + // No background task: the limiter publishes its shard inline from + // the admission path, paced by its flush interval. Once admits stop + // the published value is frozen (the unpublished tail is bounded and + // enforcement-harmless — see GlobalRateLimiter), so no timer is owed. + let global_rate_limit = GlobalRateLimiter::new(rate_limit.clone(), mesh_kv.configs()); + Arc::new(Self { + worker, + rate_limit, + global_rate_limit, + }) } /// Worker sync adapter. @@ -61,6 +74,12 @@ impl MeshAdapters { pub fn rate_limit(&self) -> &Arc { &self.rate_limit } + + /// Cluster-wide request limiter, consulted by the concurrency + /// middleware before the local token bucket. + pub fn global_rate_limit(&self) -> &Arc { + &self.global_rate_limit + } } #[cfg(test)] diff --git a/model_gateway/src/middleware/concurrency.rs b/model_gateway/src/middleware/concurrency.rs index 0625d3ffeb..eb4fe64831 100644 --- a/model_gateway/src/middleware/concurrency.rs +++ b/model_gateway/src/middleware/concurrency.rs @@ -28,6 +28,7 @@ use tracing::{debug, error, warn}; use super::token_bucket::TokenBucket; use crate::{ + mesh::global_rate_limit::unix_now_secs, observability::metrics::{metrics_labels, Metrics}, server::AppState, }; @@ -202,16 +203,34 @@ pub async fn concurrency_limit_middleware( request: Request, next: Next, ) -> Response { - // Cluster-wide rate limiting was previously enforced via the - // v1 `MeshSyncManager::check_global_rate_limit` path. That hook - // is removed in this PR. Local per-node token-bucket rate - // limiting below still applies; cluster aggregation will return - // through the v2 `RateLimitSyncAdapter` in a follow-up PR. + // Cluster-wide rate limiting: each node's admitted-request + // count gossips as an `rl:` shard; the limiter compares the + // epoch-pinned cluster aggregate against `config:rate_limit`. + // Absent mesh or config, this is a no-op. The check consumes no + // budget — only requests that ALSO clear local admission are + // recorded below, so a saturated node's locally-rejected flood + // never burns the cluster's budget. + let global_limiter = app_state + .mesh_adapters + .as_ref() + .map(|adapters| adapters.global_rate_limit().clone()); + if let Some(limiter) = &global_limiter { + if !limiter.check(unix_now_secs()) { + debug!("Cluster-wide rate limit exceeded, returning 429"); + // Distinct from the local token-bucket's RATE_LIMIT_REJECTED so a + // 429 spike is attributable to the cluster limiter vs this node. + Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_GLOBAL_REJECTED); + return StatusCode::TOO_MANY_REQUESTS.into_response(); + } + } let token_bucket = match &app_state.context.rate_limiter { Some(bucket) => bucket.clone(), None => { - // Rate limiting disabled, pass through immediately + // No local limiter: the global check above was the admission. + if let Some(limiter) = &global_limiter { + limiter.record_admit(unix_now_secs()); + } return next.run(request).await; } }; @@ -219,6 +238,9 @@ pub async fn concurrency_limit_middleware( // Try to acquire token immediately if token_bucket.try_acquire(1.0).is_ok() { debug!("Acquired token immediately"); + if let Some(limiter) = &global_limiter { + limiter.record_admit(unix_now_secs()); + } Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_ALLOWED); let response = next.run(request).await; @@ -248,6 +270,9 @@ pub async fn concurrency_limit_middleware( match permit_rx.await { Ok(Ok(())) => { debug!("Acquired token from queue"); + if let Some(limiter) = &global_limiter { + limiter.record_admit(unix_now_secs()); + } Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_ALLOWED); let response = next.run(request).await; diff --git a/model_gateway/src/middleware/scheduler/admission.rs b/model_gateway/src/middleware/scheduler/admission.rs index 9a81e1c9d9..9224034b36 100644 --- a/model_gateway/src/middleware/scheduler/admission.rs +++ b/model_gateway/src/middleware/scheduler/admission.rs @@ -29,6 +29,7 @@ use super::{ SchedulerError, SchedulerGuardBody, HEADER_X_SMG_PREEMPTED, PRIORITY_HEADER, }; use crate::{ + mesh::global_rate_limit::unix_now_secs, middleware::RouteRequestMeta, observability::metrics::{metrics_labels, Metrics}, tenant::TenantKey, @@ -112,6 +113,18 @@ pub async fn priority_admission_middleware( sched_metrics::record_clamp(resolved.requested, class, tenant.as_str()); } + // Cluster-wide rate limit, checked before the local RPS bucket and + // scheduler so a cluster-rejected request consumes neither. Mirrors the + // legacy concurrency path; the check consumes no budget (recorded only + // once the request is actually admitted below), so a rejection here + // never burns the cluster's budget. + if let Some(limiter) = &state.global_rate_limit { + if !limiter.check(unix_now_secs()) { + Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_GLOBAL_REJECTED); + return SchedulerError::QueueFull.into_response(); + } + } + // RPS sibling check (only set when an explicit per-second limit is // configured). Checked before admission so a rejected request never // consumes a slot. Tokens are not returned — refill is time-based. @@ -134,6 +147,12 @@ pub async fn priority_admission_middleware( match state.scheduler.admit(class, request_id, cancel).await { AdmitOutcome::Admitted(permit) => { + // Cleared the cluster gate, the RPS bucket, and the scheduler: + // count this admission toward the cluster window now (only + // admitted requests consume cluster budget). + if let Some(limiter) = &state.global_rate_limit { + limiter.record_admit(unix_now_secs()); + } // Hand the handler the cancel token (for preemption select!). req.extensions_mut().insert(permit.cancel_token()); let response = next.run(req).await; diff --git a/model_gateway/src/middleware/scheduler/state.rs b/model_gateway/src/middleware/scheduler/state.rs index c16cbe39d1..5ea96335bf 100644 --- a/model_gateway/src/middleware/scheduler/state.rs +++ b/model_gateway/src/middleware/scheduler/state.rs @@ -10,6 +10,7 @@ use super::{ }; use crate::{ config::types::RouterConfig, + mesh::global_rate_limit::GlobalRateLimiter, middleware::token_bucket::TokenBucket, worker::{CapacityTrackerSettings, WorkerCapacity, WorkerRegistry}, }; @@ -28,6 +29,11 @@ pub struct SchedulerState { /// it as a concurrency limiter (that would double-limit). `None` = /// no RPS limit. pub rate_limiter: Option>, + /// Cluster-wide rate limiter, checked before admission (and before the + /// RPS bucket) so the advertised `config:rate_limit` is enforced under + /// the priority scheduler too, not only on the legacy path. `None` when + /// mesh is off. + pub global_rate_limit: Option>, } /// Which admission path the protected routes use. Chosen once at startup. @@ -55,11 +61,12 @@ impl AdmissionMode { rc: &RouterConfig, registry: Arc, rate_limiter: Option>, + global_rate_limit: Option>, ) -> Self { if !rc.priority_scheduler_enabled { return Self::Legacy; } - match Self::try_build_priority(rc, registry, rate_limiter) { + match Self::try_build_priority(rc, registry, rate_limiter, global_rate_limit) { Ok(mode) => { info!("priority scheduler enabled"); mode @@ -78,6 +85,7 @@ impl AdmissionMode { rc: &RouterConfig, registry: Arc, rate_limiter: Option>, + global_rate_limit: Option>, ) -> Result { // Tier-4 fallback for WorkerCapacity comes from the legacy // --max-concurrent-requests (clamped to u16; <=0 means "disabled", @@ -120,6 +128,7 @@ impl AdmissionMode { scheduler, resolver, rate_limiter, + global_rate_limit, }))) } } diff --git a/model_gateway/src/observability/metrics.rs b/model_gateway/src/observability/metrics.rs index cdf63e2296..782870e600 100644 --- a/model_gateway/src/observability/metrics.rs +++ b/model_gateway/src/observability/metrics.rs @@ -158,7 +158,7 @@ pub(crate) fn init_metrics() { ); describe_counter!( "smg_http_rate_limit_total", - "Rate limiting decisions by result (allowed/rejected)" + "Rate limiting decisions by result (allowed/rejected/global_rejected)" ); // Layer 2: Router metrics @@ -547,6 +547,9 @@ pub mod metrics_labels { // Rate limit results pub const RATE_LIMIT_ALLOWED: &str = "allowed"; pub const RATE_LIMIT_REJECTED: &str = "rejected"; + /// Rejected by the cluster-wide limiter (vs the local token bucket), so + /// a 429 spike can be attributed to the right source. + pub const RATE_LIMIT_GLOBAL_REJECTED: &str = "global_rejected"; // Circuit breaker states pub const CB_CLOSED: &str = "closed"; diff --git a/model_gateway/src/server.rs b/model_gateway/src/server.rs index ef9a702adc..1c9edfe7de 100644 --- a/model_gateway/src/server.rs +++ b/model_gateway/src/server.rs @@ -766,6 +766,10 @@ pub fn build_app( &app_state.context.router_config, app_state.context.worker_registry.clone(), app_state.context.rate_limiter.clone(), + app_state + .mesh_adapters + .as_ref() + .map(|adapters| adapters.global_rate_limit().clone()), ); let protected_routes = with_admission_layer(