From 09b0a5c4e67fef4c2b8a98f722528db49e921690 Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Tue, 9 Jun 2026 12:01:51 -0700 Subject: [PATCH 1/2] perf(kv-index): fuse cache-aware match+insert into a single tree descent Cache-aware routing did two full top-to-bottom descents over the same request prefix per request (match then insert). Add match_and_insert / match_and_insert_with that descend once, and use them in cache_aware.rs. Benchmarked ~2x faster on the match+insert hot path, scaling with context length (131072 tokens: 58.5us -> 29.3us). Equivalence unit tests assert the fused path matches the separate match+insert on results and tenant counts. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- crates/kv_index/src/string_tree.rs | 340 ++++++++- crates/kv_index/src/token_tree.rs | 699 +++++++++++++++++- model_gateway/benches/radix_tree_benchmark.rs | 78 ++ model_gateway/src/policies/cache_aware.rs | 166 +++-- 4 files changed, 1194 insertions(+), 89 deletions(-) diff --git a/crates/kv_index/src/string_tree.rs b/crates/kv_index/src/string_tree.rs index f95c1adf14..d31dcf8e52 100644 --- a/crates/kv_index/src/string_tree.rs +++ b/crates/kv_index/src/string_tree.rs @@ -407,9 +407,23 @@ impl Tree { .entry(Arc::clone(&tenant_id)) .or_insert(0); - // Track remaining text as a slice - no allocation needed - let mut remaining = text; - let mut prev = Arc::clone(&self.root); + // Descend from the root inserting the whole text. + self.insert_from(Arc::clone(&self.root), text, tenant_id); + } + + /// Insert `remaining` for `tenant_id` starting the descent at `start` + /// (which is treated as the parent of the first edge), reusing the exact + /// node-split / tenant-attach / leaf-timestamp logic of [`Self::insert_text`]. + /// + /// `insert_text` is `insert_from(root, text, tenant_id)` after the root + /// bookkeeping. [`Self::match_and_insert_with`] reuses it to splice only the + /// *unmatched suffix* at the fall-off node, so the already-matched prefix is + /// never re-walked. The caller is responsible for the root bookkeeping + /// (`tenant_last_access_time` / `tenant_char_count` entries) and, when + /// resuming below the root, for attaching `tenant_id` to the ancestor nodes + /// on the matched path (which this method does not touch). + fn insert_from(&self, start: NodeRef, mut remaining: &str, tenant_id: TenantId) { + let mut prev = start; // Result type to carry state out of the match block // This allows the entry guard to be dropped before we update prev @@ -648,6 +662,230 @@ impl Tree { } } + /// Read-only resolution of a node's "owning" tenant, mirroring the tenant + /// pick in [`Self::match_prefix_with_counts`] **without** the cache-populate + /// or probabilistic timestamp side effects. Returns the tenant plus whether + /// the slow (iteration) path was taken — the caller replays the original's + /// single deferred side effect on the final node. + /// + /// Used by [`Self::match_and_insert`] so the match tenant is resolved from a + /// node's state **before** the fused insert adds the inserting tenant to it, + /// reproducing the original "match runs fully before insert" ordering for the + /// (otherwise non-deterministic) slow-path tenant pick. + #[expect( + clippy::unused_self, + reason = "method logically belongs to the tree; mirrors match_prefix_with_counts resolution" + )] + fn resolve_tenant_readonly(&self, node: &NodeRef) -> (TenantId, bool) { + let cached = node.last_tenant.read(); + if let Some(ref t) = *cached { + if node.tenant_last_access_time.contains_key(t.as_ref()) { + return (Arc::clone(t), false); + } + } + drop(cached); + let t = node + .tenant_last_access_time + .iter() + .next() + .map(|kv| Arc::clone(kv.key())) + .unwrap_or_else(|| Arc::from("empty")); + (t, true) + } + + /// Combined match + insert for a known `tenant` in a single descent. + /// + /// Thin wrapper over [`Self::match_and_insert_with`] for callers that already + /// know the tenant to insert for before matching (e.g. the imbalanced + /// min-load path, which routes by worker load, not cache affinity). Returns + /// the match result for any cache bookkeeping the caller needs. + pub fn match_and_insert(&self, text: &str, tenant: &str) -> PrefixMatchResult { + self.match_and_insert_with(text, move |_| Some(tenant)) + } + + /// Match in a single descent, then choose the insert tenant from the match + /// result and insert for it — still a SINGLE walk of the matched prefix. + /// + /// String analogue of `TokenTree::match_and_insert_with`. The cache-aware + /// router picks the worker it inserts for *from* the match outcome, so the + /// tenant is unknown until the match finishes. `select` runs once, after the + /// match, with the [`PrefixMatchResult`]; `Some(tenant)` inserts `text` for + /// that tenant, `None` skips the insert (the router's "no worker selected" + /// branch). + /// + /// # How the single descent is achieved + /// + /// The match phase walks the prefix exactly like + /// [`Self::match_prefix_with_counts`] (including its single deferred, + /// probabilistic timestamp touch on the resolved node — done via + /// [`Self::finish_match_and_insert`]), while recording the chain of + /// full-match nodes it descended through. After `select` yields the tenant we + /// re-attach it to those recorded ancestor nodes directly (the epoch-0 + /// intermediate attach `insert_text` performs while descending) and splice + /// only the *unmatched suffix* at the fall-off node via + /// [`Self::insert_from`]. The matched prefix is therefore compared once, not + /// twice. + /// + /// # Preserved semantics + /// + /// Identical to `match_prefix_with_counts` followed by + /// `insert_text(text, tenant)`: same matched/-input char counts, same + /// resolved tenant and probabilistic touch, same node splits, same epoch-0 + /// ancestor attaches, same final-leaf real timestamp, and same + /// `tenant_char_count` accounting. Because the match phase runs fully before + /// any insert mutation (just like the two separate calls), the tenant pick + /// observes the un-polluted tree; only the exact monotonic timestamp values + /// of insert's writes shift by a few ticks, which is immaterial to LRU. + pub fn match_and_insert_with<'t, F>(&self, text: &str, select: F) -> PrefixMatchResult + where + F: FnOnce(&PrefixMatchResult) -> Option<&'t str>, + { + // ---- Phase 1: MATCH descent (mirrors match_prefix_with_counts) ---- + // Record the full-match nodes (for ancestor re-attach) and the fall-off + // point (node + remaining slice) for the suffix splice. No mutation here + // beyond match's own deferred touch below, so the tenant pick sees the + // un-polluted tree exactly like the standalone match. + let mut remaining = text; + let mut matched_chars = 0usize; + let mut current = Arc::clone(&self.root); + // The node match resolves its tenant on (its final `curr`): the deepest + // full-match node, the partial child, or the root if nothing matched. + let mut match_curr = Arc::clone(&self.root); + // (node, char_count) for each full-match edge, in order. + let mut path: Vec<(NodeRef, usize)> = Vec::new(); + + while let Some(first_char) = remaining.chars().next() { + let child_node = current.children.get(&first_char).map(|e| e.value().clone()); + + let Some(matched_node) = child_node else { + // No child for this char: match stops at `current`. + break; + }; + + let matched_text_guard = matched_node.text.read(); + let matched_node_text_count = matched_text_guard.char_count(); + let shared_count = shared_prefix_count(remaining, matched_text_guard.as_str()); + drop(matched_text_guard); + + if shared_count == matched_node_text_count { + // Full match -> continue. Record for ancestor re-attach. + matched_chars += shared_count; + path.push((Arc::clone(&matched_node), matched_node_text_count)); + remaining = advance_by_chars(remaining, shared_count); + current = Arc::clone(&matched_node); + match_curr = matched_node; + } else { + // Partial match: match stops, resolving on this child node. + matched_chars += shared_count; + match_curr = matched_node; + // `current` stays the parent — the splice re-probes it and + // splits the partial child exactly like `insert_text`. + break; + } + } + + // ---- Match side effect + result (verbatim match_prefix_with_counts) ---- + let (match_tenant, match_slow) = self.resolve_tenant_readonly(&match_curr); + let result = self.finish_match_and_insert( + &match_curr, + &match_tenant, + match_slow, + matched_chars, + text, + ); + + // ---- Decide the insert tenant from the match result ---- + let Some(tenant) = select(&result) else { + return result; + }; + // Intern through the shared pool so the stored Arc dedups (matches + // insert_text's interning). + let tenant_id = intern_tenant(tenant); + + // ---- Phase 2: INSERT for `tenant_id` without re-walking the prefix ---- + // Root bookkeeping (mirrors insert_text). + self.root + .tenant_last_access_time + .entry(Arc::clone(&tenant_id)) + .or_insert(0); + self.tenant_char_count + .entry(Arc::clone(&tenant_id)) + .or_insert(0); + + // Re-attach the inserting tenant to every full-match ancestor node, the + // epoch-0 intermediate attach `insert_text` performs while descending. + for (node, char_count) in &path { + if !node + .tenant_last_access_time + .contains_key(tenant_id.as_ref()) + { + self.tenant_char_count + .entry(Arc::clone(&tenant_id)) + .and_modify(|count| *count += *char_count) + .or_insert(*char_count); + node.tenant_last_access_time + .insert(Arc::clone(&tenant_id), 0); + } + } + + if remaining.is_empty() { + // Loop-end: `current` is the final leaf; give it the real timestamp + // (insert_text's tail). It already received the epoch-0 attach above + // if it is on the path. + let epoch = get_epoch(); + current + .tenant_last_access_time + .insert(Arc::clone(&tenant_id), epoch); + } else { + // Fall-off: splice only the unmatched suffix below `current`, + // reusing insert_text's exact loop (handles split-continue + the + // final-leaf real timestamp). The matched prefix is not re-walked. + self.insert_from(current, remaining, tenant_id); + } + + result + } + + /// Replay `match_prefix_with_counts`'s single deferred side effect (populate + /// the `last_tenant` cache on the slow path, then a probabilistic 1/8 + /// timestamp touch) on the resolved final node, and build the match result. + /// Factored out so [`Self::match_and_insert_with`] applies it identically to + /// the standalone `match_prefix_with_counts` tail. + #[expect( + clippy::unused_self, + reason = "method logically belongs to the tree; mirrors match_prefix_with_counts tail" + )] + fn finish_match_and_insert( + &self, + match_node: &NodeRef, + tenant: &TenantId, + took_slow_path: bool, + matched_chars: usize, + text: &str, + ) -> PrefixMatchResult { + // On the slow path the original resolution populates the cache. + if took_slow_path { + *match_node.last_tenant.write() = Some(Arc::clone(tenant)); + } + + // Probabilistic (1 in 8) timestamp touch on the resolved tenant, skipping + // the synthetic "empty" tenant — identical to match_prefix_with_counts. + let epoch = get_epoch(); + if epoch & 0x7 == 0 && tenant.as_ref() != "empty" { + match_node + .tenant_last_access_time + .insert(Arc::clone(tenant), epoch); + } + + let input_char_count = text.chars().count(); + + PrefixMatchResult { + tenant: Arc::clone(tenant), + matched_char_count: matched_chars, + input_char_count, + } + } + /// Legacy prefix_match API for backward compatibility. /// Note: This computes matched_text which has allocation overhead. pub fn prefix_match_legacy(&self, text: &str) -> (String, String) { @@ -3015,4 +3253,100 @@ mod tests { paths.sort(); assert_eq!(paths, vec!["", "你好", "你好世界", "你好朋友"]); } + + /// Run the same op sequence two ways — `match_prefix_with_counts` + + /// `insert_text` (legacy pair) vs the fused `match_and_insert` — and assert + /// the match counts and resulting per-tenant char counts agree. Timestamps + /// and the probabilistic tenant cache are not compared. + fn assert_fused_matches_pair(ops: &[(&str, &str)]) { + let pair = Tree::new(); + let fused = Tree::new(); + for (text, tenant) in ops { + let r_pair = pair.match_prefix_with_counts(text); + pair.insert_text(text, tenant); + + let r_fused = fused.match_and_insert(text, tenant); + + assert_eq!( + r_pair.matched_char_count, r_fused.matched_char_count, + "matched_char_count mismatch for tenant {tenant}" + ); + assert_eq!( + r_pair.input_char_count, r_fused.input_char_count, + "input_char_count mismatch" + ); + assert_eq!( + r_pair.tenant.as_ref(), + r_fused.tenant.as_ref(), + "matched tenant mismatch for input {text:?}" + ); + } + assert_eq!( + pair.get_tenant_char_count(), + fused.get_tenant_char_count(), + "tenant char counts diverged between pair and fused" + ); + } + + #[test] + fn test_string_match_and_insert_equiv_basic() { + // Note: every query here either misses on an empty tree (resolves to the + // synthetic "empty") or falls off on a node whose `last_tenant` is + // single-valued, so the (otherwise approximate) tenant pick is + // deterministic and comparable between the two construction paths. + assert_fused_matches_pair(&[ + ("hello world", "w1"), // fresh leaf + ("hello world", "w1"), // exact re-insert (full match, loop-end) + ("hello there", "w2"), // shared "hello " prefix -> split + ("hello", "w3"), // prefix of "hello " split node + ]); + } + + #[test] + fn test_string_match_and_insert_equiv_deep_chain() { + // Build a multi-node chain ("abc" -> {"def","xyz"}), then a query that + // FULL-matches several nodes before falling off — exercising the fused + // path's ancestor re-attach (`path`) plus a deep `insert_from` splice. + assert_fused_matches_pair(&[ + ("abcdef", "w1"), // leaf "abcdef" + ("abcxyz", "w2"), // split -> "abc" + "def"/"xyz" + ("abcdefghi", "w1"), // full-match "abc"+"def", then append "ghi" + ("abcdefghi", "w3"), // full-match whole chain for a new tenant + ]); + } + + #[test] + fn test_string_match_and_insert_equiv_empty_and_unicode() { + assert_fused_matches_pair(&[ + ("", "w1"), // empty text -> tenant attached at root + ("你好世界", "w1"), // multi-byte fresh + ("你好朋友", "w2"), // multi-byte shared-prefix split + ]); + } + + #[test] + fn test_string_match_and_insert_with_select_and_skip() { + let text = "the quick brown fox"; + + let via_with = Tree::new(); + via_with.match_and_insert_with(text, |_| Some("w1")); + let via_plain = Tree::new(); + via_plain.match_and_insert(text, "w1"); + assert_eq!( + via_with.get_tenant_char_count(), + via_plain.get_tenant_char_count() + ); + + // None selection must leave the tree untouched. + let tree = Tree::new(); + tree.insert_text(text, "w1"); + let before = tree.get_tenant_char_count(); + let r = tree.match_and_insert_with(text, |_| None); + assert_eq!(r.matched_char_count, text.chars().count()); + assert_eq!( + tree.get_tenant_char_count(), + before, + "None selection must not insert" + ); + } } diff --git a/crates/kv_index/src/token_tree.rs b/crates/kv_index/src/token_tree.rs index 3593449072..3c5366cd39 100644 --- a/crates/kv_index/src/token_tree.rs +++ b/crates/kv_index/src/token_tree.rs @@ -419,13 +419,44 @@ impl TokenTree { .entry(Arc::clone(&tenant_id)) .or_insert(0); - let mut remaining = tokens; - let mut current = Arc::clone(&self.root); - let mut tokens_added = 0usize; - // Compute once: only track hit counts for LFU policy let track_lfu = self.eviction_policy == EvictionPolicy::Lfu; + // Descend from the root inserting the whole sequence. + let tokens_added = Self::insert_from( + Arc::clone(&self.root), + tokens, + Arc::clone(&tenant_id), + track_lfu, + ); + + // Update tenant token count + if tokens_added > 0 { + self.tenant_token_count + .entry(tenant_id) + .and_modify(|c| *c += tokens_added) + .or_insert(tokens_added); + } + } + + /// Insert `remaining` for `tenant_id` starting the descent at `current` + /// (treated as the parent of the first edge), returning the number of new + /// tokens to add to the tenant's count. Holds the exact node-split / + /// tenant-attach / counting logic of [`Self::insert_tokens`]; the caller + /// owns root bookkeeping and folding the returned count into + /// `tenant_token_count`. + /// + /// [`Self::match_and_insert_with`] reuses this to splice only the *unmatched + /// suffix* at the fall-off node (after re-attaching the tenant to the matched + /// ancestor nodes), so the already-matched prefix is never re-walked. + fn insert_from( + mut current: NodeRef, + mut remaining: &[TokenId], + tenant_id: TenantId, + track_lfu: bool, + ) -> usize { + let mut tokens_added = 0usize; + // Result type to carry state out of the match block // This allows the entry guard to be dropped before we update current enum InsertStep { @@ -599,13 +630,9 @@ impl TokenTree { } } - // Update tenant token count - if tokens_added > 0 { - self.tenant_token_count - .entry(tenant_id) - .and_modify(|c| *c += tokens_added) - .or_insert(tokens_added); - } + // The caller folds this into `tenant_token_count` (once, after any + // matched-prefix re-attach in match_and_insert_with). + tokens_added } /// Find longest matching prefix with detailed counts. @@ -739,6 +766,548 @@ impl TokenTree { } } + /// Combined match + insert in a SINGLE tree descent. + /// + /// Equivalent to calling [`Self::match_prefix_with_counts`] immediately + /// followed by [`Self::insert_tokens`] with the same `tokens`, but it + /// traverses the prefix only once. For long prefixes (e.g. 150K tokens) + /// this halves the number of page-key lookups, child-map probes, token + /// comparisons, and `touch_tenant` writes on the request hot path. + /// + /// # Why a single descent is correct + /// + /// `insert_tokens` descends through *full-match* children exactly the way + /// `match_prefix_with_counts` does (same page key, same page-aligned common + /// prefix, same "continue on full match" rule). Insert's descent is a + /// (depth-wise) superset of match's: match stops as soon as it hits a node + /// with no tenants (all evicted) or a partial match, whereas insert keeps + /// going on a full match and only stops when it must create or split a node. + /// So we drive the descent with insert's logic and *freeze* the match result + /// the first time match's stop condition is reached (tracked by + /// `match_frozen`). After freezing, deeper nodes only affect insert + /// accounting, never the returned match result — exactly as if the separate + /// `match` call had already returned. + /// + /// # Preserved semantics (identical to match-then-insert) + /// + /// * **Node split**: the vacant / prefix-of-child / diverge branches below + /// are copied verbatim from `insert_tokens` (intermediate node inherits the + /// child's tenant map, hit_count, creation_time, priority; child is demoted + /// to the suffix; parent pointers / page keys rewritten the same way). + /// * **Per-tenant token counting**: `tokens_added` is accumulated with the + /// same rules as `insert_tokens` (full-match Continue counts `common_len`; + /// the split branches count `common_len` only when the tenant did not + /// already own the path, plus any brand-new branch tokens) and folded into + /// `tenant_token_count` once at the end. + /// * **`touch_tenant` / timestamps**: at every node we reproduce the exact + /// touch sequence the two separate calls would have made, in the same + /// order, on the same node identities: + /// 1. The match side reads `child.get_any_tenant()` *before* any insert + /// mutation (so the all-evicted check and the routed tenant are computed + /// against the pre-insert node), then touches that tenant — exactly what + /// `match_prefix_with_counts` does, and before any split so the split's + /// tenant-map clone inherits the fresh timestamp just like the original + /// ordering (`match` ran fully before `insert`). + /// 2. The insert side then touches the inserting `tenant` on the node it + /// would have (the continued child, the newly created leaf, or the new + /// intermediate). + /// We deliberately do NOT deduplicate the two touches even when they land + /// on the same node for the same tenant: the original match-then-insert + /// pair already touched twice in that case (e.g. a full-match continuation + /// whose `get_any_tenant()` is the routed/inserting tenant), so keeping + /// both touches reproduces the LFU `hit_count` and timestamp progression + /// byte-for-byte. For the default LRU policy the second touch only advances + /// the timestamp, which is immaterial to relative ordering. + pub fn match_and_insert(&self, tokens: &[TokenId], tenant: &str) -> PrefixMatchResult { + let input_token_count = tokens.len(); + + // Align to page boundary (truncate to nearest page). Mirrors both + // `match_prefix_with_counts` and `insert_tokens`. + let aligned_len = align_to_page(tokens.len()); + if aligned_len == 0 { + // Too short to cache: `insert_tokens` is a no-op and + // `match_prefix_with_counts` returns 0 matched tokens with any + // root tenant. Reproduce the match result, skip the insert. + return PrefixMatchResult { + tenant: self + .root + .get_any_tenant() + .unwrap_or_else(|| Arc::from("empty")), + matched_token_count: 0, + input_token_count, + }; + } + let tokens = &tokens[..aligned_len]; + + let tenant_id = intern_tenant(tenant); + + // Ensure tenant exists at root (insert-side bookkeeping). + self.root + .tenant_last_access_time + .entry(Arc::clone(&tenant_id)) + .or_insert(0); + self.tenant_token_count + .entry(Arc::clone(&tenant_id)) + .or_insert(0); + + let mut remaining = tokens; + let mut current = Arc::clone(&self.root); + let mut tokens_added = 0usize; + + // Match-result accumulators. + let mut matched_tokens = 0usize; + let mut last_tenant: Option = None; + // Once the match descent would have stopped (empty node or partial + // match), stop updating the match result; insert keeps descending. + let mut match_frozen = false; + + let track_lfu = self.eviction_policy == EvictionPolicy::Lfu; + + // Carries state out of the entry match so the entry guard can be + // dropped before we advance `current` (same pattern as `insert_tokens`). + enum Step { + Done(usize), + Continue { next: NodeRef, advance: usize }, + } + + while remaining.len() >= PAGE_SIZE { + let page_key = make_page_key(remaining); + + let step = match current.children.entry(page_key) { + Entry::Vacant(entry) => { + // Match: child not found -> match stops (records nothing). + // Insert: create a new leaf node holding the remainder. + let new_node = Arc::new(Node::new(remaining.to_vec())); + new_node.set_parent(¤t, page_key); + new_node.touch_tenant(&tenant_id, track_lfu); + entry.insert(new_node); + Step::Done(remaining.len()) + } + Entry::Occupied(mut entry) => { + let child = Arc::clone(entry.get()); + let child_tokens = child.tokens.read(); + let child_len = child_tokens.len(); + + let common_len = remaining + .iter() + .zip(child_tokens.iter()) + .take_while(|(a, b)| a == b) + .count(); + let common_len = align_to_page(common_len); + + if common_len == 0 { + // Same page key but no aligned match (shouldn't happen). + // Match stops; insert adds nothing. + drop(child_tokens); + Step::Done(0) + } else if common_len == child_len { + // Full match with child -> continue traversal. + drop(child_tokens); + + // --- Match side (pre-insert node state) --- + // Read the routed tenant BEFORE insert touches the node: + // insert's touch would add `tenant_id` to the map and + // corrupt both the all-evicted check and the routed + // tenant. This reproduces `match_prefix_with_counts`'s + // Continue arm exactly (it touches `get_any_tenant()`). + if !match_frozen { + match child.get_any_tenant() { + None => { + // All tenants evicted: match stops here. + // Insert still continues (re-populates node). + match_frozen = true; + } + Some(t_match) => { + matched_tokens += common_len; + child.touch_tenant(&t_match, track_lfu); + last_tenant = Some(t_match); + } + } + } + + // --- Insert side --- + // `insert_tokens` always touches the inserting tenant on + // a full-match continuation. When the match side above + // already touched this same tenant, the original code + // ALSO touched twice (match then insert), so we keep both + // touches to preserve LFU hit_count / timestamp behavior + // byte-for-byte. + child.touch_tenant(&tenant_id, track_lfu); + Step::Continue { + next: child, + advance: common_len, + } + } else if common_len >= remaining.len() { + // Input is a prefix of the child -> split child at the + // page boundary. (Verbatim from `insert_tokens`, with the + // match-side touch interleaved before the split so the + // intermediate's tenant-map clone inherits it — exactly + // as match-then-insert ordered the writes.) + + // Match side first (touches the pre-split child). + if !match_frozen { + match child.get_any_tenant() { + None => match_frozen = true, + Some(t_match) => { + matched_tokens += common_len; + child.touch_tenant(&t_match, track_lfu); + last_tenant = Some(t_match); + } + } + } + + let common_len = align_to_page(remaining.len()); + let prefix_tokens: Vec = child_tokens[..common_len].to_vec(); + let suffix_page_key = make_page_key(&child_tokens[common_len..]); + + let tenant_already_owned = child + .tenant_last_access_time + .contains_key(tenant_id.as_ref()); + drop(child_tokens); + + let mut child_tokens_write = child.tokens.write(); + let suffix_tokens: Vec = child_tokens_write[common_len..].to_vec(); + *child_tokens_write = suffix_tokens; + drop(child_tokens_write); + + let intermediate_node = Arc::new(Node { + tokens: ParkingLotRwLock::new(prefix_tokens), + children: new_children_map(), + tenant_last_access_time: child.tenant_last_access_time.clone(), + last_tenant: ParkingLotRwLock::new(child.last_tenant.read().clone()), + parent: ParkingLotRwLock::new(Arc::downgrade(¤t)), + page_key: ParkingLotRwLock::new(Some(page_key)), + hit_count: AtomicU64::new(child.hit_count.load(Ordering::Relaxed)), + creation_time: child.creation_time, + priority: AtomicI32::new(child.priority.load(Ordering::Relaxed)), + }); + + child.set_parent(&intermediate_node, suffix_page_key); + intermediate_node + .children + .insert(suffix_page_key, Arc::clone(&child)); + + entry.insert(intermediate_node.clone()); + + intermediate_node.touch_tenant(&tenant_id, track_lfu); + + let new_tokens = if tenant_already_owned { 0 } else { common_len }; + Step::Done(new_tokens) + } else { + // Partial match -> split and add a new branch at the page + // boundary. (Verbatim from `insert_tokens`, with the + // match-side touch interleaved before the split.) + + // Match side first (touches the pre-split child). + if !match_frozen { + match child.get_any_tenant() { + None => match_frozen = true, + Some(t_match) => { + matched_tokens += common_len; + child.touch_tenant(&t_match, track_lfu); + last_tenant = Some(t_match); + } + } + } + + let prefix_tokens: Vec = child_tokens[..common_len].to_vec(); + let child_suffix_page_key = make_page_key(&child_tokens[common_len..]); + + let tenant_already_owned = child + .tenant_last_access_time + .contains_key(tenant_id.as_ref()); + drop(child_tokens); + + let mut child_tokens_write = child.tokens.write(); + let child_suffix: Vec = child_tokens_write[common_len..].to_vec(); + *child_tokens_write = child_suffix; + drop(child_tokens_write); + + let intermediate_node = Arc::new(Node { + tokens: ParkingLotRwLock::new(prefix_tokens), + children: new_children_map(), + tenant_last_access_time: child.tenant_last_access_time.clone(), + last_tenant: ParkingLotRwLock::new(child.last_tenant.read().clone()), + parent: ParkingLotRwLock::new(Arc::downgrade(¤t)), + page_key: ParkingLotRwLock::new(Some(page_key)), + hit_count: AtomicU64::new(child.hit_count.load(Ordering::Relaxed)), + creation_time: child.creation_time, + priority: AtomicI32::new(child.priority.load(Ordering::Relaxed)), + }); + + child.set_parent(&intermediate_node, child_suffix_page_key); + intermediate_node + .children + .insert(child_suffix_page_key, Arc::clone(&child)); + + let new_remaining = &remaining[common_len..]; + let new_branch_tokens = if new_remaining.len() >= PAGE_SIZE { + let new_node = Arc::new(Node::new(new_remaining.to_vec())); + let new_page_key = make_page_key(new_remaining); + new_node.set_parent(&intermediate_node, new_page_key); + new_node.touch_tenant(&tenant_id, track_lfu); + intermediate_node.children.insert(new_page_key, new_node); + new_remaining.len() + } else { + 0 + }; + + entry.insert(intermediate_node.clone()); + + intermediate_node.touch_tenant(&tenant_id, track_lfu); + + let common_tokens = if tenant_already_owned { 0 } else { common_len }; + Step::Done(new_branch_tokens + common_tokens) + } + } + }; + + match step { + Step::Done(added) => { + tokens_added += added; + break; + } + Step::Continue { next, advance } => { + tokens_added += advance; + remaining = &remaining[advance..]; + current = next; + } + } + } + + // Fold insert's token count in once (insert-side bookkeeping). + if tokens_added > 0 { + self.tenant_token_count + .entry(tenant_id) + .and_modify(|c| *c += tokens_added) + .or_insert(tokens_added); + } + + PrefixMatchResult { + tenant: last_tenant.unwrap_or_else(|| Arc::from("empty")), + matched_token_count: matched_tokens, + input_token_count, + } + } + + /// Match in a single descent, then choose the insert tenant from the match + /// result and insert for it — still a SINGLE top-to-bottom traversal of the + /// prefix. + /// + /// This is the variant the cache-aware router needs: the worker it inserts + /// for is *derived from* the match outcome (route to the matched worker on a + /// cache hit, else to the least-loaded worker), so the inserting tenant is + /// not known until the match completes. `select` is invoked exactly once, + /// after the match, with the [`PrefixMatchResult`]; returning `Some(tenant)` + /// inserts `tokens` for that tenant, and returning `None` skips the insert + /// entirely (mirroring the router's "selected worker is gone" branch, which + /// performs no insert). + /// + /// # How the single descent is achieved + /// + /// The match phase walks the prefix exactly like + /// [`Self::match_prefix_with_counts`], additionally recording the chain of + /// nodes it traverses as `path` (one `(node, advance)` per edge) and the + /// node/`remaining` slice where the walk fell off. After `select` yields the + /// tenant we *replay insert's per-node work directly on the recorded nodes* + /// (no second tree navigation): `touch_tenant` on every traversed node plus + /// the same `tokens_added` accounting, then splice the remainder at the + /// fall-off point with the very branches `insert_tokens` uses (vacant leaf / + /// prefix-of-child split / diverge split). The only tree lookups are the + /// single descent and one `entry()` re-probe at the fall-off node for the + /// splice — never a second full walk. + /// + /// # Preserved semantics + /// + /// Identical to `match_prefix_with_counts` followed by + /// `insert_tokens(tokens, tenant)`: + /// * match-side: same matched-token count, same routed tenant, same per-node + /// `touch_tenant(get_any_tenant())`, same "stop at an all-evicted node or a + /// partial match" rule; + /// * insert-side: same node-split structure, same `tenant_token_count` + /// accounting (including the existing behavior of re-counting a fully + /// matched path), same `touch_tenant(tenant)` on every node on the path and + /// on freshly created/split nodes. + /// + /// The replay reorders insert's per-node touches to *after* the match phase + /// instead of interleaving them, which only changes the exact monotonic + /// timestamp values written (immaterial to relative LRU/LFU ordering); the + /// set of touched (node, tenant) pairs and the LFU hit-count increments are + /// unchanged. Concurrency is no weaker than the original two calls — they + /// also release every guard between the separate `match` and `insert`. + pub fn match_and_insert_with<'t, F>(&self, tokens: &[TokenId], select: F) -> PrefixMatchResult + where + F: FnOnce(&PrefixMatchResult) -> Option<&'t str>, + { + let input_token_count = tokens.len(); + + let aligned_len = align_to_page(tokens.len()); + if aligned_len == 0 { + // Too short to cache: `insert_tokens` is a no-op regardless of the + // selected tenant, so just resolve + return the match result. We + // still invoke `select` so the caller's routing side effects (if + // any) run, matching a separate match-then-(skipped)-insert. + let result = PrefixMatchResult { + tenant: self + .root + .get_any_tenant() + .unwrap_or_else(|| Arc::from("empty")), + matched_token_count: 0, + input_token_count, + }; + let _ = select(&result); + return result; + } + let tokens = &tokens[..aligned_len]; + + let track_lfu = self.eviction_policy == EvictionPolicy::Lfu; + + // ---- Phase 1: MATCH descent (mirrors match_prefix_with_counts) ---- + // Additionally record every traversed edge so insert can replay its + // per-node work without re-walking, and capture the fall-off node + + // remaining slice for the splice. + let mut matched_tokens = 0usize; + let mut last_tenant: Option = None; + let mut remaining = tokens; + let mut current = Arc::clone(&self.root); + // (node, advance) for each edge we descended through, in order. + let mut path: Vec<(NodeRef, usize)> = Vec::new(); + // Once match would stop (all-evicted node / partial), freeze the match + // result but keep descending for insert (insert's reach is a superset). + let mut match_frozen = false; + + enum MatchStep { + Stop, + Continue { next: NodeRef, advance: usize }, + } + + while remaining.len() >= PAGE_SIZE { + let page_key = make_page_key(remaining); + + let step = match current.children.get(&page_key) { + None => MatchStep::Stop, + Some(child_ref) => { + let child = Arc::clone(child_ref.value()); + drop(child_ref); + + let child_tokens = child.tokens.read(); + let match_len = remaining + .iter() + .zip(child_tokens.iter()) + .take_while(|(a, b)| a == b) + .count(); + let match_len = align_to_page(match_len); + + if match_len == 0 { + drop(child_tokens); + MatchStep::Stop + } else if match_len < child_tokens.len() { + // Partial match within the node: match stops here. + if !match_frozen { + if let Some(t) = child.get_any_tenant() { + child.touch_tenant(&t, track_lfu); + matched_tokens += match_len; + last_tenant = Some(t); + } + // (If the node is all-evicted, match records nothing + // and simply stops — same as match_prefix_with_counts.) + match_frozen = true; + } + // Insert also stops descending here (it will split this + // node). Do NOT push to `path`; the splice handles it. + drop(child_tokens); + MatchStep::Stop + } else { + // Full match: match continues (if not frozen) and insert + // continues regardless. + drop(child_tokens); + if !match_frozen { + match child.get_any_tenant() { + None => { + // All-evicted: match stops, insert continues. + match_frozen = true; + } + Some(t) => { + child.touch_tenant(&t, track_lfu); + matched_tokens += match_len; + last_tenant = Some(t); + } + } + } + MatchStep::Continue { + next: child, + advance: match_len, + } + } + } + }; + + match step { + MatchStep::Stop => break, + MatchStep::Continue { next, advance } => { + path.push((Arc::clone(&next), advance)); + remaining = &remaining[advance..]; + current = next; + } + } + } + + // ---- Decide the insert tenant from the match result ---- + let result = PrefixMatchResult { + tenant: last_tenant.unwrap_or_else(|| Arc::from("empty")), + matched_token_count: matched_tokens, + input_token_count, + }; + let Some(tenant) = select(&result) else { + // No insert (router selected no worker). + return result; + }; + // Intern through the shared pool so the stored Arc dedups with every + // other call (matches insert_tokens' interning). + let tenant_id = intern_tenant(tenant); + + // ---- Phase 2: INSERT replay for `tenant_id` (no second walk) ---- + // Mirrors insert_tokens' root bookkeeping. + self.root + .tenant_last_access_time + .entry(Arc::clone(&tenant_id)) + .or_insert(0); + self.tenant_token_count + .entry(Arc::clone(&tenant_id)) + .or_insert(0); + + let mut tokens_added = 0usize; + + // Replay insert's per-node work on every edge the match descended: + // `insert_tokens` touches the inserting tenant on each full-match node + // and counts its `advance` tokens. + for (node, advance) in &path { + node.touch_tenant(&tenant_id, track_lfu); + tokens_added += *advance; + } + + // Splice only the unmatched suffix at the fall-off node (`current`), + // reusing insert_tokens' exact descent loop. It re-probes `current`'s + // children for `remaining` (a single child-map op, not a re-walk of the + // matched prefix) and handles the vacant / split / — and, under a + // concurrent split race, full-match-continue — cases identically to a + // standalone `insert_tokens`. The matched prefix above `current` was + // already re-attached by the loop above and is never re-walked. + if remaining.len() >= PAGE_SIZE { + tokens_added += + Self::insert_from(current, remaining, Arc::clone(&tenant_id), track_lfu); + } + + if tokens_added > 0 { + self.tenant_token_count + .entry(tenant_id) + .and_modify(|c| *c += tokens_added) + .or_insert(tokens_added); + } + + result + } + /// Legacy prefix_match API returning (matched_tokens, tenant_string). pub fn prefix_match_legacy(&self, tokens: &[TokenId]) -> (Vec, String) { let result = self.match_prefix_with_counts(tokens); @@ -2694,4 +3263,112 @@ mod tests { let pos_b = paths.iter().position(|p| p == &b).unwrap(); assert!(pos_a < pos_b, "page-key 0..16 < 100..116"); } + + /// Build the same sequence of operations two ways — `match_prefix_with_counts` + /// + `insert_tokens` (the legacy pair) vs the fused `match_and_insert` — and + /// assert the returned match results and the resulting per-tenant token + /// counts match step for step. Timestamps are intentionally not compared + /// (the fused path may consume a different number of monotonic ticks). + fn assert_fused_matches_pair(ops: &[(Vec, &str)]) { + let pair = TokenTree::new(); + let fused = TokenTree::new(); + for (tokens, tenant) in ops { + let r_pair = pair.match_prefix_with_counts(tokens); + pair.insert_tokens(tokens, tenant); + + let r_fused = fused.match_and_insert(tokens, tenant); + + assert_eq!( + r_pair.matched_token_count, r_fused.matched_token_count, + "matched_token_count mismatch for tenant {tenant}" + ); + assert_eq!( + r_pair.input_token_count, r_fused.input_token_count, + "input_token_count mismatch" + ); + // Tenant is approximate (get_any_tenant), but for these + // single-tenant-per-path scenarios it is deterministic. + assert_eq!( + r_pair.tenant.as_ref(), + r_fused.tenant.as_ref(), + "matched tenant mismatch" + ); + } + assert_eq!( + pair.get_tenant_token_counts(), + fused.get_tenant_token_counts(), + "tenant token counts diverged between pair and fused" + ); + } + + #[test] + fn test_match_and_insert_equiv_fresh_and_full_match() { + let a = make_tokens(1, 3); + assert_fused_matches_pair(&[ + (a.clone(), "w1"), // fresh insert (single node) + (a.clone(), "w1"), // exact re-insert (full match continue) + (a, "w2"), // same path, different tenant (full match, adds w2) + ]); + } + + #[test] + fn test_match_and_insert_equiv_prefix_of_child_split() { + let long = make_tokens(1, 3); + let short = make_tokens(1, 1); // prefix of `long` -> splits the node + assert_fused_matches_pair(&[(long, "w1"), (short, "w2")]); + } + + #[test] + fn test_match_and_insert_equiv_diverge_split() { + // Shared first page, diverging second page -> partial/diverge split. + let mut a = make_tokens(1, 1); + a.extend(make_tokens(100, 1)); + let mut b = make_tokens(1, 1); + b.extend(make_tokens(200, 1)); + assert_fused_matches_pair(&[(a, "w1"), (b, "w2")]); + } + + #[test] + fn test_match_and_insert_equiv_disjoint() { + let a = make_tokens(1, 2); + let b = make_tokens(1000, 2); + assert_fused_matches_pair(&[(a, "w1"), (b, "w2")]); + } + + #[test] + fn test_match_and_insert_equiv_short_sequence() { + // Below PAGE_SIZE: insert is a no-op, match returns 0. + assert_fused_matches_pair(&[(vec![1, 2, 3], "w1")]); + } + + /// The `_with` closure form must behave like `match_and_insert` when the + /// closure always returns the same tenant, and must skip the insert (leaving + /// the tree unchanged) when it returns `None`. + #[test] + fn test_match_and_insert_with_select_and_skip() { + let tokens = make_tokens(1, 3); + + // Always-insert closure == match_and_insert(tokens, "w1"). + let via_with = TokenTree::new(); + let r = via_with.match_and_insert_with(&tokens, |_| Some("w1")); + let via_plain = TokenTree::new(); + let r2 = via_plain.match_and_insert(&tokens, "w1"); + assert_eq!(r.matched_token_count, r2.matched_token_count); + assert_eq!( + via_with.get_tenant_token_counts(), + via_plain.get_tenant_token_counts() + ); + + // Seed a tree, then a None-returning closure must NOT mutate it. + let tree = TokenTree::new(); + tree.insert_tokens(&tokens, "w1"); + let before = tree.get_tenant_token_counts(); + let r = tree.match_and_insert_with(&tokens, |_| None); + assert_eq!(r.matched_token_count, tokens.len()); // full match observed + assert_eq!( + tree.get_tenant_token_counts(), + before, + "None selection must not insert" + ); + } } diff --git a/model_gateway/benches/radix_tree_benchmark.rs b/model_gateway/benches/radix_tree_benchmark.rs index 3d4d6caf9f..4ad7cec4ff 100644 --- a/model_gateway/benches/radix_tree_benchmark.rs +++ b/model_gateway/benches/radix_tree_benchmark.rs @@ -900,9 +900,87 @@ fn print_summary() { eprintln!("\n{}", "=".repeat(95)); } +/// Benchmark the fused `match_and_insert` against the legacy +/// `match_prefix_with_counts` + `insert_*` pair that cache-aware routing used to +/// run on every request. The win grows with prefix length because the pair walks +/// the whole prefix twice; the fused call walks it once. Includes a long-context +/// (128K-token / 512K-char) case representative of large agentic prompts where +/// the double traversal dominated routing latency. +fn bench_match_and_insert(c: &mut Criterion) { + let mut group = c.benchmark_group("match_and_insert"); + group.warm_up_time(std::time::Duration::from_millis(500)); + group.measurement_time(std::time::Duration::from_secs(2)); + group.sample_size(30); + + // Sizes span the short request path up to long-context prompts. + const TOKEN_SIZES: [usize; 4] = [1024, 16384, 65536, 131072]; + const CHAR_SIZES: [usize; 4] = [4096, 65536, 262144, 524288]; + const TENANT: &str = "grpc://worker-0.sglang.svc.cluster.local:50051"; + + // ---- TokenTree ---- + for &token_size in &TOKEN_SIZES { + // A single shared sequence: after the first insert every request is a + // full cache hit, so both the pair and the fused call traverse the whole + // prefix (worst case for the double traversal, best case for fusion). + let seq: Vec = (0..token_size as u32).collect(); + + let pair_name = format!("token_pair_{token_size}tok"); + group.bench_function(&pair_name, |b| { + let tree = TokenTree::new(); + tree.insert_tokens(&seq, TENANT); + b.iter(|| { + let r = tree.match_prefix_with_counts(black_box(&seq)); + tree.insert_tokens(black_box(&seq), TENANT); + black_box(r); + }); + }); + + // `match_and_insert_with` is the exact production hot-path entry + // (cache-aware routing picks the tenant from the match result). + let fused_name = format!("token_fused_{token_size}tok"); + group.bench_function(&fused_name, |b| { + let tree = TokenTree::new(); + tree.insert_tokens(&seq, TENANT); + b.iter(|| { + let r = tree.match_and_insert_with(black_box(&seq), |_| Some(TENANT)); + black_box(r); + }); + }); + } + + // ---- StringTree ---- + for &char_size in &CHAR_SIZES { + let text = random_ascii_string(char_size); + + let pair_name = format!("string_pair_{char_size}c"); + group.bench_function(&pair_name, |b| { + let tree = StringTree::new(); + tree.insert_text(&text, TENANT); + b.iter(|| { + let r = tree.match_prefix_with_counts(black_box(&text)); + tree.insert_text(black_box(&text), TENANT); + black_box(r); + }); + }); + + let fused_name = format!("string_fused_{char_size}c"); + group.bench_function(&fused_name, |b| { + let tree = StringTree::new(); + tree.insert_text(&text, TENANT); + b.iter(|| { + let r = tree.match_and_insert_with(black_box(&text), |_| Some(TENANT)); + black_box(r); + }); + }); + } + + group.finish(); +} + fn run_benchmarks(c: &mut Criterion) { bench_summary(c); print_summary(); + bench_match_and_insert(c); } criterion_group!(benches, run_benchmarks); diff --git a/model_gateway/src/policies/cache_aware.rs b/model_gateway/src/policies/cache_aware.rs index 34539efb4f..3b63318156 100644 --- a/model_gateway/src/policies/cache_aware.rs +++ b/model_gateway/src/policies/cache_aware.rs @@ -406,25 +406,24 @@ impl CacheAwarePolicy { .get(model_id) .map(|entry| entry.value().clone()); if let Some(tree) = tree { - // Match BEFORE insert (mirrors the string-side - // imbalanced path below). After `insert_tokens`, - // the tree contains a full path for `tokens` so a - // match returns the entire input length and we'd - // store the full sequence — at 32K tokens × 4 bytes - // × max_tree_size that's multi-GB per model. - let matched_prefix = if self.should_populate_hash_index() { - let result = tree.match_prefix_with_counts(tokens); - Some(tokens[..result.matched_token_count].to_vec()) - } else { - None - }; - tree.insert_tokens(tokens, worker_url); - if let Some(matched_prefix) = matched_prefix { + // We need the match result (the prior shared prefix) BEFORE the + // insert so the hash_index stores only that bounded prefix, not + // the full path that exists post-insert (32K tokens × 4 bytes × + // max_tree_size = multi-GB/model). `match_and_insert` resolves + // the match against the pre-insert tree and inserts in the SAME + // descent, so `result.matched_token_count` is the same prior + // prefix length the standalone match returned. When we don't + // populate the index, a plain insert (no match) suffices. + if self.should_populate_hash_index() { + let result = tree.match_and_insert(tokens, worker_url); + let matched_prefix: Vec = tokens[..result.matched_token_count].to_vec(); self.hash_index .entry(model_id.to_string()) .or_default() .token_tree .insert(kv_index::hash_token_path(tokens), matched_prefix); + } else { + tree.insert_tokens(tokens, worker_url); } } } else if let Some(text) = info.request_text { @@ -435,30 +434,24 @@ impl CacheAwarePolicy { .map(|entry| entry.value().clone()); if let Some(tree) = tree { - // Match BEFORE insert: after `insert_text`, the - // tree contains a full path for `text` so a match - // would return the entire input length and we'd - // store the full prompt — exactly the memory leak - // we're trying to avoid. The pre-insert match - // returns the prior shared prefix (~50-200 chars). - let matched_prefix = if self.should_populate_hash_index() { - let result = tree.match_prefix_with_counts(text); - Some( - text.chars() - .take(result.matched_char_count) - .collect::(), - ) - } else { - None - }; - tree.insert_text(text, worker_url); - if let Some(matched_prefix) = matched_prefix { + // Match BEFORE insert so the hash_index stores only the prior + // shared prefix (~50-200 chars), not the full prompt (20KB+) + // that exists post-insert. `match_and_insert` does both in a + // single descent; `result.matched_char_count` is the same prior + // prefix length the standalone match returned. When we don't + // populate the index, a plain insert (no match) suffices. + if self.should_populate_hash_index() { + let result = tree.match_and_insert(text, worker_url); + let matched_prefix: String = + text.chars().take(result.matched_char_count).collect(); let path_hash = kv_index::hash_node_path(text); self.hash_index .entry(model_id.to_string()) .or_default() .string_tree .insert(path_hash, matched_prefix); + } else { + tree.insert_text(text, worker_url); } } else { debug!( @@ -879,29 +872,44 @@ impl CacheAwarePolicy { .map(|entry| entry.value().clone()); if let Some(tree) = tree { - let result = tree.match_prefix_with_counts(tokens); - let match_rate = if result.input_token_count == 0 { - 0.0 - } else { - result.matched_token_count as f32 / result.input_token_count as f32 - }; - - let selected_idx = if match_rate > self.config.cache_threshold { - let tenant_url: &str = &result.tenant; - workers - .iter() - .position(|w| w.url() == tenant_url) - .filter(|&idx| workers[idx].is_healthy()) - } else { - healthy_indices - .iter() - .min_by_key(|&&idx| workers[idx].load()) - .copied() - }; + // Single tree descent: match, pick the worker from the match + // result, then insert for it — replacing the former + // match_prefix_with_counts + insert_tokens pair (two full descents + // over the same prefix). The selection closure runs once, after the + // match, mirroring the previous branch exactly: + // * cache hit (match_rate > threshold): route to the matched + // worker if it is still healthy — insert for it; + // * cache miss (match_rate <= threshold): route to the least-loaded + // worker — insert for it; + // * matched worker gone/unhealthy: select nothing and DON'T insert + // (closure returns None), falling back to first-healthy below. + let mut selected_idx: Option = None; + let result = tree.match_and_insert_with(tokens, |result| { + let match_rate = if result.input_token_count == 0 { + 0.0 + } else { + result.matched_token_count as f32 / result.input_token_count as f32 + }; - if let Some(idx) = selected_idx { - tree.insert_tokens(tokens, workers[idx].url()); + selected_idx = if match_rate > self.config.cache_threshold { + let tenant_url: &str = &result.tenant; + workers + .iter() + .position(|w| w.url() == tenant_url) + .filter(|&idx| workers[idx].is_healthy()) + } else { + healthy_indices + .iter() + .min_by_key(|&&idx| workers[idx].load()) + .copied() + }; + + // Insert for the selected worker (None => no insert, exactly + // like the old `if let Some(idx)` guard around insert_tokens). + selected_idx.map(|idx| workers[idx].url()) + }); + if let Some(idx) = selected_idx { // Record hash(full_tokens)→matched_prefix tokens. // The hash key matches what sync_tree_operation // sends on the wire (hash of full sequence). The @@ -912,7 +920,7 @@ impl CacheAwarePolicy { // incoming token delta, so maintain it alongside // the tree. Mirrors the string side at the // analogous block; reuses the match `result` - // already computed at the top of this branch. + // returned by match_and_insert_with. if self.should_populate_hash_index() { let matched_prefix: Vec = tokens[..result.matched_token_count].to_vec(); self.hash_index @@ -953,29 +961,37 @@ impl CacheAwarePolicy { .map(|entry| entry.value().clone()); if let Some(tree) = tree { - let result = tree.match_prefix_with_counts(text); - let match_rate = if result.input_char_count == 0 { - 0.0 - } else { - result.matched_char_count as f32 / result.input_char_count as f32 - }; - - let selected_idx = if match_rate > self.config.cache_threshold { - let tenant_url: &str = &result.tenant; - workers - .iter() - .position(|w| w.url() == tenant_url) - .filter(|&idx| workers[idx].is_healthy()) - } else { - healthy_indices - .iter() - .min_by_key(|&&idx| workers[idx].load()) - .copied() - }; + // Single tree descent: match, pick the worker from the match result, + // then insert for it — replacing the former match_prefix_with_counts + // + insert_text pair. Selection logic is unchanged (see the token + // path for the per-branch rationale). + let mut selected_idx: Option = None; + let result = tree.match_and_insert_with(text, |result| { + let match_rate = if result.input_char_count == 0 { + 0.0 + } else { + result.matched_char_count as f32 / result.input_char_count as f32 + }; - if let Some(idx) = selected_idx { - tree.insert_text(text, workers[idx].url()); + selected_idx = if match_rate > self.config.cache_threshold { + let tenant_url: &str = &result.tenant; + workers + .iter() + .position(|w| w.url() == tenant_url) + .filter(|&idx| workers[idx].is_healthy()) + } else { + healthy_indices + .iter() + .min_by_key(|&&idx| workers[idx].load()) + .copied() + }; + + // Insert for the selected worker (None => no insert, exactly + // like the old `if let Some(idx)` guard around insert_text). + selected_idx.map(|idx| workers[idx].url()) + }); + if let Some(idx) = selected_idx { // Record hash(full_text)→matched_prefix for mesh tenant delta // resolution. The hash key matches what sync_tree_operation sends // on the wire (hash of full text). The VALUE is only the matched From 0bc14d2ef2bdbf0c5498e14f17f7be254e53efe6 Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Tue, 9 Jun 2026 12:43:53 -0700 Subject: [PATCH 2/2] perf(kv-index): intern "empty" fallback tenant + preallocate match path Vec Address review (gemini): use the global intern pool for the rare empty-tenant fallback instead of Arc::from (no per-call heap alloc), and pre-size the match-path Vec. Re-benchmarked: fused-vs-pair speedup holds at ~2x (131072 tokens 61.0us->30.4us); kv-index 190 tests pass. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- crates/kv_index/src/string_tree.rs | 9 +++++---- crates/kv_index/src/token_tree.rs | 15 ++++++++------- 2 files changed, 13 insertions(+), 11 deletions(-) diff --git a/crates/kv_index/src/string_tree.rs b/crates/kv_index/src/string_tree.rs index d31dcf8e52..9111b377ba 100644 --- a/crates/kv_index/src/string_tree.rs +++ b/crates/kv_index/src/string_tree.rs @@ -622,7 +622,7 @@ impl Tree { .iter() .next() .map(|kv| Arc::clone(kv.key())) - .unwrap_or_else(|| Arc::from("empty")); + .unwrap_or_else(|| intern_tenant("empty")); *curr.last_tenant.write() = Some(Arc::clone(&t)); t } @@ -634,7 +634,7 @@ impl Tree { .iter() .next() .map(|kv| Arc::clone(kv.key())) - .unwrap_or_else(|| Arc::from("empty")); + .unwrap_or_else(|| intern_tenant("empty")); *curr.last_tenant.write() = Some(Arc::clone(&t)); t } @@ -689,7 +689,7 @@ impl Tree { .iter() .next() .map(|kv| Arc::clone(kv.key())) - .unwrap_or_else(|| Arc::from("empty")); + .unwrap_or_else(|| intern_tenant("empty")); (t, true) } @@ -752,7 +752,8 @@ impl Tree { // full-match node, the partial child, or the root if nothing matched. let mut match_curr = Arc::clone(&self.root); // (node, char_count) for each full-match edge, in order. - let mut path: Vec<(NodeRef, usize)> = Vec::new(); + // Pre-allocated; most matched paths are well under this depth. + let mut path: Vec<(NodeRef, usize)> = Vec::with_capacity(16); while let Some(first_char) = remaining.chars().next() { let child_node = current.children.get(&first_char).map(|e| e.value().clone()); diff --git a/crates/kv_index/src/token_tree.rs b/crates/kv_index/src/token_tree.rs index 3c5366cd39..e73231fbcc 100644 --- a/crates/kv_index/src/token_tree.rs +++ b/crates/kv_index/src/token_tree.rs @@ -650,7 +650,7 @@ impl TokenTree { tenant: self .root .get_any_tenant() - .unwrap_or_else(|| Arc::from("empty")), + .unwrap_or_else(|| intern_tenant("empty")), matched_token_count: 0, input_token_count, }; @@ -760,7 +760,7 @@ impl TokenTree { } PrefixMatchResult { - tenant: last_tenant.unwrap_or_else(|| Arc::from("empty")), + tenant: last_tenant.unwrap_or_else(|| intern_tenant("empty")), matched_token_count: matched_tokens, input_token_count, } @@ -832,7 +832,7 @@ impl TokenTree { tenant: self .root .get_any_tenant() - .unwrap_or_else(|| Arc::from("empty")), + .unwrap_or_else(|| intern_tenant("empty")), matched_token_count: 0, input_token_count, }; @@ -1084,7 +1084,7 @@ impl TokenTree { } PrefixMatchResult { - tenant: last_tenant.unwrap_or_else(|| Arc::from("empty")), + tenant: last_tenant.unwrap_or_else(|| intern_tenant("empty")), matched_token_count: matched_tokens, input_token_count, } @@ -1151,7 +1151,7 @@ impl TokenTree { tenant: self .root .get_any_tenant() - .unwrap_or_else(|| Arc::from("empty")), + .unwrap_or_else(|| intern_tenant("empty")), matched_token_count: 0, input_token_count, }; @@ -1171,7 +1171,8 @@ impl TokenTree { let mut remaining = tokens; let mut current = Arc::clone(&self.root); // (node, advance) for each edge we descended through, in order. - let mut path: Vec<(NodeRef, usize)> = Vec::new(); + // Pre-allocated; most matched paths are well under this depth. + let mut path: Vec<(NodeRef, usize)> = Vec::with_capacity(16); // Once match would stop (all-evicted node / partial), freeze the match // result but keep descending for insert (insert's reach is a superset). let mut match_frozen = false; @@ -1254,7 +1255,7 @@ impl TokenTree { // ---- Decide the insert tenant from the match result ---- let result = PrefixMatchResult { - tenant: last_tenant.unwrap_or_else(|| Arc::from("empty")), + tenant: last_tenant.unwrap_or_else(|| intern_tenant("empty")), matched_token_count: matched_tokens, input_token_count, };