[model-gateway] Add consistent hashing for ManualPolicy routing (#15907)

This commit is contained in:
Simo Lin
2025-12-26 22:53:03 -08:00
committed by GitHub
parent faecd37ed4
commit 171912a9e3
15 changed files with 1180 additions and 206 deletions
@@ -26,6 +26,7 @@ def policy_from_str(policy_str: Optional[str]) -> PolicyType:
"power_of_two": PolicyType.PowerOfTwo, "power_of_two": PolicyType.PowerOfTwo,
"bucket": PolicyType.Bucket, "bucket": PolicyType.Bucket,
"manual": PolicyType.Manual, "manual": PolicyType.Manual,
"consistent_hashing": PolicyType.ConsistentHashing,
} }
return policy_map[policy_str] return policy_map[policy_str]
@@ -13,6 +13,7 @@ pub enum PolicyType {
PowerOfTwo, PowerOfTwo,
Bucket, Bucket,
Manual, Manual,
ConsistentHashing,
} }
#[pyclass(eq)] #[pyclass(eq)]
@@ -416,6 +417,7 @@ impl Router {
bucket_adjust_interval_secs: self.bucket_adjust_interval_secs, bucket_adjust_interval_secs: self.bucket_adjust_interval_secs,
}, },
PolicyType::Manual => ConfigPolicyConfig::Manual, PolicyType::Manual => ConfigPolicyConfig::Manual,
PolicyType::ConsistentHashing => ConfigPolicyConfig::ConsistentHashing,
} }
}; };
+12 -3
View File
@@ -337,11 +337,19 @@ pub enum PolicyConfig {
bucket_adjust_interval_secs: usize, bucket_adjust_interval_secs: usize,
}, },
/// Manual routing policy supporting header-based routing: /// Manual routing policy with sticky sessions using DashMap.
/// - X-SMG-Target-Worker: Direct routing to a specific worker by URL /// - X-SMG-Routing-Key: Routes to a cached worker or assigns a new one
/// - X-SMG-Routing-Key: Consistent hash routing for session affinity /// - Provides true sticky sessions with zero key redistribution on worker add
/// - Falls back to random selection if no routing key is provided
#[serde(rename = "manual")] #[serde(rename = "manual")]
Manual, Manual,
/// Consistent hashing policy using hash ring for session affinity:
/// - X-SMG-Target-Worker: Direct routing to a specific worker by URL
/// - X-SMG-Routing-Key: Consistent hash routing for session affinity
/// - Provides O(log n) lookup with minimal redistribution (~1/N keys) on topology change
#[serde(rename = "consistent_hashing")]
ConsistentHashing,
} }
impl PolicyConfig { impl PolicyConfig {
@@ -353,6 +361,7 @@ impl PolicyConfig {
PolicyConfig::PowerOfTwo { .. } => "power_of_two", PolicyConfig::PowerOfTwo { .. } => "power_of_two",
PolicyConfig::Bucket { .. } => "bucket", PolicyConfig::Bucket { .. } => "bucket",
PolicyConfig::Manual => "manual", PolicyConfig::Manual => "manual",
PolicyConfig::ConsistentHashing => "consistent_hashing",
} }
} }
} }
+4 -1
View File
@@ -147,7 +147,10 @@ impl ConfigValidator {
fn validate_policy(policy: &PolicyConfig) -> ConfigResult<()> { fn validate_policy(policy: &PolicyConfig) -> ConfigResult<()> {
match policy { match policy {
PolicyConfig::Random | PolicyConfig::RoundRobin | PolicyConfig::Manual => {} PolicyConfig::Random
| PolicyConfig::RoundRobin
| PolicyConfig::Manual
| PolicyConfig::ConsistentHashing => {}
PolicyConfig::CacheAware { PolicyConfig::CacheAware {
cache_threshold, cache_threshold,
balance_abs_threshold: _, balance_abs_threshold: _,
+1 -1
View File
@@ -41,5 +41,5 @@ pub use worker::{
}; };
pub use worker_builder::{BasicWorkerBuilder, DPAwareWorkerBuilder}; pub use worker_builder::{BasicWorkerBuilder, DPAwareWorkerBuilder};
pub use worker_manager::{LoadMonitor, WorkerManager}; pub use worker_manager::{LoadMonitor, WorkerManager};
pub use worker_registry::{WorkerId, WorkerRegistry, WorkerRegistryStats}; pub use worker_registry::{HashRing, WorkerId, WorkerRegistry, WorkerRegistryStats};
pub use worker_service::{WorkerService, WorkerServiceError}; pub use worker_service::{WorkerService, WorkerServiceError};
+147 -2
View File
@@ -5,6 +5,11 @@
//! # Performance Optimizations //! # Performance Optimizations
//! The model index uses immutable Arc snapshots instead of RwLock for lock-free reads. //! The model index uses immutable Arc snapshots instead of RwLock for lock-free reads.
//! This is critical for high-concurrency scenarios where many requests query the same model. //! This is critical for high-concurrency scenarios where many requests query the same model.
//!
//! # Consistent Hash Ring
//! The registry maintains a pre-computed hash ring per model for O(log n) consistent hashing.
//! The ring is rebuilt only when workers are added/removed, not per-request.
//! Uses virtual nodes (150 per worker) for even distribution and blake3 for stable hashing.
use std::sync::Arc; use std::sync::Arc;
@@ -16,6 +21,118 @@ use crate::{
observability::metrics::Metrics, observability::metrics::Metrics,
}; };
/// Number of virtual nodes per physical worker for even distribution.
/// 150 is a common choice that provides good balance between memory and distribution.
const VIRTUAL_NODES_PER_WORKER: usize = 150;
/// Consistent hash ring for O(log n) worker selection.
///
/// Each worker is placed at multiple positions (virtual nodes) on the ring
/// based on hash(worker_url + vnode_index). This provides:
/// - Even key distribution across workers
/// - Minimal key redistribution when workers are added/removed (~1/N keys move)
/// - O(log n) lookup via binary search
///
/// Uses blake3 for stable, fast hashing that's consistent across Rust versions.
#[derive(Debug, Clone)]
pub struct HashRing {
/// Sorted list of (ring_position, worker_url)
/// Multiple entries per worker (virtual nodes) for even distribution.
/// Uses Arc<str> to share URL across all virtual nodes (150 refs vs 150 copies).
entries: Arc<[(u64, Arc<str>)]>,
}
impl HashRing {
/// Build a hash ring from a list of workers.
/// Creates VIRTUAL_NODES_PER_WORKER entries per worker for even distribution.
pub fn new(workers: &[Arc<dyn Worker>]) -> Self {
let mut entries: Vec<(u64, Arc<str>)> =
Vec::with_capacity(workers.len() * VIRTUAL_NODES_PER_WORKER);
for worker in workers {
// Create Arc<str> once per worker, share across all virtual nodes
let url: Arc<str> = Arc::from(worker.url());
// Create multiple virtual nodes per worker
for vnode in 0..VIRTUAL_NODES_PER_WORKER {
let vnode_key = format!("{}#{}", url, vnode);
let pos = Self::hash_position(&vnode_key);
entries.push((pos, Arc::clone(&url)));
}
}
// Sort by ring position for binary search
entries.sort_unstable_by_key(|(pos, _)| *pos);
Self {
entries: Arc::from(entries.into_boxed_slice()),
}
}
/// Hash a string to a ring position using blake3 (stable across versions).
#[inline]
fn hash_position(s: &str) -> u64 {
let hash = blake3::hash(s.as_bytes());
// Take first 8 bytes as u64
u64::from_le_bytes(hash.as_bytes()[..8].try_into().unwrap())
}
/// Find worker URL for a key using consistent hashing.
/// Returns the first healthy worker URL at or after the key's position (clockwise).
///
/// - `key`: The routing key to hash
/// - `is_healthy`: Function to check if a worker URL is healthy
pub fn find_healthy_url<F>(&self, key: &str, is_healthy: F) -> Option<&str>
where
F: Fn(&str) -> bool,
{
if self.entries.is_empty() {
return None;
}
let key_pos = Self::hash_position(key);
// Binary search to find first entry at or after key_pos
let start = self.entries.partition_point(|(pos, _)| *pos < key_pos);
// Walk clockwise from start, wrapping around
// Track visited URLs to avoid checking same worker multiple times (virtual nodes)
let mut checked_urls =
std::collections::HashSet::with_capacity(self.worker_count().min(16));
for i in 0..self.entries.len() {
let (_, url) = &self.entries[(start + i) % self.entries.len()];
let url_str: &str = url;
// Skip if we already checked this worker (from another virtual node)
if !checked_urls.insert(url_str) {
continue;
}
if is_healthy(url_str) {
return Some(url_str);
}
}
None
}
/// Check if the ring is empty
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
/// Get the number of entries in the ring (including virtual nodes)
pub fn len(&self) -> usize {
self.entries.len()
}
/// Get the number of unique workers in the ring
pub fn worker_count(&self) -> usize {
self.entries.len() / VIRTUAL_NODES_PER_WORKER.max(1)
}
}
/// Unique identifier for a worker /// Unique identifier for a worker
#[derive(Debug, Clone, Hash, Eq, PartialEq)] #[derive(Debug, Clone, Hash, Eq, PartialEq)]
pub struct WorkerId(String); pub struct WorkerId(String);
@@ -58,6 +175,10 @@ pub struct WorkerRegistry {
/// Uses Arc<[T]> instead of Arc<RwLock<Vec<T>>> for lock-free reads. /// Uses Arc<[T]> instead of Arc<RwLock<Vec<T>>> for lock-free reads.
model_index: ModelIndex, model_index: ModelIndex,
/// Consistent hash rings per model for O(log n) routing.
/// Rebuilt on worker add/remove (copy-on-write).
hash_rings: Arc<DashMap<String, Arc<HashRing>>>,
/// Workers indexed by worker type /// Workers indexed by worker type
type_workers: Arc<DashMap<WorkerType, Vec<WorkerId>>>, type_workers: Arc<DashMap<WorkerType, Vec<WorkerId>>>,
@@ -74,12 +195,29 @@ impl WorkerRegistry {
Self { Self {
workers: Arc::new(DashMap::new()), workers: Arc::new(DashMap::new()),
model_index: Arc::new(DashMap::new()), model_index: Arc::new(DashMap::new()),
hash_rings: Arc::new(DashMap::new()),
type_workers: Arc::new(DashMap::new()), type_workers: Arc::new(DashMap::new()),
connection_workers: Arc::new(DashMap::new()), connection_workers: Arc::new(DashMap::new()),
url_to_id: Arc::new(DashMap::new()), url_to_id: Arc::new(DashMap::new()),
} }
} }
/// Rebuild the hash ring for a model based on current workers in the model index
fn rebuild_hash_ring(&self, model_id: &str) {
if let Some(workers) = self.model_index.get(model_id) {
let ring = HashRing::new(&workers);
self.hash_rings.insert(model_id.to_string(), Arc::new(ring));
} else {
// No workers for this model, remove the ring
self.hash_rings.remove(model_id);
}
}
/// Get the hash ring for a model (O(1) lookup)
pub fn get_hash_ring(&self, model_id: &str) -> Option<Arc<HashRing>> {
self.hash_rings.get(model_id).map(|r| Arc::clone(&r))
}
/// Register a new worker /// Register a new worker
pub fn register(&self, worker: Arc<dyn Worker>) -> WorkerId { pub fn register(&self, worker: Arc<dyn Worker>) -> WorkerId {
let worker_id = if let Some(existing_id) = self.url_to_id.get(worker.url()) { let worker_id = if let Some(existing_id) = self.url_to_id.get(worker.url()) {
@@ -100,7 +238,7 @@ impl WorkerRegistry {
// This creates a new immutable snapshot with the added worker // This creates a new immutable snapshot with the added worker
let model_id = worker.model_id().to_string(); let model_id = worker.model_id().to_string();
self.model_index self.model_index
.entry(model_id) .entry(model_id.clone())
.and_modify(|existing| { .and_modify(|existing| {
// Create new snapshot with the additional worker // Create new snapshot with the additional worker
let mut new_workers: Vec<Arc<dyn Worker>> = existing.iter().cloned().collect(); let mut new_workers: Vec<Arc<dyn Worker>> = existing.iter().cloned().collect();
@@ -109,6 +247,9 @@ impl WorkerRegistry {
}) })
.or_insert_with(|| Arc::from(vec![worker.clone()].into_boxed_slice())); .or_insert_with(|| Arc::from(vec![worker.clone()].into_boxed_slice()));
// Rebuild hash ring for this model
self.rebuild_hash_ring(&model_id);
// Update type index (clone needed for DashMap key ownership) // Update type index (clone needed for DashMap key ownership)
self.type_workers self.type_workers
.entry(worker.worker_type().clone()) .entry(worker.worker_type().clone())
@@ -149,7 +290,8 @@ impl WorkerRegistry {
// Remove from model index using copy-on-write // Remove from model index using copy-on-write
// Create new snapshot without the removed worker // Create new snapshot without the removed worker
let worker_url = worker.url(); let worker_url = worker.url();
if let Some(mut entry) = self.model_index.get_mut(worker.model_id()) { let model_id = worker.model_id().to_string();
if let Some(mut entry) = self.model_index.get_mut(&model_id) {
let new_workers: Vec<Arc<dyn Worker>> = entry let new_workers: Vec<Arc<dyn Worker>> = entry
.iter() .iter()
.filter(|w| w.url() != worker_url) .filter(|w| w.url() != worker_url)
@@ -158,6 +300,9 @@ impl WorkerRegistry {
*entry = Arc::from(new_workers.into_boxed_slice()); *entry = Arc::from(new_workers.into_boxed_slice());
} }
// Rebuild hash ring for this model
self.rebuild_hash_ring(&model_id);
// Remove from type index // Remove from type index
if let Some(mut type_workers) = self.type_workers.get_mut(worker.worker_type()) { if let Some(mut type_workers) = self.type_workers.get_mut(worker.worker_type()) {
type_workers.retain(|id| id != worker_id); type_workers.retain(|id| id != worker_id);
@@ -810,6 +810,15 @@ impl Metrics {
.increment(1); .increment(1);
} }
/// Record consistent hashing policy execution branch for routing decisions
pub fn record_worker_consistent_hashing_policy_branch(branch: &'static str) {
counter!(
"smg_consistent_hashing_policy_branch_total",
"branch" => branch
)
.increment(1);
}
/// Set running requests per worker /// Set running requests per worker
pub fn set_worker_requests_active(worker: &str, count: usize) { pub fn set_worker_requests_active(worker: &str, count: usize) {
gauge!( gauge!(
@@ -0,0 +1,527 @@
//! Consistent hashing routing policy with header-based routing support
//!
//! Supports two routing mechanisms via HTTP headers:
//! - `X-SMG-Target-Worker`: Direct routing by worker index (0-based), returns None if unavailable
//! - `X-SMG-Routing-Key`: Consistent hash routing for session affinity
//!
//! ## Consistent Hashing
//!
//! Uses a pre-computed hash ring from WorkerRegistry where:
//! 1. Each worker is placed at a fixed position based on hash(worker_url)
//! 2. Keys are hashed to the ring, then walk clockwise to find first healthy worker
//! 3. When workers scale up/down, only keys in the affected range redistribute (~1/N keys move)
//!
//! The ring is built once when workers are added/removed, not per-request.
//! This ensures O(log n) lookup performance.
//!
//! Complexity: O(log n) binary search + O(k) walk where k = consecutive unhealthy workers.
use std::sync::Arc;
use http::header::HeaderName;
use rand::Rng as _;
use super::{LoadBalancingPolicy, SelectWorkerInfo};
use crate::{core::Worker, observability::metrics::Metrics};
/// Header for direct worker targeting by index (0-based)
static HEADER_TARGET_WORKER: HeaderName = HeaderName::from_static("x-smg-target-worker");
/// Header for consistent hash routing
static HEADER_ROUTING_KEY: HeaderName = HeaderName::from_static("x-smg-routing-key");
/// Execution branch for metrics
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Branch {
NoHealthyWorkers,
TargetWorkerHit,
TargetWorkerMiss,
RoutingKeyHit,
RandomFallback,
}
impl Branch {
#[inline]
const fn as_str(&self) -> &'static str {
match self {
Self::NoHealthyWorkers => "no_healthy_workers",
Self::TargetWorkerHit => "target_worker_hit",
Self::TargetWorkerMiss => "target_worker_miss",
Self::RoutingKeyHit => "routing_key_hit",
Self::RandomFallback => "random_fallback",
}
}
}
#[derive(Debug, Default)]
pub struct ConsistentHashingPolicy;
impl ConsistentHashingPolicy {
pub fn new() -> Self {
Self
}
/// Use consistent hashing to find a worker for the given key.
/// Uses pre-computed ring from SelectWorkerInfo if available.
///
/// The ring returns a worker URL, which we then map to an index in the workers array.
/// This correctly handles filtered worker arrays since we match by URL, not by index.
///
/// Complexity: O(n) to build healthy URL map + O(log n) ring lookup + O(k) walk
fn find_by_consistent_hash(
workers: &[Arc<dyn Worker>],
info: &SelectWorkerInfo,
key: &str,
) -> Option<usize> {
// Build URL→index map for healthy workers: O(n) once, O(1) lookups
let healthy_url_to_idx: std::collections::HashMap<&str, usize> = workers
.iter()
.enumerate()
.filter(|(_, w)| w.is_healthy())
.map(|(i, w)| (w.url(), i))
.collect();
if healthy_url_to_idx.is_empty() {
return None;
}
// Use pre-computed ring if available
if let Some(ref ring) = info.hash_ring {
// O(1) lookup per URL checked instead of O(n)
let url = ring.find_healthy_url(key, |url| healthy_url_to_idx.contains_key(url))?;
return healthy_url_to_idx.get(url).copied();
}
// Fallback: no ring provided, use simple modulo (less optimal but functional)
// This shouldn't happen in normal operation as WorkerSelectionStage provides the ring
let mut healthy_indices: Vec<usize> = healthy_url_to_idx.values().copied().collect();
healthy_indices.sort_unstable(); // Ensure deterministic order
// Use blake3 for consistent hashing in fallback too
let hash = blake3::hash(key.as_bytes());
let hash_val = u64::from_le_bytes(hash.as_bytes()[..8].try_into().unwrap());
let idx = (hash_val as usize) % healthy_indices.len();
Some(healthy_indices[idx])
}
fn select_worker_impl(
&self,
workers: &[Arc<dyn Worker>],
info: &SelectWorkerInfo,
) -> (Option<usize>, Branch) {
if workers.is_empty() {
return (None, Branch::NoHealthyWorkers);
}
// Extract routing headers - to_str() is O(1), just validates ASCII, no allocation
let target_worker = info
.headers
.and_then(|h| h.get(&HEADER_TARGET_WORKER))
.and_then(|v| v.to_str().ok())
.filter(|s| !s.is_empty());
let routing_key = info
.headers
.and_then(|h| h.get(&HEADER_ROUTING_KEY))
.and_then(|v| v.to_str().ok())
.filter(|s| !s.is_empty());
// Priority 1: X-SMG-Target-Worker - direct routing by worker index
// O(1) parse + O(1) bounds check + O(1) health check
if let Some(idx_str) = target_worker {
if let Ok(idx) = idx_str.parse::<usize>() {
if idx < workers.len() && workers[idx].is_healthy() {
return (Some(idx), Branch::TargetWorkerHit);
}
}
return (None, Branch::TargetWorkerMiss);
}
// Priority 2: X-SMG-Routing-Key - consistent hash routing (O(log n))
if let Some(key) = routing_key {
return match Self::find_by_consistent_hash(workers, info, key) {
Some(idx) => (Some(idx), Branch::RoutingKeyHit),
None => (None, Branch::NoHealthyWorkers),
};
}
// Priority 3: Implicit routing key from stable headers (session affinity)
let implicit_key = info.headers.and_then(|h| {
h.get("authorization")
.or_else(|| h.get("x-forwarded-for"))
.or_else(|| h.get("cookie"))
.and_then(|v| v.to_str().ok())
.filter(|s| !s.is_empty())
});
if let Some(key) = implicit_key {
return match Self::find_by_consistent_hash(workers, info, key) {
Some(idx) => (Some(idx), Branch::RoutingKeyHit),
None => (None, Branch::NoHealthyWorkers),
};
}
// Fallback: random selection (truly anonymous client)
let healthy_count = workers.iter().filter(|w| w.is_healthy()).count();
if healthy_count == 0 {
return (None, Branch::NoHealthyWorkers);
}
let random_healthy_idx = rand::rng().random_range(0..healthy_count);
let idx = workers
.iter()
.enumerate()
.filter(|(_, w)| w.is_healthy())
.nth(random_healthy_idx)
.map(|(i, _)| i)
.unwrap();
(Some(idx), Branch::RandomFallback)
}
}
impl LoadBalancingPolicy for ConsistentHashingPolicy {
fn select_worker(&self, workers: &[Arc<dyn Worker>], info: &SelectWorkerInfo) -> Option<usize> {
let (result, branch) = self.select_worker_impl(workers, info);
Metrics::record_worker_consistent_hashing_policy_branch(branch.as_str());
result
}
fn name(&self) -> &'static str {
"consistent_hashing"
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use super::*;
use crate::core::{BasicWorkerBuilder, HashRing, WorkerType};
fn headers_with_routing_key(key: &str) -> http::HeaderMap {
let mut headers = http::HeaderMap::new();
headers.insert("x-smg-routing-key", key.parse().unwrap());
headers
}
fn headers_with_target_worker(idx: usize) -> http::HeaderMap {
let mut headers = http::HeaderMap::new();
headers.insert("x-smg-target-worker", idx.to_string().parse().unwrap());
headers
}
fn create_workers(urls: &[&str]) -> Vec<Arc<dyn Worker>> {
urls.iter()
.map(|url| {
Arc::new(
BasicWorkerBuilder::new(*url)
.worker_type(WorkerType::Regular)
.build(),
) as Arc<dyn Worker>
})
.collect()
}
#[test]
fn test_consistent_routing() {
let policy = ConsistentHashingPolicy::new();
let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
let headers = headers_with_routing_key("user-123");
let info = SelectWorkerInfo {
headers: Some(&headers),
..Default::default()
};
let (first_result, _) = policy.select_worker_impl(&workers, &info);
let first_idx = first_result.unwrap();
// Same key should always route to same worker
for _ in 0..10 {
let (result, branch) = policy.select_worker_impl(&workers, &info);
assert_eq!(result, Some(first_idx));
assert_eq!(branch, Branch::RoutingKeyHit);
}
}
#[test]
fn test_different_keys_distribute() {
let policy = ConsistentHashingPolicy::new();
let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
let mut distribution = HashMap::new();
for i in 0..100 {
let headers = headers_with_routing_key(&format!("user-{}", i));
let info = SelectWorkerInfo {
headers: Some(&headers),
..Default::default()
};
let (result, _) = policy.select_worker_impl(&workers, &info);
*distribution.entry(result.unwrap()).or_insert(0) += 1;
}
assert!(distribution.len() > 1, "Should distribute across workers");
}
#[test]
fn test_target_worker_hit() {
let policy = ConsistentHashingPolicy::new();
let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
let headers = headers_with_target_worker(1);
let info = SelectWorkerInfo {
headers: Some(&headers),
..Default::default()
};
let (result, branch) = policy.select_worker_impl(&workers, &info);
assert_eq!(result, Some(1));
assert_eq!(branch, Branch::TargetWorkerHit);
}
#[test]
fn test_target_worker_miss_out_of_bounds() {
let policy = ConsistentHashingPolicy::new();
let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
let headers = headers_with_target_worker(5); // Out of bounds
let info = SelectWorkerInfo {
headers: Some(&headers),
..Default::default()
};
let (result, branch) = policy.select_worker_impl(&workers, &info);
assert_eq!(result, None);
assert_eq!(branch, Branch::TargetWorkerMiss);
}
#[test]
fn test_target_worker_miss_unhealthy() {
let policy = ConsistentHashingPolicy::new();
let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
workers[1].set_healthy(false);
let headers = headers_with_target_worker(1);
let info = SelectWorkerInfo {
headers: Some(&headers),
..Default::default()
};
let (result, branch) = policy.select_worker_impl(&workers, &info);
assert_eq!(result, None);
assert_eq!(branch, Branch::TargetWorkerMiss);
}
#[test]
fn test_target_worker_priority_over_routing_key() {
let policy = ConsistentHashingPolicy::new();
let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
let mut headers = http::HeaderMap::new();
headers.insert("x-smg-target-worker", "1".parse().unwrap());
headers.insert("x-smg-routing-key", "some-key".parse().unwrap());
let info = SelectWorkerInfo {
headers: Some(&headers),
..Default::default()
};
let (result, branch) = policy.select_worker_impl(&workers, &info);
assert_eq!(result, Some(1));
assert_eq!(branch, Branch::TargetWorkerHit);
}
#[test]
fn test_fallback_random_distribution() {
let policy = ConsistentHashingPolicy::new();
let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
// Without routing headers, should distribute randomly across workers
let mut distribution = HashMap::new();
for _ in 0..100 {
let info = SelectWorkerInfo::default();
let (result, branch) = policy.select_worker_impl(&workers, &info);
assert!(result.is_some());
assert_eq!(branch, Branch::RandomFallback);
*distribution.entry(result.unwrap()).or_insert(0) += 1;
}
// Should distribute across multiple workers (not always same one)
assert!(
distribution.len() > 1,
"Random fallback should distribute across workers"
);
}
#[test]
fn test_no_healthy_workers() {
let policy = ConsistentHashingPolicy::new();
let workers = create_workers(&["http://w1:8000"]);
workers[0].set_healthy(false);
let headers = headers_with_routing_key("test");
let info = SelectWorkerInfo {
headers: Some(&headers),
..Default::default()
};
let (result, branch) = policy.select_worker_impl(&workers, &info);
assert_eq!(result, None);
assert_eq!(branch, Branch::NoHealthyWorkers);
}
#[test]
fn test_empty_workers() {
let policy = ConsistentHashingPolicy::new();
let workers: Vec<Arc<dyn Worker>> = vec![];
let info = SelectWorkerInfo::default();
let (result, branch) = policy.select_worker_impl(&workers, &info);
assert_eq!(result, None);
assert_eq!(branch, Branch::NoHealthyWorkers);
}
#[test]
fn test_consistent_hash_minimal_redistribution() {
// Test that consistent hashing moves fewer keys than random redistribution
let policy = ConsistentHashingPolicy::new();
let workers = create_workers(&[
"http://w0:8000",
"http://w1:8000",
"http://w2:8000",
"http://w3:8000",
]);
let ring = Arc::new(HashRing::new(&workers));
// Record which worker each key routes to with all workers healthy
let mut key_to_worker_before: HashMap<String, usize> = HashMap::new();
for i in 0..100 {
let key = format!("user-{}", i);
let headers = headers_with_routing_key(&key);
let info = SelectWorkerInfo {
headers: Some(&headers),
hash_ring: Some(ring.clone()),
..Default::default()
};
let (result, _) = policy.select_worker_impl(&workers, &info);
key_to_worker_before.insert(key, result.unwrap());
}
// Mark worker 1 as unhealthy
workers[1].set_healthy(false);
// Record new routing and count how many keys moved
let mut moved_count = 0;
for i in 0..100 {
let key = format!("user-{}", i);
let headers = headers_with_routing_key(&key);
let info = SelectWorkerInfo {
headers: Some(&headers),
hash_ring: Some(ring.clone()),
..Default::default()
};
let (result, _) = policy.select_worker_impl(&workers, &info);
let new_worker = result.unwrap();
let old_worker = key_to_worker_before[&key];
if new_worker != old_worker {
moved_count += 1;
}
}
// With consistent hashing, approximately 1/N keys should move (N = worker count)
// Random redistribution would move approximately (N-1)/N = 75% of keys
// Verify we're significantly better than random (< 50% moved)
let keys_on_failed_worker = key_to_worker_before.values().filter(|&&w| w == 1).count();
assert!(
moved_count <= keys_on_failed_worker + 5,
"Consistent hashing should only move keys from failed worker (+small variance). \
Expected ~{}, got {}",
keys_on_failed_worker,
moved_count
);
assert!(
moved_count < 50,
"Consistent hashing should move fewer than 50% of keys (random would move ~75%), got {}%",
moved_count
);
}
#[test]
fn test_routing_key_failover_and_recovery() {
// Test that when a worker fails, keys move to another worker,
// and when it recovers, keys return to the original worker
let policy = ConsistentHashingPolicy::new();
let workers = create_workers(&["http://w0:8000", "http://w1:8000", "http://w2:8000"]);
let ring = Arc::new(HashRing::new(&workers));
// Find which worker a key routes to when all are healthy
let test_key = "session-abc-123";
let headers = headers_with_routing_key(test_key);
let info = SelectWorkerInfo {
headers: Some(&headers),
hash_ring: Some(ring.clone()),
..Default::default()
};
let (result, _) = policy.select_worker_impl(&workers, &info);
let original_idx = result.unwrap();
// Mark that worker unhealthy
workers[original_idx].set_healthy(false);
// Key should now route to a different healthy worker
let (failover_result, _) = policy.select_worker_impl(&workers, &info);
let failover_idx = failover_result.unwrap();
assert_ne!(
failover_idx, original_idx,
"Should failover to different worker"
);
assert!(
workers[failover_idx].is_healthy(),
"Failover target should be healthy"
);
// Failover should be consistent
for _ in 0..5 {
let (result, _) = policy.select_worker_impl(&workers, &info);
assert_eq!(result, Some(failover_idx), "Failover should be consistent");
}
// Recover the original worker
workers[original_idx].set_healthy(true);
// Key should route back to original worker
let (recovered_result, _) = policy.select_worker_impl(&workers, &info);
assert_eq!(
recovered_result,
Some(original_idx),
"Should return to original worker after recovery"
);
}
#[test]
fn test_empty_routing_key_uses_fallback() {
let policy = ConsistentHashingPolicy::new();
let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
let headers = headers_with_routing_key("");
let info = SelectWorkerInfo {
headers: Some(&headers),
..Default::default()
};
let (result, branch) = policy.select_worker_impl(&workers, &info);
assert!(result.is_some());
assert_eq!(branch, Branch::RandomFallback);
}
#[test]
fn test_policy_name() {
let policy = ConsistentHashingPolicy::new();
assert_eq!(policy.name(), "consistent_hashing");
}
}
+16 -2
View File
@@ -3,8 +3,8 @@
use std::sync::Arc; use std::sync::Arc;
use super::{ use super::{
BucketConfig, BucketPolicy, CacheAwareConfig, CacheAwarePolicy, LoadBalancingPolicy, BucketConfig, BucketPolicy, CacheAwareConfig, CacheAwarePolicy, ConsistentHashingPolicy,
ManualPolicy, PowerOfTwoPolicy, RandomPolicy, RoundRobinPolicy, LoadBalancingPolicy, ManualPolicy, PowerOfTwoPolicy, RandomPolicy, RoundRobinPolicy,
}; };
use crate::config::PolicyConfig; use crate::config::PolicyConfig;
@@ -47,6 +47,7 @@ impl PolicyFactory {
Arc::new(BucketPolicy::with_config(config)) Arc::new(BucketPolicy::with_config(config))
} }
PolicyConfig::Manual => Arc::new(ManualPolicy::new()), PolicyConfig::Manual => Arc::new(ManualPolicy::new()),
PolicyConfig::ConsistentHashing => Arc::new(ConsistentHashingPolicy::new()),
} }
} }
@@ -59,6 +60,9 @@ impl PolicyFactory {
"cache_aware" | "cacheaware" => Some(Arc::new(CacheAwarePolicy::new())), "cache_aware" | "cacheaware" => Some(Arc::new(CacheAwarePolicy::new())),
"bucket" => Some(Arc::new(BucketPolicy::new())), "bucket" => Some(Arc::new(BucketPolicy::new())),
"manual" => Some(Arc::new(ManualPolicy::new())), "manual" => Some(Arc::new(ManualPolicy::new())),
"consistent_hashing" | "consistenthashing" => {
Some(Arc::new(ConsistentHashingPolicy::new()))
}
_ => None, _ => None,
} }
} }
@@ -96,6 +100,12 @@ mod tests {
bucket_adjust_interval_secs: 5, bucket_adjust_interval_secs: 5,
}); });
assert_eq!(policy.name(), "bucket"); assert_eq!(policy.name(), "bucket");
let policy = PolicyFactory::create_from_config(&PolicyConfig::Manual);
assert_eq!(policy.name(), "manual");
let policy = PolicyFactory::create_from_config(&PolicyConfig::ConsistentHashing);
assert_eq!(policy.name(), "consistent_hashing");
} }
#[tokio::test] #[tokio::test]
@@ -110,6 +120,10 @@ mod tests {
assert!(PolicyFactory::create_by_name("CacheAware").is_some()); assert!(PolicyFactory::create_by_name("CacheAware").is_some());
assert!(PolicyFactory::create_by_name("bucket").is_some()); assert!(PolicyFactory::create_by_name("bucket").is_some());
assert!(PolicyFactory::create_by_name("Bucket").is_some()); assert!(PolicyFactory::create_by_name("Bucket").is_some());
assert!(PolicyFactory::create_by_name("manual").is_some());
assert!(PolicyFactory::create_by_name("Manual").is_some());
assert!(PolicyFactory::create_by_name("consistent_hashing").is_some());
assert!(PolicyFactory::create_by_name("ConsistentHashing").is_some());
assert!(PolicyFactory::create_by_name("unknown").is_none()); assert!(PolicyFactory::create_by_name("unknown").is_none());
} }
} }
+416 -188
View File
@@ -1,105 +1,159 @@
//! Manual routing policy with header-based routing support //! Manual routing policy based on routing key header
//! //!
//! Supports two routing mechanisms via HTTP headers: //! This policy provides sticky session routing where each unique routing key
//! - `X-SMG-Target-Worker`: Direct routing by worker index (0-based), returns None if unavailable //! is consistently mapped to the same worker. Unlike consistent hashing,
//! - `X-SMG-Routing-Key`: Consistent hash routing for session affinity //! this policy:
//! - Does NOT redistribute any sessions when workers are added
//! - Only remaps sessions when their assigned worker becomes unhealthy
//! - Maintains up to 2 candidate workers per routing key for fast failover
//! //!
//! Complexity: O(n) for get_healthy_worker_indices (unavoidable), O(1) for routing decisions. //! Use this when you need stronger stickiness guarantees than consistent hashing,
//! for example with stateful chat sessions where context is stored on the worker.
//!
//! ## Header
//! - `X-SMG-Routing-Key`: The routing key for sticky session routing
use std::{ use std::sync::Arc;
hash::{Hash, Hasher},
sync::Arc,
};
use dashmap::{mapref::entry::Entry, DashMap};
use http::header::HeaderName; use http::header::HeaderName;
use rand::Rng as _; use rand::Rng;
use super::{get_healthy_worker_indices, LoadBalancingPolicy, SelectWorkerInfo}; use super::{get_healthy_worker_indices, LoadBalancingPolicy, SelectWorkerInfo};
use crate::{core::Worker, observability::metrics::Metrics}; use crate::{core::Worker, observability::metrics::Metrics};
/// Header for direct worker targeting by index (0-based) /// Header for routing key based sticky sessions
static HEADER_TARGET_WORKER: HeaderName = HeaderName::from_static("x-smg-target-worker");
/// Header for consistent hash routing
static HEADER_ROUTING_KEY: HeaderName = HeaderName::from_static("x-smg-routing-key"); static HEADER_ROUTING_KEY: HeaderName = HeaderName::from_static("x-smg-routing-key");
/// Execution branch for metrics
#[derive(Debug, Clone, Copy, PartialEq, Eq)] #[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Branch { enum ExecutionBranch {
NoHealthyWorkers, NoHealthyWorkers,
TargetWorkerHit, FastPathHit,
TargetWorkerMiss, SlowPathOccupiedHit,
RoutingKeyHit, SlowPathOccupiedMiss,
RandomFallback, SlowPathVacant,
NoRoutingId,
} }
impl Branch { impl ExecutionBranch {
#[inline] fn as_str(&self) -> &'static str {
const fn as_str(&self) -> &'static str {
match self { match self {
Self::NoHealthyWorkers => "no_healthy_workers", Self::NoHealthyWorkers => "no_healthy_workers",
Self::TargetWorkerHit => "target_worker_hit", Self::FastPathHit => "fast_path_hit",
Self::TargetWorkerMiss => "target_worker_miss", Self::SlowPathOccupiedHit => "slow_path_occupied_hit",
Self::RoutingKeyHit => "routing_key_hit", Self::SlowPathOccupiedMiss => "slow_path_occupied_miss",
Self::RandomFallback => "random_fallback", Self::SlowPathVacant => "slow_path_vacant",
Self::NoRoutingId => "no_routing_id",
} }
} }
} }
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct RoutingId(String);
impl RoutingId {
fn new(id: impl Into<String>) -> Self {
Self(id.into())
}
}
const MAX_CANDIDATE_WORKERS: usize = 2;
#[derive(Debug, Clone)]
struct RoutingInfo {
candi_worker_urls: Vec<String>,
}
impl RoutingInfo {
fn push_bounded(&mut self, url: String) {
while self.candi_worker_urls.len() >= MAX_CANDIDATE_WORKERS {
self.candi_worker_urls.remove(0);
}
self.candi_worker_urls.push(url);
}
}
// TODO may optimize performance
// TODO evict old data periodically
#[derive(Debug, Default)] #[derive(Debug, Default)]
pub struct ManualPolicy; pub struct ManualPolicy {
routing_map: DashMap<RoutingId, RoutingInfo>,
}
impl ManualPolicy { impl ManualPolicy {
pub fn new() -> Self { pub fn new() -> Self {
Self Self {
routing_map: DashMap::new(),
}
}
fn select_by_routing_id(
&self,
workers: &[Arc<dyn Worker>],
routing_id: &str,
healthy_indices: &[usize],
) -> (usize, ExecutionBranch) {
let routing_id = RoutingId::new(routing_id);
// Fast path
if let Some(info) = self.routing_map.get(&routing_id) {
if let Some(idx) =
find_healthy_worker(&info.candi_worker_urls, workers, healthy_indices)
{
return (idx, ExecutionBranch::FastPathHit);
}
}
// Slow path
match self.routing_map.entry(routing_id) {
Entry::Occupied(mut entry) => {
if let Some(idx) =
find_healthy_worker(&entry.get().candi_worker_urls, workers, healthy_indices)
{
return (idx, ExecutionBranch::SlowPathOccupiedHit);
}
let selected_idx = random_select(healthy_indices);
entry
.get_mut()
.push_bounded(workers[selected_idx].url().to_string());
(selected_idx, ExecutionBranch::SlowPathOccupiedMiss)
}
Entry::Vacant(entry) => {
let selected_idx = random_select(healthy_indices);
entry.insert(RoutingInfo {
candi_worker_urls: vec![workers[selected_idx].url().to_string()],
});
(selected_idx, ExecutionBranch::SlowPathVacant)
}
}
} }
fn select_worker_impl( fn select_worker_impl(
&self, &self,
workers: &[Arc<dyn Worker>], workers: &[Arc<dyn Worker>],
info: &SelectWorkerInfo, info: &SelectWorkerInfo,
) -> (Option<usize>, Branch) { ) -> (Option<usize>, ExecutionBranch) {
// O(n) - unavoidable, need to know which workers are healthy
let healthy_indices = get_healthy_worker_indices(workers); let healthy_indices = get_healthy_worker_indices(workers);
if healthy_indices.is_empty() { if healthy_indices.is_empty() {
return (None, Branch::NoHealthyWorkers); return (None, ExecutionBranch::NoHealthyWorkers);
} }
// Extract routing headers - to_str() is O(1), just validates ASCII, no allocation // Extract routing key from header
let target_worker = info let routing_id = info
.headers
.and_then(|h| h.get(&HEADER_TARGET_WORKER))
.and_then(|v| v.to_str().ok())
.filter(|s| !s.is_empty());
let routing_key = info
.headers .headers
.and_then(|h| h.get(&HEADER_ROUTING_KEY)) .and_then(|h| h.get(&HEADER_ROUTING_KEY))
.and_then(|v| v.to_str().ok()) .and_then(|v| v.to_str().ok())
.filter(|s| !s.is_empty()); .filter(|s| !s.is_empty());
// Priority 1: X-SMG-Target-Worker - direct routing by worker index if let Some(routing_id) = routing_id {
// O(1) parse + O(1) bounds check + O(1) health check let (idx, branch) = self.select_by_routing_id(workers, routing_id, &healthy_indices);
if let Some(idx_str) = target_worker { return (Some(idx), branch);
if let Ok(idx) = idx_str.parse::<usize>() {
if idx < workers.len() && workers[idx].is_healthy() {
return (Some(idx), Branch::TargetWorkerHit);
}
}
return (None, Branch::TargetWorkerMiss);
} }
// Priority 2: X-SMG-Routing-Key - consistent hash routing (
// O(key_len) hash + O(1) modulo + O(1) index Some(random_select(&healthy_indices)),
if let Some(key) = routing_key { ExecutionBranch::NoRoutingId,
let mut hasher = std::collections::hash_map::DefaultHasher::new(); )
key.hash(&mut hasher);
let idx = (hasher.finish() as usize) % healthy_indices.len();
return (Some(healthy_indices[idx]), Branch::RoutingKeyHit);
}
// Fallback: random selection using thread-local RNG (fast, no allocation)
let idx = rand::rng().random_range(0..healthy_indices.len());
(Some(healthy_indices[idx]), Branch::RandomFallback)
} }
} }
@@ -119,6 +173,32 @@ impl LoadBalancingPolicy for ManualPolicy {
} }
} }
fn find_healthy_worker(
urls: &[String],
workers: &[Arc<dyn Worker>],
healthy_indices: &[usize],
) -> Option<usize> {
for url in urls {
if let Some(idx) = find_worker_index_by_url(workers, url) {
if healthy_indices.contains(&idx) {
return Some(idx);
}
}
}
None
}
fn find_worker_index_by_url(workers: &[Arc<dyn Worker>], url: &str) -> Option<usize> {
workers.iter().position(|w| w.url() == url)
}
// TODO: use load-aware selection later
fn random_select(healthy_indices: &[usize]) -> usize {
let mut rng = rand::rng();
let random_idx = rng.random_range(0..healthy_indices.len());
healthy_indices[random_idx]
}
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use std::collections::HashMap; use std::collections::HashMap;
@@ -126,18 +206,6 @@ mod tests {
use super::*; use super::*;
use crate::core::{BasicWorkerBuilder, WorkerType}; use crate::core::{BasicWorkerBuilder, WorkerType};
fn headers_with_routing_key(key: &str) -> http::HeaderMap {
let mut headers = http::HeaderMap::new();
headers.insert("x-smg-routing-key", key.parse().unwrap());
headers
}
fn headers_with_target_worker(idx: usize) -> http::HeaderMap {
let mut headers = http::HeaderMap::new();
headers.insert("x-smg-target-worker", idx.to_string().parse().unwrap());
headers
}
fn create_workers(urls: &[&str]) -> Vec<Arc<dyn Worker>> { fn create_workers(urls: &[&str]) -> Vec<Arc<dyn Worker>> {
urls.iter() urls.iter()
.map(|url| { .map(|url| {
@@ -150,8 +218,14 @@ mod tests {
.collect() .collect()
} }
fn headers_with_routing_key(key: &str) -> http::HeaderMap {
let mut headers = http::HeaderMap::new();
headers.insert("x-smg-routing-key", key.parse().unwrap());
headers
}
#[test] #[test]
fn test_consistent_routing() { fn test_manual_consistent_routing() {
let policy = ManualPolicy::new(); let policy = ManualPolicy::new();
let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]); let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
@@ -161,19 +235,23 @@ mod tests {
..Default::default() ..Default::default()
}; };
let (first_result, _) = policy.select_worker_impl(&workers, &info); let (first_result, branch) = policy.select_worker_impl(&workers, &info);
let first_idx = first_result.unwrap(); let first_idx = first_result.unwrap();
assert_eq!(branch, ExecutionBranch::SlowPathVacant);
// Same key should always route to same worker
for _ in 0..10 { for _ in 0..10 {
let (result, branch) = policy.select_worker_impl(&workers, &info); let (result, branch) = policy.select_worker_impl(&workers, &info);
assert_eq!(result, Some(first_idx)); assert_eq!(
assert_eq!(branch, Branch::RoutingKeyHit); result,
Some(first_idx),
"Same routing_id should route to same worker"
);
assert_eq!(branch, ExecutionBranch::FastPathHit);
} }
} }
#[test] #[test]
fn test_different_keys_distribute() { fn test_manual_different_routing_ids() {
let policy = ManualPolicy::new(); let policy = ManualPolicy::new();
let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]); let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
@@ -184,133 +262,103 @@ mod tests {
headers: Some(&headers), headers: Some(&headers),
..Default::default() ..Default::default()
}; };
let (result, _) = policy.select_worker_impl(&workers, &info); let (result, branch) = policy.select_worker_impl(&workers, &info);
assert_eq!(branch, ExecutionBranch::SlowPathVacant);
*distribution.entry(result.unwrap()).or_insert(0) += 1; *distribution.entry(result.unwrap()).or_insert(0) += 1;
} }
assert!(distribution.len() > 1, "Should distribute across workers");
}
#[test]
fn test_target_worker_hit() {
let policy = ManualPolicy::new();
let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
let headers = headers_with_target_worker(1);
let info = SelectWorkerInfo {
headers: Some(&headers),
..Default::default()
};
let (result, branch) = policy.select_worker_impl(&workers, &info);
assert_eq!(result, Some(1));
assert_eq!(branch, Branch::TargetWorkerHit);
}
#[test]
fn test_target_worker_miss_out_of_bounds() {
let policy = ManualPolicy::new();
let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
let headers = headers_with_target_worker(5); // Out of bounds
let info = SelectWorkerInfo {
headers: Some(&headers),
..Default::default()
};
let (result, branch) = policy.select_worker_impl(&workers, &info);
assert_eq!(result, None);
assert_eq!(branch, Branch::TargetWorkerMiss);
}
#[test]
fn test_target_worker_miss_unhealthy() {
let policy = ManualPolicy::new();
let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
workers[1].set_healthy(false);
let headers = headers_with_target_worker(1);
let info = SelectWorkerInfo {
headers: Some(&headers),
..Default::default()
};
let (result, branch) = policy.select_worker_impl(&workers, &info);
assert_eq!(result, None);
assert_eq!(branch, Branch::TargetWorkerMiss);
}
#[test]
fn test_target_worker_priority_over_routing_key() {
let policy = ManualPolicy::new();
let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
let mut headers = http::HeaderMap::new();
headers.insert("x-smg-target-worker", "1".parse().unwrap());
headers.insert("x-smg-routing-key", "some-key".parse().unwrap());
let info = SelectWorkerInfo {
headers: Some(&headers),
..Default::default()
};
let (result, branch) = policy.select_worker_impl(&workers, &info);
assert_eq!(result, Some(1));
assert_eq!(branch, Branch::TargetWorkerHit);
}
#[test]
fn test_fallback_random_distribution() {
let policy = ManualPolicy::new();
let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
// Without routing headers, should distribute randomly across workers
let mut distribution = HashMap::new();
for _ in 0..100 {
let info = SelectWorkerInfo::default();
let (result, branch) = policy.select_worker_impl(&workers, &info);
assert!(result.is_some());
assert_eq!(branch, Branch::RandomFallback);
*distribution.entry(result.unwrap()).or_insert(0) += 1;
}
// Should distribute across multiple workers (not always same one)
assert!( assert!(
distribution.len() > 1, distribution.len() > 1,
"Random fallback should distribute across workers" "Should distribute across multiple workers"
); );
} }
#[test] #[test]
fn test_no_healthy_workers() { fn test_manual_fallback_random() {
let policy = ManualPolicy::new(); let policy = ManualPolicy::new();
let workers = create_workers(&["http://w1:8000"]); let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
let mut counts = HashMap::new();
for _ in 0..100 {
let info = SelectWorkerInfo::default();
let (result, branch) = policy.select_worker_impl(&workers, &info);
assert_eq!(branch, ExecutionBranch::NoRoutingId);
if let Some(idx) = result {
*counts.entry(idx).or_insert(0) += 1;
}
}
assert_eq!(counts.len(), 2, "Random fallback should use all workers");
}
#[test]
fn test_manual_with_unhealthy_workers() {
let policy = ManualPolicy::new();
let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
workers[0].set_healthy(false); workers[0].set_healthy(false);
let headers = headers_with_routing_key("test-routing-id");
let info = SelectWorkerInfo {
headers: Some(&headers),
..Default::default()
};
let (result, branch) = policy.select_worker_impl(&workers, &info);
assert_eq!(result, Some(1), "Should only select healthy worker");
assert_eq!(branch, ExecutionBranch::SlowPathVacant);
for _ in 0..10 {
let (result, branch) = policy.select_worker_impl(&workers, &info);
assert_eq!(result, Some(1), "Should only select healthy worker");
assert_eq!(branch, ExecutionBranch::FastPathHit);
}
}
#[test]
fn test_manual_no_healthy_workers() {
let policy = ManualPolicy::new();
let workers = create_workers(&["http://w1:8000"]);
workers[0].set_healthy(false);
let headers = headers_with_routing_key("test"); let headers = headers_with_routing_key("test");
let info = SelectWorkerInfo { let info = SelectWorkerInfo {
headers: Some(&headers), headers: Some(&headers),
..Default::default() ..Default::default()
}; };
let (result, branch) = policy.select_worker_impl(&workers, &info); let (result, branch) = policy.select_worker_impl(&workers, &info);
assert_eq!(result, None); assert_eq!(result, None);
assert_eq!(branch, Branch::NoHealthyWorkers); assert_eq!(branch, ExecutionBranch::NoHealthyWorkers);
} }
#[test] #[test]
fn test_empty_workers() { fn test_manual_empty_routing_id() {
let policy = ManualPolicy::new(); let policy = ManualPolicy::new();
let workers: Vec<Arc<dyn Worker>> = vec![]; let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
let info = SelectWorkerInfo::default(); let mut counts = HashMap::new();
for _ in 0..100 {
let headers = headers_with_routing_key("");
let info = SelectWorkerInfo {
headers: Some(&headers),
..Default::default()
};
let (result, branch) = policy.select_worker_impl(&workers, &info); let (result, branch) = policy.select_worker_impl(&workers, &info);
assert_eq!(result, None); assert_eq!(branch, ExecutionBranch::NoRoutingId);
assert_eq!(branch, Branch::NoHealthyWorkers); if let Some(idx) = result {
*counts.entry(idx).or_insert(0) += 1;
}
}
assert_eq!(
counts.len(),
2,
"Empty routing_id should use random fallback"
);
} }
#[test] #[test]
fn test_routing_key_remaps_when_worker_unhealthy() { fn test_manual_remaps_when_worker_becomes_unhealthy() {
let policy = ManualPolicy::new(); let policy = ManualPolicy::new();
let workers = create_workers(&["http://w1:8000", "http://w2:8000"]); let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
@@ -320,37 +368,217 @@ mod tests {
..Default::default() ..Default::default()
}; };
let (first_result, _) = policy.select_worker_impl(&workers, &info); let (first_result, branch) = policy.select_worker_impl(&workers, &info);
let first_idx = first_result.unwrap(); let first_idx = first_result.unwrap();
assert_eq!(branch, ExecutionBranch::SlowPathVacant);
// Mark that worker unhealthy
workers[first_idx].set_healthy(false); workers[first_idx].set_healthy(false);
// Should now route to the other worker let (new_result, branch) = policy.select_worker_impl(&workers, &info);
let (new_result, _) = policy.select_worker_impl(&workers, &info);
let new_idx = new_result.unwrap(); let new_idx = new_result.unwrap();
assert_ne!(new_idx, first_idx); assert_ne!(new_idx, first_idx, "Should remap to healthy worker");
assert_eq!(branch, ExecutionBranch::SlowPathOccupiedMiss);
for _ in 0..10 {
let (result, branch) = policy.select_worker_impl(&workers, &info);
assert_eq!(
result,
Some(new_idx),
"Should consistently route to new worker"
);
assert_eq!(branch, ExecutionBranch::FastPathHit);
}
} }
#[test] #[test]
fn test_empty_routing_key_uses_fallback() { fn test_manual_empty_workers() {
let policy = ManualPolicy::new(); let policy = ManualPolicy::new();
let workers = create_workers(&["http://w1:8000", "http://w2:8000"]); let workers: Vec<Arc<dyn Worker>> = vec![];
let headers = headers_with_routing_key("test");
let info = SelectWorkerInfo {
headers: Some(&headers),
..Default::default()
};
let (result, branch) = policy.select_worker_impl(&workers, &info);
assert_eq!(result, None);
assert_eq!(branch, ExecutionBranch::NoHealthyWorkers);
}
let headers = headers_with_routing_key(""); #[test]
fn test_manual_single_worker() {
let policy = ManualPolicy::new();
let workers = create_workers(&["http://w1:8000"]);
let headers = headers_with_routing_key("single-test");
let info = SelectWorkerInfo { let info = SelectWorkerInfo {
headers: Some(&headers), headers: Some(&headers),
..Default::default() ..Default::default()
}; };
let (result, branch) = policy.select_worker_impl(&workers, &info); let (result, branch) = policy.select_worker_impl(&workers, &info);
assert!(result.is_some()); assert_eq!(result, Some(0));
assert_eq!(branch, Branch::RandomFallback); assert_eq!(branch, ExecutionBranch::SlowPathVacant);
for _ in 0..10 {
let (result, branch) = policy.select_worker_impl(&workers, &info);
assert_eq!(result, Some(0));
assert_eq!(branch, ExecutionBranch::FastPathHit);
}
} }
#[test] #[test]
fn test_policy_name() { fn test_manual_worker_recovery() {
let policy = ManualPolicy::new();
let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
let headers = headers_with_routing_key("recovery-test");
let info = SelectWorkerInfo {
headers: Some(&headers),
..Default::default()
};
let (first_result, branch) = policy.select_worker_impl(&workers, &info);
let first_idx = first_result.unwrap();
assert_eq!(branch, ExecutionBranch::SlowPathVacant);
workers[first_idx].set_healthy(false);
let (second_result, branch) = policy.select_worker_impl(&workers, &info);
let second_idx = second_result.unwrap();
assert_ne!(second_idx, first_idx);
assert_eq!(branch, ExecutionBranch::SlowPathOccupiedMiss);
workers[first_idx].set_healthy(true);
let (after_recovery, branch) = policy.select_worker_impl(&workers, &info);
assert_eq!(
after_recovery,
Some(first_idx),
"Should return to original worker after recovery since it's first in candidate list"
);
assert_eq!(branch, ExecutionBranch::FastPathHit);
}
#[test]
fn test_manual_max_candidate_workers_eviction() {
let policy = ManualPolicy::new();
let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
let headers = headers_with_routing_key("eviction-test");
let info = SelectWorkerInfo {
headers: Some(&headers),
..Default::default()
};
let (first_result, branch) = policy.select_worker_impl(&workers, &info);
let first_idx = first_result.unwrap();
assert_eq!(branch, ExecutionBranch::SlowPathVacant);
workers[first_idx].set_healthy(false);
let (second_result, branch) = policy.select_worker_impl(&workers, &info);
let second_idx = second_result.unwrap();
assert_ne!(second_idx, first_idx);
assert_eq!(branch, ExecutionBranch::SlowPathOccupiedMiss);
workers[second_idx].set_healthy(false);
let remaining_idx = (0..3).find(|&i| i != first_idx && i != second_idx).unwrap();
let (third_result, branch) = policy.select_worker_impl(&workers, &info);
assert_eq!(
third_result,
Some(remaining_idx),
"Should select the only remaining healthy worker"
);
assert_eq!(branch, ExecutionBranch::SlowPathOccupiedMiss);
workers[first_idx].set_healthy(true);
let (idx_after_restore, branch) = policy.select_worker_impl(&workers, &info);
assert_ne!(
idx_after_restore,
Some(first_idx),
"First worker should be evicted from candidates due to MAX_CANDIDATE_WORKERS=2"
);
assert_eq!(branch, ExecutionBranch::FastPathHit);
}
#[test]
fn test_manual_policy_name() {
let policy = ManualPolicy::new(); let policy = ManualPolicy::new();
assert_eq!(policy.name(), "manual"); assert_eq!(policy.name(), "manual");
} }
#[test]
fn test_manual_routing_info_push_bounded() {
let mut info = RoutingInfo {
candi_worker_urls: vec!["http://w1:8000".to_string()],
};
info.push_bounded("http://w2:8000".to_string());
assert_eq!(info.candi_worker_urls.len(), 2);
assert_eq!(info.candi_worker_urls[0], "http://w1:8000");
assert_eq!(info.candi_worker_urls[1], "http://w2:8000");
info.push_bounded("http://w3:8000".to_string());
assert_eq!(info.candi_worker_urls.len(), 2);
assert_eq!(
info.candi_worker_urls[0], "http://w2:8000",
"Oldest entry should be removed"
);
assert_eq!(info.candi_worker_urls[1], "http://w3:8000");
}
#[test]
fn test_manual_find_healthy_worker_priority() {
let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
let urls = vec![
"http://w1:8000".to_string(),
"http://w2:8000".to_string(),
"http://w3:8000".to_string(),
];
let healthy_indices = vec![0, 1, 2];
let result = find_healthy_worker(&urls, &workers, &healthy_indices);
assert_eq!(
result,
Some(0),
"Should return first healthy worker in urls"
);
workers[0].set_healthy(false);
let healthy_indices = vec![1, 2];
let result = find_healthy_worker(&urls, &workers, &healthy_indices);
assert_eq!(result, Some(1), "Should skip unhealthy and return next");
workers[1].set_healthy(false);
let healthy_indices = vec![2];
let result = find_healthy_worker(&urls, &workers, &healthy_indices);
assert_eq!(result, Some(2), "Should return last healthy worker");
workers[2].set_healthy(false);
let healthy_indices: Vec<usize> = vec![];
let result = find_healthy_worker(&urls, &workers, &healthy_indices);
assert_eq!(result, None, "Should return None when no healthy workers");
}
#[test]
fn test_manual_find_worker_index_by_url() {
let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
assert_eq!(
find_worker_index_by_url(&workers, "http://w1:8000"),
Some(0)
);
assert_eq!(
find_worker_index_by_url(&workers, "http://w2:8000"),
Some(1)
);
assert_eq!(
find_worker_index_by_url(&workers, "http://w3:8000"),
None,
"Should return None for unknown URL"
);
}
} }
+9 -4
View File
@@ -5,10 +5,11 @@
use std::{fmt::Debug, sync::Arc}; use std::{fmt::Debug, sync::Arc};
use crate::core::Worker; use crate::core::{HashRing, Worker};
mod bucket; mod bucket;
mod cache_aware; mod cache_aware;
mod consistent_hashing;
mod factory; mod factory;
mod manual; mod manual;
mod power_of_two; mod power_of_two;
@@ -19,6 +20,7 @@ pub mod tree;
pub use bucket::BucketPolicy; pub use bucket::BucketPolicy;
pub use cache_aware::CacheAwarePolicy; pub use cache_aware::CacheAwarePolicy;
pub use consistent_hashing::ConsistentHashingPolicy;
pub use factory::PolicyFactory; pub use factory::PolicyFactory;
pub use manual::ManualPolicy; pub use manual::ManualPolicy;
pub use power_of_two::PowerOfTwoPolicy; pub use power_of_two::PowerOfTwoPolicy;
@@ -138,15 +140,18 @@ pub(crate) fn normalize_model_key(model_id: &str) -> &str {
} }
/// Information passed to policy for worker selection /// Information passed to policy for worker selection
#[derive(Debug, Default, Clone)] #[derive(Debug, Clone, Default)]
pub struct SelectWorkerInfo<'a> { pub struct SelectWorkerInfo<'a> {
/// Request text for cache-aware routing /// Request text for cache-aware routing
pub request_text: Option<&'a str>, pub request_text: Option<&'a str>,
/// HTTP headers for header-based routing policies /// HTTP headers for header-based routing policies
/// Policies can extract routing information from headers like: /// Policies can extract routing information from headers like:
/// - X-Target-Worker: Direct routing to a specific worker by URL /// - X-SMG-Target-Worker: Direct routing to a specific worker by index
/// - X-Routing-Key: Consistent hash routing for session affinity /// - X-SMG-Routing-Key: Consistent hash routing for session affinity
pub headers: Option<&'a http::HeaderMap>, pub headers: Option<&'a http::HeaderMap>,
/// Pre-computed hash ring for O(log n) consistent hashing
/// Built and cached by WorkerRegistry, passed through to avoid per-request rebuilds
pub hash_ring: Option<Arc<HashRing>>,
} }
#[cfg(test)] #[cfg(test)]
+5 -2
View File
@@ -10,8 +10,8 @@ use tracing::{debug, info, warn};
/// All subsequent workers of the same model use the established policy. /// All subsequent workers of the same model use the established policy.
/// When the last worker of a model is removed, the policy mapping is cleaned up. /// When the last worker of a model is removed, the policy mapping is cleaned up.
use super::{ use super::{
BucketConfig, BucketPolicy, CacheAwareConfig, CacheAwarePolicy, LoadBalancingPolicy, BucketConfig, BucketPolicy, CacheAwareConfig, CacheAwarePolicy, ConsistentHashingPolicy,
ManualPolicy, PowerOfTwoPolicy, RandomPolicy, RoundRobinPolicy, LoadBalancingPolicy, ManualPolicy, PowerOfTwoPolicy, RandomPolicy, RoundRobinPolicy,
}; };
use crate::{config::types::PolicyConfig, core::Worker}; use crate::{config::types::PolicyConfig, core::Worker};
@@ -168,6 +168,8 @@ impl PolicyRegistry {
"cache_aware" => Arc::new(CacheAwarePolicy::new()), "cache_aware" => Arc::new(CacheAwarePolicy::new()),
"power_of_two" => Arc::new(PowerOfTwoPolicy::new()), "power_of_two" => Arc::new(PowerOfTwoPolicy::new()),
"bucket" => Arc::new(BucketPolicy::new()), "bucket" => Arc::new(BucketPolicy::new()),
"manual" => Arc::new(ManualPolicy::new()),
"consistent_hashing" => Arc::new(ConsistentHashingPolicy::new()),
_ => { _ => {
warn!("Unknown policy type '{}', using default", policy_type); warn!("Unknown policy type '{}', using default", policy_type);
Arc::clone(&self.default_policy) Arc::clone(&self.default_policy)
@@ -210,6 +212,7 @@ impl PolicyRegistry {
Arc::new(BucketPolicy::with_config(config)) Arc::new(BucketPolicy::with_config(config))
} }
PolicyConfig::Manual => Arc::new(ManualPolicy::new()), PolicyConfig::Manual => Arc::new(ManualPolicy::new()),
PolicyConfig::ConsistentHashing => Arc::new(ConsistentHashingPolicy::new()),
} }
} }
@@ -8,7 +8,7 @@ use tracing::{error, warn};
use super::PipelineStage; use super::PipelineStage;
use crate::{ use crate::{
core::{ConnectionMode, Worker, WorkerRegistry, WorkerType}, core::{ConnectionMode, Worker, WorkerRegistry, WorkerType, UNKNOWN_MODEL_ID},
observability::metrics::{metrics_labels, Metrics}, observability::metrics::{metrics_labels, Metrics},
policies::{PolicyRegistry, SelectWorkerInfo}, policies::{PolicyRegistry, SelectWorkerInfo},
routers::{ routers::{
@@ -148,12 +148,18 @@ impl WorkerSelectionStage {
None => self.policy_registry.get_default_policy(), None => self.policy_registry.get_default_policy(),
}; };
// Get cached hash ring for consistent hashing (O(log n) lookup)
let hash_ring = self
.worker_registry
.get_hash_ring(model_id.unwrap_or(UNKNOWN_MODEL_ID));
// Select worker using the policy // Select worker using the policy
let idx = policy.select_worker( let idx = policy.select_worker(
&available, &available,
&SelectWorkerInfo { &SelectWorkerInfo {
request_text: text, request_text: text,
headers, headers,
hash_ring,
}, },
)?; )?;
let selected = available[idx].clone(); let selected = available[idx].clone();
@@ -213,9 +219,15 @@ impl WorkerSelectionStage {
None => self.policy_registry.get_default_policy(), None => self.policy_registry.get_default_policy(),
}; };
// Get cached hash ring for consistent hashing (O(log n) lookup)
let hash_ring = self
.worker_registry
.get_hash_ring(model_id.unwrap_or(UNKNOWN_MODEL_ID));
let info = SelectWorkerInfo { let info = SelectWorkerInfo {
request_text: text, request_text: text,
headers, headers,
hash_ring,
}; };
let prefill_idx = policy.select_worker(&available_prefill, &info)?; let prefill_idx = policy.select_worker(&available_prefill, &info)?;
let decode_idx = policy.select_worker(&available_decode, &info)?; let decode_idx = policy.select_worker(&available_decode, &info)?;
@@ -19,7 +19,8 @@ use super::pd_types::api_path;
use crate::{ use crate::{
config::types::RetryConfig, config::types::RetryConfig,
core::{ core::{
is_retryable_status, RetryExecutor, Worker, WorkerLoadGuard, WorkerRegistry, WorkerType, is_retryable_status, HashRing, RetryExecutor, Worker, WorkerLoadGuard, WorkerRegistry,
WorkerType, UNKNOWN_MODEL_ID,
}, },
observability::{ observability::{
events::{self, Event}, events::{self, Event},
@@ -731,11 +732,17 @@ impl PDRouter {
let prefill_policy = self.policy_registry.get_prefill_policy(); let prefill_policy = self.policy_registry.get_prefill_policy();
let decode_policy = self.policy_registry.get_decode_policy(); let decode_policy = self.policy_registry.get_decode_policy();
// Get cached hash ring for consistent hashing
let hash_ring = self
.worker_registry
.get_hash_ring(effective_model_id.unwrap_or(UNKNOWN_MODEL_ID));
let prefill = Self::pick_worker_by_policy_arc( let prefill = Self::pick_worker_by_policy_arc(
&prefill_workers, &prefill_workers,
&*prefill_policy, &*prefill_policy,
request_text, request_text,
headers, headers,
hash_ring.clone(),
"prefill", "prefill",
)?; )?;
@@ -744,6 +751,7 @@ impl PDRouter {
&*decode_policy, &*decode_policy,
request_text, request_text,
headers, headers,
hash_ring,
"decode", "decode",
)?; )?;
@@ -770,6 +778,7 @@ impl PDRouter {
policy: &dyn LoadBalancingPolicy, policy: &dyn LoadBalancingPolicy,
request_text: Option<&str>, request_text: Option<&str>,
headers: Option<&HeaderMap>, headers: Option<&HeaderMap>,
hash_ring: Option<Arc<HashRing>>,
worker_type: &str, worker_type: &str,
) -> Result<Arc<dyn Worker>, String> { ) -> Result<Arc<dyn Worker>, String> {
if workers.is_empty() { if workers.is_empty() {
@@ -798,6 +807,7 @@ impl PDRouter {
&SelectWorkerInfo { &SelectWorkerInfo {
request_text, request_text,
headers, headers,
hash_ring,
}, },
) )
.ok_or_else(|| { .ok_or_else(|| {
+7 -1
View File
@@ -20,7 +20,7 @@ use crate::{
config::types::RetryConfig, config::types::RetryConfig,
core::{ core::{
is_retryable_status, ConnectionMode, RetryExecutor, Worker, WorkerLoadGuard, is_retryable_status, ConnectionMode, RetryExecutor, Worker, WorkerLoadGuard,
WorkerRegistry, WorkerType, WorkerRegistry, WorkerType, UNKNOWN_MODEL_ID,
}, },
observability::{ observability::{
events::{self, Event}, events::{self, Event},
@@ -169,11 +169,17 @@ impl Router {
None => self.policy_registry.get_default_policy(), None => self.policy_registry.get_default_policy(),
}; };
// Get cached hash ring for consistent hashing (O(log n) lookup)
let hash_ring = self
.worker_registry
.get_hash_ring(effective_model_id.unwrap_or(UNKNOWN_MODEL_ID));
let idx = policy.select_worker( let idx = policy.select_worker(
&available, &available,
&SelectWorkerInfo { &SelectWorkerInfo {
request_text: text, request_text: text,
headers, headers,
hash_ring,
}, },
)?; )?;