From 40277405694e33efeb06f72ae763b1bf73bd7773 Mon Sep 17 00:00:00 2001 From: Kan Wu Date: Sun, 20 Sep 2026 17:22:36 -0700 Subject: [PATCH] [sgl-router] refactor - session-aware policy (#40379) Co-authored-by: Claude Fable 5.1 --- experimental/sgl-router/POLICY_DESIGN.md | 28 +- .../sgl-router/src/policies_reorg/mod.rs | 1 + .../src/policies_reorg/session_aware.rs | 124 ++++++ .../sgl-router/tests/component/main.rs | 1 + .../component/policies_reorg_session_aware.rs | 410 ++++++++++++++++++ .../tests/proxy/chat_routing/reorg.rs | 2 + .../proxy/chat_routing/reorg/session_aware.rs | 119 +++++ 7 files changed, 681 insertions(+), 4 deletions(-) create mode 100644 experimental/sgl-router/src/policies_reorg/session_aware.rs create mode 100644 experimental/sgl-router/tests/component/policies_reorg_session_aware.rs create mode 100644 experimental/sgl-router/tests/proxy/chat_routing/reorg/session_aware.rs diff --git a/experimental/sgl-router/POLICY_DESIGN.md b/experimental/sgl-router/POLICY_DESIGN.md index 53cec7bc5..de1f79ad6 100644 --- a/experimental/sgl-router/POLICY_DESIGN.md +++ b/experimental/sgl-router/POLICY_DESIGN.md @@ -172,7 +172,7 @@ bucket, both groups share this one request-length decision. Policy fallback on a cache/affinity miss stays within that group's candidates. There is no second pass with relaxed admission and no post-policy substitution. -SLO ordering, cache lookup, and session/sticky policies are follow-ups. Their +SLO ordering, cache lookup, global session modes, and sticky policies are follow-ups. Their integration must preserve bucket-first selection and the same-bucket PD rule. Cross-bucket affinity probing is not part of this interface. Session/routing keys still pass through `PickRequest` for policies operating inside the selected @@ -294,6 +294,18 @@ separate mechanism. | `StickyPolicy` | Reuse an admitted routing-key binding; use the configured fallback for new or missing keys | | `CacheAwarePolicy` | Prefer a usable prefix under cache and pressure rules; use a load-based fallback on a miss | +Session assignments are scoped by model, bucket ID, stage, and session key. +`SessionAwarePolicy::new(store, engine_load)` receives shared state; the caller +owns the store's idle timeout and eviction task. Missing or empty session keys +use power-of-two without creating assignments. A new or out-of-group binding +uses power-of-two with `AllowAll`, then the session policy checks its selected +engine before binding. A concurrent live assignment wins, but is checked before +returning it; rejection ends that attempt without rewriting the binding or +retrying another engine. Existing bindings are reused regardless of pressure +when admitted. Session policies can be attached independently to each role. +Programmatic reorg callers configure `model.affinity.session_id_header` for HTTP +header extraction; this does not enable legacy global modes or backup escape. + Session and sticky policies do not create assignments for missing keys. A binding outside the candidates cannot win. A missing binding may invoke policy fallback within the group; hard admission rejection remains an error. @@ -557,8 +569,15 @@ Implemented here: - `AppContext::chat_routing` configures legacy versus reorg routing on the same endpoint and carries the reorg model-resolver map. -Follow-up order: concrete admission (#40271), then bucket SLO ordering -in a separate PR, followed by remaining policies and production configuration. +- `SessionAwarePolicy` reuses admitted model/bucket/role-scoped bindings from a + shared `AffinityStore`, falling back to power-of-two for new or keyless sessions. + Assignments follow admission; concurrent binding winners are rechecked. + Rejection preserves existing bindings and advances to the next bucket. + The caller owns expiry and sweeper lifecycle. A binding may remain after a + later PD group fails, because it records placement rather than dispatch. + +Follow-up order: cache-aware selection (#40366), concrete admission (#40271), +then bucket SLO ordering, remaining policies, and production configuration. Not yet implemented in the reorg path: @@ -566,7 +585,8 @@ Not yet implemented in the reorg path: - SLO estimates, targets, and bucket preference ordering. - CLI/configuration parsing, validation, and model-specific construction. The YAML above is illustrative; reorg resolvers are installed in code. -- Session modes, prefix memoization, and cache-aware selection. +- Global session modes and sticky routing-key affinity. +- Prefix memoization and cache-aware selection. - Shared load interpretation, dispatch correction, and policy-specific dispatch-timestamp requirements. - PD compatibility filtering, retry integration, and legacy-route switchover. diff --git a/experimental/sgl-router/src/policies_reorg/mod.rs b/experimental/sgl-router/src/policies_reorg/mod.rs index 34e5ca265..1fa209bdf 100644 --- a/experimental/sgl-router/src/policies_reorg/mod.rs +++ b/experimental/sgl-router/src/policies_reorg/mod.rs @@ -6,6 +6,7 @@ pub mod admission; pub mod power_of_two; +pub mod session_aware; use std::fmt::Debug; use std::sync::Arc; diff --git a/experimental/sgl-router/src/policies_reorg/session_aware.rs b/experimental/sgl-router/src/policies_reorg/session_aware.rs new file mode 100644 index 000000000..f56273e0d --- /dev/null +++ b/experimental/sgl-router/src/policies_reorg/session_aware.rs @@ -0,0 +1,124 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Bucket-scoped session placement. Admission rejection preserves the binding +//! and returns to the bucket loop; it never selects a backup inside the group. + +use std::sync::Arc; +use std::time::Instant; + +use futures::future::BoxFuture; + +use crate::state::load_monitor::engine_reported_load::{ + EngineReportedLoadSnapshot, EngineReportedLoadTable, +}; +use crate::state::AffinityStore; +use crate::workers::Worker; + +use super::admission::{AllowAll, Decision, EngineAdmission}; +use super::power_of_two::PowerOfTwoPolicy; +use super::{Pick, PickError, PickRequest, Policy, Rejection}; + +#[derive(Debug)] +pub struct SessionAwarePolicy { + store: Arc, + engine_load: Arc, + fallback: PowerOfTwoPolicy, + pub admission: Arc, +} + +impl SessionAwarePolicy { + /// The caller owns the shared store's idle timeout and sweeper lifecycle. + /// This policy implements bucket-scoped affinity, without legacy global modes + /// or primary/backup pressure escape. + pub fn new(store: Arc, engine_load: Arc) -> Self { + Self { + store, + fallback: PowerOfTwoPolicy::new(Arc::clone(&engine_load)), + engine_load, + admission: Arc::new(AllowAll), + } + } + + fn assignment_key(request: &PickRequest<'_>) -> Option { + let session = request.session_key.filter(|key| !key.is_empty())?; + // Length prefixes keep arbitrary model, bucket and session strings + // unambiguous, including embedded delimiters. Roles never share bindings. + Some(format!( + "session:{:?}:{}:{}{}:{}{}", + request.stage, + request.model.0.len(), + request.model.0, + request.bucket.len(), + request.bucket, + session + )) + } + + fn check( + &self, + engine: &Worker, + request: &PickRequest<'_>, + load: &EngineReportedLoadSnapshot, + ) -> Result<(), PickError> { + match self + .admission + .check(engine, request, load.fresh_load_for_url(&engine.url))? + { + Decision::Allow => Ok(()), + Decision::Reject(reason) => Err(PickError::AdmissionRejected(Rejection { + engine: engine.id.clone(), + reason, + })), + } + } +} + +impl Policy for SessionAwarePolicy { + fn pick<'a>( + &'a self, + engines: &'a [Arc], + request: &'a PickRequest<'a>, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + if engines.is_empty() { + return Err(PickError::NoCandidates); + } + let key = Self::assignment_key(request); + if let Some(bound) = key.as_ref().and_then(|key| self.store.bound(key, engines)) { + let load = self.engine_load.capture_snapshot(Instant::now()); + self.check(bound, request, &load)?; + return Ok(Pick { + engine: Arc::clone(bound), + reason: "session_primary", + }); + } + + // The nested power-of-two policy uses AllowAll. The session owner + // checks its chosen engine before creating or replacing a binding. + let mut pick = self.pick_fallback(engines, request).await?; + let load = self.engine_load.capture_snapshot(Instant::now()); + self.check(&pick.engine, request, &load)?; + let Some(key) = key else { + pick.reason = "no_session"; + return Ok(pick); + }; + + let effective = self.store.bind(key, &pick.engine, engines); + if !Arc::ptr_eq(effective, &pick.engine) { + // A racing first assignment wins. Check it once, without + // rewriting a rejected binding or retrying another engine. + self.check(effective, request, &load)?; + pick.reason = "session_primary"; + } else { + pick.reason = "assigned"; + } + pick.engine = Arc::clone(effective); + Ok(pick) + }) + } + + fn fallback(&self) -> Option<&dyn Policy> { + Some(&self.fallback) + } +} diff --git a/experimental/sgl-router/tests/component/main.rs b/experimental/sgl-router/tests/component/main.rs index 1a5accf8b..24d4dc09d 100644 --- a/experimental/sgl-router/tests/component/main.rs +++ b/experimental/sgl-router/tests/component/main.rs @@ -13,5 +13,6 @@ mod policies; mod policies_reorg; mod policies_reorg_load; mod policies_reorg_power_of_two; +mod policies_reorg_session_aware; mod tokenizer; mod workers; diff --git a/experimental/sgl-router/tests/component/policies_reorg_session_aware.rs b/experimental/sgl-router/tests/component/policies_reorg_session_aware.rs new file mode 100644 index 000000000..0cae2dabf --- /dev/null +++ b/experimental/sgl-router/tests/component/policies_reorg_session_aware.rs @@ -0,0 +1,410 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use sgl_router::discovery::{ModelId, WorkerId, WorkerSpec}; +use sgl_router::policies_reorg::admission::{Decision, EngineAdmission}; +use sgl_router::policies_reorg::session_aware::SessionAwarePolicy; +use sgl_router::policies_reorg::{PickError, PickRequest, Policy, Stage}; +use sgl_router::state::load_monitor::engine_reported_load::{ + EngineReportedLoadTable, EngineReportedWorkerLoad, LoadStat, +}; +use sgl_router::state::load_monitor::router_inflight_load::MockClock; +use sgl_router::state::AffinityStore; +use sgl_router::workers::Worker; + +fn engine(id: &str, active: usize) -> Arc { + let engine = Arc::new(Worker::new(WorkerSpec { + id: WorkerId(id.into()), + url: format!("http://{id}"), + mode: Stage::Plain, + model_ids: vec![ModelId("m".into())], + bootstrap_port: None, + })); + engine.active_requests.store(active, Ordering::Relaxed); + engine +} + +fn request(model: &ModelId) -> PickRequest<'_> { + PickRequest { + bucket: "bucket", + session_key: Some("session"), + ..PickRequest::new(model, Stage::Plain, 8) + } +} + +fn policy() -> (SessionAwarePolicy, Arc) { + let store = AffinityStore::new(Duration::from_secs(60)); + ( + SessionAwarePolicy::new(store.clone(), EngineReportedLoadTable::new()), + store, + ) +} + +#[derive(Debug, Default)] +struct Admission { + reject: AtomicBool, + invalid: AtomicBool, + calls: Mutex)>>, +} + +impl EngineAdmission for Admission { + fn check( + &self, + engine: &Worker, + _: &PickRequest<'_>, + load: Option<&EngineReportedWorkerLoad>, + ) -> Result { + self.calls + .lock() + .unwrap() + .push((engine.id.0.clone(), load.map(|load| load.num_waiting_reqs))); + if self.invalid.load(Ordering::Relaxed) { + return Err(PickError::InvalidSignal("invalid admission input".into())); + } + Ok(if self.reject.load(Ordering::Relaxed) { + Decision::Reject("full".into()) + } else { + Decision::Allow + }) + } +} + +#[tokio::test] +async fn new_sessions_choose_lower_pressure_then_reuse_binding() { + let (policy, store) = policy(); + let engines = [engine("a", 0), engine("b", 9)]; + let model = ModelId("m".into()); + let request = request(&model); + let first = policy.pick(&engines, &request).await.unwrap(); + assert_eq!(first.engine.id.0, "a"); + assert_eq!(first.reason, "assigned"); + engines[0].active_requests.store(100, Ordering::Relaxed); + for _ in 0..5 { + let pick = policy.pick(&engines, &request).await.unwrap(); + assert!(Arc::ptr_eq(&pick.engine, &engines[0])); + assert_eq!(pick.reason, "session_primary"); + } + assert_eq!(store.len(), 1); +} + +#[tokio::test] +async fn missing_and_empty_keys_use_admitted_power_of_two_without_binding() { + let (mut policy, store) = policy(); + let admission = Arc::new(Admission::default()); + policy.admission = admission.clone(); + let model = ModelId("m".into()); + let engines = [engine("a", 9), engine("b", 0)]; + for key in [None, Some("")] { + let request = PickRequest { + session_key: key, + ..request(&model) + }; + let pick = policy.pick(&engines, &request).await.unwrap(); + assert_eq!(pick.engine.id.0, "b"); + assert_eq!(pick.reason, "no_session"); + admission.reject.store(true, Ordering::Relaxed); + assert!(matches!( + policy.pick(&engines, &request).await, + Err(PickError::AdmissionRejected(_)) + )); + admission.reject.store(false, Ordering::Relaxed); + } + assert!(store.is_empty()); + assert_eq!(admission.calls.lock().unwrap().len(), 4); +} + +#[tokio::test] +async fn rejected_new_and_existing_sessions_never_rebind_or_try_another_engine() { + let (mut policy, store) = policy(); + let admission = Arc::new(Admission::default()); + policy.admission = admission.clone(); + let engines = [engine("a", 0), engine("b", 9)]; + let model = ModelId("m".into()); + let request = request(&model); + admission.reject.store(true, Ordering::Relaxed); + assert!(matches!( + policy.pick(&engines, &request).await, + Err(PickError::AdmissionRejected(_)) + )); + assert!(store.is_empty()); + admission.reject.store(false, Ordering::Relaxed); + policy.pick(&engines, &request).await.unwrap(); + engines[0].active_requests.store(100, Ordering::Relaxed); + admission.reject.store(true, Ordering::Relaxed); + assert!(matches!( + policy.pick(&engines, &request).await, + Err(PickError::AdmissionRejected(_)) + )); + admission.reject.store(false, Ordering::Relaxed); + assert_eq!( + policy.pick(&engines, &request).await.unwrap().engine.id.0, + "a" + ); + assert_eq!(store.len(), 1); + assert_eq!(admission.calls.lock().unwrap().len(), 4); +} + +#[tokio::test] +async fn removed_binding_is_replaced_only_after_admission() { + let (mut policy, store) = policy(); + let admission = Arc::new(Admission::default()); + policy.admission = admission.clone(); + let model = ModelId("m".into()); + let request = request(&model); + let original = [engine("a", 0)]; + let replacement = [engine("b", 0)]; + policy.pick(&original, &request).await.unwrap(); + admission.reject.store(true, Ordering::Relaxed); + assert!(matches!( + policy.pick(&replacement, &request).await, + Err(PickError::AdmissionRejected(_)) + )); + admission.reject.store(false, Ordering::Relaxed); + assert_eq!( + policy.pick(&original, &request).await.unwrap().reason, + "session_primary" + ); + assert_eq!( + policy.pick(&replacement, &request).await.unwrap().reason, + "assigned" + ); + assert_eq!(store.len(), 1); +} + +#[tokio::test] +async fn bindings_are_scoped_by_model_bucket_role_and_session() { + let (policy, store) = policy(); + let models = [ModelId("m".into()), ModelId("other".into())]; + let original = [engine("a", 0)]; + let fleet = [original[0].clone(), engine("b", 0)]; + policy.pick(&original, &request(&models[0])).await.unwrap(); + original[0].active_requests.store(100, Ordering::Relaxed); + let requests = [ + request(&models[1]), + PickRequest { + bucket: "other", + ..request(&models[0]) + }, + PickRequest { + stage: Stage::Prefill, + ..request(&models[0]) + }, + PickRequest { + stage: Stage::Decode, + ..request(&models[0]) + }, + PickRequest { + session_key: Some("other"), + ..request(&models[0]) + }, + ]; + for request in requests { + let pick = policy.pick(&fleet, &request).await.unwrap(); + assert_eq!(pick.engine.id.0, "b"); + assert_eq!(pick.reason, "assigned"); + } + assert_eq!( + policy + .pick(&fleet, &request(&models[0])) + .await + .unwrap() + .engine + .id + .0, + "a" + ); + assert_eq!(store.len(), 6); +} + +#[tokio::test] +async fn embedded_delimiters_do_not_alias_scopes() { + let (policy, store) = policy(); + let models = [ModelId("a\0b".into()), ModelId("a".into())]; + let first = PickRequest { + bucket: "c", + session_key: Some("d\0e"), + ..request(&models[0]) + }; + let second = PickRequest { + bucket: "b\0c", + session_key: Some("d\0e"), + ..request(&models[1]) + }; + let third = PickRequest { + bucket: "c\0d", + session_key: Some("e"), + ..request(&models[0]) + }; + let fleet = [engine("a", 0), engine("b", 9)]; + policy.pick(&fleet, &first).await.unwrap(); + fleet[0].active_requests.store(100, Ordering::Relaxed); + for request in [second, third] { + assert_eq!( + policy.pick(&fleet, &request).await.unwrap().engine.id.0, + "b" + ); + } + assert_eq!(store.len(), 3); +} + +#[tokio::test] +async fn same_id_replacement_returns_the_live_candidate_instance() { + let (policy, _) = policy(); + let model = ModelId("m".into()); + let request = request(&model); + policy.pick(&[engine("a", 0)], &request).await.unwrap(); + let replacement = [engine("a", 0)]; + let pick = policy.pick(&replacement, &request).await.unwrap(); + assert!(Arc::ptr_eq(&pick.engine, &replacement[0])); + assert_eq!(pick.reason, "session_primary"); +} + +#[tokio::test] +async fn empty_candidates_skip_admission_and_invalid_signals_never_bind() { + let (mut policy, store) = policy(); + let admission = Arc::new(Admission::default()); + policy.admission = admission.clone(); + let model = ModelId("m".into()); + let request = request(&model); + assert!(matches!( + policy.pick(&[], &request).await, + Err(PickError::NoCandidates) + )); + assert!(admission.calls.lock().unwrap().is_empty()); + admission.invalid.store(true, Ordering::Relaxed); + assert!(matches!( + policy.pick(&[engine("a", 0)], &request).await, + Err(PickError::InvalidSignal(_)) + )); + assert!(store.is_empty()); +} + +#[tokio::test] +async fn admission_receives_fresh_load_on_assignment_and_reuse() { + let store = AffinityStore::new(Duration::from_secs(60)); + let table = EngineReportedLoadTable::new(); + let mut policy = SessionAwarePolicy::new(store, table.clone()); + let admission = Arc::new(Admission::default()); + policy.admission = admission.clone(); + let engines = [engine("a", 0)]; + let model = ModelId("m".into()); + let request = request(&model); + for waiting in [3, 7] { + table.set( + &engines[0].url, + 0, + LoadStat { + num_waiting_reqs: waiting, + num_running_reqs: 1, + num_tokens: 10, + max_total_num_tokens: 100, + native_cache: None, + }, + Instant::now(), + ); + policy.pick(&engines, &request).await.unwrap(); + } + assert_eq!( + *admission.calls.lock().unwrap(), + vec![("a".into(), Some(3)), ("a".into(), Some(7))] + ); +} + +#[tokio::test] +async fn shared_store_refreshes_active_sessions_and_expires_idle_ones() { + let clock = Arc::new(MockClock::new(Instant::now())); + let store = AffinityStore::with_clock(Duration::from_secs(10), clock.clone()); + let policy = SessionAwarePolicy::new(store.clone(), EngineReportedLoadTable::new()); + let model = ModelId("m".into()); + let fleet = [engine("a", 0), engine("b", 9)]; + let hot = request(&model); + let cold = PickRequest { + session_key: Some("cold"), + ..hot + }; + policy.pick(&fleet, &hot).await.unwrap(); + policy.pick(&fleet, &cold).await.unwrap(); + clock.advance(Duration::from_secs(8)); + fleet[0].active_requests.store(100, Ordering::Relaxed); + policy.pick(&fleet, &hot).await.unwrap(); + clock.advance(Duration::from_secs(8)); + assert_eq!(store.sweep_expired(), 1); + assert_eq!(policy.pick(&fleet, &hot).await.unwrap().engine.id.0, "a"); + assert_eq!(policy.pick(&fleet, &cold).await.unwrap().engine.id.0, "b"); +} + +#[derive(Debug)] +struct RacingAdmission { + competitor: SessionAwarePolicy, + winner: Arc, + reject_winner: bool, + calls: Mutex>, +} + +impl EngineAdmission for RacingAdmission { + fn check( + &self, + engine: &Worker, + request: &PickRequest<'_>, + _: Option<&EngineReportedWorkerLoad>, + ) -> Result { + self.calls.lock().unwrap().push(engine.id.0.clone()); + if engine.id.0 == "a" { + // Complete a competing first request after this request selected a, + // but before it can commit. The effective binding is now b. + futures::executor::block_on( + self.competitor + .pick(std::slice::from_ref(&self.winner), request), + )?; + } + Ok(if self.reject_winner && engine.id == self.winner.id { + Decision::Reject("racing winner full".into()) + } else { + Decision::Allow + }) + } +} + +#[tokio::test] +async fn concurrent_assignment_winner_is_checked_and_preserved_on_rejection() { + for reject_winner in [false, true] { + let (mut policy, store) = policy(); + let engines = [engine("a", 0), engine("b", 9)]; + let admission = Arc::new(RacingAdmission { + competitor: SessionAwarePolicy::new(store.clone(), EngineReportedLoadTable::new()), + winner: engines[1].clone(), + reject_winner, + calls: Mutex::new(Vec::new()), + }); + policy.admission = admission.clone(); + let model = ModelId("m".into()); + let request = request(&model); + let result = policy.pick(&engines, &request).await; + if reject_winner { + assert!( + matches!(result, Err(PickError::AdmissionRejected(rejection)) if rejection.engine.0 == "b") + ); + } else { + let pick = result.unwrap(); + assert_eq!(pick.engine.id.0, "b"); + assert_eq!(pick.reason, "session_primary"); + } + assert_eq!(*admission.calls.lock().unwrap(), vec!["a", "b"]); + assert_eq!( + admission + .competitor + .pick(&engines, &request) + .await + .unwrap() + .engine + .id + .0, + "b" + ); + assert_eq!(store.len(), 1); + } +} diff --git a/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs b/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs index 013ea6d6d..b9040a2c6 100644 --- a/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs +++ b/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs @@ -12,6 +12,8 @@ use sgl_router::server::app_context::ChatRouting; use sgl_router::state::load_monitor::engine_reported_load::EngineReportedWorkerLoad; use std::sync::Mutex; +mod session_aware; + type PickCall = (String, Stage, u64, Option); #[derive(Debug)] diff --git a/experimental/sgl-router/tests/proxy/chat_routing/reorg/session_aware.rs b/experimental/sgl-router/tests/proxy/chat_routing/reorg/session_aware.rs new file mode 100644 index 000000000..a398e9fea --- /dev/null +++ b/experimental/sgl-router/tests/proxy/chat_routing/reorg/session_aware.rs @@ -0,0 +1,119 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +use super::*; + +#[tokio::test] +async fn session_aware_reuses_custom_header_binding_after_load_changes() { + use sgl_router::config::AffinityConfig; + use sgl_router::policies_reorg::session_aware::SessionAwarePolicy; + use sgl_router::state::load_monitor::engine_reported_load::EngineReportedLoadTable; + use sgl_router::state::AffinityStore; + use std::sync::atomic::Ordering; + use std::time::Duration; + + let primary = MockWorker::start(vec![]).await; + let other = MockWorker::start(vec![]).await; + let store = AffinityStore::new(Duration::from_secs(60)); + let policy = Arc::new(SessionAwarePolicy::new( + store.clone(), + EngineReportedLoadTable::new(), + )); + let mut ctx = context( + &[ + ("primary", Stage::Plain, &primary), + ("other", Stage::Plain, &other), + ], + vec![Bucket::new( + "session", + BucketGroups::Plain(EngineGroup::new(policy)), + )], + ); + Arc::get_mut(&mut ctx).unwrap().config.model.affinity = Some(AffinityConfig { + session_id_header: "x-test-session".into(), + ..Default::default() + }); + let primary_worker = ctx.registry.get(&WorkerId("primary".into())).unwrap(); + let other_worker = ctx.registry.get(&WorkerId("other".into())).unwrap(); + other_worker.active_requests.store(10, Ordering::Relaxed); + let app = build_router(ctx); + for _ in 0..2 { + let mut req = request(body("hello")); + req.headers_mut() + .insert("x-test-session", "same-session".parse().unwrap()); + let response = app.clone().oneshot(req).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let _ = response.into_body().collect().await.unwrap(); + assert!(primary.captured.lock().unwrap().last_body.is_some()); + assert!(other.captured.lock().unwrap().last_body.is_none()); + primary_worker.active_requests.store(100, Ordering::Relaxed); + other_worker.active_requests.store(0, Ordering::Relaxed); + } + assert_eq!(store.len(), 1); +} + +#[tokio::test] +async fn rejected_session_binding_advances_buckets_without_reassignment_or_dispatch() { + use sgl_router::config::AffinityConfig; + use sgl_router::policies_reorg::session_aware::SessionAwarePolicy; + use sgl_router::state::load_monitor::engine_reported_load::EngineReportedLoadTable; + use sgl_router::state::AffinityStore; + use std::time::Duration; + + let primary = MockWorker::start(vec![]).await; + let backup = MockWorker::start(vec![]).await; + let store = AffinityStore::new(Duration::from_secs(60)); + let table = EngineReportedLoadTable::new(); + let mut rejected = SessionAwarePolicy::new(store.clone(), table.clone()); + rejected.admission = Arc::new(RejectAll); + let accepted = SessionAwarePolicy::new(store.clone(), table.clone()); + let mut ctx = context( + &[ + ("primary", Stage::Plain, &primary), + ("backup", Stage::Plain, &backup), + ], + vec![ + Bucket::new( + "a-primary", + BucketGroups::Plain(EngineGroup { + worker_ids: Some([WorkerId("primary".into())].into_iter().collect()), + policy: Arc::new(rejected), + }), + ), + Bucket::new( + "b-backup", + BucketGroups::Plain(EngineGroup { + worker_ids: Some([WorkerId("backup".into())].into_iter().collect()), + policy: Arc::new(accepted), + }), + ), + ], + ); + Arc::get_mut(&mut ctx).unwrap().config.model.affinity = Some(AffinityConfig::default()); + let seed = SessionAwarePolicy::new(store.clone(), table); + let model = ModelId("tiny".into()); + let pick_request = PickRequest { + bucket: "a-primary", + session_key: Some("same-session"), + ..PickRequest::new(&model, Stage::Plain, 1) + }; + let engine = ctx.registry.get(&WorkerId("primary".into())).unwrap(); + seed.pick(std::slice::from_ref(&engine), &pick_request) + .await + .unwrap(); + let mut req = request(body("hello")); + req.headers_mut() + .insert("x-session-id", "same-session".parse().unwrap()); + let response = build_router(ctx).oneshot(req).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let _ = response.into_body().collect().await.unwrap(); + assert!(primary.captured.lock().unwrap().last_body.is_none()); + assert!(backup.captured.lock().unwrap().last_body.is_some()); + assert_eq!(store.len(), 2); + let retained = seed + .pick(std::slice::from_ref(&engine), &pick_request) + .await + .unwrap(); + assert_eq!(retained.reason, "session_primary"); + assert_eq!(retained.engine.id.0, "primary"); +}