diff --git a/crates/kv_index/benches/throughput_bench.rs b/crates/kv_index/benches/throughput_bench.rs index fc5caadb10..2f8c5a12c7 100644 --- a/crates/kv_index/benches/throughput_bench.rs +++ b/crates/kv_index/benches/throughput_bench.rs @@ -667,7 +667,9 @@ async fn run_benchmark(args: &Args, traces: Vec>) -> BenchmarkRe let num_total_workers = args.num_workers * args.duplication_factor; for w in 0..num_total_workers { - indexer.intern_worker(&format!("worker-{w}")); + indexer + .intern_worker(&format!("worker-{w}")) + .expect("worker id space exhausted"); } let traces: Vec>> = traces.into_iter().map(Arc::new).collect(); diff --git a/crates/kv_index/src/event_tree.rs b/crates/kv_index/src/event_tree.rs index d503547064..919ea6cbf4 100644 --- a/crates/kv_index/src/event_tree.rs +++ b/crates/kv_index/src/event_tree.rs @@ -23,8 +23,8 @@ use std::{ fmt, sync::{ - atomic::{AtomicU32, AtomicUsize, Ordering}, - Arc, + atomic::{AtomicU32, AtomicU64, AtomicUsize, Ordering}, + Arc, OnceLock, }, }; @@ -43,9 +43,16 @@ const INDEX_SHARD_COUNT: usize = 1024; /// These maps hold at most ~500 entries (one per worker), so 8 shards is sufficient. const WORKER_SHARD_COUNT: usize = 8; -/// Maximum number of workers supported. tree_sizes is a flat Vec indexed by worker_id, -/// giving lock-free reads on the query hot path (array index vs DashMap hash+lock+probe). -const MAX_WORKERS: usize = 2048; +/// Length of the first `TreeSizes` segment (covers worker ids 0..2048). +/// Segment `s` doubles to `FIRST_SEGMENT_LEN << s` entries, so worker count is +/// unbounded while reads stay a lock-free array index on the query hot path. +const FIRST_SEGMENT_LEN: usize = 2048; + +/// log2 of [`FIRST_SEGMENT_LEN`]. +const FIRST_SEGMENT_BITS: u32 = FIRST_SEGMENT_LEN.trailing_zeros(); + +/// Number of doubling segments needed to cover the entire u32 worker-id space. +const SEGMENT_COUNT: usize = (u32::BITS - FIRST_SEGMENT_BITS + 1) as usize; /// Position-independent content hash of tokens within a single block. /// Computed via XXH3-64 from token IDs. Same tokens always produce the same hash @@ -105,6 +112,21 @@ impl fmt::Display for ApplyError { impl std::error::Error for ApplyError {} +/// Error returned by [`PositionalIndexer::intern_worker`] when the u32 worker-id +/// space is exhausted. Ids are assigned monotonically and never recycled, so this +/// can only be reached after `u32::MAX + 1` distinct worker URLs have been interned +/// over the indexer's lifetime. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct WorkerIdExhausted; + +impl fmt::Display for WorkerIdExhausted { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "worker id space exhausted (u32::MAX workers interned)") + } +} + +impl std::error::Error for WorkerIdExhausted {} + /// Overlap scores: how many consecutive blocks each worker has cached. /// /// Keys are internal `u32` worker IDs. Use [`PositionalIndexer::worker_id`] to @@ -235,6 +257,92 @@ impl SeqEntry { } } +// --------------------------------------------------------------------------- +// TreeSizes: growable, lock-free per-worker block counters +// --------------------------------------------------------------------------- + +/// Per-worker block counters indexed by worker id, with no worker-count cap. +/// +/// Storage is segmented with doubling sizes: segment `s` covers ids +/// `[FIRST_SEGMENT_LEN * (2^s - 1), FIRST_SEGMENT_LEN * (2^(s+1) - 1))`, so +/// [`SEGMENT_COUNT`] segments cover every possible u32 id. Segments are +/// allocated on first write and never moved, which keeps reads a lock-free +/// array index (~1ns) on the query hot path — the property the previous fixed +/// 2048-slot Vec was built for — while supporting unbounded worker counts. +/// +/// Memory: 8 bytes per worker id, at most ~2x the high-water id due to +/// doubling; segments are never freed. Ids are never recycled, so cost grows +/// with workers ever interned — same lifecycle as the `worker_to_id` map, +/// whose per-URL entries are an order of magnitude larger. +struct TreeSizes { + segments: [OnceLock>; SEGMENT_COUNT], +} + +impl TreeSizes { + fn new() -> Self { + Self { + segments: std::array::from_fn(|_| OnceLock::new()), + } + } + + /// Map a worker id to (segment index, offset within segment). + /// + /// Computed in u64: `id + FIRST_SEGMENT_LEN` overflows u32 for ids near + /// `u32::MAX`, and the last segment's length (`2^32`) overflows 32-bit usize. + #[inline] + fn locate(id: u32) -> (usize, usize) { + let virtual_idx = id as u64 + FIRST_SEGMENT_LEN as u64; + let msb = 63 - virtual_idx.leading_zeros(); + let segment = (msb - FIRST_SEGMENT_BITS) as usize; + let offset = (virtual_idx - (1u64 << msb)) as usize; + (segment, offset) + } + + #[inline] + fn segment_len(segment: usize) -> usize { + FIRST_SEGMENT_LEN << segment + } + + /// Counter slot for a worker id, allocating its segment on first use. + fn slot(&self, id: u32) -> &AtomicUsize { + let (segment, offset) = Self::locate(id); + let entries = self.segments[segment].get_or_init(|| { + (0..Self::segment_len(segment)) + .map(|_| AtomicUsize::new(0)) + .collect() + }); + &entries[offset] + } + + /// Lock-free read. Ids whose segment was never written read as 0. + #[inline] + fn load(&self, id: u32) -> usize { + let (segment, offset) = Self::locate(id); + match self.segments[segment].get() { + Some(entries) => entries[offset].load(Ordering::Relaxed), + None => 0, + } + } + + /// Reset a worker's count to 0 without allocating its segment if absent. + fn reset(&self, id: u32) { + let (segment, offset) = Self::locate(id); + if let Some(entries) = self.segments[segment].get() { + entries[offset].store(0, Ordering::Relaxed); + } + } + + /// Sum of all counters across allocated segments. + fn total(&self) -> usize { + self.segments + .iter() + .filter_map(OnceLock::get) + .flat_map(|entries| entries.iter()) + .map(|size| size.load(Ordering::Relaxed)) + .sum() + } +} + // --------------------------------------------------------------------------- // PositionalIndexer // --------------------------------------------------------------------------- @@ -259,17 +367,29 @@ pub struct PositionalIndexer { /// No capacity limit — grows as blocks are stored. index: DashMap<(usize, ContentHash), SeqEntry, FxBuildHasher>, /// Per-worker block counts, tracked atomically for O(1) reads during queries. - /// Flat Vec indexed by worker_id — lock-free reads on the query hot path - /// (array index ~1ns vs DashMap hash+lock+probe ~25ns per access). - tree_sizes: Vec, + /// Segmented array indexed by worker_id — lock-free reads on the query hot + /// path (array index ~1ns vs DashMap hash+lock+probe ~25ns per access), + /// growable so the worker count is unbounded. + tree_sizes: TreeSizes, /// Worker URL → internal u32 ID (fast path: DashMap shard read). worker_to_id: DashMap, u32, FxBuildHasher>, - /// Monotonic counter for assigning new worker IDs. - next_worker_id: AtomicU32, + /// Monotonic counter for assigning new worker IDs. Never recycled; u64 so + /// exhaustion of the u32 id space is detected instead of wrapping. + next_worker_id: AtomicU64, + /// Process-unique, non-zero id for this indexer instance ("epoch"). + /// Consumers cache `worker URL → worker id` resolutions on long-lived + /// worker objects; comparing this id makes such caches self-invalidating + /// when an indexer is dropped and a fresh one is created for the same + /// model (worker ids from different instances are unrelated). + instance_id: u32, /// Jump size for search optimization (default 64). jump_size: usize, } +/// Process-global counter for [`PositionalIndexer`] instance ids. Starts at 1: +/// 0 is reserved as the "no cached id" sentinel in consumer-side caches. +static NEXT_INSTANCE_ID: AtomicU32 = AtomicU32::new(1); + impl PositionalIndexer { /// Create a new PositionalIndexer with the given jump size. /// @@ -278,15 +398,35 @@ impl PositionalIndexer { /// when workers drain. Default: 64. pub fn new(jump_size: usize) -> Self { assert!(jump_size > 0, "jump_size must be greater than 0"); + let instance_id = loop { + let id = NEXT_INSTANCE_ID.fetch_add(1, Ordering::Relaxed); + // Skip the reserved 0 sentinel on the (unreachable in practice) + // u32 wrap after 2^32 indexer constructions. + if id != 0 { + break id; + } + }; Self { index: DashMap::with_hasher_and_shard_amount(FxBuildHasher, INDEX_SHARD_COUNT), - tree_sizes: (0..MAX_WORKERS).map(|_| AtomicUsize::new(0)).collect(), + tree_sizes: TreeSizes::new(), worker_to_id: DashMap::with_hasher_and_shard_amount(FxBuildHasher, WORKER_SHARD_COUNT), - next_worker_id: AtomicU32::new(0), + next_worker_id: AtomicU64::new(0), + instance_id, jump_size, } } + /// Process-unique, non-zero id for this indexer instance ("epoch"). + /// + /// Two `PositionalIndexer` instances never share an id, so a cached + /// `(instance_id, worker_id)` pair is valid for an indexer if and only if + /// the instance id matches. Used by routing policies to cache + /// [`intern_worker`](Self::intern_worker) results on worker objects + /// instead of probing the URL-keyed map on every request. + pub fn instance_id(&self) -> u32 { + self.instance_id + } + /// Get the internal u32 ID for a worker URL, if it has been interned. /// /// Used by consumers to look up scores in [`OverlapScores`] by worker URL. @@ -359,7 +499,9 @@ impl PositionalIndexer { // Atomically update tree_sizes — lock-free array index. if num_new_blocks > 0 { - self.tree_sizes[worker_id as usize].fetch_add(num_new_blocks, Ordering::Relaxed); + self.tree_sizes + .slot(worker_id) + .fetch_add(num_new_blocks, Ordering::Relaxed); } Ok(()) @@ -399,7 +541,9 @@ impl PositionalIndexer { } if num_removed > 0 { - self.tree_sizes[worker_id as usize].fetch_sub(num_removed, Ordering::Relaxed); + self.tree_sizes + .slot(worker_id) + .fetch_sub(num_removed, Ordering::Relaxed); } } @@ -416,7 +560,7 @@ impl PositionalIndexer { } } } - self.tree_sizes[worker_id as usize].store(0, Ordering::Relaxed); + self.tree_sizes.reset(worker_id); } /// Remove a worker entirely — takes ownership of blocks, cleans index, worker is gone. @@ -431,16 +575,12 @@ impl PositionalIndexer { } } } - self.tree_sizes[worker_id as usize].store(0, Ordering::Relaxed); + self.tree_sizes.reset(worker_id); } /// Get total number of blocks across all workers. pub fn current_size(&self) -> usize { - let n = self.next_worker_id.load(Ordering::Relaxed) as usize; - self.tree_sizes[..n] - .iter() - .map(|size| size.load(Ordering::Relaxed)) - .sum() + self.tree_sizes.total() } /// Find overlap scores for a request's content hash sequence. @@ -506,22 +646,28 @@ impl PositionalIndexer { /// Intern a worker URL to an internal u32 ID. /// Fast path: DashMap shard read (no lock). Slow path: DashMap entry API (once per worker). - pub fn intern_worker(&self, worker: &str) -> u32 { + /// + /// Ids are assigned monotonically and never recycled: a URL keeps its id for + /// the indexer's lifetime (including across [`remove_worker`](Self::remove_worker)), + /// and new URLs always get a fresh id. Must never panic — it runs inside + /// per-worker subscription tasks where a panic would silently stop KV event + /// indexing for that worker. The only error is u32 id-space exhaustion. + pub fn intern_worker(&self, worker: &str) -> Result { // Fast path: already interned if let Some(entry) = self.worker_to_id.get(worker) { - return *entry.value(); + return Ok(*entry.value()); + } + // Slow path: the entry API holds the shard lock, so the vacant arm runs + // at most once per URL. Nothing is inserted on the error path. + match self.worker_to_id.entry(Arc::from(worker)) { + Entry::Occupied(entry) => Ok(*entry.get()), + Entry::Vacant(entry) => { + let id = self.next_worker_id.fetch_add(1, Ordering::Relaxed); + let id = u32::try_from(id).map_err(|_| WorkerIdExhausted)?; + entry.insert(id); + Ok(id) + } } - // Slow path: DashMap entry API handles the race — or_insert_with runs at most once. - let id = *self - .worker_to_id - .entry(Arc::from(worker)) - .or_insert_with(|| self.next_worker_id.fetch_add(1, Ordering::Relaxed)) - .value(); - assert!( - (id as usize) < MAX_WORKERS, - "worker count {id} exceeds MAX_WORKERS ({MAX_WORKERS})" - ); - id } // ----------------------------------------------------------------------- @@ -694,10 +840,9 @@ impl PositionalIndexer { } scores.scores = internal_scores; for &int_id in scores.scores.keys() { - scores.tree_sizes.insert( - int_id, - self.tree_sizes[int_id as usize].load(Ordering::Relaxed), - ); + scores + .tree_sizes + .insert(int_id, self.tree_sizes.load(int_id)); } return scores; } @@ -743,10 +888,9 @@ impl PositionalIndexer { // Populate tree_sizes from atomic counters — lock-free array index. for &int_id in scores.scores.keys() { - scores.tree_sizes.insert( - int_id, - self.tree_sizes[int_id as usize].load(Ordering::Relaxed), - ); + scores + .tree_sizes + .insert(int_id, self.tree_sizes.load(int_id)); } scores @@ -810,7 +954,7 @@ mod tests { fn test_store_and_find_single_worker() { let indexer = PositionalIndexer::new(64); let blocks = make_blocks(&[10, 20, 30]); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); indexer.apply_stored(w1, &blocks, None, &mut wb1).unwrap(); @@ -823,7 +967,7 @@ mod tests { fn test_store_partial_prefix_match() { let indexer = PositionalIndexer::new(64); let blocks = make_blocks(&[10, 20, 30]); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); indexer.apply_stored(w1, &blocks, None, &mut wb1).unwrap(); @@ -836,7 +980,7 @@ mod tests { fn test_store_no_match() { let indexer = PositionalIndexer::new(64); let blocks = make_blocks(&[10, 20, 30]); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); indexer.apply_stored(w1, &blocks, None, &mut wb1).unwrap(); @@ -849,8 +993,8 @@ mod tests { let indexer = PositionalIndexer::new(64); let blocks_w1 = make_blocks(&[10, 20, 30]); let blocks_w2 = make_blocks(&[10, 20]); - let w1 = indexer.intern_worker("http://w1:8000"); - let w2 = indexer.intern_worker("http://w2:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); + let w2 = indexer.intern_worker("http://w2:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); let mut wb2 = WorkerBlockMap::default(); indexer @@ -870,7 +1014,7 @@ mod tests { let indexer = PositionalIndexer::new(64); let blocks = make_blocks(&[10, 20, 30]); let seq_hash_of_30 = blocks[2].seq_hash; - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); indexer.apply_stored(w1, &blocks, None, &mut wb1).unwrap(); indexer.apply_removed(w1, &[seq_hash_of_30], &mut wb1); @@ -886,8 +1030,8 @@ mod tests { let indexer = PositionalIndexer::new(64); let blocks_w1 = make_blocks(&[10, 20, 30]); let blocks_w2 = make_blocks(&[10, 20]); - let w1 = indexer.intern_worker("http://w1:8000"); - let w2 = indexer.intern_worker("http://w2:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); + let w2 = indexer.intern_worker("http://w2:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); let mut wb2 = WorkerBlockMap::default(); indexer @@ -909,8 +1053,8 @@ mod tests { let indexer = PositionalIndexer::new(64); let blocks_w1 = make_blocks(&[10, 20, 30]); let blocks_w2 = make_blocks(&[10, 20]); - let w1 = indexer.intern_worker("http://w1:8000"); - let w2 = indexer.intern_worker("http://w2:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); + let w2 = indexer.intern_worker("http://w2:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); let mut wb2 = WorkerBlockMap::default(); indexer @@ -931,7 +1075,7 @@ mod tests { // First store: blocks at positions 0, 1 let blocks1 = make_blocks(&[10, 20]); let parent_seq_hash = blocks1[1].seq_hash; - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); indexer.apply_stored(w1, &blocks1, None, &mut wb1).unwrap(); @@ -959,7 +1103,7 @@ mod tests { fn test_store_with_parent_error_worker_not_tracked() { let indexer = PositionalIndexer::new(64); let blocks = make_blocks(&[10, 20]); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); let result = indexer.apply_stored(w1, &blocks, Some(SequenceHash(999)), &mut wb1); assert!(matches!(result, Err(ApplyError::WorkerNotTracked))); @@ -969,7 +1113,7 @@ mod tests { fn test_store_with_parent_error_parent_not_found() { let indexer = PositionalIndexer::new(64); let blocks1 = make_blocks(&[10, 20]); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); indexer.apply_stored(w1, &blocks1, None, &mut wb1).unwrap(); @@ -982,7 +1126,7 @@ mod tests { fn test_remove_missing_block_is_noop() { let indexer = PositionalIndexer::new(64); let blocks = make_blocks(&[10, 20, 30]); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); indexer.apply_stored(w1, &blocks, None, &mut wb1).unwrap(); @@ -993,7 +1137,7 @@ mod tests { #[test] fn test_remove_unknown_worker_is_noop() { let indexer = PositionalIndexer::new(64); - let w1 = indexer.intern_worker("http://unknown:8000"); + let w1 = indexer.intern_worker("http://unknown:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); indexer.apply_removed(w1, &[SequenceHash(1)], &mut wb1); } @@ -1002,7 +1146,7 @@ mod tests { fn test_remove_worker() { let indexer = PositionalIndexer::new(64); let blocks = make_blocks(&[10, 20, 30]); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); indexer.apply_stored(w1, &blocks, None, &mut wb1).unwrap(); indexer.remove_worker(w1, wb1); @@ -1015,9 +1159,9 @@ mod tests { #[test] fn test_multiple_workers_same_position() { let indexer = PositionalIndexer::new(64); - let w1 = indexer.intern_worker("http://w1:8000"); - let w2 = indexer.intern_worker("http://w2:8000"); - let w3 = indexer.intern_worker("http://w3:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); + let w2 = indexer.intern_worker("http://w2:8000").unwrap(); + let w3 = indexer.intern_worker("http://w3:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); let mut wb2 = WorkerBlockMap::default(); let mut wb3 = WorkerBlockMap::default(); @@ -1040,7 +1184,7 @@ mod tests { #[test] fn test_empty_blocks_is_noop() { let indexer = PositionalIndexer::new(64); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); indexer.apply_stored(w1, &[], None, &mut wb1).unwrap(); assert_eq!(indexer.current_size(), 0); @@ -1050,7 +1194,7 @@ mod tests { fn test_single_block_sequence() { let indexer = PositionalIndexer::new(64); let blocks = make_blocks(&[42]); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); indexer.apply_stored(w1, &blocks, None, &mut wb1).unwrap(); @@ -1081,7 +1225,7 @@ mod tests { let indexer = PositionalIndexer::new(4); // small jump_size to exercise jump logic let values: Vec = (1..=20).collect(); let blocks = make_blocks(&values); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); indexer.apply_stored(w1, &blocks, None, &mut wb1).unwrap(); @@ -1095,8 +1239,8 @@ mod tests { // w1 has 10 blocks, w2 has 6 let values_w1: Vec = (1..=10).collect(); let values_w2: Vec = (1..=6).collect(); - let w1 = indexer.intern_worker("http://w1:8000"); - let w2 = indexer.intern_worker("http://w2:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); + let w2 = indexer.intern_worker("http://w2:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); let mut wb2 = WorkerBlockMap::default(); indexer @@ -1119,9 +1263,9 @@ mod tests { let v1: Vec = (1..=12).collect(); let v2: Vec = (1..=7).collect(); let v3: Vec = (1..=4).collect(); - let w1 = indexer.intern_worker("http://w1:8000"); - let w2 = indexer.intern_worker("http://w2:8000"); - let w3 = indexer.intern_worker("http://w3:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); + let w2 = indexer.intern_worker("http://w2:8000").unwrap(); + let w3 = indexer.intern_worker("http://w3:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); let mut wb2 = WorkerBlockMap::default(); let mut wb3 = WorkerBlockMap::default(); @@ -1152,7 +1296,9 @@ mod tests { let writer = thread::spawn(move || { for i in 0..100u64 { let blocks = make_blocks(&[i * 10, i * 10 + 1, i * 10 + 2]); - let wid = indexer_writer.intern_worker(&format!("http://w{i}:8000")); + let wid = indexer_writer + .intern_worker(&format!("http://w{i}:8000")) + .unwrap(); let mut wb = WorkerBlockMap::default(); let _ = indexer_writer.apply_stored(wid, &blocks, None, &mut wb); } @@ -1181,8 +1327,8 @@ mod tests { seq_hash: SequenceHash(100), content_hash: ContentHash(10), }]; - let w1 = indexer.intern_worker("http://w1:8000"); - let w2 = indexer.intern_worker("http://w2:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); + let w2 = indexer.intern_worker("http://w2:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); let mut wb2 = WorkerBlockMap::default(); indexer @@ -1212,8 +1358,8 @@ mod tests { // Worker 1: position 0 = content 10, position 1 = content 99 // Prefix at pos 1 = XXH3(10 || 99) let blocks_w1 = make_blocks(&[10, 99]); - let w1 = indexer.intern_worker("http://w1:8000"); - let w2 = indexer.intern_worker("http://w2:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); + let w2 = indexer.intern_worker("http://w2:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); let mut wb2 = WorkerBlockMap::default(); indexer @@ -1245,7 +1391,7 @@ mod tests { fn test_early_exit_returns_score_one() { let indexer = PositionalIndexer::new(64); let blocks = make_blocks(&[10, 20, 30]); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); indexer.apply_stored(w1, &blocks, None, &mut wb1).unwrap(); @@ -1260,7 +1406,7 @@ mod tests { fn test_early_exit_no_match() { let indexer = PositionalIndexer::new(64); let blocks = make_blocks(&[10, 20, 30]); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); indexer.apply_stored(w1, &blocks, None, &mut wb1).unwrap(); @@ -1281,7 +1427,7 @@ mod tests { #[test] fn test_worker_id_after_store() { let indexer = PositionalIndexer::default(); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); indexer .apply_stored(w1, &make_blocks(&[10]), None, &mut wb1) @@ -1297,7 +1443,7 @@ mod tests { fn test_tree_sizes_after_store_and_remove() { let indexer = PositionalIndexer::new(64); let blocks = make_blocks(&[10, 20, 30, 40, 50]); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); indexer.apply_stored(w1, &blocks, None, &mut wb1).unwrap(); assert_eq!(indexer.current_size(), 5); @@ -1315,7 +1461,7 @@ mod tests { fn test_duplicate_store_does_not_inflate_tree_size() { let indexer = PositionalIndexer::new(64); let blocks = make_blocks(&[10, 20, 30]); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); // First store: 3 new blocks @@ -1339,7 +1485,7 @@ mod tests { #[test] fn test_remove_worker_nonexistent_is_noop() { let indexer = PositionalIndexer::default(); - let w = indexer.intern_worker("http://ghost:8000"); + let w = indexer.intern_worker("http://ghost:8000").unwrap(); indexer.remove_worker(w, WorkerBlockMap::default()); // no-op, no panic assert_eq!(indexer.current_size(), 0); } @@ -1349,7 +1495,7 @@ mod tests { let indexer = Arc::new(PositionalIndexer::new(4)); let content: Vec = (1..=20).collect(); let blocks = make_blocks(&content); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); indexer.apply_stored(w1, &blocks, None, &mut wb1).unwrap(); @@ -1374,7 +1520,7 @@ mod tests { let worker_content: Vec = (1..=5).collect(); handles.push(std::thread::spawn(move || { let worker = format!("http://writer{i}:8000"); - let wid = idx.intern_worker(&worker); + let wid = idx.intern_worker(&worker).unwrap(); let mut wb = WorkerBlockMap::default(); let blks = make_blocks(&worker_content); for _ in 0..50 { @@ -1396,8 +1542,8 @@ mod tests { fn test_dashmap_cleanup_no_memory_leak() { let indexer = PositionalIndexer::default(); let blocks = make_blocks(&[10, 20, 30]); - let w1 = indexer.intern_worker("http://w1:8000"); - let w2 = indexer.intern_worker("http://w2:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); + let w2 = indexer.intern_worker("http://w2:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); let mut wb2 = WorkerBlockMap::default(); indexer.apply_stored(w1, &blocks, None, &mut wb1).unwrap(); @@ -1446,7 +1592,7 @@ mod tests { fn test_query_prefix_of_stored() { let indexer = PositionalIndexer::default(); let blocks = make_blocks(&[10, 20, 30, 40, 50]); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); indexer.apply_stored(w1, &blocks, None, &mut wb1).unwrap(); @@ -1460,8 +1606,8 @@ mod tests { let indexer = PositionalIndexer::default(); let blocks_w1 = make_blocks(&[10, 20, 30]); let blocks_w2 = make_blocks(&[99, 88, 77]); - let w1 = indexer.intern_worker("http://w1:8000"); - let w2 = indexer.intern_worker("http://w2:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); + let w2 = indexer.intern_worker("http://w2:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); let mut wb2 = WorkerBlockMap::default(); indexer @@ -1492,8 +1638,8 @@ mod tests { assert_eq!(indexer.current_size(), 0); let blocks = make_blocks(&[10, 20, 30]); - let w1 = indexer.intern_worker("http://w1:8000"); - let w2 = indexer.intern_worker("http://w2:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); + let w2 = indexer.intern_worker("http://w2:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); let mut wb2 = WorkerBlockMap::default(); indexer.apply_stored(w1, &blocks, None, &mut wb1).unwrap(); @@ -1591,7 +1737,7 @@ mod tests { }) .collect(); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); indexer.apply_stored(w1, &blocks, None, &mut wb1).unwrap(); @@ -1614,7 +1760,7 @@ mod tests { content_hash: compute_content_hash(chunk), }) .collect(); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); indexer.apply_stored(w1, &blocks, None, &mut wb1).unwrap(); @@ -1653,8 +1799,8 @@ mod tests { }) .collect(); - let sglang = indexer.intern_worker("http://sglang:8000"); - let vllm = indexer.intern_worker("http://vllm:8000"); + let sglang = indexer.intern_worker("http://sglang:8000").unwrap(); + let vllm = indexer.intern_worker("http://vllm:8000").unwrap(); let mut wb_sg = WorkerBlockMap::default(); let mut wb_vl = WorkerBlockMap::default(); indexer @@ -1682,7 +1828,7 @@ mod tests { chunk_size: usize, worker_blocks: &mut WorkerBlockMap, ) { - let worker_id = indexer.intern_worker(worker); + let worker_id = indexer.intern_worker(worker).unwrap(); let all_blocks = make_blocks(content); let mut offset = 0; let mut parent: Option = None; @@ -1702,7 +1848,7 @@ mod tests { let indexer = PositionalIndexer::new(32); let full: Vec = (1..=128).collect(); let full_blocks = make_blocks(&full); - let full_id = indexer.intern_worker("http://full:8000"); + let full_id = indexer.intern_worker("http://full:8000").unwrap(); let mut wb_full = WorkerBlockMap::default(); indexer .apply_stored(full_id, &full_blocks, None, &mut wb_full) @@ -1711,7 +1857,7 @@ mod tests { for &depth in &[31, 32, 33] { let partial_blocks = make_blocks(&full[..depth]); let worker = format!("http://depth{depth}:8000"); - let wid = indexer.intern_worker(&worker); + let wid = indexer.intern_worker(&worker).unwrap(); let mut wb = WorkerBlockMap::default(); indexer .apply_stored(wid, &partial_blocks, None, &mut wb) @@ -1721,7 +1867,7 @@ mod tests { for &depth in &[63, 64, 65] { let partial_blocks = make_blocks(&full[..depth]); let worker = format!("http://depth{depth}:8000"); - let wid = indexer.intern_worker(&worker); + let wid = indexer.intern_worker(&worker).unwrap(); let mut wb = WorkerBlockMap::default(); indexer .apply_stored(wid, &partial_blocks, None, &mut wb) @@ -1745,7 +1891,7 @@ mod tests { let content: Vec = (1..=len as u64).collect(); let blocks = make_blocks(&content); let worker = format!("http://len{len}:8000"); - let wid = indexer.intern_worker(&worker); + let wid = indexer.intern_worker(&worker).unwrap(); let mut wb = WorkerBlockMap::default(); indexer.apply_stored(wid, &blocks, None, &mut wb).unwrap(); @@ -1767,7 +1913,7 @@ mod tests { let content = &full[..len]; let blocks = make_blocks(content); let worker = format!("http://len{len}:8000"); - let wid = indexer.intern_worker(&worker); + let wid = indexer.intern_worker(&worker).unwrap(); let mut wb = WorkerBlockMap::default(); indexer.apply_stored(wid, &blocks, None, &mut wb).unwrap(); @@ -1789,7 +1935,7 @@ mod tests { for &depth in &depths { let blocks = make_blocks(&full[..depth]); let worker = format!("http://w{depth}:8000"); - let wid = indexer.intern_worker(&worker); + let wid = indexer.intern_worker(&worker).unwrap(); let mut wb = WorkerBlockMap::default(); indexer.apply_stored(wid, &blocks, None, &mut wb).unwrap(); } @@ -1814,9 +1960,9 @@ mod tests { let mut content_w1 = shared.clone(); content_w1.extend(1001..=1060); let blocks_w1 = make_blocks(&content_w1); - let w1 = indexer.intern_worker("http://w1:8000"); - let w2 = indexer.intern_worker("http://w2:8000"); - let w3 = indexer.intern_worker("http://w3:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); + let w2 = indexer.intern_worker("http://w2:8000").unwrap(); + let w3 = indexer.intern_worker("http://w3:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); let mut wb2 = WorkerBlockMap::default(); let mut wb3 = WorkerBlockMap::default(); @@ -1847,7 +1993,7 @@ mod tests { let indexer = PositionalIndexer::new(64); let content: Vec = (1..=1000).collect(); let blocks = make_blocks(&content); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); indexer.apply_stored(w1, &blocks, None, &mut wb1).unwrap(); @@ -1904,7 +2050,7 @@ mod tests { #[test] fn test_multiple_disjoint_sequences_per_worker() { let indexer = PositionalIndexer::new(64); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); let blocks1 = make_blocks(&[10, 20, 30]); @@ -1929,7 +2075,7 @@ mod tests { let indexer = PositionalIndexer::new(32); let content: Vec = (1..=100).collect(); let blocks = make_blocks(&content); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); indexer.apply_stored(w1, &blocks, None, &mut wb1).unwrap(); @@ -1949,7 +2095,7 @@ mod tests { fn test_remove_parent_does_not_cascade() { let indexer = PositionalIndexer::new(1); let blocks = make_blocks(&[10, 20, 30, 40, 50]); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); indexer.apply_stored(w1, &blocks, None, &mut wb1).unwrap(); @@ -1964,7 +2110,7 @@ mod tests { #[test] fn test_long_sequence_clear_and_rebuild() { let indexer = PositionalIndexer::new(32); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); let original: Vec = (1..=100).collect(); @@ -1996,7 +2142,7 @@ mod tests { for &depth in &depths { let blocks = make_blocks(&content[..depth]); let worker = format!("http://w{depth}:8000"); - let wid = indexer.intern_worker(&worker); + let wid = indexer.intern_worker(&worker).unwrap(); let mut wb = WorkerBlockMap::default(); indexer.apply_stored(wid, &blocks, None, &mut wb).unwrap(); } @@ -2017,4 +2163,173 @@ mod tests { ); } } + + // ----------------------------------------------------------------------- + // Worker interning: growth past 2048, id lifecycle, exhaustion + // ----------------------------------------------------------------------- + + #[test] + fn test_tree_sizes_locate_covers_u32_id_space() { + // First segment: ids 0..2048. + assert_eq!(TreeSizes::locate(0), (0, 0)); + assert_eq!(TreeSizes::locate(2047), (0, 2047)); + // Doubling segments: 2048..6144, 6144..14336, ... + assert_eq!(TreeSizes::locate(2048), (1, 0)); + assert_eq!(TreeSizes::locate(6143), (1, 4095)); + assert_eq!(TreeSizes::locate(6144), (2, 0)); + // The largest possible id maps inside the last segment. + let (segment, offset) = TreeSizes::locate(u32::MAX); + assert_eq!(segment, SEGMENT_COUNT - 1); + assert!(offset < TreeSizes::segment_len(segment)); + // Segment starts are contiguous: each boundary id maps to offset 0. + let mut start = 0u64; + for segment in 0..SEGMENT_COUNT { + if start > u32::MAX as u64 { + break; + } + assert_eq!(TreeSizes::locate(start as u32), (segment, 0)); + start += TreeSizes::segment_len(segment) as u64; + } + } + + #[test] + fn test_intern_worker_past_2048_workers() { + let indexer = PositionalIndexer::new(64); + // Ids must be dense and uncapped — 2049+ used to panic. + let ids: Vec = (0..5000u32) + .map(|w| indexer.intern_worker(&format!("http://w{w}:8000")).unwrap()) + .collect(); + for (expected, &id) in ids.iter().enumerate() { + assert_eq!(id, expected as u32); + } + + // Shared 3-block prefix plus a distinct tail per worker, for workers on + // both sides of the old 2048 cap (2048 is also a segment boundary). + let shared: Vec = vec![10, 20, 30]; + let probe_ids = [0u32, 1, 2047, 2048, 2049, 4999]; + for &wid in &probe_ids { + let mut content = shared.clone(); + content.push(1_000_000 + wid as u64); + let blocks = make_blocks(&content); + let mut wb = WorkerBlockMap::default(); + indexer.apply_stored(wid, &blocks, None, &mut wb).unwrap(); + } + + // All probed workers share the 3-block prefix. + let scores = indexer.find_matches(&hashes(&shared), false); + for &wid in &probe_ids { + assert_eq!(scores.scores.get(&wid), Some(&3), "worker {wid} score"); + assert_eq!( + scores.tree_sizes.get(&wid), + Some(&4), + "worker {wid} tree_size" + ); + } + + // Only worker 2049 has the tail block — the rest drain at depth 3. + let mut full = shared.clone(); + full.push(1_000_000 + 2049); + let scores = indexer.find_matches(&hashes(&full), false); + assert_eq!(scores.scores.get(&2049), Some(&4)); + assert_eq!(scores.scores.get(&2048), Some(&3)); + + assert_eq!(indexer.current_size(), 4 * probe_ids.len()); + } + + #[test] + fn test_worker_ids_not_recycled_after_removal() { + let indexer = PositionalIndexer::new(64); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); + let w2 = indexer.intern_worker("http://w2:8000").unwrap(); + assert_eq!((w1, w2), (0, 1)); + + let blocks = make_blocks(&[10, 20, 30]); + let mut wb1 = WorkerBlockMap::default(); + indexer.apply_stored(w1, &blocks, None, &mut wb1).unwrap(); + indexer.remove_worker(w1, wb1); + + // The removed worker's URL keeps its id; it is not freed for reuse. + assert_eq!(indexer.intern_worker("http://w1:8000").unwrap(), w1); + // New URLs continue monotonically — removal never recycles ids. + assert_eq!(indexer.intern_worker("http://w3:8000").unwrap(), 2); + + // The removed worker no longer matches; re-storing under its id works. + let scores = indexer.find_matches(&hashes(&[10, 20, 30]), false); + assert!(scores.scores.is_empty()); + let mut wb1 = WorkerBlockMap::default(); + indexer.apply_stored(w1, &blocks, None, &mut wb1).unwrap(); + let scores = indexer.find_matches(&hashes(&[10, 20, 30]), false); + assert_eq!(scores.scores.get(&w1), Some(&3)); + assert_eq!(indexer.current_size(), 3); + } + + #[test] + fn test_remove_worker_in_unallocated_segment_is_noop() { + let indexer = PositionalIndexer::new(64); + // High ids whose tree_sizes segments were never written: removal and + // clear must not panic (and must not allocate the segment). + indexer.remove_worker(123_456, WorkerBlockMap::default()); + indexer.apply_cleared(654_321, &mut WorkerBlockMap::default()); + assert_eq!(indexer.current_size(), 0); + } + + #[test] + fn test_intern_worker_exhaustion_returns_error() { + let indexer = PositionalIndexer::new(64); + // Jump the counter to the last valid id. Interning never touches + // tree_sizes, so no segment is allocated for this id. + indexer + .next_worker_id + .store(u32::MAX as u64, Ordering::Relaxed); + assert_eq!(indexer.intern_worker("http://last:8000").unwrap(), u32::MAX); + // The id space is now exhausted: new URLs error, known URLs still resolve. + assert_eq!( + indexer.intern_worker("http://one-too-many:8000"), + Err(WorkerIdExhausted) + ); + assert_eq!(indexer.intern_worker("http://last:8000").unwrap(), u32::MAX); + assert_eq!(indexer.worker_id("http://last:8000"), Some(u32::MAX)); + } + + #[test] + fn test_concurrent_interning_across_segment_boundary() { + let indexer = Arc::new(PositionalIndexer::new(64)); + // Pre-assign ids up to just below the 2048 segment boundary, so the + // concurrent writers race across it. + for w in 0..2040 { + indexer + .intern_worker(&format!("http://pre{w}:8000")) + .unwrap(); + } + + let mut handles = Vec::new(); + for t in 0..4u32 { + let idx = Arc::clone(&indexer); + handles.push(std::thread::spawn(move || { + for i in 0..8u32 { + let wid = idx.intern_worker(&format!("http://t{t}-{i}:8000")).unwrap(); + let mut wb = WorkerBlockMap::default(); + let blocks = make_blocks(&[wid as u64 * 100 + 1, wid as u64 * 100 + 2]); + idx.apply_stored(wid, &blocks, None, &mut wb).unwrap(); + } + })); + } + for handle in handles { + handle.join().unwrap(); + } + + // 4 threads x 8 workers x 2 blocks each, ids 2040..2072 straddling the + // segment boundary; every worker's blocks must be matchable. + assert_eq!(indexer.current_size(), 64); + for t in 0..4u32 { + for i in 0..8u32 { + let wid = indexer.worker_id(&format!("http://t{t}-{i}:8000")).unwrap(); + let scores = indexer.find_matches( + &hashes(&[wid as u64 * 100 + 1, wid as u64 * 100 + 2]), + false, + ); + assert_eq!(scores.scores.get(&wid), Some(&2)); + } + } + } } diff --git a/crates/kv_index/src/lib.rs b/crates/kv_index/src/lib.rs index 3fbd854c20..443a008727 100644 --- a/crates/kv_index/src/lib.rs +++ b/crates/kv_index/src/lib.rs @@ -21,7 +21,7 @@ mod token_tree; pub use common::{MatchResult, TenantId}; pub use event_tree::{ compute_content_hash, compute_request_content_hashes, ApplyError, ContentHash, OverlapScores, - PositionalIndexer, SequenceHash, StoredBlock, WorkerBlockMap, WorkerId, + PositionalIndexer, SequenceHash, StoredBlock, WorkerBlockMap, WorkerId, WorkerIdExhausted, }; pub use path_hash::{hash_node_path, hash_token_path, GLOBAL_EVICTION_HASH}; // Re-export under names matching old tree.rs API for easier migration diff --git a/model_gateway/benches/radix_tree_benchmark.rs b/model_gateway/benches/radix_tree_benchmark.rs index 4ad7cec4ff..45296a05df 100644 --- a/model_gateway/benches/radix_tree_benchmark.rs +++ b/model_gateway/benches/radix_tree_benchmark.rs @@ -445,7 +445,7 @@ fn build_populated_indexer( let mut all_worker_chunks = Vec::with_capacity(workers.len()); for worker in workers { - let worker_id = indexer.intern_worker(worker); + let worker_id = indexer.intern_worker(worker).unwrap(); let mut wb = WorkerBlockMap::default(); indexer .apply_stored(worker_id, &shared_blocks, None, &mut wb) @@ -500,7 +500,7 @@ macro_rules! bench_indexer_store { for _ in 0..iters { let indexer = PositionalIndexer::new(32); for worker in &workers { - let worker_id = indexer.intern_worker(worker); + let worker_id = indexer.intern_worker(worker).unwrap(); let mut wb = WorkerBlockMap::default(); let chunks = generate_token_chunks($blocks_per_worker, $block_size); let blocks = chunks_to_stored_blocks(&chunks); @@ -596,7 +596,7 @@ macro_rules! bench_indexer_concurrent { (0..$num_threads) .map(|t| { let chunks = &worker_chunks[t % workers.len()]; - let worker_id = indexer.intern_worker(&workers[t % workers.len()]); + let worker_id = indexer.intern_worker(&workers[t % workers.len()]).unwrap(); // Read data: pre-computed content hashes let query_tokens = flatten_tokens(chunks); diff --git a/model_gateway/src/observability/metrics.rs b/model_gateway/src/observability/metrics.rs index 5c7d079517..13aebb3828 100644 --- a/model_gateway/src/observability/metrics.rs +++ b/model_gateway/src/observability/metrics.rs @@ -230,6 +230,11 @@ pub(crate) fn init_metrics() { "smg_worker_errors_total", "Worker-level errors by worker_type, connection_mode, error_type" ); + describe_counter!( + "smg_kv_event_subscription_failures_total", + "KV event subscription task failures by worker and reason \ + (panic, join_error, intern_failed)" + ); describe_gauge!( "smg_manual_policy_cache_entries", "Number of routing entries in manual policy cache" @@ -944,6 +949,18 @@ impl Metrics { .set(if healthy { 1.0 } else { 0.0 }); } + /// Record a KV event subscription task failure (panic, join error, or + /// worker-id intern failure) + pub fn record_kv_event_subscription_failure(worker_url: &str, reason: &'static str) { + let worker_interned = intern_string(worker_url); + counter!( + "smg_kv_event_subscription_failures_total", + "worker" => worker_interned, + "reason" => reason + ) + .increment(1); + } + // ======================================================================== // Layer 3: Worker resilience metrics (circuit breaker) // ======================================================================== diff --git a/model_gateway/src/policies/cache_aware.rs b/model_gateway/src/policies/cache_aware.rs index 137b3c2e73..6f84368264 100644 --- a/model_gateway/src/policies/cache_aware.rs +++ b/model_gateway/src/policies/cache_aware.rs @@ -42,6 +42,7 @@ */ use std::{ + cmp::Reverse, collections::HashMap, sync::{ atomic::{AtomicBool, Ordering}, @@ -50,7 +51,7 @@ use std::{ }; use dashmap::DashMap; -use kv_index::{compute_request_content_hashes, PositionalIndexer, TokenTree, Tree}; +use kv_index::{ContentHash, PositionalIndexer, TokenTree, Tree, WorkerId}; use openai_protocol::worker::WorkerLoadResponse; use parking_lot::RwLock; use rand::Rng; @@ -789,7 +790,7 @@ impl LoadBalancingPolicy for CacheAwarePolicy { // 3. Approximate string tree: Tree prefix matching (HTTP) if let Some(tokens) = request_tokens { if self.has_event_indexer(model_id) { - self.select_worker_event_driven(workers, tokens, &healthy_indices, model_id) + self.select_worker_event_driven(workers, tokens, info, &healthy_indices, model_id) } else { self.select_worker_with_tokens(workers, tokens, &healthy_indices, model_id) } @@ -847,6 +848,7 @@ impl CacheAwarePolicy { &self, workers: &[Arc], tokens: &[u32], + info: &SelectWorkerInfo, healthy_indices: &[usize], model_id: &str, ) -> Option { @@ -859,8 +861,13 @@ impl CacheAwarePolicy { .block_size(model_id) .unwrap_or(self.config.block_size); - if let Some(idx) = - Self::score_overlap(workers, tokens, healthy_indices, &indexer, block_size) + // Full-prompt hashing is memoized on the request: PD mode runs + // selection twice (prefill + decode) and retries re-run it, all with + // the same `SelectWorkerInfo`, so the O(tokens) XXH3 pass happens at + // most once per request instead of once per `select_worker` call. + let content_hashes = info.request_content_hashes(tokens, block_size); + + if let Some(idx) = Self::score_overlap(workers, &content_hashes, healthy_indices, &indexer) { return Some(idx); } @@ -880,66 +887,84 @@ impl CacheAwarePolicy { /// Score healthy workers by PositionalIndexer overlap and select the best. /// - /// Returns `Some(idx)` if at least one worker has cached blocks matching the - /// request. Returns `None` if the request is too short for a full block or - /// no workers have matching data. + /// `content_hashes` is the request's per-block hash sequence, computed (and + /// memoized) by the caller. Returns `Some(idx)` if at least one worker has + /// cached blocks matching the request. Returns `None` if the request is + /// too short for a full block or no workers have matching data. + /// + /// The scoring loop is integer-keyed: each worker's interned u32 id comes + /// from a per-worker cache ([`Worker::cached_kv_worker_id`], one relaxed + /// atomic load) and the overlap maps are probed by id. The URL-keyed map + /// is touched at most once per (worker, indexer instance) — previously it + /// was probed 2-3 times per worker on every request. fn score_overlap( workers: &[Arc], - tokens: &[u32], + content_hashes: &[ContentHash], healthy_indices: &[usize], indexer: &PositionalIndexer, - block_size: usize, ) -> Option { - let content_hashes = compute_request_content_hashes(tokens, block_size); if content_hashes.is_empty() { return None; } - let overlap = indexer.find_matches(&content_hashes, false); + let overlap = indexer.find_matches(content_hashes, false); if overlap.scores.is_empty() { return None; } // Select worker with best overlap among those that actually match. - // Tie-break: lower load, then smaller tree size. - let best_idx = healthy_indices - .iter() - .copied() - .filter(|&idx| { - indexer - .worker_id(workers[idx].url()) - .and_then(|id| overlap.scores.get(&id)) - .copied() - .unwrap_or(0) - > 0 - }) - .max_by_key(|&idx| { - let wid = indexer.worker_id(workers[idx].url()); - let score = wid - .and_then(|id| overlap.scores.get(&id)) - .copied() - .unwrap_or(0); - let load = workers[idx].load(); - let tree_size = wid - .and_then(|id| overlap.tree_sizes.get(&id)) - .copied() - .unwrap_or(0); - (score, std::cmp::Reverse(load), std::cmp::Reverse(tree_size)) - })?; + // Tie-break: lower load, then smaller tree size. On fully equal keys + // the later worker wins, matching the previous `max_by_key` behavior. + type ScoreKey = (u32, Reverse, Reverse); + let mut best: Option<(usize, ScoreKey)> = None; + for &idx in healthy_indices { + let Some(wid) = Self::resolve_worker_id(&workers[idx], indexer) else { + continue; + }; + let score = overlap.scores.get(&wid).copied().unwrap_or(0); + if score == 0 { + continue; + } + let tree_size = overlap.tree_sizes.get(&wid).copied().unwrap_or(0); + let key = (score, Reverse(workers[idx].load()), Reverse(tree_size)); + if best.as_ref().is_none_or(|(_, best_key)| key >= *best_key) { + best = Some((idx, key)); + } + } + let (best_idx, (best_score, _, _)) = best?; debug!( worker = workers[best_idx].url(), - score = indexer - .worker_id(workers[best_idx].url()) - .and_then(|id| overlap.scores.get(&id)) - .copied() - .unwrap_or(0), + score = best_score, "Event-driven routing: overlap match" ); workers[best_idx].increment_processed(); Some(best_idx) } + /// Interned kv_index id for a worker, cached on the worker object across + /// requests and validated against the indexer instance ("epoch") so a + /// replaced indexer never sees another instance's ids. + /// + /// First contact per (worker, indexer instance) interns the URL — ids are + /// assigned once and never recycled within an instance, so this is the + /// same id the KV-event path uses — and caches it. Afterwards resolution + /// is one relaxed atomic load. Returns `None` only on u32 id-space + /// exhaustion (the worker then scores as unmatched, exactly like a worker + /// the indexer has never seen). + fn resolve_worker_id( + worker: &Arc, + indexer: &PositionalIndexer, + ) -> Option { + let epoch = indexer.instance_id(); + if let Some(wid) = worker.cached_kv_worker_id(epoch) { + return Some(wid); + } + let wid = indexer.intern_worker(worker.url()).ok()?; + worker.cache_kv_worker_id(epoch, wid); + Some(wid) + } + /// Select worker using token-based tree (gRPC path) fn select_worker_with_tokens( &self, @@ -1015,9 +1040,15 @@ impl CacheAwarePolicy { return Some(idx); } - // Selected worker no longer exists or unhealthy - fall back to first healthy - // Stale entries will be cleaned up by LRU eviction - healthy_indices.first().copied() + // Matched worker no longer exists or unhealthy — fall back to the + // least-loaded healthy worker. Always taking the *first* healthy + // index herded every fallback request (e.g. all tenants of a dead + // worker) onto one replica. Stale entries will be cleaned up by + // LRU eviction. + healthy_indices + .iter() + .min_by_key(|&&idx| workers[idx].load()) + .copied() } else { debug!( "Warning: No token tree found for model '{}', using random worker selection", @@ -1096,9 +1127,14 @@ impl CacheAwarePolicy { return Some(idx); } - // Selected worker no longer exists or unhealthy - fall back to first healthy - // Stale entries will be cleaned up by LRU eviction - healthy_indices.first().copied() + // Matched worker no longer exists or unhealthy — fall back to the + // least-loaded healthy worker (see the token path above for the + // anti-herding rationale). Stale entries will be cleaned up by + // LRU eviction. + healthy_indices + .iter() + .min_by_key(|&&idx| workers[idx].load()) + .copied() } else { debug!( "Warning: No string tree found for model '{}', using random worker selection", @@ -1119,7 +1155,10 @@ impl Default for CacheAwarePolicy { #[cfg(test)] mod tests { - use kv_index::{compute_content_hash, SequenceHash, StoredBlock, WorkerBlockMap}; + use kv_index::{ + compute_content_hash, compute_request_content_hashes, SequenceHash, StoredBlock, + WorkerBlockMap, + }; use openai_protocol::worker::{HealthCheckConfig, SchedulerLoadSnapshot, WorkerStatus}; use super::*; @@ -1688,7 +1727,7 @@ mod tests { jump_size: usize, ) -> Arc { let indexer = Arc::new(PositionalIndexer::new(jump_size)); - let worker_id = indexer.intern_worker(worker_url); + let worker_id = indexer.intern_worker(worker_url).unwrap(); let mut wb = WorkerBlockMap::default(); let blocks: Vec = token_chunks .iter() @@ -1748,10 +1787,12 @@ mod tests { // Query with matching tokens — should select w1 let result = CacheAwarePolicy::score_overlap( &workers, - &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16], + &compute_request_content_hashes( + &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16], + 4, + ), &[0, 1], &indexer, - 4, ); assert_eq!(result, Some(0)); // w1 } @@ -1773,10 +1814,9 @@ mod tests { // Completely different tokens — no overlap → None let result = CacheAwarePolicy::score_overlap( &workers, - &[100, 200, 300, 400, 500, 600, 700, 800], + &compute_request_content_hashes(&[100, 200, 300, 400, 500, 600, 700, 800], 4), &[0], &indexer, - 4, ); assert_eq!(result, None); } @@ -1804,8 +1844,8 @@ mod tests { // Store same blocks for both workers (equal overlap) let indexer = Arc::new(PositionalIndexer::new(4)); - let w1_id = indexer.intern_worker("http://w1:8000"); - let w2_id = indexer.intern_worker("http://w2:8000"); + let w1_id = indexer.intern_worker("http://w1:8000").unwrap(); + let w2_id = indexer.intern_worker("http://w2:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); let mut wb2 = WorkerBlockMap::default(); let blocks = vec![StoredBlock { @@ -1824,7 +1864,12 @@ mod tests { .unwrap(); // Equal overlap → tie-break by load → w2 wins (lower load) - let result = CacheAwarePolicy::score_overlap(&workers, &[1, 2, 3, 4], &[0, 1], &indexer, 4); + let result = CacheAwarePolicy::score_overlap( + &workers, + &compute_request_content_hashes(&[1, 2, 3, 4], 4), + &[0, 1], + &indexer, + ); assert_eq!(result, Some(1)); // w2 (lower load) } @@ -1848,8 +1893,8 @@ mod tests { policy.init_workers(&workers); let indexer = Arc::new(PositionalIndexer::new(4)); - let w1_id = indexer.intern_worker("http://w1:8000"); - let w2_id = indexer.intern_worker("http://w2:8000"); + let w1_id = indexer.intern_worker("http://w1:8000").unwrap(); + let w2_id = indexer.intern_worker("http://w2:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); let mut wb2 = WorkerBlockMap::default(); @@ -1877,7 +1922,12 @@ mod tests { .unwrap(); // Equal overlap, equal load → tie-break by tree size → w1 wins (smaller) - let result = CacheAwarePolicy::score_overlap(&workers, &[1, 2, 3, 4], &[0, 1], &indexer, 4); + let result = CacheAwarePolicy::score_overlap( + &workers, + &compute_request_content_hashes(&[1, 2, 3, 4], 4), + &[0, 1], + &indexer, + ); assert_eq!(result, Some(0)); // w1 (smaller tree) } @@ -1893,7 +1943,12 @@ mod tests { let indexer = setup_indexer_with_blocks("http://w1:8000", &[&[1, 2, 3, 4]], 4); // Request shorter than block_size → no full blocks → None - let result = CacheAwarePolicy::score_overlap(&workers, &[1, 2, 3], &[0], &indexer, 4); + let result = CacheAwarePolicy::score_overlap( + &workers, + &compute_request_content_hashes(&[1, 2, 3], 4), + &[0], + &indexer, + ); assert_eq!(result, None); } @@ -1917,8 +1972,8 @@ mod tests { policy.init_workers(&workers); let indexer = Arc::new(PositionalIndexer::new(4)); - let w1_id = indexer.intern_worker("http://w1:8000"); - let w2_id = indexer.intern_worker("http://w2:8000"); + let w1_id = indexer.intern_worker("http://w1:8000").unwrap(); + let w2_id = indexer.intern_worker("http://w2:8000").unwrap(); let mut wb1 = WorkerBlockMap::default(); let mut wb2 = WorkerBlockMap::default(); @@ -1957,14 +2012,172 @@ mod tests { // Query with all 4 blocks worth of tokens → w1 wins (higher overlap: 4 vs 2) let result = CacheAwarePolicy::score_overlap( &workers, - &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16], + &compute_request_content_hashes( + &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16], + 4, + ), &[0, 1], &indexer, - 4, ); assert_eq!(result, Some(0)); // w1 (higher overlap) } + /// The pre-K3 `score_overlap` (URL-keyed `worker_id` probes per worker + + /// `max_by_key`), kept verbatim as a reference oracle for the equivalence + /// test below. + fn reference_score_overlap( + workers: &[Arc], + content_hashes: &[ContentHash], + healthy_indices: &[usize], + indexer: &PositionalIndexer, + ) -> Option { + if content_hashes.is_empty() { + return None; + } + let overlap = indexer.find_matches(content_hashes, false); + if overlap.scores.is_empty() { + return None; + } + healthy_indices + .iter() + .copied() + .filter(|&idx| { + indexer + .worker_id(workers[idx].url()) + .and_then(|id| overlap.scores.get(&id)) + .copied() + .unwrap_or(0) + > 0 + }) + .max_by_key(|&idx| { + let wid = indexer.worker_id(workers[idx].url()); + let score = wid + .and_then(|id| overlap.scores.get(&id)) + .copied() + .unwrap_or(0); + let load = workers[idx].load(); + let tree_size = wid + .and_then(|id| overlap.tree_sizes.get(&id)) + .copied() + .unwrap_or(0); + (score, Reverse(load), Reverse(tree_size)) + }) + } + + /// Equivalence guard for the integer-keyed scoring loop: given the same + /// indexer state, loads, and health, the new `score_overlap` must select + /// exactly the worker the previous URL-probing implementation selected — + /// including ties (equal score/load/tree-size keeps the later worker) and + /// workers unknown to the indexer. + #[test] + fn test_score_overlap_matches_reference_implementation() { + let mk = |url: &str| { + BasicWorkerBuilder::new(url) + .worker_type(WorkerType::Regular) + .health_config(no_health_check()) + .build() + }; + let w1 = mk("http://w1:8000"); + let w2 = mk("http://w2:8000"); + let w3 = mk("http://w3:8000"); + let w4 = mk("http://w4:8000"); // never feeds the indexer + for _ in 0..5 { + w3.increment_load(); + } + let workers: Vec> = + vec![Arc::new(w1), Arc::new(w2), Arc::new(w3), Arc::new(w4)]; + let healthy = [0, 1, 2, 3]; + + // w1: 4 blocks; w2: first 2 blocks; w3: same 4 blocks (higher load). + let indexer = Arc::new(PositionalIndexer::new(4)); + let chunk = |i: u32| [i * 4 + 1, i * 4 + 2, i * 4 + 3, i * 4 + 4]; + for (url, n_blocks) in [ + ("http://w1:8000", 4u32), + ("http://w2:8000", 2), + ("http://w3:8000", 4), + ] { + let wid = indexer.intern_worker(url).unwrap(); + let mut wb = WorkerBlockMap::default(); + let blocks: Vec = (0..n_blocks) + .map(|i| StoredBlock { + seq_hash: SequenceHash(u64::from(i) + 1), + content_hash: compute_content_hash(&chunk(i)), + }) + .collect(); + indexer.apply_stored(wid, &blocks, None, &mut wb).unwrap(); + } + + let full: Vec = (1..=16).collect(); + let queries: [&[u32]; 4] = [ + &full, // full 4-block match: w1 vs w3 tie on score + &[1, 2, 3, 4], // 1-block match: 3-way tie on score + &[1, 2, 3], // too short: no full block + &[9, 9, 9, 9, 9], // no overlap + ]; + for tokens in queries { + let hashes = compute_request_content_hashes(tokens, 4); + let expected = reference_score_overlap(&workers, &hashes, &healthy, &indexer); + let got = CacheAwarePolicy::score_overlap(&workers, &hashes, &healthy, &indexer); + assert_eq!(got, expected, "selection diverged for tokens {tokens:?}"); + // Repeat with warm per-worker id caches — must still agree. + let got_warm = CacheAwarePolicy::score_overlap(&workers, &hashes, &healthy, &indexer); + assert_eq!(got_warm, expected, "warm-cache selection diverged"); + } + + // Fully equal keys (same score, load, and tree size for w1 and w3): + // both implementations must keep the same (later) worker. + for _ in 0..5 { + workers[0].increment_load(); + } + let hashes = compute_request_content_hashes(&full, 4); + let expected = reference_score_overlap(&workers, &hashes, &healthy, &indexer); + let got = CacheAwarePolicy::score_overlap(&workers, &hashes, &healthy, &indexer); + assert_eq!(got, expected, "selection diverged on a full tie"); + assert_eq!(got, Some(2), "full tie must keep the later worker"); + } + + /// Worker-side cached ids are validated against the indexer instance + /// ("epoch"): replacing a model's indexer must not let ids interned in the + /// old instance leak into scoring against the new one, even when the raw + /// u32 ids collide across instances. + #[test] + fn test_score_overlap_invalidates_ids_across_indexer_instances() { + let workers: Vec> = vec![ + Arc::new( + BasicWorkerBuilder::new("http://w1:8000") + .worker_type(WorkerType::Regular) + .health_config(no_health_check()) + .build(), + ), + Arc::new( + BasicWorkerBuilder::new("http://w2:8000") + .worker_type(WorkerType::Regular) + .health_config(no_health_check()) + .build(), + ), + ]; + + // First indexer: only w1 has the block. Scoring caches per-worker ids. + let indexer1 = setup_indexer_with_blocks("http://w1:8000", &[&[1, 2, 3, 4]], 4); + let hashes = compute_request_content_hashes(&[1, 2, 3, 4], 4); + assert_eq!( + CacheAwarePolicy::score_overlap(&workers, &hashes, &[0, 1], &indexer1), + Some(0) + ); + + // Replacement indexer (same model in production): only w2 has the + // block, and w2 interns FIRST so it gets the same raw id (0) that w1 + // had in the old instance. A stale cache would credit w1 with w2's + // score; epoch validation must force re-interning instead. + let indexer2 = setup_indexer_with_blocks("http://w2:8000", &[&[1, 2, 3, 4]], 4); + assert_ne!(indexer1.instance_id(), indexer2.instance_id()); + assert_eq!( + CacheAwarePolicy::score_overlap(&workers, &hashes, &[0, 1], &indexer2), + Some(1), + "stale worker id from the old indexer instance leaked into scoring" + ); + } + // -- select_worker_event_driven integration tests -- #[test] @@ -2167,7 +2380,7 @@ mod tests { // Store blocks using block_size=8 (tokens chunked in groups of 8) let indexer = Arc::new(PositionalIndexer::new(4)); - let w1_id = indexer.intern_worker("http://w1:8000"); + let w1_id = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb = WorkerBlockMap::default(); let block = vec![StoredBlock { seq_hash: SequenceHash(1), @@ -2300,4 +2513,152 @@ mod tests { .unwrap(); assert_eq!(idx, idx2); // token tree cache affinity preserved } + + /// PD mode runs `select_worker` twice (prefill + decode) with one + /// `SelectWorkerInfo`: the full-prompt content-hash pass must run once, + /// be reused on the second selection, and both selections must agree. + #[test] + fn test_event_driven_content_hashes_memoized_across_selections() { + let policy = CacheAwarePolicy::with_config(test_config()); // block_size=4 + let workers: Vec> = vec![ + Arc::new( + BasicWorkerBuilder::new("http://w1:8000") + .worker_type(WorkerType::Regular) + .health_config(no_health_check()) + .build(), + ), + Arc::new( + BasicWorkerBuilder::new("http://w2:8000") + .worker_type(WorkerType::Regular) + .health_config(no_health_check()) + .build(), + ), + ]; + policy.init_workers(&workers); + + let monitor = Arc::new(KvEventMonitor::new(Some(4))); + let indexer = + setup_indexer_with_blocks("http://w1:8000", &[&[1, 2, 3, 4], &[5, 6, 7, 8]], 4); + monitor.indexers.insert("unknown".to_string(), indexer); + policy.set_kv_event_monitor(Some(monitor)); + + let tokens: [u32; 8] = [1, 2, 3, 4, 5, 6, 7, 8]; + let info = SelectWorkerInfo { + tokens: Some(&tokens), + ..Default::default() + }; + assert!(info.content_hashes.get().is_none()); + + let first = policy.select_worker(&workers, &info); + assert_eq!(first, Some(0)); + let (memo_block_size, memo_hashes) = info + .content_hashes + .get() + .expect("event-driven scoring must populate the request hash memo"); + assert_eq!(*memo_block_size, 4); + assert_eq!(memo_hashes.len(), 2); // 8 tokens at block_size 4 + + // Second selection (the PD decode pass) reuses the memo and agrees. + let second = policy.select_worker(&workers, &info); + assert_eq!(second, first); + } + + /// A cache hit on a tenant that is now unhealthy must fall back to the + /// least-loaded healthy worker, not herd onto the first healthy index + /// (token-tree path). + #[test] + fn test_token_tree_dead_tenant_falls_back_to_min_load() { + let policy = CacheAwarePolicy::with_config(test_config()); + + let w1 = BasicWorkerBuilder::new("http://w1:8000") + .worker_type(WorkerType::Regular) + .health_config(no_health_check()) + .build(); + let w2 = BasicWorkerBuilder::new("http://w2:8000") + .worker_type(WorkerType::Regular) + .health_config(no_health_check()) + .build(); + let w3 = BasicWorkerBuilder::new("http://w3:8000") + .worker_type(WorkerType::Regular) + .health_config(no_health_check()) + .build(); + // w1 (the first healthy index) is busier than w2. + for _ in 0..5 { + w1.increment_load(); + } + + let workers: Vec> = vec![Arc::new(w1), Arc::new(w2), Arc::new(w3)]; + policy.init_workers(&workers); + + // The tree maps these tokens to w3 (full match > cache_threshold)... + let tokens: [u32; 8] = [1, 2, 3, 4, 5, 6, 7, 8]; + policy + .token_trees + .get("unknown") + .unwrap() + .insert_tokens(&tokens, "http://w3:8000"); + // ...but w3 is no longer healthy. + workers[2].set_status(WorkerStatus::NotReady); + + let idx = policy + .select_worker( + &workers, + &SelectWorkerInfo { + tokens: Some(&tokens), + ..Default::default() + }, + ) + .unwrap(); + assert_eq!( + idx, 1, + "fallback must pick min-load (w2), not first healthy (w1)" + ); + } + + /// Same as above for the string-tree (HTTP) path. + #[test] + fn test_string_tree_dead_tenant_falls_back_to_min_load() { + let policy = CacheAwarePolicy::with_config(test_config()); + + let w1 = BasicWorkerBuilder::new("http://w1:8000") + .worker_type(WorkerType::Regular) + .health_config(no_health_check()) + .build(); + let w2 = BasicWorkerBuilder::new("http://w2:8000") + .worker_type(WorkerType::Regular) + .health_config(no_health_check()) + .build(); + let w3 = BasicWorkerBuilder::new("http://w3:8000") + .worker_type(WorkerType::Regular) + .health_config(no_health_check()) + .build(); + for _ in 0..5 { + w1.increment_load(); + } + + let workers: Vec> = vec![Arc::new(w1), Arc::new(w2), Arc::new(w3)]; + policy.init_workers(&workers); + + let text = "the quick brown fox jumps over the lazy dog"; + policy + .string_trees + .get("unknown") + .unwrap() + .insert_text(text, "http://w3:8000"); + workers[2].set_status(WorkerStatus::NotReady); + + let idx = policy + .select_worker( + &workers, + &SelectWorkerInfo { + request_text: Some(text), + ..Default::default() + }, + ) + .unwrap(); + assert_eq!( + idx, 1, + "fallback must pick min-load (w2), not first healthy (w1)" + ); + } } diff --git a/model_gateway/src/policies/least_load.rs b/model_gateway/src/policies/least_load.rs index 4067edec01..675af0d7af 100644 --- a/model_gateway/src/policies/least_load.rs +++ b/model_gateway/src/policies/least_load.rs @@ -1,6 +1,9 @@ use std::{ collections::HashMap, - sync::{Arc, RwLock}, + sync::{ + atomic::{AtomicU64, Ordering}, + Arc, PoisonError, RwLock, + }, }; use openai_protocol::worker::WorkerLoadResponse; @@ -79,11 +82,16 @@ pub const DEFAULT_THROUGHPUT: f64 = 2000.0; /// in-flight correction absorbs staleness between polls. #[derive(Debug)] pub struct LeastLoadPolicy { - /// Cached load reports from the worker monitor (keyed by worker URL). - cached_loads: RwLock>, - /// In-flight token-work dispatched per worker since its last load poll - /// (keyed by worker URL); reset when a fresh report arrives. - inflight_tokens: RwLock>, + /// Per-worker load-scoring state keyed by worker URL. + /// + /// Read-mostly: written only by `update_loads` (poll interval) and + /// `remove_worker` (churn). Per-request selection takes a **read** guard — + /// concurrent selections never serialize on each other — and credits the + /// chosen worker's in-flight counter with a relaxed `fetch_add` on the + /// entry's atomic. (Previously a *write* guard was held across the whole + /// O(workers) scoring scan of every request, serializing all routing for + /// the model on one std `RwLock`.) + loads: RwLock>, /// KV-pressure weight `λ_t` (seconds). kv_pressure_weight: f64, /// Mean prefill length (tokens) for estimating in-flight token-work when a @@ -94,6 +102,28 @@ pub struct LeastLoadPolicy { default_throughput: f64, } +/// Per-worker load-scoring state: the latest poll snapshot plus the token-work +/// dispatched since that snapshot was taken. +#[derive(Debug)] +struct WorkerLoadState { + /// Latest load report from the worker monitor. + report: WorkerLoadResponse, + /// In-flight token-work dispatched since `report` was taken. Credited with + /// a relaxed `fetch_add` at selection time (under the map's read guard, so + /// concurrent credits are never lost) and reset to 0 when a fresh report + /// arrives. + inflight_tokens: AtomicU64, +} + +impl WorkerLoadState { + fn new(report: WorkerLoadResponse) -> Self { + Self { + report, + inflight_tokens: AtomicU64::new(0), + } + } +} + impl LeastLoadPolicy { pub fn new() -> Self { Self::with_params( @@ -117,8 +147,7 @@ impl LeastLoadPolicy { default_throughput: f64, ) -> Self { Self { - cached_loads: RwLock::new(HashMap::new()), - inflight_tokens: RwLock::new(HashMap::new()), + loads: RwLock::new(HashMap::new()), kv_pressure_weight: if kv_pressure_weight.is_finite() && kv_pressure_weight >= 0.0 { kv_pressure_weight } else { @@ -135,24 +164,24 @@ impl LeastLoadPolicy { /// Expected-wait score for a worker (lower is better). /// - /// `inflight` maps worker URL -> token-work dispatched since its last poll. - /// `nominal_throughput` (a peer-derived mean) estimates drain rate for a - /// worker missing a fresh snapshot; `fleet_has_loads` is false only when no - /// worker reports at all, in which case we fall back to join-shortest-queue - /// on the live in-flight count (which, unlike the since-poll estimate, - /// reflects completions and so suits backends that never report loads). + /// `state` is the worker's load entry (latest report + since-poll in-flight + /// token-work), resolved once by the caller. `nominal_throughput` (a + /// peer-derived mean) estimates drain rate for a worker missing a fresh + /// snapshot; `fleet_has_loads` is false only when no worker reports at all, + /// in which case we fall back to join-shortest-queue on the live in-flight + /// count (which, unlike the since-poll estimate, reflects completions and + /// so suits backends that never report loads). fn score( &self, worker: &Arc, - loads: Option<&HashMap>, - inflight: &HashMap, + state: Option<&WorkerLoadState>, nominal_throughput: f64, fleet_has_loads: bool, ) -> f64 { - let url = worker.url(); - match loads.and_then(|m| m.get(url)) { - Some(load) => { - let inflight_tokens = inflight.get(url).copied().unwrap_or(0) as f64; + match state { + Some(state) => { + let load = &state.report; + let inflight_tokens = state.inflight_tokens.load(Ordering::Relaxed) as f64; let queued_tokens = load.total_waiting_uncached_tokens() as f64; let live_throughput = load.total_gen_throughput(); let throughput = if live_throughput > 0.0 { @@ -195,62 +224,59 @@ impl LoadBalancingPolicy for LeastLoadPolicy { return Some(healthy[0]); } - let loads_guard = self.cached_loads.read().ok(); - let loads = loads_guard.as_deref(); - - // Fleet-nominal throughput (mean of positive reports) stands in for a - // worker missing a fresh snapshot; `fleet_has_loads` distinguishes a - // partial gap (estimate that worker's drain time at the nominal rate) - // from a fully dark fleet (fall back to join-shortest-queue). - let (tp_sum, tp_count) = healthy - .iter() - .filter_map(|&i| loads.and_then(|m| m.get(workers[i].url()))) - .map(|l| l.total_gen_throughput()) - .filter(|t| *t > 0.0) - .fold((0.0, 0u32), |(s, n), t| (s + t, n + 1)); - let nominal_throughput = if tp_count > 0 { - tp_sum / tp_count as f64 - } else { - self.default_throughput - }; - let fleet_has_loads = loads - .map(|m| healthy.iter().any(|&i| m.contains_key(workers[i].url()))) - .unwrap_or(false); - - // Held across selection so the in-flight estimate stays consistent and - // the chosen worker can be credited before the guard is released. - let mut inflight = self - .inflight_tokens - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - - let mut best = healthy[0]; - let mut best_score = self.score( - &workers[best], - loads, - &inflight, - nominal_throughput, - fleet_has_loads, - ); - for &idx in &healthy[1..] { - let s = self.score( - &workers[idx], - loads, - &inflight, - nominal_throughput, - fleet_has_loads, - ); - if s < best_score { - best = idx; - best_score = s; + // Selection holds only a shared read guard: concurrent requests score + // in parallel, and the chosen worker's in-flight credit is a relaxed + // `fetch_add` on its entry's atomic (never lost, no exclusive lock). + let (best, best_score) = { + let loads = self.loads.read().unwrap_or_else(PoisonError::into_inner); + + // Resolve each healthy worker's load entry once — a single + // URL-keyed probe per worker per request; the throughput and + // scoring passes below are string-free. + let entries: Vec<(usize, Option<&WorkerLoadState>)> = healthy + .iter() + .map(|&i| (i, loads.get(workers[i].url()))) + .collect(); + + // Fleet-nominal throughput (mean of positive reports) stands in for + // a worker missing a fresh snapshot; `fleet_has_loads` distinguishes + // a partial gap (estimate that worker's drain time at the nominal + // rate) from a fully dark fleet (fall back to join-shortest-queue). + let (tp_sum, tp_count) = entries + .iter() + .filter_map(|(_, state)| state.map(|s| s.report.total_gen_throughput())) + .filter(|t| *t > 0.0) + .fold((0.0, 0u32), |(s, n), t| (s + t, n + 1)); + let nominal_throughput = if tp_count > 0 { + tp_sum / tp_count as f64 + } else { + self.default_throughput + }; + let fleet_has_loads = entries.iter().any(|(_, state)| state.is_some()); + + let mut best = (healthy[0], None); + let mut best_score = f64::INFINITY; + for &(idx, state) in &entries { + let s = self.score(&workers[idx], state, nominal_throughput, fleet_has_loads); + if s < best_score { + best = (idx, state); + best_score = s; + } } - } - - // In-flight correction: credit the chosen worker with this request's - // token-work until its next poll refreshes the snapshot. - let req_tokens = self.request_tokens(info); - *inflight.entry(workers[best].url().to_string()).or_insert(0) += req_tokens; - drop(inflight); + let (best, best_state) = best; + + // In-flight correction: credit the chosen worker with this + // request's token-work until its next poll refreshes the snapshot. + // A worker without a load entry takes no credit: the estimate is + // only ever read for workers that have a report, and the + // `update_loads` that installs one resets the estimate anyway. + if let Some(state) = best_state { + state + .inflight_tokens + .fetch_add(self.request_tokens(info), Ordering::Relaxed); + } + (best, best_score) + }; debug!( "least_load selected {} (score {:.4}, in_flight {})", @@ -267,25 +293,24 @@ impl LoadBalancingPolicy for LeastLoadPolicy { } fn update_loads(&self, loads: &HashMap) { - if let Ok(mut cached) = self.cached_loads.write() { - cached.extend(loads.iter().map(|(k, v)| (k.clone(), v.clone()))); - } - // A fresh snapshot already reflects work up to the poll, so reset the - // since-poll in-flight estimate for the workers it covers. - if let Ok(mut inflight) = self.inflight_tokens.write() { - for url in loads.keys() { - inflight.insert(url.clone(), 0); + let mut map = self.loads.write().unwrap_or_else(PoisonError::into_inner); + for (url, report) in loads { + // A fresh snapshot already reflects work up to the poll, so reset + // the since-poll in-flight estimate for the workers it covers. + if let Some(state) = map.get_mut(url) { + state.report = report.clone(); + state.inflight_tokens.store(0, Ordering::Relaxed); + } else { + map.insert(url.clone(), WorkerLoadState::new(report.clone())); } } } fn remove_worker(&self, url: &str) { - if let Ok(mut cached) = self.cached_loads.write() { - cached.remove(url); - } - if let Ok(mut inflight) = self.inflight_tokens.write() { - inflight.remove(url); - } + self.loads + .write() + .unwrap_or_else(PoisonError::into_inner) + .remove(url); } fn as_any(&self) -> &dyn std::any::Any { @@ -483,20 +508,20 @@ mod tests { policy.select_worker(&workers, &info); } assert!(policy - .inflight_tokens + .loads .read() .unwrap() .values() - .any(|&v| v > 0)); + .any(|s| s.inflight_tokens.load(Ordering::Relaxed) > 0)); // A fresh poll clears the since-poll estimate. policy.update_loads(&loads); assert!(policy - .inflight_tokens + .loads .read() .unwrap() .values() - .all(|&v| v == 0)); + .all(|s| s.inflight_tokens.load(Ordering::Relaxed) == 0)); } #[test] @@ -516,13 +541,94 @@ mod tests { loads.insert("http://a:8000".to_string(), make_load(0, 0.5, 100.0)); loads.insert("http://b:8000".to_string(), make_load(0, 0.3, 100.0)); policy.update_loads(&loads); - assert_eq!(policy.cached_loads.read().unwrap().len(), 2); + assert_eq!(policy.loads.read().unwrap().len(), 2); // Removing a worker drops only its entry (no unbounded growth on churn). policy.remove_worker("http://a:8000"); - let cached = policy.cached_loads.read().unwrap(); + let cached = policy.loads.read().unwrap(); assert_eq!(cached.len(), 1); assert!(!cached.contains_key("http://a:8000")); assert!(cached.contains_key("http://b:8000")); } + + /// K4 regression: parallel selections must not lose in-flight credits. + /// + /// The exclusive write guard formerly held across the scoring scan made + /// scan+credit one atomic transaction; with lock-free credits the + /// invariant that must survive is that every request's token-work lands + /// exactly once — the total credited in-flight equals the sum dispatched. + #[test] + fn concurrent_selection_loses_no_inflight_credits() { + const THREADS: usize = 8; + const PER_THREAD: usize = 250; + const TOKENS_PER_REQUEST: usize = 7; + + let policy = LeastLoadPolicy::new(); + let workers = vec![mk("http://a:8000"), mk("http://b:8000")]; + let mut loads = HashMap::new(); + loads.insert("http://a:8000".to_string(), make_load(0, 0.1, 100.0)); + loads.insert("http://b:8000".to_string(), make_load(0, 0.1, 100.0)); + policy.update_loads(&loads); + + let tokens: Vec = (0..TOKENS_PER_REQUEST as u32).collect(); + std::thread::scope(|s| { + for _ in 0..THREADS { + s.spawn(|| { + for _ in 0..PER_THREAD { + let info = SelectWorkerInfo { + tokens: Some(&tokens), + ..Default::default() + }; + policy.select_worker(&workers, &info).unwrap(); + } + }); + } + }); + + let map = policy.loads.read().unwrap(); + let total: u64 = map + .values() + .map(|s| s.inflight_tokens.load(Ordering::Relaxed)) + .sum(); + assert_eq!( + total, + (THREADS * PER_THREAD * TOKENS_PER_REQUEST) as u64, + "in-flight token credits were lost under concurrent selection" + ); + } + + /// Selections concurrent with `update_loads`/`remove_worker` must not + /// deadlock or panic, and the map must reflect the final write afterwards + /// (the read-guard scan tolerates writers between requests). + #[test] + fn concurrent_selection_with_load_updates_is_safe() { + let policy = LeastLoadPolicy::new(); + let workers = vec![mk("http://a:8000"), mk("http://b:8000")]; + let mut loads = HashMap::new(); + loads.insert("http://a:8000".to_string(), make_load(0, 0.1, 100.0)); + loads.insert("http://b:8000".to_string(), make_load(0, 0.1, 100.0)); + policy.update_loads(&loads); + + std::thread::scope(|s| { + for _ in 0..4 { + s.spawn(|| { + for _ in 0..200 { + let info = SelectWorkerInfo::default(); + assert!(policy.select_worker(&workers, &info).is_some()); + } + }); + } + s.spawn(|| { + for _ in 0..50 { + policy.update_loads(&loads); + policy.remove_worker("http://b:8000"); + let mut one = HashMap::new(); + one.insert("http://b:8000".to_string(), make_load(0, 0.1, 100.0)); + policy.update_loads(&one); + } + }); + }); + + assert_eq!(policy.loads.read().unwrap().len(), 2); + } } diff --git a/model_gateway/src/policies/mod.rs b/model_gateway/src/policies/mod.rs index 00b2560fc1..16f6a13b62 100644 --- a/model_gateway/src/policies/mod.rs +++ b/model_gateway/src/policies/mod.rs @@ -3,8 +3,13 @@ //! This module provides a unified abstraction for routing policies that work //! across both regular and prefill-decode (PD) routing modes. -use std::{fmt::Debug, sync::Arc}; +use std::{ + borrow::Cow, + fmt::Debug, + sync::{Arc, OnceLock}, +}; +use kv_index::{compute_request_content_hashes, ContentHash}; use openai_protocol::worker::WorkerLoadResponse; use crate::worker::{HashRing, Worker}; @@ -201,6 +206,42 @@ pub struct SelectWorkerInfo<'a> { /// 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>, + /// Request-scoped memo of the per-block content hashes derived from + /// `tokens` (see [`SelectWorkerInfo::request_content_hashes`]). Stored as + /// `(block_size, hashes)` so the memo is only reused for the block size it + /// was computed with. Filled lazily by the first policy that needs it; + /// callers that run selection more than once per request (PD prefill + + /// decode, retries) reuse the same `SelectWorkerInfo` so the full-prompt + /// hash pass runs at most once. + pub content_hashes: OnceLock<(usize, Vec)>, +} + +impl SelectWorkerInfo<'_> { + /// Per-block content hashes of `tokens` for `block_size`, memoized on this + /// request. + /// + /// The memo assumes one `tokens` value per `SelectWorkerInfo` (construction + /// sites build the struct once per request). The block size is recorded + /// with the memo: a call with a different block size recomputes instead of + /// reusing hashes chunked at the wrong width (returns `Cow::Owned` in that + /// case, without disturbing the memo). + pub fn request_content_hashes( + &self, + tokens: &[u32], + block_size: usize, + ) -> Cow<'_, [ContentHash]> { + let (memo_block_size, hashes) = self.content_hashes.get_or_init(|| { + ( + block_size, + compute_request_content_hashes(tokens, block_size), + ) + }); + if *memo_block_size == block_size { + Cow::Borrowed(hashes.as_slice()) + } else { + Cow::Owned(compute_request_content_hashes(tokens, block_size)) + } + } } #[cfg(test)] @@ -217,6 +258,29 @@ mod tests { } } + #[test] + fn request_content_hashes_memoized_per_block_size() { + let tokens: [u32; 8] = [1, 2, 3, 4, 5, 6, 7, 8]; + let info = SelectWorkerInfo { + tokens: Some(&tokens), + ..Default::default() + }; + + let first = info.request_content_hashes(&tokens, 4); + assert_eq!(first.len(), 2); + // Same block size: served from the memo (same allocation, no rehash). + let again = info.request_content_hashes(&tokens, 4); + assert_eq!(first.as_ptr(), again.as_ptr()); + assert_eq!(*first, *again); + + // Different block size: recomputed at the right width; the memo keeps + // the original block size. + let other = info.request_content_hashes(&tokens, 8); + assert_eq!(other.len(), 1); + assert_eq!(other[0], kv_index::compute_content_hash(&tokens)); + assert_eq!(info.content_hashes.get().unwrap().0, 4); + } + #[test] fn test_get_healthy_worker_indices() { let workers: Vec> = vec![ diff --git a/model_gateway/src/routers/grpc/common/stages/worker_selection.rs b/model_gateway/src/routers/grpc/common/stages/worker_selection.rs index a1b270de2e..7cf197b5f1 100644 --- a/model_gateway/src/routers/grpc/common/stages/worker_selection.rs +++ b/model_gateway/src/routers/grpc/common/stages/worker_selection.rs @@ -161,6 +161,7 @@ impl WorkerSelectionStage { tokens, headers, hash_ring, + content_hashes: Default::default(), }, )?; let selected = available[idx].clone(); @@ -267,11 +268,15 @@ impl WorkerSelectionStage { // Get cached hash ring for consistent hashing (O(log n) lookup) let hash_ring = self.worker_registry.get_hash_ring(model_id); + // One SelectWorkerInfo for both selections: request-scoped memos (the + // per-block content hashes for cache-aware scoring) are computed once + // and shared between the prefill and decode passes. let info = SelectWorkerInfo { request_text: text, tokens, headers, hash_ring, + content_hashes: Default::default(), }; let prefill_idx = policy.select_worker(&available_prefill, &info)?; let decode_idx = policy.select_worker(&available_decode, &info)?; diff --git a/model_gateway/src/routers/http/pd_router.rs b/model_gateway/src/routers/http/pd_router.rs index d17b6d3468..9fbf975166 100644 --- a/model_gateway/src/routers/http/pd_router.rs +++ b/model_gateway/src/routers/http/pd_router.rs @@ -876,6 +876,7 @@ impl PDRouter { tokens: None, // HTTP doesn't have tokens, use gRPC for PrefixHash headers, hash_ring, + content_hashes: Default::default(), }, ) .ok_or_else(|| { diff --git a/model_gateway/src/routers/http/router.rs b/model_gateway/src/routers/http/router.rs index 0799467f1e..04f6475f55 100644 --- a/model_gateway/src/routers/http/router.rs +++ b/model_gateway/src/routers/http/router.rs @@ -180,6 +180,7 @@ impl Router { tokens: None, // HTTP doesn't have tokens, use gRPC for PrefixHash headers, hash_ring, + content_hashes: Default::default(), }, )?; @@ -573,6 +574,7 @@ impl Router { tokens: None, headers, hash_ring, + content_hashes: Default::default(), }, ) { Some(i) => i, diff --git a/model_gateway/src/worker/kv_event_monitor.rs b/model_gateway/src/worker/kv_event_monitor.rs index f5cc043de8..19b7d9acec 100644 --- a/model_gateway/src/worker/kv_event_monitor.rs +++ b/model_gateway/src/worker/kv_event_monitor.rs @@ -13,6 +13,7 @@ use std::{collections::HashMap, fmt, sync::Arc, time::Duration}; use dashmap::DashMap; +use futures::FutureExt as _; use kv_index::{ compute_content_hash, ApplyError, PositionalIndexer, SequenceHash, StoredBlock, WorkerBlockMap, }; @@ -23,9 +24,12 @@ use tokio::{ sync::{oneshot, Mutex}, task::JoinHandle, }; -use tracing::{debug, info, warn}; +use tracing::{debug, error, info, warn}; -use crate::worker::{ConnectionMode, Worker, UNKNOWN_MODEL_ID}; +use crate::{ + observability::metrics::Metrics, + worker::{ConnectionMode, Worker, UNKNOWN_MODEL_ID}, +}; /// Default jump size for new `PositionalIndexer` instances. const DEFAULT_JUMP_SIZE: usize = 64; @@ -138,6 +142,7 @@ impl KvEventMonitor { let (shutdown_tx, shutdown_rx) = oneshot::channel(); let loop_model_id = model_id.clone(); + let task_url = url.clone(); #[expect( clippy::disallowed_methods, @@ -145,15 +150,34 @@ impl KvEventMonitor { handle is stored and graceful shutdown is sent on removal" )] let handle = tokio::spawn(async move { - Self::subscription_loop( + // Catch panics here so they surface when they happen — a bare + // JoinError would only be observed at worker removal, leaving the + // index silently frozen for this worker until then. + let result = std::panic::AssertUnwindSafe(Self::subscription_loop( worker, worker_url, indexer, block_sizes, loop_model_id, shutdown_rx, - ) + )) + .catch_unwind() .await; + if let Err(payload) = result { + let msg = payload + .downcast_ref::<&str>() + .copied() + .map(String::from) + .or_else(|| payload.downcast_ref::().cloned()) + .unwrap_or_else(|| "(non-string panic)".into()); + error!( + worker_url = %task_url, + panic.message = %msg, + "KV event subscription task panicked; KV events from this \ + worker no longer feed cache-aware routing" + ); + Metrics::record_kv_event_subscription_failure(&task_url, "panic"); + } }); handles.insert( @@ -183,7 +207,16 @@ impl KvEventMonitor { info!(worker_url = %worker_url, "Stopping KV event subscription"); // Signal graceful shutdown — task cleans up its worker_blocks in the indexer. let _ = sub.shutdown_tx.send(()); - let _ = sub.handle.await; + // Panics are caught inside the task; a JoinError here (abort or a + // panic that escaped the guard) must still be surfaced, not discarded. + if let Err(e) = sub.handle.await { + error!( + worker_url = %worker_url, + error = %e, + "KV event subscription task failed" + ); + Metrics::record_kv_event_subscription_failure(worker_url, "join_error"); + } // Re-check under lock whether this was the last worker for the model. // Must re-acquire lock after shutdown to avoid TOCTOU with concurrent @@ -215,7 +248,14 @@ impl KvEventMonitor { for (url, sub) in subscriptions { debug!(worker_url = %url, "Stopping KV event subscription"); let _ = sub.shutdown_tx.send(()); - let _ = sub.handle.await; + if let Err(e) = sub.handle.await { + error!( + worker_url = %url, + error = %e, + "KV event subscription task failed" + ); + Metrics::record_kv_event_subscription_failure(&url, "join_error"); + } } } @@ -308,7 +348,19 @@ impl KvEventMonitor { model_id: String, mut shutdown_rx: oneshot::Receiver<()>, ) { - let worker_id = indexer.intern_worker(&worker_url); + let worker_id = match indexer.intern_worker(&worker_url) { + Ok(id) => id, + Err(e) => { + error!( + worker_url = %worker_url, + error = %e, + "Failed to intern worker; KV events from this worker will \ + not feed cache-aware routing" + ); + Metrics::record_kv_event_subscription_failure(&worker_url, "intern_failed"); + return; + } + }; let mut worker_blocks = WorkerBlockMap::default(); let mut last_seq: u64 = 0; let mut reconnect_delay_ms = INITIAL_RECONNECT_DELAY_MS; @@ -673,7 +725,7 @@ mod tests { #[test] fn test_apply_stored_no_parent() { let indexer = PositionalIndexer::new(64); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb = WorkerBlockMap::default(); let stored = KvBlocksStored { blocks: vec![ @@ -702,7 +754,7 @@ mod tests { #[test] fn test_apply_stored_with_parent() { let indexer = PositionalIndexer::new(64); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb = WorkerBlockMap::default(); let stored1 = KvBlocksStored { @@ -734,7 +786,7 @@ mod tests { #[test] fn test_apply_stored_fallback_on_worker_not_tracked() { let indexer = PositionalIndexer::new(64); - let w1 = indexer.intern_worker("http://new-worker:8000"); + let w1 = indexer.intern_worker("http://new-worker:8000").unwrap(); let mut wb = WorkerBlockMap::default(); // Pass parent_block_hash for an untracked worker — should fallback to no parent. @@ -755,7 +807,7 @@ mod tests { #[test] fn test_apply_removed() { let indexer = PositionalIndexer::new(64); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb = WorkerBlockMap::default(); let stored = KvBlocksStored { @@ -790,7 +842,7 @@ mod tests { #[test] fn test_apply_cleared_event() { let indexer = PositionalIndexer::new(64); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb = WorkerBlockMap::default(); let stored = KvBlocksStored { @@ -813,7 +865,7 @@ mod tests { #[test] fn test_apply_event_dispatch_stored() { let indexer = PositionalIndexer::new(64); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb = WorkerBlockMap::default(); let event = KvCacheEvent { event_id: 1, @@ -836,7 +888,7 @@ mod tests { #[test] fn test_apply_event_dispatch_removed() { let indexer = PositionalIndexer::new(64); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb = WorkerBlockMap::default(); let stored_event = KvCacheEvent { @@ -868,7 +920,7 @@ mod tests { #[test] fn test_apply_event_dispatch_cleared() { let indexer = PositionalIndexer::new(64); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb = WorkerBlockMap::default(); KvEventMonitor::apply_event( @@ -908,7 +960,7 @@ mod tests { #[test] fn test_apply_event_no_data() { let indexer = PositionalIndexer::new(64); - let w1 = indexer.intern_worker("http://w1:8000"); + let w1 = indexer.intern_worker("http://w1:8000").unwrap(); let mut wb = WorkerBlockMap::default(); let event = KvCacheEvent { event_id: 1, diff --git a/model_gateway/src/worker/worker.rs b/model_gateway/src/worker/worker.rs index 7d520903d5..019d9b08ed 100644 --- a/model_gateway/src/worker/worker.rs +++ b/model_gateway/src/worker/worker.rs @@ -291,6 +291,25 @@ pub trait Worker: Send + Sync + fmt::Debug + 'static { /// Increment the processed requests counter fn increment_processed(&self); + /// Cached kv_index worker id for the indexer instance ("epoch") given, if + /// one was stored via [`cache_kv_worker_id`](Self::cache_kv_worker_id). + /// + /// Event-driven cache-aware scoring resolves each worker's interned u32 id + /// at most once per (worker, indexer instance) and rereads it as a single + /// atomic load, instead of probing a URL-keyed map on every request. + /// Default implementation has no cache and always returns `None` (callers + /// then resolve through the indexer directly). + fn cached_kv_worker_id(&self, indexer_epoch: u32) -> Option { + let _ = indexer_epoch; + None + } + + /// Cache this worker's interned kv_index id for the indexer instance + /// ("epoch") given. Default implementation is a no-op. + fn cache_kv_worker_id(&self, indexer_epoch: u32, worker_id: u32) { + let _ = (indexer_epoch, worker_id); + } + /// Get worker-specific metadata fn metadata(&self) -> &WorkerMetadata; @@ -784,6 +803,14 @@ pub struct WorkerRuntime { processed_counter: AtomicUsize, worker_routing_key_load: WorkerRoutingKeyLoad, revision: AtomicU64, + /// Cached kv_index worker id, packed as `(indexer_epoch << 32) | worker_id` + /// with `0` meaning "unset" (indexer epochs are non-zero). Written by + /// cache-aware event-driven scoring on first contact with an indexer + /// instance so the per-request scoring loop reads one atomic instead of + /// probing a URL-keyed map. Lives in the shared runtime so a same-URL + /// `replace()` keeps the cache (the id is a pure function of the URL + /// within one indexer instance). + kv_worker_id: AtomicU64, } impl WorkerRuntime { @@ -797,6 +824,7 @@ impl WorkerRuntime { processed_counter: AtomicUsize::new(0), worker_routing_key_load: WorkerRoutingKeyLoad::new(url), revision: AtomicU64::new(0), + kv_worker_id: AtomicU64::new(0), } } @@ -891,6 +919,27 @@ impl WorkerRuntime { pub fn increment_processed(&self) { self.processed_counter.fetch_add(1, Ordering::Relaxed); } + + // ── Cached kv_index worker id ─────────────────────────────────── + + /// Cached kv_index worker id for the indexer instance `indexer_epoch`, + /// or `None` when unset or cached for a different instance. + pub fn cached_kv_worker_id(&self, indexer_epoch: u32) -> Option { + if indexer_epoch == 0 { + return None; + } + let packed = self.kv_worker_id.load(Ordering::Relaxed); + ((packed >> 32) as u32 == indexer_epoch).then_some(packed as u32) + } + + /// Cache this worker's interned kv_index id for the indexer instance + /// `indexer_epoch` (must be non-zero; kv_index guarantees that). + pub fn cache_kv_worker_id(&self, indexer_epoch: u32, worker_id: u32) { + self.kv_worker_id.store( + (u64::from(indexer_epoch) << 32) | u64::from(worker_id), + Ordering::Relaxed, + ); + } } /// Basic worker implementation @@ -1093,6 +1142,16 @@ impl Worker for BasicWorker { self.runtime.load().increment_processed(); } + fn cached_kv_worker_id(&self, indexer_epoch: u32) -> Option { + self.runtime.load().cached_kv_worker_id(indexer_epoch) + } + + fn cache_kv_worker_id(&self, indexer_epoch: u32, worker_id: u32) { + self.runtime + .load() + .cache_kv_worker_id(indexer_epoch, worker_id); + } + fn metadata(&self) -> &WorkerMetadata { &self.metadata }