diff --git a/crates/aisix-core/src/snapshot.rs b/crates/aisix-core/src/snapshot.rs index 75079e81..78f891d4 100644 --- a/crates/aisix-core/src/snapshot.rs +++ b/crates/aisix-core/src/snapshot.rs @@ -18,17 +18,38 @@ use crate::resource::{Resource, ResourceEntry}; use arc_swap::ArcSwap; use dashmap::DashMap; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::Arc; /// Per-kind table with primary id-index and secondary name-index. /// /// Both indices point at the same `Arc>` so there is no /// duplicate storage — the name map just holds ids. -#[derive(Debug, Clone)] +#[derive(Debug)] pub struct ResourceTable { by_id: DashMap>>, by_name: DashMap, + /// Cached entry count, maintained by [`ResourceTable::insert`] / + /// [`ResourceTable::remove`]. DashMap's own `len()` / `is_empty()` + /// visit every shard (a CAS pair per shard), so per-request + /// emptiness checks on the hot path go through this counter + /// instead — one relaxed load, O(1) regardless of shard count. + count: AtomicUsize, +} + +/// Manual impl: `AtomicUsize` is not `Clone`. The count is re-seeded +/// from the cloned map's length, which the etcd watch supervisor's +/// clone-then-mutate update cycle relies on being exact. +impl Clone for ResourceTable { + fn clone(&self) -> Self { + let by_id = self.by_id.clone(); + let count = AtomicUsize::new(by_id.len()); + Self { + by_id, + by_name: self.by_name.clone(), + count, + } + } } impl Default for ResourceTable { @@ -36,6 +57,7 @@ impl Default for ResourceTable { Self { by_id: DashMap::new(), by_name: DashMap::new(), + count: AtomicUsize::new(0), } } } @@ -46,11 +68,11 @@ impl ResourceTable { } pub fn len(&self) -> usize { - self.by_id.len() + self.count.load(Ordering::Relaxed) } pub fn is_empty(&self) -> bool { - self.by_id.is_empty() + self.len() == 0 } /// Insert or replace an entry, updating both indices. @@ -70,12 +92,21 @@ impl ResourceTable { } self.by_name.insert(name, id.clone()); - self.by_id.insert(id, Arc::new(entry)); + // Provisional increment BEFORE the map insert, corrected after a + // replace. Orders the count so it can only ever read high during + // a mutation window, never low: the empty fast paths may take + // one redundant full scan, but can never skip an entry that is + // already visible in the map. + self.count.fetch_add(1, Ordering::Relaxed); + if self.by_id.insert(id, Arc::new(entry)).is_some() { + self.count.fetch_sub(1, Ordering::Relaxed); + } } /// Remove by id; also removes the matching name index entry. pub fn remove(&self, id: &str) -> Option>> { let (_, entry) = self.by_id.remove(id)?; + self.count.fetch_sub(1, Ordering::Relaxed); let name = entry.value.name().to_string(); self.by_name.remove_if(&name, |_, v| v == id); Some(entry) @@ -104,8 +135,13 @@ impl ResourceTable { } /// Snapshot of all entries. Callers get owned `Arc` clones, so iteration - /// does not hold DashMap shards. + /// does not hold DashMap shards. O(1) when the table is empty — the + /// per-request callers (exporter fan-out, policy scans) skip the + /// all-shards walk on unconfigured deployments. pub fn entries(&self) -> Vec>> { + if self.is_empty() { + return Vec::new(); + } self.by_id.iter().map(|kv| kv.value().clone()).collect() } @@ -115,7 +151,7 @@ impl ResourceTable { /// guard is held during the scan, so `pred` must not call back into /// this table. pub fn any(&self, pred: impl Fn(&ResourceEntry) -> bool) -> bool { - self.by_id.iter().any(|kv| pred(kv.value())) + !self.is_empty() && self.by_id.iter().any(|kv| pred(kv.value())) } /// The single entry satisfying `pred`, without materialising the @@ -128,6 +164,9 @@ impl ResourceTable { &self, pred: impl Fn(&ResourceEntry) -> bool, ) -> (Option>>, bool) { + if self.is_empty() { + return (None, false); + } let mut found: Option>> = None; for kv in self.by_id.iter() { if !pred(kv.value()) { @@ -290,6 +329,42 @@ mod tests { assert!(t.get_by_name("alpha").is_none()); } + /// The cached count must stay exact through every mutation shape: + /// fresh insert, same-id replace, remove, remove-miss, and clone. + #[test] + fn cached_count_tracks_all_mutations() { + let t = ResourceTable::::new(); + assert_eq!(t.len(), 0); + assert!(t.is_empty()); + + t.insert(entry("a-1", "alpha")); + t.insert(entry("b-2", "beta")); + assert_eq!(t.len(), 2); + assert!(!t.is_empty()); + + // Same-id replace (update, incl. rename) must not double-count. + t.insert(entry("a-1", "aleph")); + assert_eq!(t.len(), 2); + + // Remove-miss must not decrement. + assert!(t.remove("missing").is_none()); + assert_eq!(t.len(), 2); + + assert!(t.remove("a-1").is_some()); + assert_eq!(t.len(), 1); + + // Clone re-seeds the counter from the cloned map. + let c = t.clone(); + assert_eq!(c.len(), 1); + c.insert(entry("c-3", "gamma")); + assert_eq!(c.len(), 2); + assert_eq!(t.len(), 1); // original untouched + + assert!(t.remove("b-2").is_some()); + assert_eq!(t.len(), 0); + assert!(t.is_empty()); + } + #[test] fn snapshot_handle_atomic_swap() { let handle: SnapshotHandle = SnapshotHandle::new(0); diff --git a/crates/aisix-core/src/wildcard.rs b/crates/aisix-core/src/wildcard.rs index fd82e2ed..47e3d5d1 100644 --- a/crates/aisix-core/src/wildcard.rs +++ b/crates/aisix-core/src/wildcard.rs @@ -33,12 +33,23 @@ pub fn wildcard_capture(pattern: &str, candidate: &str) -> Option { /// Whether `pattern` matches `candidate`. A pattern with a single `*` is /// glob-matched; any other pattern matches only an exactly-equal candidate. +/// +/// Same decision as `wildcard_capture(..).is_some()` but without +/// materialising the capture — this runs on the per-request authz path, +/// where a bare `"*"` allowlist entry would otherwise copy the whole +/// candidate just to discard it. pub fn wildcard_matches(pattern: &str, candidate: &str) -> bool { - if pattern.contains('*') { - wildcard_capture(pattern, candidate).is_some() - } else { - pattern == candidate + let Some(star) = pattern.find('*') else { + return pattern == candidate; + }; + if pattern[star + 1..].contains('*') { + return false; // only a single '*' is supported } + let prefix = &pattern[..star]; + let suffix = &pattern[star + 1..]; + candidate.len() >= prefix.len() + suffix.len() + && candidate.starts_with(prefix) + && candidate.ends_with(suffix) } #[cfg(test)] @@ -82,6 +93,48 @@ mod tests { assert_eq!(wildcard_capture("a/*/*", "a/b/c"), None); } + /// Differential oracle: `wildcard_matches` must agree with + /// `wildcard_capture(..).is_some()` for every pattern shape — the + /// two are duplicate decision procedures guarding the authz + /// allowlists, and a future edit to either must not let them + /// diverge silently. + #[test] + fn matches_agrees_with_capture_oracle() { + let pats = [ + "*", + "openai/*", + "*-sfx", + "gpt-*-preview", + "aa*aa", + "a/*/*", + "**", + "lit", + "", + ]; + let cands = [ + "", + "openai/", + "openai/gpt-4o", + "aaa", + "aaaa", + "gpt-4o-preview", + "gpt-4o-final", + "a/b/c", + "-sfx", + "lit", + ]; + for p in pats { + for c in cands { + let expected = if p.contains('*') { + wildcard_capture(p, c).is_some() + } else { + p == c + }; + assert_eq!(wildcard_matches(p, c), expected, "p={p:?} c={c:?}"); + } + } + } + #[test] fn matches_handles_literals_and_globs() { assert!(wildcard_matches("*", "anything")); diff --git a/crates/aisix-guardrails/src/build.rs b/crates/aisix-guardrails/src/build.rs index a6f889f4..248c4941 100644 --- a/crates/aisix-guardrails/src/build.rs +++ b/crates/aisix-guardrails/src/build.rs @@ -1174,8 +1174,19 @@ impl LiveGuardrailIndex { /// Cheap on the cache-hit path (one lock acquire + version compare + /// arc clone + `O(n)` linear walk over attachment rows). Rebuilds only /// on snapshot version change. + /// + /// An empty index (no guardrails configured — the default) resolves + /// to the same empty chain for every request, so that case returns + /// early: no resolve walk, no applied-set copy, no sink attach (a + /// chain with no members never reports to the sink). This is the one + /// chokepoint every endpoint family resolves through, so the fast + /// path covers them all. pub fn resolve(&self, ctx: &RequestContext<'_>) -> GuardrailChain { - self.current() + let index = self.current(); + if index.is_empty() { + return GuardrailChain::empty(); + } + index .resolve(ctx) .with_metrics_sink(self.metrics_sink.clone()) } @@ -2174,6 +2185,22 @@ mod tests { assert!(!live.is_empty()); } + #[tokio::test] + async fn live_index_empty_fast_path_resolves_empty_chain() { + // Zero-config fast path: an empty index resolves to an empty + // chain with an empty applied set — the same observable shape + // the full resolve walk produces on an empty index. + let live = LiveGuardrailIndex::new(SnapshotHandle::new(AisixSnapshot::new()), None); + let chain = live.resolve(&RequestContext { + model_id: "m", + api_key_id: "k", + team_id: None, + }); + assert!(chain.is_empty()); + assert!(chain.applied().is_empty()); + assert!(!chain.check_input(&req("anything")).await.is_block()); + } + // ----------------------------------------------------------------------- // per-execution metrics sink (AISIX-Cloud#1076) // ----------------------------------------------------------------------- diff --git a/crates/aisix-proxy/src/auth.rs b/crates/aisix-proxy/src/auth.rs index 1477c823..c43cef0f 100644 --- a/crates/aisix-proxy/src/auth.rs +++ b/crates/aisix-proxy/src/auth.rs @@ -48,12 +48,39 @@ pub struct JwtIdentity { /// metric answers "how many" but not "who, when, against what" — so the /// denial log line, which an operator turns on precisely when investigating, /// carries the identifying detail instead. -#[derive(Clone, Copy, Default)] +#[derive(Clone, Copy)] pub(crate) struct DenialContext<'a> { pub method: &'a str, pub path: &'a str, pub request_id: &'a str, - pub source_ip: &'a str, + pub source_ip: LazySourceIp<'a>, +} + +/// Source IP for denial logs, resolved only when a denial actually logs +/// it. The successful-auth path has no consumer for it, so it never pays +/// the forwarded-header scan + IP formatting — `ClientContext` resolves +/// the same value once for the handlers instead. +#[derive(Clone, Copy)] +pub(crate) enum LazySourceIp<'a> { + /// Resolve from the request parts + real-ip config on first use. + Deferred( + &'a axum::http::request::Parts, + &'a crate::client_ip::ResolvedRealIp, + ), + /// Already resolved by the caller (WebSocket subprotocol auth, which + /// has a `ClientContext` in hand). + Ready(&'a str), +} + +impl LazySourceIp<'_> { + pub(crate) fn resolve(&self) -> std::borrow::Cow<'_, str> { + match self { + Self::Deferred(parts, cfg) => { + std::borrow::Cow::Owned(crate::client_ip::source_ip_from_parts(parts, cfg)) + } + Self::Ready(s) => std::borrow::Cow::Borrowed(s), + } + } } impl AuthenticatedKey { @@ -72,7 +99,6 @@ where async fn from_request_parts(parts: &mut Parts, state: &S) -> Result { let proxy_state = ProxyState::from_ref(state); - let source_ip = crate::client_ip::source_ip_from_parts(parts, &proxy_state.real_ip); let request_id = parts .extensions .get::() @@ -82,7 +108,8 @@ where method: parts.method.as_str(), path: parts.uri.path(), request_id, - source_ip: &source_ip, + // Deferred: only a denial resolves the source IP. + source_ip: LazySourceIp::Deferred(&*parts, proxy_state.real_ip.as_ref()), }; let token = match extract_bearer(parts) { Ok(t) => t, @@ -100,7 +127,7 @@ where http_method = %ctx.method, path = %ctx.path, request_id = %ctx.request_id, - source_ip = %ctx.source_ip, + source_ip = %ctx.source_ip.resolve(), "rejected inbound request without a credential", ); return Err(e); @@ -137,7 +164,10 @@ pub(crate) async fn authenticate_token( ctx: DenialContext<'_>, ) -> Result { let snapshot = state.snapshot.load(); - if crate::jwt::looks_like_jwt(token) && crate::jwt::any_enabled_provider(&snapshot) { + // Provider gate first: it is O(1) when no trust provider is + // configured, so key-only deployments never pay the structural JWT + // probe (base64 + JSON decode of the header segment). + if crate::jwt::any_enabled_provider(&snapshot) && crate::jwt::looks_like_jwt(token) { return crate::jwt::authenticate_jwt(state, &snapshot, token, ctx).await; // No enabled trust provider: fall through to the key path so a // custom-imported key that happens to look like a JWT keeps @@ -172,7 +202,9 @@ pub(crate) async fn authenticate_token( ProxyError::ApiKeyDisabled, )); } - if entry.value.is_expired_at(chrono::Utc::now()) { + // Read the wall clock only for keys that actually carry a deadline; + // enforcement is unchanged (`is_expired_at` is false for `None`). + if entry.value.expires_at.is_some() && entry.value.is_expired_at(chrono::Utc::now()) { return Err(deny_key( state, "key_expired", @@ -217,7 +249,7 @@ fn deny_key( http_method = %ctx.method, path = %ctx.path, request_id = %ctx.request_id, - source_ip = %ctx.source_ip, + source_ip = %ctx.source_ip.resolve(), "rejected inbound credential", ); } else { @@ -229,7 +261,7 @@ fn deny_key( http_method = %ctx.method, path = %ctx.path, request_id = %ctx.request_id, - source_ip = %ctx.source_ip, + source_ip = %ctx.source_ip.resolve(), "rejected inbound credential", ); } diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index dbd76ffa..e16c4fb8 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -4224,6 +4224,13 @@ pub(crate) fn sanitize_tag(s: String) -> String { if s.is_empty() { return s; } + // Clean-input fast path: within the cap (≤256 bytes ⇒ ≤256 chars) + // and no control characters means the filtered copy would be + // byte-identical — return the input without re-collecting. This + // runs several times per request (tags, user-agent, PK names). + if s.len() <= 256 && !s.chars().any(|c| c.is_control()) { + return s; + } s.chars().filter(|c| !c.is_control()).take(256).collect() } diff --git a/crates/aisix-proxy/src/health.rs b/crates/aisix-proxy/src/health.rs index c763354e..d33d5bbc 100644 --- a/crates/aisix-proxy/src/health.rs +++ b/crates/aisix-proxy/src/health.rs @@ -15,12 +15,12 @@ use dashmap::DashMap; use std::sync::atomic::AtomicBool; -use std::sync::atomic::{AtomicU32, AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicU32, AtomicU64, AtomicUsize, Ordering}; use std::sync::Arc; use std::time::{Duration, SystemTime}; use aisix_core::snapshot::SnapshotHandle; -use aisix_core::AisixSnapshot; +use aisix_core::{AisixSnapshot, RoutingStrategy}; use aisix_obs::{DeploymentLabels, DeploymentState, Metrics}; use axum::http::header::{HeaderName, HeaderValue, CONTENT_TYPE}; use axum::http::StatusCode; @@ -366,10 +366,87 @@ impl Entry { } } +/// Version-gated answer to "does any configured consumer depend on the +/// exact per-request bookkeeping write path?" +/// +/// Three predicates, all derived from the live snapshot: +/// - a Model routes with `least_busy` (reads the in-flight counters, +/// `crate::routing::order_attempts_by_metric`) +/// - a Model routes with `least_latency` (reads the latency EWMA) +/// - a Model has background health checks enabled (the background +/// checker and the health surface observe tracker state) +/// +/// While **any** predicate holds, every bookkeeping method below runs its +/// historical write path unchanged, so configured deployments keep +/// byte-identical behavior. Only when none holds do the trackers take the +/// cheap read-first paths — writes whose consumers provably don't exist. +/// +/// The predicate set is recomputed at most once per snapshot version +/// (packed with the version into one atomic so the pair can never be +/// observed torn). A racing store between the version read and the table +/// walk can cache bits against a stale version; the next call detects the +/// mismatch and recomputes, so the value converges immediately. +#[derive(Debug)] +pub struct BookkeepingFlags { + snapshot: SnapshotHandle, + /// `(snapshot version << 3) | predicate bits`, or [`UNCOMPUTED`]. + packed: AtomicU64, +} + +const FLAG_LEAST_BUSY: u64 = 1; +const FLAG_LEAST_LATENCY: u64 = 1 << 1; +const FLAG_HEALTH_CHECKS: u64 = 1 << 2; +const FLAG_BITS: u64 = 0b111; +const UNCOMPUTED: u64 = u64::MAX; + +impl BookkeepingFlags { + pub fn new(snapshot: SnapshotHandle) -> Arc { + Arc::new(Self { + snapshot, + packed: AtomicU64::new(UNCOMPUTED), + }) + } + + /// True when any predicate holds — the trackers then use their + /// historical write paths. + pub fn any_active(&self) -> bool { + self.bits() != 0 + } + + fn bits(&self) -> u64 { + let ver = self.snapshot.version(); + let packed = self.packed.load(Ordering::Relaxed); + if packed != UNCOMPUTED && packed >> 3 == ver { + return packed & FLAG_BITS; + } + let snap = self.snapshot.load(); + let mut bits = 0; + for entry in snap.models.entries() { + let m = &entry.value; + if let Some(routing) = &m.routing { + match routing.strategy { + RoutingStrategy::LeastBusy => bits |= FLAG_LEAST_BUSY, + RoutingStrategy::LeastLatency => bits |= FLAG_LEAST_LATENCY, + _ => {} + } + } + if m.background_model_check.as_ref().is_some_and(|c| c.enabled) { + bits |= FLAG_HEALTH_CHECKS; + } + } + self.packed.store((ver << 3) | bits, Ordering::Relaxed); + bits + } +} + /// Shared tracker — one per `ProxyState`, cloned cheaply via `Arc`. #[derive(Default, Debug)] pub struct HealthTracker { entries: DashMap, + /// `None` (tests, lightweight constructors) means "assume active": + /// the historical write path always runs. The production bootstrap + /// wires the shared [`BookkeepingFlags`]. + flags: Option>, } /// Smoothing factor for the per-target latency EWMA. Higher = more weight on @@ -394,6 +471,10 @@ pub struct ModelRuntimeStatusTracker { /// only on a cooldown transition. `None` falls back to model-id-only /// labels. snapshot: Option>, + /// `None` (tests, lightweight constructors) means "assume active": + /// every method runs its historical write path. See + /// [`BookkeepingFlags`]. + flags: Option>, } /// RAII guard that decrements a target's in-flight counter when dropped. @@ -402,12 +483,19 @@ pub struct ModelRuntimeStatusTracker { /// stream body so the count stays raised until the stream ends or is /// cancelled, matching the request's true lifetime. pub struct InFlightGuard { - counter: Arc, + /// `None` is the no-op guard handed out while bookkeeping is + /// inactive (no configured consumer); dropping it does nothing. A + /// guard armed before a config change that deactivates bookkeeping + /// still decrements the counter it incremented, so the count can + /// never go negative. + counter: Option>, } impl Drop for InFlightGuard { fn drop(&mut self) { - self.counter.fetch_sub(1, Ordering::Relaxed); + if let Some(counter) = &self.counter { + counter.fetch_sub(1, Ordering::Relaxed); + } } } @@ -416,12 +504,39 @@ impl HealthTracker { Self::default() } + /// Production constructor: consults the shared [`BookkeepingFlags`] + /// so the per-request success write can take the read-first path + /// when no configured consumer exists. + pub fn with_flags(flags: Arc) -> Self { + Self { + entries: DashMap::new(), + flags: Some(flags), + } + } + + fn bookkeeping_active(&self) -> bool { + self.flags.as_ref().is_none_or(|f| f.any_active()) + } + /// Record a successful upstream response for `model`. pub fn record_success(&self, model: &str) { - self.entries - .entry(model.to_string()) - .or_default() - .on_success(); + if self.bookkeeping_active() { + self.entries + .entry(model.to_string()) + .or_default() + .on_success(); + return; + } + // Read-first path: a model with no tracked failures is already + // Healthy — skip the key allocation and the shard write lock the + // `entry()` API pays on every call. The counter is atomic, so + // the reset happens under the read guard; a miss means the model + // never failed and there is nothing to reset. + if let Some(e) = self.entries.get(model) { + if e.consecutive_failures.load(Ordering::Relaxed) != 0 { + e.consecutive_failures.store(0, Ordering::Relaxed); + } + } } /// Record a failed upstream call (any non-4xx bridge error) for `model`. @@ -465,14 +580,20 @@ impl ModelRuntimeStatusTracker { pub fn with_observability( metrics: Arc, snapshot: SnapshotHandle, + flags: Arc, ) -> Self { Self { entries: DashMap::new(), metrics: Some(metrics), snapshot: Some(snapshot), + flags: Some(flags), } } + fn bookkeeping_active(&self) -> bool { + self.flags.as_ref().is_none_or(|f| f.any_active()) + } + pub fn mark_cooldown(&self, model_id: &str, ttl: Duration, reason: impl Into) { let now = SystemTime::now(); let until = now + ttl; @@ -497,6 +618,38 @@ impl ModelRuntimeStatusTracker { } pub fn mark_healthy(&self, model_id: &str) { + if self.bookkeeping_active() { + if let Some(mut entry) = self.entries.get_mut(model_id) { + entry.unhealthy = false; + entry.cooldown_until = None; + entry.status_reason = None; + self.sync_deployment_state(model_id, &mut entry, SystemTime::now()); + } + return; + } + // Read-first path (no configured bookkeeping consumer): in the + // steady state — entry clean, gauge already Healthy — a read + // guard and a few field loads replace the per-request shard + // write lock + wall-clock read. The write path below still runs + // whenever there is anything to do (cooldown early-recovery, + // the first-success Healthy publish on an entry begin_in_flight + // created), so the `aisix_deployment_state` series behaves + // exactly as before. A MISSING entry stays a no-op, exactly as + // on the historical path: the single-attempt handlers call + // mark_healthy without begin_in_flight, and their targets must + // not grow a series they never had. + let needs_write = match self.entries.get(model_id) { + Some(e) => { + e.unhealthy + || e.cooldown_until.is_some() + || e.status_reason.is_some() + || e.emitted_state != Some(DeploymentState::Healthy) + } + None => false, + }; + if !needs_write { + return; + } if let Some(mut entry) = self.entries.get_mut(model_id) { entry.unhealthy = false; entry.cooldown_until = None; @@ -640,6 +793,13 @@ impl ModelRuntimeStatusTracker { /// each successful upstream attempt; drives the `least_latency` routing /// strategy. Independent of health/cooldown state. pub fn record_latency(&self, model_id: &str, latency_ms: u32) { + // The EWMA's only reader is `least_latency` target ordering; with + // no such strategy configured the sample has no consumer. When + // the strategy is (re)configured the EWMA cold-starts, exactly as + // it does on process start. + if !self.bookkeeping_active() { + return; + } let sample = f64::from(latency_ms); self.entries .entry(model_id.to_string()) @@ -663,6 +823,25 @@ impl ModelRuntimeStatusTracker { /// Mark one request as in flight to `model_id` and return a guard that /// decrements the count when dropped. Drives the `least_busy` strategy. pub fn begin_in_flight(&self, model_id: &str) -> InFlightGuard { + // The counter's only reader is `least_busy` target ordering; with + // no such strategy configured, hand out a no-op guard instead of + // paying the counter RMWs and guard refcount per request. When + // the strategy is (re)configured, counting resumes for new + // requests; requests already in flight hold no-op guards, so the + // count transiently underreads until they drain — the same + // cold-start the counter has on process start. + // + // The ENTRY-CREATION side effect is preserved: `mark_healthy`'s + // first-success Healthy publish keys off the entry this method + // creates, and only the endpoints that call begin_in_flight may + // publish that series (the single-attempt handlers never do). + // Steady state downgrades to a read-guard existence check. + if !self.bookkeeping_active() { + if self.entries.get(model_id).is_none() { + self.entries.entry(model_id.to_string()).or_default(); + } + return InFlightGuard { counter: None }; + } let counter = Arc::clone( &self .entries @@ -671,7 +850,9 @@ impl ModelRuntimeStatusTracker { .in_flight, ); counter.fetch_add(1, Ordering::Relaxed); - InFlightGuard { counter } + InFlightGuard { + counter: Some(counter), + } } /// Current in-flight request count for `model_id`. @@ -764,6 +945,175 @@ mod tests { assert_eq!(t.all_levels().len(), 1); } + // ------------------------------------------------------------------- + // On-demand bookkeeping (BookkeepingFlags) + // ------------------------------------------------------------------- + + fn model_json(routing_strategy: Option<&str>, health_check: bool) -> aisix_core::Model { + let mut v = serde_json::json!({ "name": "vg", "display_name": "vg" }); + if let Some(s) = routing_strategy { + v["routing"] = serde_json::json!({ + "strategy": s, + "targets": [{ "model": "d1" }], + }); + } + if health_check { + v["background_model_check"] = serde_json::json!({ + "enabled": true, + "interval_seconds": 5, + "timeout_seconds": 1, + "prompt": "ping", + "max_tokens": 1, + "stale_after_seconds": 60, + }); + } + serde_json::from_value(v).expect("test model json") + } + + fn snapshot_with(model: Option) -> AisixSnapshot { + let snap = AisixSnapshot::new(); + if let Some(m) = model { + snap.models + .insert(aisix_core::ResourceEntry::new("m-1", m, 1)); + } + snap + } + + fn inactive_tracker() -> (SnapshotHandle, ModelRuntimeStatusTracker) { + let handle = SnapshotHandle::new(snapshot_with(None)); + let flags = BookkeepingFlags::new(handle.clone()); + let t = ModelRuntimeStatusTracker { + entries: DashMap::new(), + metrics: None, + snapshot: None, + flags: Some(flags), + }; + (handle, t) + } + + #[test] + fn bookkeeping_flags_derive_from_snapshot() { + for (model, expect) in [ + (None, false), + (Some(model_json(Some("round_robin"), false)), false), + (Some(model_json(Some("weighted"), false)), false), + (Some(model_json(Some("failover"), false)), false), + // least_cost ranks by static configured cost, not runtime + // bookkeeping — it must NOT activate the write paths. + (Some(model_json(Some("least_cost"), false)), false), + (Some(model_json(Some("least_busy"), false)), true), + (Some(model_json(Some("least_latency"), false)), true), + (Some(model_json(None, true)), true), + ] { + let described = format!("{model:?}"); + let flags = BookkeepingFlags::new(SnapshotHandle::new(snapshot_with(model))); + assert_eq!(flags.any_active(), expect, "for {described}"); + } + } + + #[test] + fn inactive_bookkeeping_skips_inflight_and_latency() { + let (_handle, t) = inactive_tracker(); + let g = t.begin_in_flight("d1"); + assert_eq!(t.in_flight("d1"), 0, "no-op guard must not count"); + drop(g); + assert_eq!(t.in_flight("d1"), 0, "no-op guard must not underflow"); + t.record_latency("d1", 100); + assert_eq!(t.latency_ewma_ms("d1"), None); + } + + #[test] + fn bookkeeping_reactivates_on_snapshot_swap() { + let (handle, t) = inactive_tracker(); + let g = t.begin_in_flight("d1"); + assert_eq!(t.in_flight("d1"), 0); + drop(g); + + // Config change introduces a least_busy router → counting resumes + // (version-gated recompute, no restart needed). + handle.store(snapshot_with(Some(model_json(Some("least_busy"), false)))); + let g = t.begin_in_flight("d1"); + assert_eq!(t.in_flight("d1"), 1); + // Any active predicate takes the WHOLE family back to the old + // path — least_busy alone re-enables the EWMA write too. + t.record_latency("d1", 100); + assert_eq!(t.latency_ewma_ms("d1"), Some(100.0)); + drop(g); + assert_eq!(t.in_flight("d1"), 0); + } + + #[test] + fn inactive_mark_healthy_publishes_healthy_once_then_reads() { + let (_handle, t) = inactive_tracker(); + // Real multi-attempt-endpoint sequence: begin_in_flight creates + // the entry (no-op guard, but the side effect is preserved), + // then the first success publishes Healthy. + drop(t.begin_in_flight("d1")); + t.mark_healthy("d1"); + { + let e = t + .entries + .get("d1") + .expect("entry created by begin_in_flight"); + assert_eq!(e.emitted_state, Some(DeploymentState::Healthy)); + } + // Steady state: read-only, entry untouched. + t.mark_healthy("d1"); + assert_eq!(t.status("d1").status, RuntimeStatus::Healthy); + } + + /// The single-attempt handlers (embeddings, images, completions, + /// audio, rerank, count_tokens) call mark_healthy WITHOUT + /// begin_in_flight. On the historical path their targets never got + /// an entry — and therefore never published `aisix_deployment_state` + /// — so the inactive fast path must stay a no-op for a missing + /// entry, or zero-config deployments grow a series main never had. + #[test] + fn inactive_mark_healthy_without_begin_in_flight_stays_noop() { + let (_handle, t) = inactive_tracker(); + t.mark_healthy("embeddings-only-target"); + assert!( + t.entries.get("embeddings-only-target").is_none(), + "mark_healthy on a never-seen id must not create an entry" + ); + } + + #[test] + fn inactive_mark_healthy_still_recovers_cooldown_early() { + let (_handle, t) = inactive_tracker(); + t.mark_cooldown("d1", Duration::from_secs(3600), "upstream failure"); + assert_eq!(t.status("d1").status, RuntimeStatus::Cooldown); + // A success during cooldown still recovers immediately — the + // read-first path detects the dirty entry and takes the full + // write path. + t.mark_healthy("d1"); + assert_eq!(t.status("d1").status, RuntimeStatus::Healthy); + assert_eq!( + t.entries.get("d1").unwrap().emitted_state, + Some(DeploymentState::Healthy) + ); + } + + #[test] + fn inactive_record_success_creates_no_entry_but_still_resets() { + let flags = BookkeepingFlags::new(SnapshotHandle::new(snapshot_with(None))); + let t = HealthTracker::with_flags(flags); + // Happy path: no failures → no entry, no allocation. + t.record_success("m"); + assert!( + t.all_levels().is_empty(), + "no entry for a never-failed model" + ); + assert_eq!(t.level("m"), HealthLevel::Healthy); + // Reset still works through the read guard. + for _ in 0..10 { + t.record_failure("m"); + } + assert_eq!(t.level("m"), HealthLevel::Down); + t.record_success("m"); + assert_eq!(t.level("m"), HealthLevel::Healthy); + } + #[tokio::test] async fn livez_default_success_is_plain_ok() { let state = LivezState::new(); @@ -901,9 +1251,11 @@ mod tests { .insert(ResourceEntry::new("m-cool", model, 1)); let metrics = Arc::new(Metrics::new(false)); + let handle = SnapshotHandle::new(snapshot); let tracker = ModelRuntimeStatusTracker::with_observability( metrics.clone(), - SnapshotHandle::new(snapshot), + handle.clone(), + BookkeepingFlags::new(handle), ); // First mark = a fresh transition (counter++, gauge → Down). The @@ -949,9 +1301,11 @@ mod tests { #[test] fn gauge_returns_to_healthy_after_a_cooldown_expires_naturally() { let metrics = Arc::new(Metrics::new(false)); + let handle = SnapshotHandle::new(AisixSnapshot::new()); let tracker = ModelRuntimeStatusTracker::with_observability( metrics.clone(), - SnapshotHandle::new(AisixSnapshot::new()), + handle.clone(), + BookkeepingFlags::new(handle), ); tracker.mark_cooldown( @@ -982,9 +1336,11 @@ mod tests { #[test] fn gauge_tracks_background_check_failures_and_recovery() { let metrics = Arc::new(Metrics::new(false)); + let handle = SnapshotHandle::new(AisixSnapshot::new()); let tracker = ModelRuntimeStatusTracker::with_observability( metrics.clone(), - SnapshotHandle::new(AisixSnapshot::new()), + handle.clone(), + BookkeepingFlags::new(handle), ); tracker.mark_unhealthy("m-bg", Some(503), "background_check_failed"); @@ -1000,12 +1356,16 @@ mod tests { #[test] fn repeated_success_neither_churns_the_gauge_nor_the_cooldown_counter() { let metrics = Arc::new(Metrics::new(false)); + let handle = SnapshotHandle::new(AisixSnapshot::new()); let tracker = ModelRuntimeStatusTracker::with_observability( metrics.clone(), - SnapshotHandle::new(AisixSnapshot::new()), + handle.clone(), + BookkeepingFlags::new(handle), ); - // begin_in_flight is what creates the entry on the request path. + // With bookkeeping active begin_in_flight creates the entry; with + // it inactive (this tracker: empty snapshot) the first mark_healthy + // does. Either way the assertions below must hold. drop(tracker.begin_in_flight("m-ok")); tracker.mark_healthy("m-ok"); tracker.mark_healthy("m-ok"); diff --git a/crates/aisix-proxy/src/jwt.rs b/crates/aisix-proxy/src/jwt.rs index ed09e19e..1976390c 100644 --- a/crates/aisix-proxy/src/jwt.rs +++ b/crates/aisix-proxy/src/jwt.rs @@ -96,17 +96,25 @@ pub(crate) fn looks_like_jwt(token: &str) -> bool { if token.len() > MAX_JWT_BYTES { return false; } - let parts: Vec<&str> = token.splitn(4, '.').collect(); - if parts.len() != 3 || parts.iter().any(|p| p.is_empty()) { + // Exactly three non-empty segments — checked on the iterator so the + // per-request path allocates nothing. + let mut parts = token.splitn(4, '.'); + let (Some(header), Some(payload), Some(sig), None) = + (parts.next(), parts.next(), parts.next(), parts.next()) + else { + return false; + }; + if header.is_empty() || payload.is_empty() || sig.is_empty() { return false; } - matches!(b64url_json(parts[0]), Some(v) if v.get("alg").is_some()) + matches!(b64url_json(header), Some(v) if v.get("alg").is_some()) } /// True when the snapshot has at least one enabled trust provider — the -/// gate for entering the JWT path at all. +/// gate for entering the JWT path at all. O(1) on deployments with no +/// providers configured (the common case). pub(crate) fn any_enabled_provider(snapshot: &AisixSnapshot) -> bool { - snapshot.oidc_providers.any(|e| e.value.enabled) + !snapshot.oidc_providers.is_empty() && snapshot.oidc_providers.any(|e| e.value.enabled) } fn b64url_json(segment: &str) -> Option { @@ -525,7 +533,7 @@ fn deny( http_method = %ctx.method, path = %ctx.path, request_id = %ctx.request_id, - source_ip = %ctx.source_ip, + source_ip = %ctx.source_ip.resolve(), "rejected inbound JWT (pre-verification)", ); } else { @@ -538,7 +546,7 @@ fn deny( http_method = %ctx.method, path = %ctx.path, request_id = %ctx.request_id, - source_ip = %ctx.source_ip, + source_ip = %ctx.source_ip.resolve(), "rejected inbound JWT", ); } diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index 374c42e0..4403212c 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -197,19 +197,15 @@ pub fn build_router(state: ProxyState) -> Router { state.clone(), enforce_request_body_limit, )) + // One layer for both per-request telemetry guards: the in-flight + // gauge and the client-cancel recorder (see + // `record_request_telemetry`). Sits outside the body-limit layers + // so a hang-up during body upload is captured too, and inside + // `ensure_request_id` so the emitted line carries the same + // request id the caller was handed. .layer(middleware::from_fn_with_state( state.clone(), - record_in_flight_request, - )) - // Record requests the caller abandoned before any response head - // existed. Sits outside `record_in_flight_request` so a hang-up - // during body upload (which the body-limit layers above are - // awaiting) is captured too, and inside `ensure_request_id` so - // the emitted line carries the same request id the caller was - // handed. See `record_client_cancel`. - .layer(middleware::from_fn_with_state( - state.clone(), - record_client_cancel, + record_request_telemetry, )) // Identify the data plane on every response, including error // envelopes and short-circuited responses from the layers @@ -256,27 +252,6 @@ pub fn build_router(state: ProxyState) -> Router { )) } -async fn record_in_flight_request( - State(state): State, - request: Request, - next: Next, -) -> Response { - // Normalize to a bounded route template BEFORE using the path as a - // metric label. This middleware runs before authentication and before - // route matching, so the raw `request.uri().path()` is fully - // attacker-controlled — the `/passthrough/:provider/*rest` wildcard - // suffix (or any 404 path) would otherwise let an unauthenticated - // caller mint unbounded Prometheus time series (#451). - let endpoint = normalize_endpoint_label(request.uri().path()); - let inbound_protocol = inbound_protocol_for_endpoint(endpoint).to_string(); - let _guard = InFlightGuard::new( - state.metrics.clone(), - endpoint.to_string(), - inbound_protocol, - ); - next.run(request).await -} - /// Collapse a raw request path to a fixed route template so metric labels /// stay bounded regardless of caller-supplied path segments. Keep this /// allowlist in sync with the routes registered in `build_router`; any @@ -336,17 +311,21 @@ fn inbound_protocol_for_endpoint(endpoint: &str) -> &'static str { struct InFlightGuard { metrics: std::sync::Arc, - endpoint: String, - inbound_protocol: String, + /// Bounded route template + protocol family — both `'static` by + /// construction (`normalize_endpoint_label` / + /// `inbound_protocol_for_endpoint`), so the guard owns no + /// allocations. + endpoint: &'static str, + inbound_protocol: &'static str, } impl InFlightGuard { fn new( metrics: std::sync::Arc, - endpoint: String, - inbound_protocol: String, + endpoint: &'static str, + inbound_protocol: &'static str, ) -> Self { - metrics.increment_proxy_in_flight(&endpoint, &inbound_protocol); + metrics.increment_proxy_in_flight(endpoint, inbound_protocol); Self { metrics, endpoint, @@ -358,7 +337,7 @@ impl InFlightGuard { impl Drop for InFlightGuard { fn drop(&mut self) { self.metrics - .decrement_proxy_in_flight(&self.endpoint, &self.inbound_protocol); + .decrement_proxy_in_flight(self.endpoint, self.inbound_protocol); } } @@ -372,22 +351,32 @@ pub(crate) const CLIENT_CLOSED_REQUEST: u16 = 499; /// `ClientDisconnected` error class. const CLIENT_DISCONNECTED_KIND: &str = "client_disconnected"; -/// Record a request whose caller hung up before the response head was -/// written. +/// One middleware for both per-request telemetry guards: the in-flight +/// gauge and the client-cancel recorder. The two used to be separate +/// layers; they sit at the same position in the stack with nothing +/// between them, so a single layer arms both and the per-request boxed +/// service hop (and one of two route-normalize calls) disappears. +/// +/// In-flight gauge: incremented before the inner service runs, +/// decremented on guard drop — including cancellation. /// -/// Every endpoint logs and meters itself at the end of its own handler — -/// 29 `emit_access_log` call sites across 12 modules. When the client -/// disconnects first, axum drops the handler future and *none* of that -/// code runs: the request leaves no access-log line, no usage event and -/// no metric. It is invisible exactly where an operator most needs it, -/// because the usual reason a caller gives up is a long -/// time-to-first-token. +/// Client-cancel: records a request whose caller hung up before the +/// response head was written. Every endpoint logs and meters itself at +/// the end of its own handler — 29 `emit_access_log` call sites across +/// 12 modules. When the client disconnects first, axum drops the +/// handler future and *none* of that code runs: the request leaves no +/// access-log line, no usage event and no metric. It is invisible +/// exactly where an operator most needs it, because the usual reason a +/// caller gives up is a long time-to-first-token. /// /// A cancelled future is only observable from `Drop`, so arm a guard, /// disarm it once the inner service yields a response, and emit from /// `Drop` when it is still armed. Doing it in one layer rather than in /// each handler also keeps the endpoint family from drifting the way the /// request-id header did before `ensure_request_id` (see request_id.rs). +/// On cancellation the in-flight guard (declared later) drops first, +/// then the cancel guard emits — the same order the nested layers +/// produced. /// /// This is NOT the streaming-disconnect path: once SSE bytes flow the /// response head is already committed, so the handler has logged and the @@ -395,15 +384,22 @@ const CLIENT_DISCONNECTED_KIND: &str = "client_disconnected"; /// `chat::build_sse_stream`). Response bodies are polled after this /// middleware has returned, so a mid-stream hang-up leaves the guard /// disarmed and is not double-counted here. -async fn record_client_cancel( +async fn record_request_telemetry( State(state): State, request: Request, next: Next, ) -> Response { + // Normalize to a bounded route template BEFORE using the path as a + // metric label. This middleware runs before authentication and before + // route matching, so the raw `request.uri().path()` is fully + // attacker-controlled — the `/passthrough/:provider/*rest` wildcard + // suffix (or any 404 path) would otherwise let an unauthenticated + // caller mint unbounded Prometheus time series (#451). + let endpoint = normalize_endpoint_label(request.uri().path()); let mut guard = ClientCancelGuard { armed: true, metrics: state.metrics.clone(), - endpoint: normalize_endpoint_label(request.uri().path()), + endpoint, method: request.method().clone(), uri: request.uri().clone(), request_id: request @@ -413,6 +409,11 @@ async fn record_client_cancel( .unwrap_or_default(), started: std::time::Instant::now(), }; + let _in_flight = InFlightGuard::new( + state.metrics.clone(), + endpoint, + inbound_protocol_for_endpoint(endpoint), + ); let response = next.run(request).await; guard.armed = false; response diff --git a/crates/aisix-proxy/src/quota.rs b/crates/aisix-proxy/src/quota.rs index 1bf1544a..8253a7ff 100644 --- a/crates/aisix-proxy/src/quota.rs +++ b/crates/aisix-proxy/src/quota.rs @@ -288,7 +288,9 @@ async fn reserve_layers( model_rl: Option<&ModelRateLimit>, mcp_server: Option<&str>, ) -> Result { - let mut reservations = Vec::with_capacity(8); + // Starts empty so the common no-limits request never allocates; + // the first reservation (if any) grows it on demand. + let mut reservations = Vec::new(); // Layer 1: API key inline rate limit. let key_limits = auth.key().rate_limit.clone().unwrap_or_default(); @@ -354,6 +356,13 @@ async fn reserve_policy_layers( reservations: &mut Vec, ) -> Result<(), ProxyError> { let snap = state.snapshot.load(); + // O(1) empty check before anything else: deployments with no + // rate-limit policies (the default) skip the wall-clock read and + // the per-shard table scan below entirely. Covers both callers — + // the request gate and the per-target gate. + if snap.rate_limit_policies.is_empty() { + return Ok(()); + } let now = chrono::Utc::now(); for entry in snap.rate_limit_policies.entries() { let policy = &entry.value; diff --git a/crates/aisix-proxy/src/realtime.rs b/crates/aisix-proxy/src/realtime.rs index 16855614..a9030b67 100644 --- a/crates/aisix-proxy/src/realtime.rs +++ b/crates/aisix-proxy/src/realtime.rs @@ -328,7 +328,7 @@ async fn authenticate( method: "GET", path: "/v1/realtime", request_id: &client.request_id, - source_ip: &client.source_ip, + source_ip: crate::auth::LazySourceIp::Ready(&client.source_ip), }; if let Some(auth) = headers.get(axum::http::header::AUTHORIZATION) { let s = auth.to_str().map_err(|_| ProxyError::MissingAuth)?; diff --git a/crates/aisix-proxy/src/state.rs b/crates/aisix-proxy/src/state.rs index 6431f811..fe209afa 100644 --- a/crates/aisix-proxy/src/state.rs +++ b/crates/aisix-proxy/src/state.rs @@ -371,10 +371,14 @@ impl ProxyState { // The bootstrap constructor is the one place the tracker gets a // metrics sink + snapshot handle, so cooldown transitions emit // `aisix_deployment_*`. Clone both before they are moved into the - // struct below. + // struct below. Both trackers consult one shared BookkeepingFlags + // so the "does any configured consumer read this?" answer can't + // drift between them. + let bookkeeping_flags = crate::health::BookkeepingFlags::new(snapshot.clone()); let runtime_status = Arc::new(ModelRuntimeStatusTracker::with_observability( metrics.clone(), snapshot.clone(), + Arc::clone(&bookkeeping_flags), )); Self::from_inner(ProxyStateInner { snapshot, @@ -386,7 +390,7 @@ impl ProxyState { semantic_cache: Arc::new(crate::semantic::SemanticVectorCache::default()), guardrail_index, budgets: Arc::new(BudgetClient::disabled()), - health: Arc::new(HealthTracker::new()), + health: Arc::new(HealthTracker::with_flags(bookkeeping_flags)), livez: Arc::new(LivezState::new()), config_apply_age: None, runtime_status,