From 8492a8527fc0f683eb9baeac0d7f182c107f0dfc Mon Sep 17 00:00:00 2001 From: Jarvis Date: Wed, 19 Aug 2026 10:30:54 +0000 Subject: [PATCH] feat(routing): consistent_hash strategy + priority tiers; fold weighted into round_robin Load-balancing taxonomy rework (AISIX-Cloud#1206), aligned with the APISIX balancer model: weight and priority are per-target attributes available under every strategy, and consistent hashing is a strategy of its own rather than a flag on another one. - targets[].priority (i32, default 0, higher preferred, negatives for backups): targets partition into tiers; the strategy orders each tier independently (per-tier balancer state); tiers concatenate, so a lower tier only receives traffic when every higher-tier target failed or was health-filtered. max_fallbacks caps the concatenated walk. - strategy consistent_hash: ketama-style ring (160 points x weight units, FNV-1a + splitmix64 finalizer), keyed by the hash_on source chain (header / cookie / api_key / client_ip; default = x-aisix-routing-key header, then the caller's API key id). A key's ring order doubles as its failover order, so a failed target's keys spread to their ring successors and every other key keeps its mapping. - round_robin is now smooth WEIGHTED round-robin (the nginx algorithm); equal weights keep the old declaration-order cycle exactly. - least_busy now scores (1 + in-flight) / weight. - REMOVED: the weighted strategy (fold into round_robin) and the sticky flag (superseded by consistent_hash). The enum value now fails row-level in the lenient loader; declarative configs must migrate (weighted -> round_robin keeping weights; weighted+sticky -> consistent_hash). The control plane migrates stored rows in the companion PR. Fixes api7/AISIX-Cloud#1206 --- README.md | 9 +- crates/aisix-admin/src/openapi.rs | 6 +- crates/aisix-core/src/lib.rs | 13 +- crates/aisix-core/src/models/mod.rs | 5 +- crates/aisix-core/src/models/routing.rs | 224 ++++- .../tests/model_schema_characterization.rs | 8 +- crates/aisix-proxy/src/chat.rs | 8 +- crates/aisix-proxy/src/client_ip.rs | 17 - crates/aisix-proxy/src/count_tokens.rs | 8 +- crates/aisix-proxy/src/health.rs | 2 +- crates/aisix-proxy/src/messages.rs | 8 +- crates/aisix-proxy/src/responses.rs | 8 +- crates/aisix-proxy/src/routing.rs | 888 ++++++++++++------ crates/aisix-proxy/src/semantic.rs | 2 + schemas/resources/model.schema.json | 83 +- schemas/resources/routing.schema.json | 95 +- .../e2e/src/cases/canary-routing-e2e.test.ts | 146 --- .../cases/consistent-hash-routing-e2e.test.ts | 418 +++++++++ ...t.ts => routing-priority-edit-e2e.test.ts} | 60 +- .../src/cases/routing-strategies-e2e.test.ts | 78 +- .../weighted-routing-distribution-e2e.test.ts | 77 +- 21 files changed, 1510 insertions(+), 653 deletions(-) delete mode 100644 tests/e2e/src/cases/canary-routing-e2e.test.ts create mode 100644 tests/e2e/src/cases/consistent-hash-routing-e2e.test.ts rename tests/e2e/src/cases/{weighted-routing-edit-e2e.test.ts => routing-priority-edit-e2e.test.ts} (73%) diff --git a/README.md b/README.md index 758aaf5e..4c3a2b87 100644 --- a/README.md +++ b/README.md @@ -154,10 +154,11 @@ Covered by 183 end-to-end scenario files (496 cases) that run against real gatew - **Anthropic Messages API** — `POST /v1/messages` as a first-class route, working against **any** configured upstream: requests and responses (including streaming) are translated both ways when a model points at a non-Anthropic provider. -- **Routing & failover** — virtual/routing models with six strategies: `round_robin`, - `weighted` (with sticky/canary hashing), `failover`, plus metric-based `least_cost`, - `least_latency`, and `least_busy`. Retry budgets, cooldowns, tag-conditional targets, - and per-attempt timeouts. +- **Routing & failover** — virtual/routing models with six strategies: `round_robin` + (smooth weighted round-robin), `consistent_hash` (session affinity keyed by header / + cookie / API key / client IP), `failover`, plus metric-based `least_cost`, + `least_latency`, and `least_busy`. Per-target `priority` tiers (active/backup pools), + retry budgets, cooldowns, tag-conditional targets, and per-attempt timeouts. - **Ensemble models** — fan one request out to a panel of models concurrently, then have a judge model synthesize a single answer, with a minimum-successful-responses threshold. - **Semantic routing** — one virtual model that dispatches by the *meaning* of each diff --git a/crates/aisix-admin/src/openapi.rs b/crates/aisix-admin/src/openapi.rs index c383fa03..a280f4b4 100644 --- a/crates/aisix-admin/src/openapi.rs +++ b/crates/aisix-admin/src/openapi.rs @@ -2234,13 +2234,17 @@ fn add_variant_titles(doc: &mut Value) { "/components/schemas/RoutingStrategy/oneOf", &[ "Round robin", - "Weighted", + "Consistent hash", "Failover", "Least cost", "Least latency", "Least busy", ], ), + ( + "/components/schemas/HashOnType/oneOf", + &["Header", "Cookie", "API key", "Client IP"], + ), ( "/components/schemas/SlsContentMode/oneOf", &["Metadata only", "Full content"], diff --git a/crates/aisix-core/src/lib.rs b/crates/aisix-core/src/lib.rs index 5ca08d34..868f7220 100644 --- a/crates/aisix-core/src/lib.rs +++ b/crates/aisix-core/src/lib.rs @@ -47,12 +47,13 @@ pub use models::{ validate_rate_limit_policy, A2aAgent, A2aAuthType, A2aProtocolVersion, Adapter, AisixSnapshot, ApiKey, AppliedGuardrail, CachePolicy, CooldownConfig, ExporterKind, Guardrail, GuardrailExecution, GuardrailHookPoint, GuardrailKind, GuardrailMetricsSink, - GuardrailMonitorHit, KeywordConfig, KeywordPattern, McpAuthType, McpProtocolVersion, - McpRateLimit, McpServer, McpServerType, McpTransport, Model, ObservabilityExporter, - ParamConstraints, PassthroughAuthMode, PassthroughCredentialMode, PassthroughRoute, - PolicyScope, PolicyWindow, ProviderKey, RateLimit, RateLimitPolicy, RequestOverrides, - ResponseOverrides, Routing, RoutingStrategy, RoutingTarget, SchemaError, StreamDoneMarker, - TelemetryKind, TelemetryTags, WhenAllUnavailablePolicy, DEFAULT_COOLDOWN_TRIGGER_STATUSES, + GuardrailMonitorHit, HashOnSource, HashOnType, KeywordConfig, KeywordPattern, McpAuthType, + McpProtocolVersion, McpRateLimit, McpServer, McpServerType, McpTransport, Model, + ObservabilityExporter, ParamConstraints, PassthroughAuthMode, PassthroughCredentialMode, + PassthroughRoute, PolicyScope, PolicyWindow, ProviderKey, RateLimit, RateLimitPolicy, + RequestOverrides, ResponseOverrides, Routing, RoutingStrategy, RoutingTarget, SchemaError, + StreamDoneMarker, TelemetryKind, TelemetryTags, WhenAllUnavailablePolicy, + DEFAULT_COOLDOWN_TRIGGER_STATUSES, }; pub use resource::{Resource, ResourceEntry}; pub use snapshot::{ResourceTable, SnapshotHandle}; diff --git a/crates/aisix-core/src/models/mod.rs b/crates/aisix-core/src/models/mod.rs index 86686c98..5c13d627 100644 --- a/crates/aisix-core/src/models/mod.rs +++ b/crates/aisix-core/src/models/mod.rs @@ -76,7 +76,10 @@ pub use provider_key::{ }; pub use rate_limit::{McpRateLimit, RateLimit}; pub use rate_limit_policy::{PolicyScope, PolicyWindow, RateLimitPolicy}; -pub use routing::{Routing, RoutingStrategy, RoutingTarget, WhenAllUnavailablePolicy}; +pub use routing::{ + default_hash_on, HashOnSource, HashOnType, Routing, RoutingStrategy, RoutingTarget, + WhenAllUnavailablePolicy, +}; pub use schema::{ validate_a2a_agent, validate_a2a_agent_lenient, validate_apikey, validate_apikey_lenient, validate_cache_policy, validate_cache_policy_lenient, validate_claim_mapping, diff --git a/crates/aisix-core/src/models/routing.rs b/crates/aisix-core/src/models/routing.rs index 2c9f579b..25e3251f 100644 --- a/crates/aisix-core/src/models/routing.rs +++ b/crates/aisix-core/src/models/routing.rs @@ -6,24 +6,35 @@ //! Failures may retry the current target and then fall back to later //! targets. //! -//! Positional strategies (spec §3) pick a *starting* target, then walk +//! Targets partition into **priority tiers** first (`priority`, higher +//! value preferred — the APISIX node-priority convention): the strategy +//! orders targets *within* each tier, tiers concatenate best-first, and a +//! lower tier is only reached when every higher-tier target failed or is +//! unavailable. All targets default to priority `0`, so priority is inert +//! unless configured. +//! +//! Positional strategies pick a *starting* target per tier, then walk //! forward on failure: -//! - `round_robin`: cycle through targets in declaration order. -//! - `weighted`: pick a target with probability proportional to its -//! `weight`; falls back to round-robin when weights are missing. +//! - `round_robin`: smooth weighted round-robin over target `weight`s +//! (equal weights degrade to a plain declaration-order cycle). +//! - `consistent_hash`: ketama-style consistent hashing of the request's +//! hash key (see [`HashOnSource`]) over the tier's targets, `weight` +//! scaling each target's share of the ring. The same key keeps landing +//! on the same target; on failure the walk follows the ring, so only +//! the failed target's keys move. //! - `failover`: always start at the first target; only move down the -//! list on failure. +//! list on failure. Declaration order is the priority order. //! -//! Metric-ordered strategies rank *all* targets by a runtime signal and -//! attempt them best-first, falling forward down the ranked order: +//! Metric-ordered strategies rank targets by a runtime signal within each +//! tier and attempt them best-first, falling forward down the ranked order: //! - `least_cost`: cheapest target first, by the target model's `cost` //! (combined input+output per-1K price). Targets without a `cost` rank //! last. //! - `least_latency`: fastest target first, by a moving average of recent //! observed upstream latency (time-to-first-token for streaming). Targets //! with no latency samples yet rank first so they get probed. -//! - `least_busy`: least-loaded target first, by the number of in-flight -//! requests currently dispatched to each target. +//! - `least_busy`: least-loaded target first, by in-flight requests +//! divided by target `weight` (the APISIX least_conn score). //! //! See [`RoutingStrategy::is_metric_based`]. @@ -34,10 +45,15 @@ use serde::{Deserialize, Serialize}; )] #[serde(rename_all = "snake_case")] pub enum RoutingStrategy { - /// Cycle through targets in declaration order. + /// Smooth weighted round-robin over target `weight`s. Equal (or absent) + /// weights degrade to a plain declaration-order cycle. RoundRobin, - /// Pick targets by configured weight. Missing target weights fall back to 1. - Weighted, + /// Ketama-style consistent hashing of the request's hash key (see + /// `hash_on`) over the targets, `weight` scaling each target's share of + /// the ring. The same key keeps landing on the same target while it is + /// healthy; on failure the walk follows the ring so only the failed + /// target's keys move. + ConsistentHash, /// Always start with the first target and move to later targets only /// after failure. #[default] @@ -50,11 +66,80 @@ pub enum RoutingStrategy { /// upstream latency (time-to-first-token for streaming), then fall /// forward. Targets with no samples yet rank first so they get probed. LeastLatency, - /// Rank targets least-loaded-first by the number of in-flight requests - /// currently dispatched to each target, then fall forward. + /// Rank targets least-loaded-first by in-flight requests divided by + /// target `weight` (the APISIX least_conn score), then fall forward. LeastBusy, } +/// Which request attribute a [`HashOnSource`] reads the hash key from. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)] +#[serde(rename_all = "snake_case")] +pub enum HashOnType { + /// A request header, named by `name`. + Header, + /// A cookie from the request's `Cookie` header, named by `name`. + Cookie, + /// The caller's API key id. + ApiKey, + /// The caller's resolved client IP (honouring the trusted-proxy + /// configuration). + ClientIp, +} + +/// One source for the `consistent_hash` hash key. Sources are tried in +/// order; the first one that yields a non-empty value wins. +#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema, PartialEq, Eq)] +pub struct HashOnSource { + /// Which request attribute supplies the hash key. + #[serde(rename = "type")] + pub source_type: HashOnType, + /// The header or cookie name to read. Required for `header` and + /// `cookie` sources; not accepted for `api_key` or `client_ip`. + #[serde(default, skip_serializing_if = "Option::is_none")] + #[schemars(length(min = 1))] + pub name: Option, +} + +impl HashOnSource { + pub fn header(name: impl Into) -> Self { + Self { + source_type: HashOnType::Header, + name: Some(name.into()), + } + } + + pub fn cookie(name: impl Into) -> Self { + Self { + source_type: HashOnType::Cookie, + name: Some(name.into()), + } + } + + pub fn api_key() -> Self { + Self { + source_type: HashOnType::ApiKey, + name: None, + } + } + + pub fn client_ip() -> Self { + Self { + source_type: HashOnType::ClientIp, + name: None, + } + } +} + +/// Default hash-key chain when `hash_on` is not configured: the +/// `x-aisix-routing-key` request header, falling back to the caller's API +/// key id. +pub fn default_hash_on() -> Vec { + vec![ + HashOnSource::header("x-aisix-routing-key"), + HashOnSource::api_key(), + ] +} + impl RoutingStrategy { /// Whether the strategy ranks the full target set by a runtime metric /// (rather than picking a start index and walking positionally). These @@ -74,9 +159,18 @@ pub struct RoutingTarget { /// Model alias for a direct model that can receive routed traffic. #[schemars(length(min = 1))] pub model: String, - /// Target weight for `weighted` routing. Other strategies ignore this field. + /// Target weight, default `1`. Used by `round_robin` (rotation share), + /// `consistent_hash` (share of the hash ring), and `least_busy` + /// (in-flight divided by weight). `failover`, `least_cost`, and + /// `least_latency` accept the field but do not use it. #[serde(default, skip_serializing_if = "Option::is_none")] pub weight: Option, + /// Priority tier, default `0`; a higher value is preferred (the APISIX + /// node-priority convention — give backup targets `-1`). The strategy + /// orders targets within each tier; a lower tier is only tried when + /// every higher-tier target failed or is unavailable. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub priority: Option, /// Tags for tag/metadata-conditional routing. When a request carries /// routing tags, only targets whose tags intersect the request's are /// eligible; a target tagged `"default"` is the fallback used when nothing @@ -97,6 +191,7 @@ impl RoutingTarget { Self { model: model.into(), weight: None, + priority: None, tags: None, } } @@ -106,6 +201,11 @@ impl RoutingTarget { self } + pub fn with_priority(mut self, priority: i32) -> Self { + self.priority = Some(priority); + self + } + pub fn with_tags(mut self, tags: Vec) -> Self { self.tags = Some(tags); self @@ -115,6 +215,10 @@ impl RoutingTarget { self.weight.unwrap_or(1) } + pub fn priority_or_default(&self) -> i32 { + self.priority.unwrap_or(0) + } + /// True if this target carries at least one tag. pub fn has_tags(&self) -> bool { self.tags.as_ref().is_some_and(|t| !t.is_empty()) @@ -175,15 +279,13 @@ pub struct Routing { /// Policy to apply when every target is unavailable because of runtime health or cooldown state. #[serde(default, skip_serializing_if = "Option::is_none")] pub when_all_unavailable: Option, - /// Sticky (deterministic) target selection for `weighted` routing — the - /// A/B / canary knob. When `true`, a request's target is chosen by hashing a - /// stability key (the `x-aisix-routing-key` header, else the caller's API - /// key) into the weight distribution, so the same key consistently lands on - /// the same target while the aggregate split still honors the weights. When - /// absent/`false`, `weighted` samples independently per request (the - /// default). Ignored by non-`weighted` strategies. + /// Where the `consistent_hash` hash key comes from: an ordered chain of + /// sources, the first non-empty value winning. Defaults to the + /// `x-aisix-routing-key` request header, falling back to the caller's + /// API key id. Only valid with `strategy: consistent_hash`. #[serde(default, skip_serializing_if = "Option::is_none")] - pub sticky: Option, + #[schemars(length(min = 1))] + pub hash_on: Option>, } impl Routing { @@ -192,8 +294,12 @@ impl Routing { // Resolving that needs the target Model and the DP config, so it lives // in `aisix_proxy::routing::effective_retries`. - pub fn sticky_or_default(&self) -> bool { - self.sticky.unwrap_or(false) + /// The effective hash-key source chain for `consistent_hash`. + pub fn hash_on_or_default(&self) -> Vec { + match &self.hash_on { + Some(chain) if !chain.is_empty() => chain.clone(), + _ => default_hash_on(), + } } pub fn max_fallbacks_or_default(&self) -> usize { @@ -230,20 +336,22 @@ mod tests { #[test] fn deserialises_full_routing_block() { let json = r#"{ - "strategy": "weighted", + "strategy": "round_robin", "targets": [ {"model": "primary", "weight": 90}, - {"model": "backup", "weight": 10} + {"model": "backup", "weight": 10, "priority": -1} ], "retries": 2, "max_fallbacks": 1, "retry_on_429": true }"#; let r: Routing = serde_json::from_str(json).unwrap(); - assert_eq!(r.strategy, RoutingStrategy::Weighted); + assert_eq!(r.strategy, RoutingStrategy::RoundRobin); assert_eq!(r.targets.len(), 2); assert_eq!(r.targets[0].model, "primary"); assert_eq!(r.targets[0].weight_or_default(), 90); + assert_eq!(r.targets[0].priority_or_default(), 0); + assert_eq!(r.targets[1].priority_or_default(), -1); assert_eq!(r.retries, Some(2)); assert_eq!(r.max_fallbacks_or_default(), 1); assert!(r.retry_on_429_or_default()); @@ -270,7 +378,7 @@ mod tests { retry_on_429: None, fallback_on_statuses: None, when_all_unavailable: None, - sticky: None, + hash_on: None, }; assert_eq!(r.max_fallbacks_or_default(), 0); } @@ -285,7 +393,7 @@ mod tests { retry_on_429: None, fallback_on_statuses: None, when_all_unavailable: None, - sticky: None, + hash_on: None, }; assert_eq!(r.max_fallbacks_or_default(), 0); } @@ -325,14 +433,56 @@ mod tests { } #[test] - fn sticky_parses_and_defaults_false() { - let off: Routing = serde_json::from_str(r#"{"targets":[{"model":"a"}]}"#).unwrap(); - assert!(!off.sticky_or_default()); - let on: Routing = serde_json::from_str( - r#"{"strategy":"weighted","sticky":true,"targets":[{"model":"a"},{"model":"b"}]}"#, + fn parses_consistent_hash_with_hash_on_chain() { + let r: Routing = serde_json::from_str( + r#"{ + "strategy": "consistent_hash", + "hash_on": [ + {"type": "header", "name": "x-session-id"}, + {"type": "cookie", "name": "sid"}, + {"type": "api_key"}, + {"type": "client_ip"} + ], + "targets": [{"model": "a"}, {"model": "b"}] + }"#, + ) + .unwrap(); + assert_eq!(r.strategy, RoutingStrategy::ConsistentHash); + let chain = r.hash_on_or_default(); + assert_eq!(chain.len(), 4); + assert_eq!(chain[0], HashOnSource::header("x-session-id")); + assert_eq!(chain[1], HashOnSource::cookie("sid")); + assert_eq!(chain[2], HashOnSource::api_key()); + assert_eq!(chain[3], HashOnSource::client_ip()); + } + + #[test] + fn hash_on_defaults_to_routing_key_header_then_api_key() { + let r: Routing = + serde_json::from_str(r#"{"strategy":"consistent_hash","targets":[{"model":"a"}]}"#) + .unwrap(); + assert_eq!(r.hash_on_or_default(), default_hash_on()); + assert_eq!( + default_hash_on()[0], + HashOnSource::header("x-aisix-routing-key") + ); + } + + #[test] + fn removed_weighted_strategy_and_sticky_flag_are_rejected() { + // `weighted` merged into `round_robin` and `sticky` was replaced by + // `strategy: consistent_hash` (AISIX-Cloud#1206). The enum value must + // fail row-level so a stale kine row cannot silently change meaning. + let weighted: Result = + serde_json::from_str(r#"{"strategy":"weighted","targets":[{"model":"a"}]}"#); + assert!(weighted.is_err()); + // `sticky` is now just an unknown field: lenient serde tolerates it + // (forward/backward compat), the strict write-path schema rejects it. + let sticky: Routing = serde_json::from_str( + r#"{"strategy":"round_robin","sticky":true,"targets":[{"model":"a"}]}"#, ) .unwrap(); - assert!(on.sticky_or_default()); + assert_eq!(sticky.strategy, RoutingStrategy::RoundRobin); } #[test] @@ -376,7 +526,7 @@ mod tests { assert!(RoutingStrategy::LeastBusy.is_metric_based()); assert!(!RoutingStrategy::Failover.is_metric_based()); assert!(!RoutingStrategy::RoundRobin.is_metric_based()); - assert!(!RoutingStrategy::Weighted.is_metric_based()); + assert!(!RoutingStrategy::ConsistentHash.is_metric_based()); } #[test] diff --git a/crates/aisix-core/tests/model_schema_characterization.rs b/crates/aisix-core/tests/model_schema_characterization.rs index 469553a0..166260eb 100644 --- a/crates/aisix-core/tests/model_schema_characterization.rs +++ b/crates/aisix-core/tests/model_schema_characterization.rs @@ -115,8 +115,12 @@ fn accept_routing_full() { json!({ "display_name": "r", "routing": { - "strategy": "weighted", - "targets": [{"model": "a", "weight": 3}, {"model": "b", "weight": 1}], + "strategy": "consistent_hash", + "hash_on": [{"type": "header", "name": "x-session-id"}, {"type": "api_key"}], + "targets": [ + {"model": "a", "weight": 3}, + {"model": "b", "weight": 1, "priority": -1} + ], "retries": 2, "max_fallbacks": 1, "retry_on_429": true, diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index 3360f97b..4446d7e8 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -1366,12 +1366,8 @@ async fn dispatch( &virtual_entry.value, RoutingRequest { tags: &client.routing_tags, - stability_key: Some( - client - .routing_key - .as_deref() - .unwrap_or(auth.entry.id.as_str()), - ), + headers: Some(&client.headers), + api_key_id: auth.entry.id.as_str(), source_ip: &client.source_ip, }, ) diff --git a/crates/aisix-proxy/src/client_ip.rs b/crates/aisix-proxy/src/client_ip.rs index 1771d4aa..c81c0b7d 100644 --- a/crates/aisix-proxy/src/client_ip.rs +++ b/crates/aisix-proxy/src/client_ip.rs @@ -121,11 +121,6 @@ fn parse_forwarded_token(tok: &str) -> Option { /// tags never reach the upstream request body. pub const ROUTING_TAGS_HEADER: &str = "x-aisix-routing-tags"; -/// Header carrying the stability key for sticky (A/B / canary) weighted -/// routing. When present, a request consistently maps to the same weighted -/// target; absent, the caller's API key is used as the key instead. -pub const ROUTING_KEY_HEADER: &str = "x-aisix-routing-key"; - /// Per-request client attribution. Resolved once via the extractor and /// threaded into the usage event by each handler's emit fn. #[derive(Debug, Clone, Default)] @@ -135,9 +130,6 @@ pub struct ClientContext { /// Routing tags from [`ROUTING_TAGS_HEADER`], used to select among a /// routing model's tagged targets. Empty when the header is absent. pub routing_tags: Vec, - /// Stability key from [`ROUTING_KEY_HEADER`] for sticky weighted routing. - /// `None` when the header is absent (the caller's API key is used instead). - pub routing_key: Option, /// Per-request correlation id, resolved from the [`RequestId`] the /// `ensure_request_id` middleware stamped into the request extensions. /// Handlers use it for both the usage event and the response header, so @@ -211,14 +203,6 @@ where .map(parse_routing_tags) .unwrap_or_default(); - let routing_key = parts - .headers - .get(ROUTING_KEY_HEADER) - .and_then(|v| v.to_str().ok()) - .map(str::trim) - .filter(|s| !s.is_empty()) - .map(str::to_owned); - let request_id = parts .extensions .get::() @@ -229,7 +213,6 @@ where source_ip, user_agent, routing_tags, - routing_key, request_id, headers: Arc::new(parts.headers.clone()), caller: parts diff --git a/crates/aisix-proxy/src/count_tokens.rs b/crates/aisix-proxy/src/count_tokens.rs index f4f1c421..3783b7a1 100644 --- a/crates/aisix-proxy/src/count_tokens.rs +++ b/crates/aisix-proxy/src/count_tokens.rs @@ -215,12 +215,8 @@ async fn dispatch( &model_entry.value, crate::routing::RoutingRequest { tags: &client.routing_tags, - stability_key: Some( - client - .routing_key - .as_deref() - .unwrap_or(auth.entry.id.as_str()), - ), + headers: Some(&client.headers), + api_key_id: auth.entry.id.as_str(), source_ip: &client.source_ip, }, )?; diff --git a/crates/aisix-proxy/src/health.rs b/crates/aisix-proxy/src/health.rs index 5975bbfe..4ec783b5 100644 --- a/crates/aisix-proxy/src/health.rs +++ b/crates/aisix-proxy/src/health.rs @@ -1010,7 +1010,7 @@ mod tests { 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("consistent_hash"), 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. diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index a6cd5ed2..21cf8878 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -649,12 +649,8 @@ async fn dispatch( &model_entry.value, crate::routing::RoutingRequest { tags: &client.routing_tags, - stability_key: Some( - client - .routing_key - .as_deref() - .unwrap_or(auth.entry.id.as_str()), - ), + headers: Some(&client.headers), + api_key_id: auth.entry.id.as_str(), source_ip: &client.source_ip, }, )?; diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index dc7bb5cf..65ec2648 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -618,12 +618,8 @@ async fn dispatch( &model_entry.value, crate::routing::RoutingRequest { tags: &client.routing_tags, - stability_key: Some( - client - .routing_key - .as_deref() - .unwrap_or(auth.entry.id.as_str()), - ), + headers: Some(&client.headers), + api_key_id: auth.entry.id.as_str(), source_ip: &client.source_ip, }, )?; diff --git a/crates/aisix-proxy/src/routing.rs b/crates/aisix-proxy/src/routing.rs index c1b47634..cc5bfcc2 100644 --- a/crates/aisix-proxy/src/routing.rs +++ b/crates/aisix-proxy/src/routing.rs @@ -3,35 +3,46 @@ //! When a request lands on a Model with `routing` configured, the proxy //! asks the [`RoutingRegistry`] for an iterator of underlying target //! Model names in attempt-order. The registry owns the per-virtual- -//! model state (round-robin counter, weighted PRNG seed); selection +//! model state (smooth-WRR counters, consistent-hash rings); selection //! itself is pure given that state. //! -//! Positional strategies (spec §3.5) pick a starting target, then walk +//! Targets partition into **priority tiers** first (`priority`, higher +//! value preferred, APISIX-style — each tier gets its own balancing +//! state, mirroring APISIX's per-priority pickers): the strategy orders +//! targets within each tier, tiers concatenate best-first, and a lower +//! tier is only reached after every higher-tier target failed or was +//! filtered as unavailable. +//! +//! Positional strategies pick a starting target per tier, then walk //! forward on failure: -//! - **failover**: always start at `targets[0]`, walk forward on failure. -//! - **round_robin**: each *new* request advances a per-model counter -//! so callers spread evenly across targets. -//! - **weighted**: pick a starting target with probability proportional -//! to `weight`, then walk forward on failure (weights only affect the -//! *first* target choice — once we're falling back, order is positional). +//! - **failover**: always start at the tier's first target, walk forward. +//! - **round_robin**: smooth weighted round-robin over target `weight`s +//! (equal weights degrade to a plain cycle). +//! - **consistent_hash**: ketama-style hashing of the request's hash key +//! (the `hash_on` chain) over the tier's ring; the walk follows the +//! ring, so a failed target's keys spread to their ring successors and +//! every other key keeps its mapping. //! -//! Metric-ordered strategies rank the whole target set by a runtime signal -//! (attempted best-first, then falling forward). They can't be ordered from -//! `pick_targets` because the ranking key lives on the resolved target +//! Metric-ordered strategies rank targets by a runtime signal within each +//! tier (attempted best-first, then falling forward). They can't be ordered +//! from `pick_targets` because the ranking key lives on the resolved target //! Models / runtime state, so `resolve_attempt_models` ranks them instead: //! - **least_cost**: cheapest target first, by combined input+output per-1K //! price; targets without a `cost` rank last. //! - **least_latency**: fastest target first, by an EWMA of observed upstream //! latency; targets with no samples yet rank first (probe, then exploit). -//! - **least_busy**: least-loaded target first, by in-flight request count. +//! - **least_busy**: least-loaded target first, by in-flight requests +//! divided by target `weight` (the APISIX least_conn score). use aisix_core::{ - AisixSnapshot, Model, Routing, RoutingStrategy, RoutingTarget, WhenAllUnavailablePolicy, + AisixSnapshot, HashOnType, Model, Routing, RoutingStrategy, RoutingTarget, + WhenAllUnavailablePolicy, }; use aisix_gateway::BridgeError; +use axum::http::HeaderMap; use dashmap::DashMap; use rand::Rng; -use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Mutex; use std::time::Duration; use crate::error::ProxyError; @@ -441,10 +452,111 @@ where Err(last_err.unwrap_or_else(|| BridgeError::Config("retry loop produced no error".into()))) } +/// Balancing state is keyed per (virtual model, tier priority). Priority is +/// part of the identity — the APISIX convention: tiers hold disjoint target +/// sets and must not share rotation state. +type TierKey = (String, i32); + +/// Smooth weighted round-robin state for one (virtual model, tier). +struct WrrState { + /// Fingerprint of the (model, weight) list this state was built for. A + /// config change — or a different tag/IP-filtered subset — resets the + /// rotation rather than letting stale counters index a different list. + fingerprint: u64, + current: Vec, +} + +const RING_POINTS_PER_UNIT: usize = 160; +/// Weights are reduced (gcd, then proportional scaling) to at most this many +/// total units before being multiplied by [`RING_POINTS_PER_UNIT`], bounding +/// ring memory per tier regardless of the configured weight magnitudes. +const RING_MAX_UNITS: u64 = 64; + +/// One tier's ketama-style consistent-hash ring. Every target owns +/// `160 × weight-units` pseudo-random points on the u64 hash circle; a key +/// maps to the first point at or after its own hash. Removing a target +/// (health filtering) moves only that target's keys — each lands on its +/// ring successor — and every other key keeps its mapping. +struct HashRing { + fingerprint: u64, + /// (hash point, index into the tier's target list), sorted by point. + points: Vec<(u64, u32)>, +} + +impl HashRing { + fn build(fingerprint: u64, targets: &[RoutingTarget]) -> Self { + fn gcd(mut a: u64, mut b: u64) -> u64 { + while b != 0 { + (a, b) = (b, a % b); + } + a + } + // An explicit weight of 0 still gets one unit: a target with no ring + // points would be silently unreachable in its tier, turning a weight + // typo into a dropped target. + let weights: Vec = targets + .iter() + .map(|t| u64::from(t.weight_or_default().max(1))) + .collect(); + let g = weights.iter().fold(0, |acc, w| gcd(acc, *w)).max(1); + let mut units: Vec = weights.iter().map(|w| w / g).collect(); + let sum: u64 = units.iter().sum(); + if sum > RING_MAX_UNITS { + units = units + .iter() + .map(|u| ((u * RING_MAX_UNITS) / sum).max(1)) + .collect(); + } + let total_points = units.iter().sum::() as usize * RING_POINTS_PER_UNIT; + let mut points = Vec::with_capacity(total_points); + for (idx, (target, units)) in targets.iter().zip(&units).enumerate() { + for i in 0..(*units as usize) * RING_POINTS_PER_UNIT { + let mut h = fnv1a_extend(FNV_OFFSET_BASIS, target.model.as_bytes()); + h = fnv1a_extend(h, &[0]); + h = fnv1a_extend(h, &(i as u32).to_le_bytes()); + points.push((mix64(h), idx as u32)); + } + } + points.sort_unstable(); + Self { + fingerprint, + points, + } + } + + /// The tier's targets in this key's deterministic preference order: + /// the key's own point first, then successive ring positions. This + /// doubles as the failover order within the tier. + fn preference_order(&self, key_hash: u64, n_targets: usize) -> Vec { + let mut order = Vec::with_capacity(n_targets); + let mut seen = vec![false; n_targets]; + if !self.points.is_empty() { + let start = self.points.partition_point(|(p, _)| *p < key_hash); + for off in 0..self.points.len() { + let (_, idx) = self.points[(start + off) % self.points.len()]; + if !seen[idx as usize] { + seen[idx as usize] = true; + order.push(idx); + if order.len() == n_targets { + return order; + } + } + } + } + // Degenerate rings (no points) still yield every target. + for (i, present) in seen.iter().enumerate() { + if !present { + order.push(i as u32); + } + } + order + } +} + #[derive(Default)] pub struct RoutingRegistry { - // virtual model name → atomic round-robin cursor - cursors: DashMap, + wrr: DashMap>, + rings: DashMap>, } impl RoutingRegistry { @@ -453,14 +565,17 @@ impl RoutingRegistry { } /// Pick the target order for one request. The first element is the - /// initial target; subsequent elements are later fallback targets (in - /// declaration order, wrapping if needed). Length is bounded by the - /// initial target plus `routing.max_fallbacks_or_default()`. + /// initial target; subsequent elements are later fallback targets. + /// Targets partition into priority tiers (higher value first); the + /// strategy orders each tier independently and tiers concatenate, so a + /// lower tier is reached only after every higher-tier target failed or + /// was filtered. Length is bounded by the initial target plus + /// `routing.max_fallbacks_or_default()`. pub fn pick_targets( &self, virtual_name: &str, routing: &Routing, - stability_key: Option<&str>, + hash_key: &str, ) -> Vec { if routing.targets.is_empty() { return Vec::new(); @@ -468,61 +583,145 @@ impl RoutingRegistry { // Metric-ordered strategies (least_cost, …) can't be ranked here: // the ranking key lives on the resolved target Models / runtime // state, which `resolve_attempt_models` has and this does not. Hand - // back the full declaration-order list; ranking and `max_fallbacks` - // truncation happen there instead. + // back the full declaration-order list; tier-aware ranking and + // `max_fallbacks` truncation happen there instead. if routing.strategy.is_metric_based() { return routing.targets.iter().map(|t| t.model.clone()).collect(); } - let start = self.starting_index(virtual_name, routing, stability_key); - attempt_order( - &routing.targets, - start, - routing.max_fallbacks_or_default() + 1, - ) + let mut order = Vec::with_capacity(routing.targets.len()); + for tier in partition_by_priority(&routing.targets) { + let priority = tier[0].priority_or_default(); + match routing.strategy { + RoutingStrategy::Failover => { + order.extend(tier.iter().map(|t| t.model.clone())); + } + RoutingStrategy::RoundRobin => { + let start = self.wrr_pick(virtual_name, priority, &tier); + order.extend(attempt_order(&tier, start, tier.len())); + } + RoutingStrategy::ConsistentHash => { + let ring = self.ring_for(virtual_name, priority, &tier); + for idx in ring.preference_order(stable_hash(hash_key), tier.len()) { + order.push(tier[idx as usize].model.clone()); + } + } + RoutingStrategy::LeastCost + | RoutingStrategy::LeastLatency + | RoutingStrategy::LeastBusy => { + unreachable!("metric strategies short-circuit above") + } + } + } + order.truncate(routing.max_fallbacks_or_default() + 1); + order } - fn starting_index( - &self, - virtual_name: &str, - routing: &Routing, - stability_key: Option<&str>, - ) -> usize { - match routing.strategy { - RoutingStrategy::Failover => 0, - RoutingStrategy::RoundRobin => self.advance_cursor(virtual_name, routing.targets.len()), - RoutingStrategy::Weighted => { - // Sticky (A/B / canary) routing makes the weighted pick - // deterministic in the request's stability key; otherwise each - // request samples the weight distribution independently. - let sticky_key = routing - .sticky_or_default() - .then_some(stability_key) - .flatten(); - weighted_pick(&routing.targets, sticky_key) + /// Smooth weighted round-robin (the nginx algorithm): every pick adds + /// each target's weight to its running counter, takes the max, and + /// subtracts the weight total from the winner. Proportional AND + /// interleaved; equal weights degrade to a declaration-order cycle. + fn wrr_pick(&self, virtual_name: &str, priority: i32, tier: &[RoutingTarget]) -> usize { + let weights: Vec = tier + .iter() + .map(|t| i64::from(t.weight_or_default().max(1))) + .collect(); + let total: i64 = weights.iter().sum(); + let fingerprint = tier_fingerprint(tier); + let entry = self + .wrr + .entry((virtual_name.to_string(), priority)) + .or_insert_with(|| { + Mutex::new(WrrState { + fingerprint, + current: vec![0; weights.len()], + }) + }); + let mut state = entry + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state.fingerprint != fingerprint || state.current.len() != weights.len() { + *state = WrrState { + fingerprint, + current: vec![0; weights.len()], + }; + } + let mut best = 0; + for (i, w) in weights.iter().enumerate() { + state.current[i] += w; + if state.current[i] > state.current[best] { + best = i; } - // Metric-ordered strategies never reach here — `pick_targets` - // short-circuits them before computing a start index. - RoutingStrategy::LeastCost - | RoutingStrategy::LeastLatency - | RoutingStrategy::LeastBusy => 0, } + state.current[best] -= total; + best } - fn advance_cursor(&self, virtual_name: &str, modulo: usize) -> usize { - let entry = self.cursors.entry(virtual_name.to_string()).or_default(); - let prev = entry.fetch_add(1, Ordering::Relaxed); - prev % modulo + /// The cached ring for one (virtual model, tier), rebuilt when the + /// tier's (model, weight) fingerprint changes — a config edit, or a + /// different tag/IP-filtered subset. One entry per key: alternating + /// subsets rebuild rather than accumulate, keeping the map bounded by + /// the number of configured (group, tier) pairs. + fn ring_for( + &self, + virtual_name: &str, + priority: i32, + tier: &[RoutingTarget], + ) -> std::sync::Arc { + let fingerprint = tier_fingerprint(tier); + let key = (virtual_name.to_string(), priority); + if let Some(ring) = self.rings.get(&key) { + if ring.fingerprint == fingerprint { + return ring.clone(); + } + } + let ring = std::sync::Arc::new(HashRing::build(fingerprint, tier)); + self.rings.insert(key, ring.clone()); + ring } } impl std::fmt::Debug for RoutingRegistry { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("RoutingRegistry") - .field("virtual_models_seen", &self.cursors.len()) + .field("wrr_tiers", &self.wrr.len()) + .field("rings", &self.rings.len()) .finish() } } +/// Split targets into priority tiers, highest priority first, declaration +/// order preserved within each tier. Every target defaults to priority 0, +/// so an unconfigured group is a single tier and this is a no-op shape. +fn partition_by_priority(targets: &[RoutingTarget]) -> Vec> { + let mut priorities: Vec = targets.iter().map(|t| t.priority_or_default()).collect(); + priorities.sort_unstable_by(|a, b| b.cmp(a)); + priorities.dedup(); + priorities + .into_iter() + .map(|p| { + targets + .iter() + .filter(|t| t.priority_or_default() == p) + .cloned() + .collect() + }) + .collect() +} + +/// Fingerprint of a tier's identity for balancing-state reuse: the ordered +/// (model, weight) pairs. Priorities are already part of the state key and +/// tags do not affect selection within a surviving subset. +fn tier_fingerprint(tier: &[RoutingTarget]) -> u64 { + let mut h = FNV_OFFSET_BASIS; + for t in tier { + h = fnv1a_extend(h, t.model.as_bytes()); + h = fnv1a_extend(h, &[0]); + h = fnv1a_extend(h, &t.weight_or_default().to_le_bytes()); + h = fnv1a_extend(h, &[1]); + } + h +} + /// Narrow a routing model's targets to those eligible for this request's /// routing tags, mirroring LiteLLM's tag-based routing: /// * No target is tagged → tag routing isn't in use; every target eligible. @@ -572,59 +771,95 @@ fn attempt_order(targets: &[RoutingTarget], start_idx: usize, limit: usize) -> V order } -/// Pick an index by weighted-random. Ignores zero weights; a fully-zero -/// list falls back to index 0 deterministically. -/// -/// Per #197: each call must draw an INDEPENDENT sample from the weight -/// distribution. The prior implementation used -/// `SystemTime::now().subsec_nanos() + Instant::now().elapsed().as_nanos()` -/// as entropy, which has two correctness bugs that compound: -/// 1. `Instant::now().elapsed()` always returns ~0 (the Instant was -/// just created), so the mix is effectively just subsec_nanos. -/// 2. Under rapid-fire requests (e2e fires N=100 in tight loop), -/// consecutive subsec_nanos values differ by a near-constant -/// step (≈1 µs of wall-clock per request). Modular reduction -/// `entropy() % total_weight` against that step pattern aliases -/// to a single bin — every request lands on the same target. -/// Empirical observation: 200/0 split on a configured 70/30. -/// -/// Use `rand::thread_rng()` instead. The thread-local PRNG is seeded -/// from OS entropy on first use and is independent across calls; the -/// distribution converges to the configured weights over a finite -/// sample (per the spec the e2e pins). -/// -/// With a `sticky_key` (A/B / canary routing) the pick is instead a -/// deterministic function of that key, so the same key always resolves to the -/// same target while the aggregate split still honors the weights. -fn weighted_pick(targets: &[RoutingTarget], sticky_key: Option<&str>) -> usize { - let total: u64 = targets.iter().map(|t| t.weight_or_default() as u64).sum(); - if total == 0 { - return 0; - } - let pick = match sticky_key { - Some(key) => stable_hash(key) % total, - None => rand::thread_rng().gen_range(0..total), - }; - let mut acc: u64 = 0; - for (i, t) in targets.iter().enumerate() { - acc += t.weight_or_default() as u64; - if pick < acc { - return i; - } +const FNV_OFFSET_BASIS: u64 = 0xcbf2_9ce4_8422_2325; +const FNV_PRIME: u64 = 0x0000_0100_0000_01b3; + +/// Fold `bytes` into a running 64-bit FNV-1a hash. Deterministic across +/// processes, replicas, and toolchains by design (the std hasher is not) — +/// every consistent-hash artifact (ring points, key hashes, fingerprints) +/// must agree everywhere, and MUST NOT change across DP versions: changing +/// this function remaps every session's target. +fn fnv1a_extend(mut h: u64, bytes: &[u8]) -> u64 { + for b in bytes { + h ^= u64::from(*b); + h = h.wrapping_mul(FNV_PRIME); } - targets.len() - 1 + h +} + +/// The splitmix64 finalizer. FNV-1a alone has weak avalanche on short, +/// structured inputs — ring points hashed from `name\0index` cluster into +/// narrow bands, handing one target most of the circle (observed: 32 +/// distinct keys all mapping to one of two equal-weight targets). Ketama +/// implementations use MD5/CRC32 for exactly this reason; a strong final +/// mix restores uniform dispersion while keeping the pipeline dependency- +/// free and deterministic. Same stability contract as `fnv1a_extend`. +fn mix64(mut x: u64) -> u64 { + x ^= x >> 30; + x = x.wrapping_mul(0xbf58_476d_1ce4_e5b9); + x ^= x >> 27; + x = x.wrapping_mul(0x94d0_49bb_1331_11eb); + x ^= x >> 31; + x } -/// Stable 64-bit FNV-1a hash used to map a sticky-routing key into the weight -/// distribution. Deterministic across processes and toolchains by design (the -/// std hasher is not), so a given key always resolves to the same target. +/// Stable, well-dispersed 64-bit hash of a consistent-hash key. fn stable_hash(s: &str) -> u64 { - let mut h: u64 = 0xcbf2_9ce4_8422_2325; // FNV-1a offset basis - for b in s.as_bytes() { - h ^= *b as u64; - h = h.wrapping_mul(0x0000_0100_0000_01b3); // FNV-1a prime + mix64(fnv1a_extend(FNV_OFFSET_BASIS, s.as_bytes())) +} + +/// Resolve the `consistent_hash` key for one request by walking the +/// routing model's `hash_on` chain (default: the `x-aisix-routing-key` +/// header, then the caller's API key id). The first source yielding a +/// non-empty value wins; when nothing yields, the empty string keeps the +/// pick deterministic rather than random. +fn resolve_hash_key(routing: &Routing, req: &RoutingRequest<'_>) -> String { + for source in routing.hash_on_or_default() { + let value = match source.source_type { + HashOnType::Header => source.name.as_deref().and_then(|name| { + req.headers + .and_then(|h| h.get(name)) + .and_then(|v| v.to_str().ok()) + .map(str::trim) + .filter(|s| !s.is_empty()) + .map(str::to_owned) + }), + HashOnType::Cookie => source + .name + .as_deref() + .and_then(|name| req.headers.and_then(|h| cookie_value(h, name))), + HashOnType::ApiKey => Some(req.api_key_id.trim()) + .filter(|s| !s.is_empty()) + .map(str::to_owned), + HashOnType::ClientIp => Some(req.source_ip.trim()) + .filter(|s| !s.is_empty()) + .map(str::to_owned), + }; + if let Some(value) = value { + return value; + } } - h + String::new() +} + +/// Extract a cookie's value from the request's `Cookie` header(s): +/// `name=value` pairs separated by `;`, first match wins. Values are taken +/// verbatim (no unquoting) — the key only needs to be stable, not parsed. +fn cookie_value(headers: &HeaderMap, name: &str) -> Option { + for header in headers.get_all(axum::http::header::COOKIE) { + let Ok(raw) = header.to_str() else { continue }; + for pair in raw.split(';') { + let mut it = pair.splitn(2, '='); + let key = it.next().unwrap_or("").trim(); + if key == name { + let value = it.next().unwrap_or("").trim(); + if !value.is_empty() { + return Some(value.to_owned()); + } + } + } + } + None } /// Combined per-1K unit price used to rank `least_cost` targets. A target @@ -666,18 +901,30 @@ fn order_attempts_by_metric( }); } RoutingStrategy::LeastBusy => { - attempts.sort_by_key(|a| runtime_status.in_flight(&a.id)); + // The APISIX least_conn score: in-flight scaled by 1/weight, so a + // heavier target absorbs proportionally more concurrency. `+1` + // keeps idle targets ranked by weight instead of tying at 0. + let score = |a: &AttemptModel| { + (runtime_status.in_flight(&a.id) as f64 + 1.0) / f64::from(a.weight.max(1)) + }; + attempts.sort_by(|a, b| score(a).total_cmp(&score(b))); } - RoutingStrategy::Failover | RoutingStrategy::RoundRobin | RoutingStrategy::Weighted => {} + RoutingStrategy::Failover + | RoutingStrategy::RoundRobin + | RoutingStrategy::ConsistentHash => {} } } /// One concrete (non-routing) Model the dispatch loop will attempt, paired -/// with its snapshot id so health/cooldown tracking can key on it. +/// with its snapshot id so health/cooldown tracking can key on it, and the +/// routing-target attributes selection still needs after resolution +/// (priority for tier-aware metric ranking, weight for `least_busy`). #[derive(Clone)] pub(crate) struct AttemptModel { pub id: String, pub model: Model, + pub priority: i32, + pub weight: u32, } /// Outcome of routing-candidate filtering. Lifts the "all candidates @@ -757,11 +1004,11 @@ pub(crate) fn filter_attempt_models( } /// Per-request routing inputs threaded into [`resolve_attempt_models`]: the -/// tags that gate tag/metadata routing, the stability key for sticky -/// (A/B / canary) weighted selection, and the caller's resolved source IP -/// for the per-target client-IP allowlist. Tags come from request headers; -/// the stability key is the routing-key header when present, otherwise the -/// caller's API key id. +/// tags that gate tag/metadata routing, the raw material for the +/// `consistent_hash` hash key (the inbound headers plus the caller's API key +/// id — the configured `hash_on` chain is evaluated against them at +/// resolution time), and the caller's resolved source IP, used both for the +/// per-target client-IP allowlist and as the `client_ip` hash source. /// /// `source_ip` defaults to the empty string, which /// [`aisix_core::Model::ip_allowed`] treats as "not in range" — so a caller @@ -770,7 +1017,8 @@ pub(crate) fn filter_attempt_models( #[derive(Clone, Copy, Default)] pub(crate) struct RoutingRequest<'a> { pub tags: &'a [String], - pub stability_key: Option<&'a str>, + pub headers: Option<&'a HeaderMap>, + pub api_key_id: &'a str, pub source_ip: &'a str, } @@ -821,6 +1069,8 @@ pub(crate) fn resolve_attempt_models( return Ok(vec![AttemptModel { id: virtual_id.to_string(), model: virtual_model.clone(), + priority: 0, + weight: 1, }]); }; @@ -853,7 +1103,12 @@ pub(crate) fn resolve_attempt_models( }; let routing = &filtered_routing; - let names = routing_registry.pick_targets(virtual_name, routing, req.stability_key); + let hash_key = if routing.strategy == RoutingStrategy::ConsistentHash { + resolve_hash_key(routing, &req) + } else { + String::new() + }; + let names = routing_registry.pick_targets(virtual_name, routing, &hash_key); if names.is_empty() { return Err(ProxyError::InvalidRequest( "routing model has no targets".into(), @@ -866,16 +1121,28 @@ pub(crate) fn resolve_attempt_models( "routing target {name:?} does not resolve to a Model" )) })?; + // Duplicate target models are rejected at the write path, so the + // first match is the only match. + let target = routing + .targets + .iter() + .find(|t| t.model == *name) + .expect("picked name comes from routing.targets"); resolved.push(AttemptModel { id: target_entry.id.clone(), model: target_entry.value.clone(), + priority: target.priority_or_default(), + weight: target.weight_or_default(), }); } // Metric-ordered strategies get the full target set from `pick_targets`; // rank it best-first here (target Models are now resolved) and cap it to - // the same attempt budget the positional strategies apply upstream. + // the same attempt budget the positional strategies apply upstream. The + // metric sort runs first, then a stable sort on priority — so tiers + // concatenate highest-first with the metric order preserved inside each. if routing.strategy.is_metric_based() { order_attempts_by_metric(routing.strategy, &mut resolved, runtime_status); + resolved.sort_by_key(|a| std::cmp::Reverse(a.priority)); resolved.truncate(routing.max_fallbacks_or_default() + 1); } match filter_attempt_models( @@ -913,7 +1180,7 @@ mod tests { retry_on_429: None, fallback_on_statuses: None, when_all_unavailable: None, - sticky: None, + hash_on: None, } } @@ -932,59 +1199,96 @@ mod tests { } #[test] - fn sticky_weighted_pick_is_deterministic_per_key() { + fn chash_preference_order_is_deterministic_per_key() { let targets = vec![ RoutingTarget::new("a").with_weight(50), RoutingTarget::new("b").with_weight(50), ]; - let first = weighted_pick(&targets, Some("session-1")); + let ring = HashRing::build(tier_fingerprint(&targets), &targets); + let first = ring.preference_order(stable_hash("session-1"), targets.len()); + assert_eq!(first.len(), 2); for _ in 0..50 { - assert_eq!(weighted_pick(&targets, Some("session-1")), first); + assert_eq!( + ring.preference_order(stable_hash("session-1"), targets.len()), + first + ); } } #[test] - fn sticky_weighted_pick_spreads_distinct_keys() { + fn chash_spreads_distinct_keys() { // Distinct keys shouldn't all funnel to one target. let targets = vec![ RoutingTarget::new("a").with_weight(50), RoutingTarget::new("b").with_weight(50), ]; + let ring = HashRing::build(tier_fingerprint(&targets), &targets); let mut seen = [false; 2]; for i in 0..200 { - seen[weighted_pick(&targets, Some(&format!("k{i}")))] = true; + seen[ring.preference_order(stable_hash(&format!("k{i}")), 2)[0] as usize] = true; } assert!(seen[0] && seen[1]); } #[test] - fn sticky_weighted_pick_honors_extreme_weights() { - // A 100/0 canary split lands every key on the weighted target. + fn chash_weight_scales_a_targets_share_of_keys() { let targets = vec![ - RoutingTarget::new("stable").with_weight(100), - RoutingTarget::new("canary").with_weight(0), + RoutingTarget::new("heavy").with_weight(90), + RoutingTarget::new("light").with_weight(10), ]; - for i in 0..50 { - assert_eq!(weighted_pick(&targets, Some(&format!("k{i}"))), 0); + let ring = HashRing::build(tier_fingerprint(&targets), &targets); + let mut heavy = 0; + for i in 0..1000 { + if ring.preference_order(stable_hash(&format!("k{i}")), 2)[0] == 0 { + heavy += 1; + } } + // 90/10 configured; allow generous slack for hash variance. + assert!( + (800..=980).contains(&heavy), + "expected ~900/1000 keys on the heavy target, got {heavy}" + ); } #[test] - fn sticky_routing_pins_a_key_to_one_target() { + fn chash_removing_a_target_only_moves_its_own_keys() { + // The consistent-hash property the whole feature hangs on: dropping + // one target must not remap keys whose first choice survives. + let full = vec![ + RoutingTarget::new("a"), + RoutingTarget::new("b"), + RoutingTarget::new("c"), + ]; + let ring = HashRing::build(tier_fingerprint(&full), &full); + let shrunk: Vec = vec![full[0].clone(), full[2].clone()]; // drop "b" + let shrunk_ring = HashRing::build(tier_fingerprint(&shrunk), &shrunk); + for i in 0..500 { + let h = stable_hash(&format!("k{i}")); + let before = full[ring.preference_order(h, 3)[0] as usize].model.clone(); + let after = shrunk[shrunk_ring.preference_order(h, 2)[0] as usize] + .model + .clone(); + if before != "b" { + assert_eq!(before, after, "key k{i} moved although its target survived"); + } + } + } + + #[test] + fn chash_pick_targets_pins_a_key_and_walks_the_ring() { let reg = RoutingRegistry::new(); - let mut routing = r( - RoutingStrategy::Weighted, + let routing = r( + RoutingStrategy::ConsistentHash, vec![ RoutingTarget::new("stable").with_weight(90), RoutingTarget::new("canary").with_weight(10), ], - Some(0), // only the chosen start target + None, ); - routing.sticky = Some(true); - let first = reg.pick_targets("v", &routing, Some("user-42")); - assert_eq!(first.len(), 1); + let first = reg.pick_targets("v", &routing, "user-42"); + assert_eq!(first.len(), 2, "walk covers the whole tier"); for _ in 0..20 { - assert_eq!(reg.pick_targets("v", &routing, Some("user-42")), first); + assert_eq!(reg.pick_targets("v", &routing, "user-42"), first); } } @@ -1136,7 +1440,7 @@ mod tests { None, ); for _ in 0..5 { - let order = reg.pick_targets("v", &routing, None); + let order = reg.pick_targets("v", &routing, ""); assert_eq!(order, vec!["primary", "secondary", "tertiary"]); } } @@ -1155,7 +1459,7 @@ mod tests { ); let mut firsts = Vec::new(); for _ in 0..6 { - let order = reg.pick_targets("v", &routing, None); + let order = reg.pick_targets("v", &routing, ""); firsts.push(order[0].clone()); } // Two full cycles of a→b→c. @@ -1171,10 +1475,10 @@ mod tests { Some(1), ); // Two distinct virtual models advance independently. - assert_eq!(reg.pick_targets("v1", &routing, None)[0], "a"); - assert_eq!(reg.pick_targets("v2", &routing, None)[0], "a"); - assert_eq!(reg.pick_targets("v1", &routing, None)[0], "b"); - assert_eq!(reg.pick_targets("v2", &routing, None)[0], "b"); + assert_eq!(reg.pick_targets("v1", &routing, "")[0], "a"); + assert_eq!(reg.pick_targets("v2", &routing, "")[0], "a"); + assert_eq!(reg.pick_targets("v1", &routing, "")[0], "b"); + assert_eq!(reg.pick_targets("v2", &routing, "")[0], "b"); } #[test] @@ -1190,178 +1494,200 @@ mod tests { Some(2), ); // First call starts at a → a, b, c - assert_eq!(reg.pick_targets("v", &routing, None), vec!["a", "b", "c"]); + assert_eq!(reg.pick_targets("v", &routing, ""), vec!["a", "b", "c"]); // Second call starts at b → b, c, a - assert_eq!(reg.pick_targets("v", &routing, None), vec!["b", "c", "a"]); + assert_eq!(reg.pick_targets("v", &routing, ""), vec!["b", "c", "a"]); } #[test] - fn weighted_picks_from_targets_and_falls_back_in_order() { + fn wrr_first_pick_prefers_the_heavier_weight_and_walks_forward() { let reg = RoutingRegistry::new(); let routing = r( - RoutingStrategy::Weighted, + RoutingStrategy::RoundRobin, vec![ RoutingTarget::new("a").with_weight(99), RoutingTarget::new("b").with_weight(1), ], Some(1), ); - // We just assert correctness of the *order* shape: - // exactly two attempts, distinct targets, both targets covered. - // (Aggregate distribution is pinned by the dedicated tests - // below.) - let order = reg.pick_targets("v", &routing, None); - assert_eq!(order.len(), 2); - assert!(order.iter().any(|t| t == "a")); - assert!(order.iter().any(|t| t == "b")); + // Smooth WRR is deterministic: the first pick is the heavy target, + // the walk continues in declaration order. + assert_eq!(reg.pick_targets("v", &routing, ""), vec!["a", "b"]); } #[test] - fn weighted_with_all_zero_weights_picks_index_zero_deterministically() { - let targets = vec![ - RoutingTarget::new("a").with_weight(0), - RoutingTarget::new("b").with_weight(0), - ]; - assert_eq!(weighted_pick(&targets, None), 0); + fn wrr_distribution_matches_weights_exactly() { + // Smooth WRR is exact, not stochastic: over one full cycle of + // total-weight picks, each target is chosen exactly `weight` times. + let reg = RoutingRegistry::new(); + let routing = r( + RoutingStrategy::RoundRobin, + vec![ + RoutingTarget::new("a").with_weight(70), + RoutingTarget::new("b").with_weight(30), + ], + Some(0), + ); + let mut counts = [0usize; 2]; + for _ in 0..100 { + match reg.pick_targets("v", &routing, "")[0].as_str() { + "a" => counts[0] += 1, + _ => counts[1] += 1, + } + } + assert_eq!(counts, [70, 30]); } - /// Aggregate-distribution property: across many trials, a 100/1 - /// weight bias must converge to ~99% on the heavy target. Pre-#197 - /// the threshold sat at ≥ 60% to absorb the weak nanos-clock entropy - /// — that gate would also pass a weight-half-sensitivity regression - /// (~75% would slip through). With proper PRNG entropy in - /// `weighted_pick`, the empirical bin should land within ~1% of - /// the analytic 100/(100+1) = 99.0% expectation; we assert ≥ 95% - /// (≈4σ band for n=5000, rejects half-sensitivity AND weight-blind). #[test] - fn weighted_pick_aggregate_distribution_favors_heavier_weight() { - let targets = vec![ - RoutingTarget::new("a").with_weight(100), - RoutingTarget::new("b").with_weight(1), - ]; - let n = 5_000; - let a_count = (0..n) - .filter(|_| weighted_pick(&targets, None) == 0) - .count(); - // Uniform 50/50 → ~2500. Weighted 100/1 → ~4950 in theory. - // 95% threshold (4750) rejects both a weight-blind impl - // (~50%) AND a half-sensitivity regression (~75% would also - // fail). With proper PRNG entropy this gate has ~5σ margin; - // CI-flake risk is negligible. - assert!( - a_count * 100 / n >= 95, - "weight=100 target should dominate aggregate picks; got {a_count}/{n}", + fn wrr_interleaves_rather_than_bursting() { + // The nginx smooth-WRR property: 2/1 yields a,b,a per cycle, not + // a,a,b — heavier targets spread across the cycle. + let reg = RoutingRegistry::new(); + let routing = r( + RoutingStrategy::RoundRobin, + vec![ + RoutingTarget::new("a").with_weight(2), + RoutingTarget::new("b").with_weight(1), + ], + Some(0), ); + let picks: Vec = (0..6) + .map(|_| reg.pick_targets("v", &routing, "").remove(0)) + .collect(); + assert_eq!(picks, vec!["a", "b", "a", "a", "b", "a"]); } - /// Companion to the above: that test passes both for a correctly - /// weighted impl AND for an "always pick index 0" regression (since - /// the heavy weight is at index 0). Swap the weights so the heavy - /// target sits at index 1 — a weight-blind impl that always picks - /// the first target would now fail this test, while a correct - /// weighted impl still favors index 1. #[test] - fn weighted_pick_aggregate_distribution_respects_index_swap() { - let targets = vec![ - RoutingTarget::new("a").with_weight(1), - RoutingTarget::new("b").with_weight(100), - ]; - let n = 5_000; - let b_count = (0..n) - .filter(|_| weighted_pick(&targets, None) == 1) - .count(); - assert!( - b_count * 100 / n >= 95, - "weight=100 target at index 1 should dominate aggregate picks; got {b_count}/{n}", + fn wrr_zero_weights_clamp_to_one() { + // weight: 0 clamps to 1 (the write path forbids 0; clamping keeps a + // hand-written 0 reachable instead of silently dropping the target). + let reg = RoutingRegistry::new(); + let routing = r( + RoutingStrategy::RoundRobin, + vec![ + RoutingTarget::new("a").with_weight(0), + RoutingTarget::new("b").with_weight(0), + ], + Some(0), ); + let picks: Vec = (0..4) + .map(|_| reg.pick_targets("v", &routing, "").remove(0)) + .collect(); + assert_eq!(picks, vec!["a", "b", "a", "b"]); } - /// Issue #197 regression: a 70/30 weighted split must land near - /// 70/30 over a finite sample. The pre-fix nanos-clock entropy - /// collapsed to a single bin under rapid-fire calls (observed - /// 200/0 in e2e on a configured 70/30); a proper PRNG converges - /// to the analytic distribution. - /// - /// Tolerance: n=1000 with p=0.7 has σ=√(np(1-p))=√210≈14.49. A ±50 - /// absolute window is ~3.45σ → P(false positive) ≈ 0.056%. The - /// pre-fix collapse-to-one-bin failure produces 1000/0 which is - /// ~33σ outside the window — caught with overwhelming margin. #[test] - fn weighted_pick_70_30_split_converges_to_configured_ratio() { - let targets = vec![ - RoutingTarget::new("a").with_weight(70), - RoutingTarget::new("b").with_weight(30), - ]; - let n = 1_000; - let a_count = (0..n) - .filter(|_| weighted_pick(&targets, None) == 0) - .count(); - // Expected ~700; tolerance window [650, 750] (≈±3.45σ). - assert!( - (650..=750).contains(&a_count), - "70/30 weighted split must land near 700/1000; got {a_count}/{n} on heavy target", + fn wrr_state_resets_when_the_tier_config_changes() { + let reg = RoutingRegistry::new(); + let before = r( + RoutingStrategy::RoundRobin, + vec![RoutingTarget::new("a"), RoutingTarget::new("b")], + Some(0), + ); + assert_eq!(reg.pick_targets("v", &before, "")[0], "a"); + assert_eq!(reg.pick_targets("v", &before, "")[0], "b"); + // New target list → fingerprint mismatch → rotation restarts. + let after = r( + RoutingStrategy::RoundRobin, + vec![RoutingTarget::new("x"), RoutingTarget::new("y")], + Some(0), ); + assert_eq!(reg.pick_targets("v", &after, "")[0], "x"); } - /// 3-target coverage: a weight-blind impl that only ever picks - /// `targets[0]` if `pick < sum/n` (and `targets[1]` otherwise) - /// would pass every 2-target test in this module but fail with - /// 3+ targets — the third bin would starve. Pin a 50/30/20 split - /// and assert each bin lands within a generous tolerance window. - /// - /// n=2000 chosen so the smallest bin (20% → ~400) has σ ≈ 17.9; - /// ±100 window ≈ 5.6σ for that bin, larger margins for the other - /// two. #[test] - fn weighted_pick_50_30_20_split_distributes_to_all_three_bins() { - let targets = vec![ - RoutingTarget::new("a").with_weight(50), - RoutingTarget::new("b").with_weight(30), - RoutingTarget::new("c").with_weight(20), - ]; - let n = 2_000; - let mut counts = [0_usize; 3]; - for _ in 0..n { - counts[weighted_pick(&targets, None)] += 1; - } - // Expected 1000/600/400. ±100 window catches a weight-blind - // 2-target collapse (where the 3rd bin would be 0) AND - // sample noise. - assert!( - (900..=1100).contains(&counts[0]), - "50%-weighted bin should land near 1000/2000; got {counts:?}", - ); - assert!( - (500..=700).contains(&counts[1]), - "30%-weighted bin should land near 600/2000; got {counts:?}", + fn priority_tiers_concatenate_highest_first() { + // The A/B two-pool shape from AISIX-Cloud#1206: priority 0 is the + // active pool, priority -1 the backup; the walk exhausts the whole + // active tier before any backup target. + let reg = RoutingRegistry::new(); + let routing = r( + RoutingStrategy::Failover, + vec![ + RoutingTarget::new("b1").with_priority(-1), + RoutingTarget::new("a1"), + RoutingTarget::new("a2"), + RoutingTarget::new("b2").with_priority(-1), + ], + None, ); - assert!( - (300..=500).contains(&counts[2]), - "20%-weighted bin should land near 400/2000; got {counts:?}", + assert_eq!( + reg.pick_targets("v", &routing, ""), + vec!["a1", "a2", "b1", "b2"] ); } - /// Zero-weight-in-the-middle: a weight=0 target between two - /// non-zero targets must NEVER be picked. The CDF predicate - /// `pick < acc` (strict less-than) is what enforces this — a - /// weight-0 segment doesn't widen `acc` so the predicate skips - /// past it. A regression that used `<=` would incidentally pick - /// the zero-weight bin on the boundary value of `pick`. #[test] - fn weighted_pick_zero_weight_target_in_middle_is_never_picked() { - let targets = vec![ - RoutingTarget::new("a").with_weight(10), - RoutingTarget::new("b").with_weight(0), - RoutingTarget::new("c").with_weight(10), - ]; - let n = 2_000; - let b_count = (0..n) - .filter(|_| weighted_pick(&targets, None) == 1) - .count(); + fn priority_tiers_run_the_strategy_per_tier() { + // Each tier owns its own WRR rotation (the APISIX per-priority + // picker rule): the backup tier's order is its own strategy pick, + // not a continuation of the active tier's. + let reg = RoutingRegistry::new(); + let routing = r( + RoutingStrategy::RoundRobin, + vec![ + RoutingTarget::new("a1"), + RoutingTarget::new("a2"), + RoutingTarget::new("b1").with_priority(-1), + RoutingTarget::new("b2").with_priority(-1), + ], + None, + ); assert_eq!( - b_count, 0, - "weight=0 target must never be picked; got {b_count}/{n}", + reg.pick_targets("v", &routing, ""), + vec!["a1", "a2", "b1", "b2"] + ); + // Second call: both tiers advanced their own rotation. + assert_eq!( + reg.pick_targets("v", &routing, ""), + vec!["a2", "a1", "b2", "b1"] + ); + } + + #[test] + fn priority_tiers_chash_hashes_within_each_tier() { + let reg = RoutingRegistry::new(); + let routing = r( + RoutingStrategy::ConsistentHash, + vec![ + RoutingTarget::new("a1"), + RoutingTarget::new("a2"), + RoutingTarget::new("a3"), + RoutingTarget::new("b1").with_priority(-1), + RoutingTarget::new("b2").with_priority(-1), + ], + None, + ); + let order = reg.pick_targets("v", &routing, "session-7"); + assert_eq!(order.len(), 5); + // Every active-tier target precedes every backup target. + let a_positions: Vec = order + .iter() + .enumerate() + .filter(|(_, t)| t.starts_with('a')) + .map(|(i, _)| i) + .collect(); + assert_eq!(a_positions, vec![0, 1, 2]); + // Deterministic per key. + assert_eq!(reg.pick_targets("v", &routing, "session-7"), order); + // A different key may start elsewhere but keeps the tier boundary. + let other = reg.pick_targets("v", &routing, "session-8"); + assert!(other[..3].iter().all(|t| t.starts_with('a'))); + } + + #[test] + fn max_fallbacks_caps_across_tiers() { + let reg = RoutingRegistry::new(); + let routing = r( + RoutingStrategy::Failover, + vec![ + RoutingTarget::new("a1"), + RoutingTarget::new("a2"), + RoutingTarget::new("b1").with_priority(-1), + ], + Some(1), ); + assert_eq!(reg.pick_targets("v", &routing, ""), vec!["a1", "a2"]); } #[test] @@ -1372,7 +1698,7 @@ mod tests { vec![RoutingTarget::new("a"), RoutingTarget::new("b")], Some(0), ); - let order = reg.pick_targets("v", &routing, None); + let order = reg.pick_targets("v", &routing, ""); assert_eq!(order, vec!["a"]); } @@ -1380,7 +1706,7 @@ mod tests { fn empty_targets_yields_empty_order() { let reg = RoutingRegistry::new(); let routing = r(RoutingStrategy::Failover, vec![], None); - assert!(reg.pick_targets("v", &routing, None).is_empty()); + assert!(reg.pick_targets("v", &routing, "").is_empty()); } #[test] @@ -1921,6 +2247,8 @@ mod tests { AttemptModel { id: id.to_string(), model, + priority: 0, + weight: 1, } } @@ -1939,6 +2267,8 @@ mod tests { AttemptModel { id: id.to_string(), model, + priority: 0, + weight: 1, } } @@ -2067,7 +2397,7 @@ mod tests { ); // Ranking needs resolved Models, so pick_targets hands back every // target untouched regardless of max_fallbacks. - assert_eq!(reg.pick_targets("v", &routing, None), vec!["a", "b", "c"]); + assert_eq!(reg.pick_targets("v", &routing, ""), vec!["a", "b", "c"]); } #[test] diff --git a/crates/aisix-proxy/src/semantic.rs b/crates/aisix-proxy/src/semantic.rs index 976f9007..32ee735e 100644 --- a/crates/aisix-proxy/src/semantic.rs +++ b/crates/aisix-proxy/src/semantic.rs @@ -322,6 +322,8 @@ fn attempt_for_target(snapshot: &AisixSnapshot, alias: &str) -> Result { - let app: SpawnedApp | undefined; - let seed: SeedClient | undefined; - let etcdReachable = false; - const upstreams: OpenAiUpstream[] = []; - - beforeAll(async () => { - const etcd = new EtcdClient(); - etcdReachable = await etcd.ping(); - if (!etcdReachable) return; - - app = await spawnApp(); - seed = new SeedClient(etcd, app.etcdPrefix); - await seed.createApiKey({ - key_hash: CALLER_KEY_HASH, - allowed_models: ["*"], - }); - }); - - afterAll(async () => { - await app?.exit(); - await Promise.all(upstreams.map((u) => u.close())); - }); - - async function createOpenAiModel( - displayName: string, - upstream: OpenAiUpstream, - ): Promise { - if (!seed) throw new Error("seed client not initialized"); - const providerKey = await seed.createProviderKey({ - display_name: `${displayName}-pk`, - api_key: "sk-mock", - api_base: `${upstream.baseUrl}/v1`, - }); - await seed.createModel({ - display_name: displayName, - provider: "openai", - model_name: "gpt-4o-mini", - provider_key_id: providerKey.id, - }); - } - - function client(): OpenAI { - return new OpenAI({ - apiKey: CALLER_PLAINTEXT, - baseURL: `${app?.proxyUrl}/v1`, - maxRetries: 0, - }); - } - - async function askWithKey(key: string): Promise { - const completion = await client().chat.completions.create( - { model: "canary-router", messages: [{ role: "user", content: "hi" }] }, - { headers: { "x-aisix-routing-key": key } }, - ); - return completion.choices[0]?.message.content ?? null; - } - - test("pins a stability key to one target while splitting across keys", async (ctx) => { - if (!etcdReachable || !app || !seed) { - ctx.skip(); - return; - } - - const stable = await startOpenAiUpstream({ nonStreamBody: okBody("stable-served") }); - const canary = await startOpenAiUpstream({ nonStreamBody: okBody("canary-served") }); - upstreams.push(stable, canary); - // Router BEFORE its targets: watch events apply in revision order, so - // once /v1/models lists both targets the router is in the snapshot too. - await seed.createModel({ - display_name: "canary-router", - routing: { - strategy: "weighted", - sticky: true, - targets: [ - { model: "canary-stable", weight: 50 }, - { model: "canary-new", weight: 50 }, - ], - }, - }); - await createOpenAiModel("canary-stable", stable); - await createOpenAiModel("canary-new", canary); - - // Gate on the DP snapshot via /v1/models — authenticates only once the - // caller key has propagated, lists the targets only once the snapshot - // has them, and dispatches to no target (which would skew the - // per-target counts below). - await waitConfigPropagation(async () => { - const res = await fetch(`${app!.proxyUrl}/v1/models`, { - headers: { authorization: `Bearer ${CALLER_PLAINTEXT}` }, - }); - if (res.status !== 200) return false; - const ids = ((await res.json()) as { data?: Array<{ id?: string }> }).data?.map((m) => m.id) ?? []; - return ids.includes("canary-stable") && ids.includes("canary-new"); - }); - - // Same key → same target on every request (sticky). - const repeated = await Promise.all( - Array.from({ length: 6 }, () => askWithKey("user-A")), - ); - expect(new Set(repeated).size).toBe(1); - - // Distinct keys spread across both targets (the split is honored, not a - // single funnel). Deterministic hashing keeps this stable across runs. - const served = new Set(); - for (let i = 0; i < 32; i++) { - served.add(await askWithKey(`user-${i}`)); - } - expect(served).toEqual(new Set(["stable-served", "canary-served"])); - }); -}); diff --git a/tests/e2e/src/cases/consistent-hash-routing-e2e.test.ts b/tests/e2e/src/cases/consistent-hash-routing-e2e.test.ts new file mode 100644 index 00000000..a65a6abc --- /dev/null +++ b/tests/e2e/src/cases/consistent-hash-routing-e2e.test.ts @@ -0,0 +1,418 @@ +import { createHash } from "node:crypto"; +import OpenAI from "openai"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + EtcdClient, + SeedClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +// E2E: `strategy: consistent_hash` + per-target `priority` tiers — the +// AISIX-Cloud#1206 two-pool shape. Contracts pinned here: +// +// 1. The same hash key keeps landing on the same target; distinct keys +// spread across the tier (ketama ring, weight-scaled). +// 2. The `hash_on` chain is honored in order (cookie first here), with +// the caller's API key as the configured fallback. +// 3. Priority tiers: the backup tier receives ZERO traffic while the +// active tier has a healthy target; a fully-down active tier shifts +// traffic to the backup within the SAME request (in-request spill); +// an active target leaving cooldown takes its traffic back. +// 4. A single failed member redistributes within its own tier — never +// to the backup tier. +// +// Reference: OpenAI Chat Completions shape the caller sees +// (https://platform.openai.com/docs/api-reference/chat). + +const CALLER_PLAINTEXT = "sk-chash-routing-e2e-caller"; +const CALLER_KEY_HASH = createHash("sha256") + .update(CALLER_PLAINTEXT) + .digest("hex"); + +function okBody(content: string) { + return { + id: `cmpl-${content}`, + object: "chat.completion", + created: Math.floor(Date.now() / 1000), + model: "gpt-4o-mini", + choices: [ + { + index: 0, + message: { role: "assistant", content }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }; +} + +const err503 = { + status: 503, + errorBody: { error: { message: "instance down", type: "server_error" } }, +}; + +describe("consistent-hash routing + priority tiers e2e", () => { + let app: SpawnedApp | undefined; + let seed: SeedClient | undefined; + let etcdReachable = false; + const upstreams: OpenAiUpstream[] = []; + + beforeAll(async () => { + const etcd = new EtcdClient(); + etcdReachable = await etcd.ping(); + if (!etcdReachable) return; + + app = await spawnApp(); + seed = new SeedClient(etcd, app.etcdPrefix); + await seed.createApiKey({ + key_hash: CALLER_KEY_HASH, + allowed_models: ["*"], + }); + }); + + afterAll(async () => { + await app?.exit(); + await Promise.all(upstreams.map((u) => u.close())); + }); + + async function createMember( + displayName: string, + upstream: OpenAiUpstream, + extra: Record = {}, + ): Promise { + if (!seed) throw new Error("seed client not initialized"); + const providerKey = await seed.createProviderKey({ + display_name: `${displayName}-pk`, + api_key: "sk-mock", + api_base: `${upstream.baseUrl}/v1`, + }); + await seed.createModel({ + display_name: displayName, + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: providerKey.id, + ...extra, + }); + } + + function client(): OpenAI { + return new OpenAI({ + apiKey: CALLER_PLAINTEXT, + baseURL: `${app?.proxyUrl}/v1`, + maxRetries: 0, + }); + } + + async function ask( + model: string, + headers: Record, + ): Promise { + const completion = await client().chat.completions.create( + { model, messages: [{ role: "user", content: "hi" }] }, + { headers }, + ); + return completion.choices[0]?.message.content ?? null; + } + + // Gate on the DP snapshot via /v1/models — authenticates only once the + // caller key has propagated, lists the members only once the snapshot + // has them, and dispatches to no target (which would warm cooldowns and + // skew the per-target counts the assertions rely on). + async function waitMembersVisible(members: string[]): Promise { + await waitConfigPropagation(async () => { + const res = await fetch(`${app!.proxyUrl}/v1/models`, { + headers: { authorization: `Bearer ${CALLER_PLAINTEXT}` }, + }); + if (res.status !== 200) return false; + const ids = + ((await res.json()) as { data?: Array<{ id?: string }> }).data?.map( + (m) => m.id, + ) ?? []; + return members.every((m) => ids.includes(m)); + }); + } + + test("same key sticks to one target while distinct keys spread", async (ctx) => { + if (!etcdReachable || !app || !seed) { + ctx.skip(); + return; + } + + const one = await startOpenAiUpstream({ nonStreamBody: okBody("one-served") }); + const two = await startOpenAiUpstream({ nonStreamBody: okBody("two-served") }); + upstreams.push(one, two); + // Router BEFORE its targets: watch events apply in revision order, so + // once /v1/models lists both targets the router is in the snapshot too. + await seed.createModel({ + display_name: "ch-basic", + routing: { + strategy: "consistent_hash", + targets: [ + { model: "ch-basic-one", weight: 50 }, + { model: "ch-basic-two", weight: 50 }, + ], + }, + }); + await createMember("ch-basic-one", one); + await createMember("ch-basic-two", two); + await waitMembersVisible(["ch-basic-one", "ch-basic-two"]); + + // Same key → same target on every request. + const repeated = await Promise.all( + Array.from({ length: 6 }, () => + ask("ch-basic", { "x-aisix-routing-key": "user-A" }), + ), + ); + expect(new Set(repeated).size).toBe(1); + + // Distinct keys spread across both targets. Deterministic hashing + // keeps this stable across runs. + const served = new Set(); + for (let i = 0; i < 32; i++) { + served.add(await ask("ch-basic", { "x-aisix-routing-key": `user-${i}` })); + } + expect(served).toEqual(new Set(["one-served", "two-served"])); + }); + + test("hash_on chain reads the cookie first and falls back to the api key", async (ctx) => { + if (!etcdReachable || !app || !seed) { + ctx.skip(); + return; + } + + const one = await startOpenAiUpstream({ nonStreamBody: okBody("ck-one") }); + const two = await startOpenAiUpstream({ nonStreamBody: okBody("ck-two") }); + upstreams.push(one, two); + await seed.createModel({ + display_name: "ch-cookie", + routing: { + strategy: "consistent_hash", + hash_on: [ + { type: "cookie", name: "sid" }, + { type: "api_key" }, + ], + targets: [{ model: "ch-cookie-one" }, { model: "ch-cookie-two" }], + }, + }); + await createMember("ch-cookie-one", one); + await createMember("ch-cookie-two", two); + await waitMembersVisible(["ch-cookie-one", "ch-cookie-two"]); + + // A cookie-keyed session is stable across requests… + const viaCookie = await Promise.all( + Array.from({ length: 5 }, () => + ask("ch-cookie", { cookie: "theme=dark; sid=sess-42" }), + ), + ); + expect(new Set(viaCookie).size).toBe(1); + + // …and distinct cookie values are what spreads the traffic — proving + // the cookie (not the shared caller key) is the operative source. + const spread = new Set(); + for (let i = 0; i < 32; i++) { + spread.add(await ask("ch-cookie", { cookie: `sid=sess-${i}` })); + } + expect(spread).toEqual(new Set(["ck-one", "ck-two"])); + + // Without the cookie every request falls back to the caller's API + // key — one shared key, one consistent target. + const viaApiKey = await Promise.all( + Array.from({ length: 5 }, () => ask("ch-cookie", {})), + ); + expect(new Set(viaApiKey).size).toBe(1); + }); + + test("backup tier idles while the active tier is healthy", async (ctx) => { + if (!etcdReachable || !app || !seed) { + ctx.skip(); + return; + } + + const a1 = await startOpenAiUpstream({ nonStreamBody: okBody("a1-served") }); + const a2 = await startOpenAiUpstream({ nonStreamBody: okBody("a2-served") }); + const b1 = await startOpenAiUpstream({ nonStreamBody: okBody("b1-served") }); + upstreams.push(a1, a2, b1); + await seed.createModel({ + display_name: "ch-pools", + routing: { + strategy: "consistent_hash", + targets: [ + { model: "ch-pools-a1" }, + { model: "ch-pools-a2" }, + { model: "ch-pools-b1", priority: -1 }, + ], + }, + }); + await createMember("ch-pools-a1", a1); + await createMember("ch-pools-a2", a2); + await createMember("ch-pools-b1", b1); + await waitMembersVisible(["ch-pools-a1", "ch-pools-a2", "ch-pools-b1"]); + + const b1Baseline = b1.receivedRequests.length; + const served = new Set(); + for (let i = 0; i < 16; i++) { + served.add(await ask("ch-pools", { "x-aisix-routing-key": `user-${i}` })); + } + // Every response came from the active tier, and the backup upstream + // never saw a single request. + expect([...served].every((s) => s === "a1-served" || s === "a2-served")).toBe( + true, + ); + expect(served.size).toBe(2); // both active members participate + expect(b1.receivedRequests.length - b1Baseline).toBe(0); + }); + + test("a fully-down active tier spills to the backup tier within one request, with hash affinity there", async (ctx) => { + if (!etcdReachable || !app || !seed) { + ctx.skip(); + return; + } + + const a1 = await startOpenAiUpstream(err503); + const a2 = await startOpenAiUpstream(err503); + const b1 = await startOpenAiUpstream({ nonStreamBody: okBody("bk-one") }); + const b2 = await startOpenAiUpstream({ nonStreamBody: okBody("bk-two") }); + upstreams.push(a1, a2, b1, b2); + await seed.createModel({ + display_name: "ch-down", + routing: { + strategy: "consistent_hash", + targets: [ + { model: "ch-down-a1" }, + { model: "ch-down-a2" }, + { model: "ch-down-b1", priority: -1 }, + { model: "ch-down-b2", priority: -1 }, + ], + }, + }); + await createMember("ch-down-a1", a1); + await createMember("ch-down-a2", a2); + await createMember("ch-down-b1", b1); + await createMember("ch-down-b2", b2); + await waitMembersVisible([ + "ch-down-a1", + "ch-down-a2", + "ch-down-b1", + "ch-down-b2", + ]); + + // The FIRST request discovers both active members down and still + // succeeds — the walk crosses the tier boundary inside one request. + const first = await ask("ch-down", { "x-aisix-routing-key": "user-A" }); + expect(first === "bk-one" || first === "bk-two").toBe(true); + expect( + a1.receivedRequests.length + a2.receivedRequests.length, + ).toBeGreaterThan(0); + + // While the active tier cools down, the same key stays on the same + // backup target (hash affinity holds inside the backup tier too). + const repeated = await Promise.all( + Array.from({ length: 6 }, () => + ask("ch-down", { "x-aisix-routing-key": "user-A" }), + ), + ); + expect(new Set(repeated)).toEqual(new Set([first])); + + // Distinct keys spread across BOTH backup members. + const served = new Set(); + for (let i = 0; i < 32; i++) { + served.add(await ask("ch-down", { "x-aisix-routing-key": `u-${i}` })); + } + expect(served).toEqual(new Set(["bk-one", "bk-two"])); + }); + + test("an active target leaving cooldown takes its traffic back from the backup", async (ctx) => { + if (!etcdReachable || !app || !seed) { + ctx.skip(); + return; + } + + // The active member fails exactly once, then serves; its cooldown is + // 1s so recovery is observable within the test. + const a1 = await startOpenAiUpstream({ + scriptedResponses: [{ ...err503 }], + nonStreamBody: okBody("active-served"), + }); + const b1 = await startOpenAiUpstream({ nonStreamBody: okBody("backup-served") }); + upstreams.push(a1, b1); + await seed.createModel({ + display_name: "ch-recover", + routing: { + strategy: "consistent_hash", + targets: [ + { model: "ch-recover-a1" }, + { model: "ch-recover-b1", priority: -1 }, + ], + }, + }); + await createMember("ch-recover-a1", a1, { + cooldown: { default_seconds: 1 }, + }); + await createMember("ch-recover-b1", b1); + await waitMembersVisible(["ch-recover-a1", "ch-recover-b1"]); + + // 1st request: the active member fails (scripted 503), the backup + // absorbs it in-request. + expect(await ask("ch-recover", { "x-aisix-routing-key": "s1" })).toBe( + "backup-served", + ); + const a1AfterFailure = a1.receivedRequests.length; + + // While the active member cools down, traffic goes STRAIGHT to the + // backup — the cooled member is not even attempted. + expect(await ask("ch-recover", { "x-aisix-routing-key": "s1" })).toBe( + "backup-served", + ); + expect(a1.receivedRequests.length).toBe(a1AfterFailure); + + // After the 1s cooldown expires the recovered member takes back its + // traffic (poll rather than sleep a fixed amount — config watches and + // timers make exact timing environment-dependent). + await waitConfigPropagation(async () => { + return (await ask("ch-recover", { "x-aisix-routing-key": "s1" })) === + "active-served"; + }); + }); + + test("a single failed member redistributes within its tier, never to the backup", async (ctx) => { + if (!etcdReachable || !app || !seed) { + ctx.skip(); + return; + } + + const a1 = await startOpenAiUpstream(err503); + const a2 = await startOpenAiUpstream({ nonStreamBody: okBody("peer-served") }); + const b1 = await startOpenAiUpstream({ nonStreamBody: okBody("backup-served") }); + upstreams.push(a1, a2, b1); + await seed.createModel({ + display_name: "ch-partial", + routing: { + strategy: "consistent_hash", + targets: [ + { model: "ch-partial-a1" }, + { model: "ch-partial-a2" }, + { model: "ch-partial-b1", priority: -1 }, + ], + }, + }); + await createMember("ch-partial-a1", a1); + await createMember("ch-partial-a2", a2); + await createMember("ch-partial-b1", b1); + await waitMembersVisible(["ch-partial-a1", "ch-partial-a2", "ch-partial-b1"]); + + const b1Baseline = b1.receivedRequests.length; + for (let i = 0; i < 12; i++) { + // Keys whose first choice is the dead member fail over to its tier + // peer; keys mapped to the healthy peer are untouched. Either way + // the answer comes from inside the active tier. + expect(await ask("ch-partial", { "x-aisix-routing-key": `k-${i}` })).toBe( + "peer-served", + ); + } + expect(b1.receivedRequests.length - b1Baseline).toBe(0); + }); +}); diff --git a/tests/e2e/src/cases/weighted-routing-edit-e2e.test.ts b/tests/e2e/src/cases/routing-priority-edit-e2e.test.ts similarity index 73% rename from tests/e2e/src/cases/weighted-routing-edit-e2e.test.ts rename to tests/e2e/src/cases/routing-priority-edit-e2e.test.ts index 4fb16d44..0fd04095 100644 --- a/tests/e2e/src/cases/weighted-routing-edit-e2e.test.ts +++ b/tests/e2e/src/cases/routing-priority-edit-e2e.test.ts @@ -11,24 +11,25 @@ import { type SpawnedApp, } from "../harness/index.js"; -// E2E: a LIVE edit to a weighted routing model's weights re-takes +// E2E: a LIVE edit to a routing model's target priorities re-takes // effect on the dispatch path (#196 L1, ai-gateway #127 L1). // // The sibling weighted-routing-distribution-e2e pins that the INITIAL // weights are honored. The gap this closes: after the model is live -// and serving, an operator rewrites the weights on the stored model, and -// the change must propagate through the etcd watch and the weighted -// scheduler must REBUILD — a scheduler that cached its weight wheel on -// first dispatch and never rebuilt on config update would keep serving -// the old split, silently ignoring the operator's change. +// and serving, an operator rewrites the target priorities on the stored +// model, and the change must propagate through the etcd watch and the +// scheduler must REBUILD — a scheduler that cached its tier partition +// (or WRR wheel) on first dispatch and never rebuilt on config update +// would keep serving the old layout, silently ignoring the operator's +// change. // -// Design is deterministic (no statistics): weight 0 = excluded (see -// routing-strategies-e2e "weighted picks the positive-weight target"). -// - Start [wr-edit-a: 100, wr-edit-b: 0] → every dispatch hits A. -// - Edit to [wr-edit-a: 0, wr-edit-b: 100] → every dispatch hits B. +// Design is deterministic (no statistics): the active tier takes ALL +// traffic while it is healthy, the backup tier none (AISIX-Cloud#1206). +// - Start [wr-edit-a: priority 0, wr-edit-b: priority -1] → all A. +// - Edit to [wr-edit-a: priority -1, wr-edit-b: priority 0] → all B. // The propagation signal is unambiguous: a probe through the virtual -// model returning "served by B" is IMPOSSIBLE under the old [100,0] -// config, so it proves the edit is live before we count. +// model returning "served by B" is IMPOSSIBLE under the old layout, so +// it proves the edit is live before we count. // // Reference: OpenAI Chat Completions shape the caller sees // (https://platform.openai.com/docs/api-reference/chat). @@ -53,7 +54,7 @@ function upstreamBody(content: string, id: string): Record { }; } -describe("weighted routing live-edit: changing weights shifts real traffic (#196 L1)", () => { +describe("routing live-edit: swapping target priorities shifts real traffic (#196 L1)", () => { let app: SpawnedApp | undefined; let upstreamA: OpenAiUpstream | undefined; let upstreamB: OpenAiUpstream | undefined; @@ -98,15 +99,15 @@ describe("weighted routing live-edit: changing weights shifts real traffic (#196 model_name: "gpt-4o-mini", provider_key_id: pkB.id, }); - // Virtual model: weighted, ALL traffic to A initially (B excluded - // via weight 0). Capture the generated id so we can PUT it below. + // Virtual model: round_robin with B parked in a backup tier — ALL + // traffic to A initially. Capture the generated id for the PUT below. const virtual = await seed.createModel({ display_name: "wr-edit-virtual", routing: { - strategy: "weighted", + strategy: "round_robin", targets: [ - { model: "wr-edit-a", weight: 100 }, - { model: "wr-edit-b", weight: 0 }, + { model: "wr-edit-a" }, + { model: "wr-edit-b", priority: -1 }, ], }, }); @@ -124,7 +125,7 @@ describe("weighted routing live-edit: changing weights shifts real traffic (#196 await upstreamB?.close(); }); - test("editing weights [100,0] → [0,100] flips the served upstream", async (ctx) => { + test("swapping tier priorities flips the served upstream", async (ctx) => { if (!etcdReachable || !app || !upstreamA || !upstreamB || !seed || !virtualId) { ctx.skip(); return; @@ -149,7 +150,7 @@ describe("weighted routing live-edit: changing weights shifts real traffic (#196 }; // Readiness: both leaves registered, then the virtual model serving - // A under the initial [100,0] weights. + // A under the initial tier layout (A active, B backup). await waitConfigPropagation(async () => { try { const a = await client.chat.completions.create({ @@ -174,7 +175,7 @@ describe("weighted routing live-edit: changing weights shifts real traffic (#196 }); await waitConfigPropagation(async () => (await callVirtual("ready-virtual")) === "served by A"); - // --- Phase 1: under [100,0], every dispatch must hit A. --- + // --- Phase 1: A is the active tier — every dispatch must hit A. --- const aBase1 = upstreamA.receivedRequests.length; const bBase1 = upstreamB.receivedRequests.length; for (let i = 0; i < BATCH; i++) { @@ -183,25 +184,26 @@ describe("weighted routing live-edit: changing weights shifts real traffic (#196 expect(upstreamA.receivedRequests.length - aBase1).toBe(BATCH); expect(upstreamB.receivedRequests.length - bBase1).toBe(0); - // --- Edit: invert the weights to [0,100] by rewriting the document. --- + // --- Edit: swap the tiers by rewriting the document. --- await seed.update("models", virtualId, { display_name: "wr-edit-virtual", routing: { - strategy: "weighted", + strategy: "round_robin", targets: [ - { model: "wr-edit-a", weight: 0 }, - { model: "wr-edit-b", weight: 100 }, + { model: "wr-edit-a", priority: -1 }, + { model: "wr-edit-b" }, ], }, }); // Propagation signal: a virtual dispatch returning "served by B" is - // impossible under the old [100,0] config, so it proves the edit is - // live + the scheduler rebuilt. If the scheduler never rebuilds on a - // config edit (the regression this test targets), this times out. + // impossible under the old tier layout (B was backup behind a healthy + // A), so it proves the edit is live + the scheduler rebuilt. If the + // scheduler never rebuilds on a config edit (the regression this test + // targets), this times out. await waitConfigPropagation(async () => (await callVirtual("post-edit-probe")) === "served by B"); - // --- Phase 2: under [0,100], every dispatch must hit B. --- + // --- Phase 2: after the swap B is the active tier — every dispatch must hit B. --- const aBase2 = upstreamA.receivedRequests.length; const bBase2 = upstreamB.receivedRequests.length; for (let i = 0; i < BATCH; i++) { diff --git a/tests/e2e/src/cases/routing-strategies-e2e.test.ts b/tests/e2e/src/cases/routing-strategies-e2e.test.ts index 615ca3a7..bc607b78 100644 --- a/tests/e2e/src/cases/routing-strategies-e2e.test.ts +++ b/tests/e2e/src/cases/routing-strategies-e2e.test.ts @@ -350,63 +350,66 @@ describe("routing strategies and retry behavior e2e", () => { expect(second.receivedRequests.length - secondBaseline).toBe(2); }); - test("weighted picks the positive-weight target first and falls forward from there", async (ctx) => { + test("priority tiers: the active tier is tried first, the backup tier absorbs its failure", async (ctx) => { if (!etcdReachable || !app || !seed) { ctx.skip(); return; } - const zeroWeightBefore = await startOpenAiUpstream({ + // Active tier (priority 0) is down; first backup tier (priority -1) + // serves; a second backup (priority -2) exists but max_fallbacks: 1 + // caps the walk before it — pinning both the tier order and the cap. + const primaryDown = await startOpenAiUpstream({ + status: 503, + errorBody: { error: { message: "active tier down", type: "server_error" } }, + }); + const backup = await startOpenAiUpstream({ nonStreamBody: { - id: "cmpl-routing-weighted-before", + id: "cmpl-routing-priority-backup", object: "chat.completion", created: Math.floor(Date.now() / 1000), model: "gpt-4o-mini", choices: [ { index: 0, - message: { role: "assistant", content: "should-not-run" }, + message: { role: "assistant", content: "backup tier served" }, finish_reason: "stop", }, ], usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, }, }); - const weightedPrimary = await startOpenAiUpstream({ - status: 503, - errorBody: { error: { message: "weighted primary down", type: "server_error" } }, - }); - const forwardFallback = await startOpenAiUpstream({ + const lastResort = await startOpenAiUpstream({ nonStreamBody: { - id: "cmpl-routing-weighted-fallback", + id: "cmpl-routing-priority-last", object: "chat.completion", created: Math.floor(Date.now() / 1000), model: "gpt-4o-mini", choices: [ { index: 0, - message: { role: "assistant", content: "weighted fallback worked" }, + message: { role: "assistant", content: "should-not-run" }, finish_reason: "stop", }, ], usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, }, }); - upstreams.push(zeroWeightBefore, weightedPrimary, forwardFallback); + upstreams.push(primaryDown, backup, lastResort); - await createOpenAiModel("routing-weighted-before", zeroWeightBefore); - await createOpenAiModel("routing-weighted-primary", weightedPrimary); - await createOpenAiModel("routing-weighted-fallback", forwardFallback); - await waitUntilModelResponds("routing-weighted-before", "should-not-run"); - await waitUntilModelResponds("routing-weighted-fallback", "weighted fallback worked"); + await createOpenAiModel("routing-priority-primary", primaryDown); + await createOpenAiModel("routing-priority-backup", backup); + await createOpenAiModel("routing-priority-last", lastResort); + await waitUntilModelResponds("routing-priority-backup", "backup tier served"); + await waitUntilModelResponds("routing-priority-last", "should-not-run"); await seed.createModel({ - display_name: "routing-weighted-virtual", + display_name: "routing-priority-virtual", routing: { - strategy: "weighted", + strategy: "failover", targets: [ - { model: "routing-weighted-before", weight: 0 }, - { model: "routing-weighted-primary", weight: 1 }, - { model: "routing-weighted-fallback", weight: 0 }, + { model: "routing-priority-last", priority: -2 }, + { model: "routing-priority-primary" }, + { model: "routing-priority-backup", priority: -1 }, ], max_fallbacks: 1, }, @@ -418,25 +421,26 @@ describe("routing strategies and retry behavior e2e", () => { maxRetries: 0, }); - // Gate on DP-snapshot presence rather than probing the - // virtual — probe would warm the weighted primary's 502 cooldown - // and skew per-target hit counts. Both direct models' readiness - // was already established above; this confirms the routing record - // has reached the DP snapshot. - await waitSeedApplied("routing-weighted"); + // Gate on DP-snapshot presence rather than probing the virtual — + // a probe would warm the primary's cooldown and skew per-target + // hit counts. Both healthy models' readiness was established + // above; this confirms the routing record reached the DP snapshot. + await waitSeedApplied("routing-priority"); - const beforeBaseline = zeroWeightBefore.receivedRequests.length; - const primaryBaseline = weightedPrimary.receivedRequests.length; - const fallbackBaseline = forwardFallback.receivedRequests.length; + const lastBaseline = lastResort.receivedRequests.length; + const primaryBaseline = primaryDown.receivedRequests.length; + const backupBaseline = backup.receivedRequests.length; const completion = await client.chat.completions.create({ - model: "routing-weighted-virtual", - messages: [{ role: "user", content: "weighted routing request" }], + model: "routing-priority-virtual", + messages: [{ role: "user", content: "priority routing request" }], }); - expect(completion.choices[0]?.message.content).toBe("weighted fallback worked"); - expect(zeroWeightBefore.receivedRequests.length - beforeBaseline).toBe(0); - expect(weightedPrimary.receivedRequests.length - primaryBaseline).toBe(1); - expect(forwardFallback.receivedRequests.length - fallbackBaseline).toBe(1); + // Declaration order puts the priority -2 target FIRST — tier order + // must out-rank declaration order for this to pass. + expect(completion.choices[0]?.message.content).toBe("backup tier served"); + expect(primaryDown.receivedRequests.length - primaryBaseline).toBe(1); + expect(backup.receivedRequests.length - backupBaseline).toBe(1); + expect(lastResort.receivedRequests.length - lastBaseline).toBe(0); }); }); diff --git a/tests/e2e/src/cases/weighted-routing-distribution-e2e.test.ts b/tests/e2e/src/cases/weighted-routing-distribution-e2e.test.ts index f9809545..657d9aa7 100644 --- a/tests/e2e/src/cases/weighted-routing-distribution-e2e.test.ts +++ b/tests/e2e/src/cases/weighted-routing-distribution-e2e.test.ts @@ -11,35 +11,24 @@ import { type SpawnedApp, } from "../harness/index.js"; -// E2E: weighted routing distribution. A virtual Model carries a -// Routing block with `strategy: "weighted"` and two targets — `wr-a` -// (weight 70) and `wr-b` (weight 30). Per docs `api-admin.md` §4.1 -// (direct vs routing two-mode split) and the routing schema enum -// `["round_robin", "weighted", "failover"]`, the gateway is expected -// to dispatch incoming traffic in a 70:30 ratio across the two -// targets. +// E2E: weighted round-robin distribution. A virtual Model carries a +// Routing block with `strategy: "round_robin"` and two targets — `wr-a` +// (weight 70) and `wr-b` (weight 30). round_robin is smooth WEIGHTED +// round-robin (the nginx algorithm; AISIX-Cloud#1206 merged the former +// `weighted` strategy into it), so the gateway must dispatch traffic in +// an exact 70:30 ratio across the two targets. // // One contract pinned here: // -// - Weighted strategy honours the integer `weight` field per -// target. After N requests the observed split lands inside a -// statistically reasonable tolerance window around the declared -// ratio. A regression that ignored `weight` and round-robined -// instead would fail (each side would land ~50%, well outside -// [55, 85] / [15, 45]). +// - round_robin honours the integer `weight` field per target, and is +// deterministic: smooth WRR is periodic with period = total weight +// (100 here), and ANY window of one full period contains each +// target exactly `weight` times — so 100 sequential requests land +// at exactly 70/30 regardless of how many warm-up probes ran +// before the counted batch. // // Reference: OpenAI Chat Completions API spec for the shape the // caller sees (https://platform.openai.com/docs/api-reference/chat). -// -// The 100-request count and the [55, 85] / [15, 45] tolerance are -// chosen so a scheduler that completely ignores weight (e.g. -// round-robins or pins to one target) cannot pass — round-robin -// lands at 50/50, well outside [55, 85] for the heavy side — while -// the legitimate 70/30 path stays comfortably inside. Two -// independent binomial windows: 70±15 over n=100 with σ≈4.58 puts -// the gate at ~3.3σ, P(false positive) ≈ 0.1%. The previous ±10 -// gate sat at ~2.2σ (≈2.8%) and tripped roughly once per ~36 CI -// runs — wide enough to be a steady CI flake. const CALLER_PLAINTEXT = "sk-wr-e2e-caller"; const CALLER_KEY_HASH = createHash("sha256") @@ -49,14 +38,8 @@ const CALLER_KEY_HASH = createHash("sha256") const TOTAL_REQUESTS = 100; const HEAVY_WEIGHT = 70; const LIGHT_WEIGHT = 30; -// Tolerance: weight ±15 absolute on a 100-sample window. See header -// comment for the statistical-power tradeoff vs the previous ±10. -const HEAVY_LO = 55; -const HEAVY_HI = 85; -const LIGHT_LO = 15; -const LIGHT_HI = 45; - -describe("weighted routing distribution e2e: 70/30 split lands inside [55,85] / [15,45]", () => { + +describe("weighted round-robin distribution e2e: 70/30 split is exact over one WRR period", () => { let app: SpawnedApp | undefined; let upstreamA: OpenAiUpstream | undefined; let upstreamB: OpenAiUpstream | undefined; @@ -129,13 +112,12 @@ describe("weighted routing distribution e2e: 70/30 split lands inside [55,85] / model_name: "gpt-4o-mini", provider_key_id: pkB.id, }); - // Virtual Model: routing-only, weighted strategy. Per the schema - // enum the gateway publishes (round_robin / weighted / failover), - // `weighted` should honour each target's `weight` integer. + // Virtual Model: routing-only, weighted round-robin. round_robin + // honours each target's `weight` integer exactly (smooth WRR). await seed.createModel({ display_name: "wr-virtual", routing: { - strategy: "weighted", + strategy: "round_robin", targets: [ { model: "wr-a", weight: HEAVY_WEIGHT }, { model: "wr-b", weight: LIGHT_WEIGHT }, @@ -143,7 +125,7 @@ describe("weighted routing distribution e2e: 70/30 split lands inside [55,85] / }, }); // Caller is allowed all three Models so the readiness probes can - // hit the leaves directly without firing the weighted dispatcher. + // hit the leaves directly without firing the WRR dispatcher. await seed.createApiKey({ key_hash: CALLER_KEY_HASH, allowed_models: ["wr-virtual", "wr-a", "wr-b"], @@ -171,7 +153,7 @@ describe("weighted routing distribution e2e: 70/30 split lands inside [55,85] / // Two-stage readiness gate: probe each leaf Model directly so // both ProviderKey registrations are observed by the proxy // before we exercise the virtual router. Probing through - // `wr-virtual` here would fire the weighted dispatcher and + // `wr-virtual` here would fire the WRR dispatcher and // pollute the count baseline. await waitConfigPropagation(async () => { try { @@ -195,10 +177,9 @@ describe("weighted routing distribution e2e: 70/30 split lands inside [55,85] / return false; } }); - // One probe through the virtual Model so the weighted - // dispatcher's lazy state (if any — schedulers often build the - // weight wheel on first dispatch) is constructed before we start - // counting. + // One probe through the virtual Model so the WRR dispatcher's + // lazy state is constructed before we start counting. Periodicity + // makes the exact assertion below offset-independent. await waitConfigPropagation(async () => { try { const probe = await client.chat.completions.create({ @@ -237,13 +218,13 @@ describe("weighted routing distribution e2e: 70/30 split lands inside [55,85] / // upstreams could still appear "balanced" by ratio. expect(aDelta + bDelta).toBe(TOTAL_REQUESTS); - // Distribution assertion: heavy side ~70, light side ~30, both - // inside ±15. A round-robin regression (50/50) fails both gates; - // a pin-to-one regression (100/0) fails both gates. - expect(aDelta).toBeGreaterThanOrEqual(HEAVY_LO); - expect(aDelta).toBeLessThanOrEqual(HEAVY_HI); - expect(bDelta).toBeGreaterThanOrEqual(LIGHT_LO); - expect(bDelta).toBeLessThanOrEqual(LIGHT_HI); + // Distribution assertion: smooth WRR is exact over one full period + // (total weight = 100 = TOTAL_REQUESTS), at any starting offset. An + // equal-rotation regression lands 50/50; a pin-to-one regression + // lands 100/0; a random-sampling regression flakes around 70 — all + // fail an exact gate. + expect(aDelta).toBe(HEAVY_WEIGHT); + expect(bDelta).toBe(LIGHT_WEIGHT); // Per-test timeout lifted to 90s. The default suite timeout // (60s, vitest.config.ts) is tight for 100 sequential round-trips // when upstream latency drifts above ~50ms/call; 90s leaves