Repository navigation
perf(kv-index): fuse cache-aware match+insert into a single tree descent - #1615
Conversation
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>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (2)
📝 WalkthroughWalkthroughAdds fused single-descent ChangesRadix Tree Fused Match+Insert API
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
github-actions
Bot
added
benchmarks
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces fused match_and_insert and match_and_insert_with methods to both StringTree and TokenTree to perform prefix matching and insertion in a single tree descent, significantly reducing traversal overhead on the request hot path. It also updates the cache-aware routing policy to utilize these fused operations and adds comprehensive tests and benchmarks. The review feedback suggests minor performance optimizations on the hot path, specifically pre-allocating the path vector to avoid re-allocations and using intern_tenant("empty") instead of Arc::from("empty") to prevent unnecessary heap allocations.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| 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(); | ||
|
|
There was a problem hiding this comment.
On the request hot path, we can avoid re-allocating the Vec for path on every call by pre-allocating it with a small capacity (e.g., 16), which is sufficient for most tree depths.
| let mut path: Vec<(NodeRef, usize)> = Vec::with_capacity(16); |
| .iter() | ||
| .next() | ||
| .map(|kv| Arc::clone(kv.key())) | ||
| .unwrap_or_else(|| Arc::from("empty")); |
There was a problem hiding this comment.
Using Arc::from("empty") allocates a new heap-allocated str inside Arc on every fallback. We can use intern_tenant("empty") instead, which leverages the global static intern pool to completely avoid heap allocations.
| .unwrap_or_else(|| Arc::from("empty")); | |
| .unwrap_or_else(|| intern_tenant("empty")); |
References
- For types that are frequently cloned on hot paths and represent a small, repeated set of values (e.g., worker IDs or tenant IDs), use an interned string type like
Arc<str>to improve performance by making clones cheap (atomic reference count increments). Implement an interning mechanism where methods take&strand intern internally.
| 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(); |
There was a problem hiding this comment.
On the request hot path, we can avoid re-allocating the Vec for path on every call by pre-allocating it with a small capacity (e.g., 16), which is sufficient for most tree depths.
| let mut path: Vec<(NodeRef, usize)> = Vec::new(); | |
| let mut path: Vec<(NodeRef, usize)> = Vec::with_capacity(16); |
| tenant: self | ||
| .root | ||
| .get_any_tenant() | ||
| .unwrap_or_else(|| Arc::from("empty")), |
There was a problem hiding this comment.
Using Arc::from("empty") allocates a new heap-allocated str inside Arc on every fallback. We can use intern_tenant("empty") instead, which leverages the global static intern pool to completely avoid heap allocations.
| .unwrap_or_else(|| Arc::from("empty")), | |
| .unwrap_or_else(|| intern_tenant("empty")), |
References
- For types that are frequently cloned on hot paths and represent a small, repeated set of values (e.g., worker IDs or tenant IDs), use an interned string type like
Arc<str>to improve performance by making clones cheap (atomic reference count increments). Implement an interning mechanism where methods take&strand intern internally.
| } | ||
|
|
||
| PrefixMatchResult { | ||
| tenant: last_tenant.unwrap_or_else(|| Arc::from("empty")), |
There was a problem hiding this comment.
Using Arc::from("empty") allocates a new heap-allocated str inside Arc on every fallback. We can use intern_tenant("empty") instead, which leverages the global static intern pool to completely avoid heap allocations.
| tenant: last_tenant.unwrap_or_else(|| Arc::from("empty")), | |
| tenant: last_tenant.unwrap_or_else(|| intern_tenant("empty")), |
References
- For types that are frequently cloned on hot paths and represent a small, repeated set of values (e.g., worker IDs or tenant IDs), use an interned string type like
Arc<str>to improve performance by making clones cheap (atomic reference count increments). Implement an interning mechanism where methods take&strand intern internally.
| tenant: self | ||
| .root | ||
| .get_any_tenant() | ||
| .unwrap_or_else(|| Arc::from("empty")), |
There was a problem hiding this comment.
Using Arc::from("empty") allocates a new heap-allocated str inside Arc on every fallback. We can use intern_tenant("empty") instead, which leverages the global static intern pool to completely avoid heap allocations.
| .unwrap_or_else(|| Arc::from("empty")), | |
| .unwrap_or_else(|| intern_tenant("empty")), |
References
- For types that are frequently cloned on hot paths and represent a small, repeated set of values (e.g., worker IDs or tenant IDs), use an interned string type like
Arc<str>to improve performance by making clones cheap (atomic reference count increments). Implement an interning mechanism where methods take&strand intern internally.
|
|
||
| // ---- Decide the insert tenant from the match result ---- | ||
| let result = PrefixMatchResult { | ||
| tenant: last_tenant.unwrap_or_else(|| Arc::from("empty")), |
There was a problem hiding this comment.
Using Arc::from("empty") allocates a new heap-allocated str inside Arc on every fallback. We can use intern_tenant("empty") instead, which leverages the global static intern pool to completely avoid heap allocations.
| tenant: last_tenant.unwrap_or_else(|| Arc::from("empty")), | |
| tenant: last_tenant.unwrap_or_else(|| intern_tenant("empty")), |
References
- For types that are frequently cloned on hot paths and represent a small, repeated set of values (e.g., worker IDs or tenant IDs), use an interned string type like
Arc<str>to improve performance by making clones cheap (atomic reference count increments). Implement an interning mechanism where methods take&strand intern internally.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/kv_index/src/string_tree.rs`:
- Around line 815-843: Recorded ancestor nodes in path must be revalidated
before folding their cached char_count into tenant_char_count because
intervening splits can change which node represents the matched edge; update the
loop that iterates path (the block that checks node.tenant_last_access_time and
updates self.tenant_char_count) to first confirm each NodeRef still corresponds
to the same matched edge/length (compare the node's current edge length or
parent/child linkage against the recorded char_count or expected child on the
path), and if any mismatch is detected abort the replay and fall back to a fresh
insertion (call self.insert_from(current, remaining, tenant_id) or otherwise
re-run insert_text path discovery) instead of applying stale char_count or
attaching the tenant to a stale node; ensure tenant_last_access_time updates
only occur for validated nodes.
In `@crates/kv_index/src/token_tree.rs`:
- Around line 907-939: The current match_and_insert interleaves match-side and
insert-side touches per node (child.touch_tenant calls), changing ancestor
last_access_time; instead, traverse the full-match chain performing only
match-side logic (get_any_tenant, match_frozen, matched_tokens,
child.touch_tenant for the match-side only) and record the nodes/tenants that
need insert-side touches (including duplicates when match and insert are same
tenant) into a temporary list; after the traversal finishes (i.e., when you know
the full sequence like match_prefix_with_counts would), iterate that recorded
list and call child.touch_tenant(&tenant_id, track_lfu) in the original
ancestor-to-leaf order to preserve legacy access-order semantics; keep existing
behavior around match_frozen, last_tenant, common_len, and Step::Continue
semantics.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: a6a308cb-c9cc-4e7b-8ffc-8aa98ceeedad
📒 Files selected for processing (4)
crates/kv_index/src/string_tree.rscrates/kv_index/src/token_tree.rsmodel_gateway/benches/radix_tree_benchmark.rsmodel_gateway/src/policies/cache_aware.rs
| // 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); |
There was a problem hiding this comment.
Revalidate replayed nodes before folding their cached char_count into tenant_char_count.
This replay assumes every recorded NodeRef still represents the same matched edge it did in Phase 1. If another insert splits one of those nodes before Lines 815-843 run, that ref can now be the suffix child while char_count still reflects the old unsplit edge. The replay then attaches the tenant to the stale node, increments tenant_char_count by the stale full-edge length, and skips the new intermediate. The route may “self-heal” on a later request, but the inflated size accounting does not, so size-based eviction will act on corrupted data. Please validate that each recorded node is still on the matched path or fall back to a fresh insert when the path has changed.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@crates/kv_index/src/string_tree.rs` around lines 815 - 843, Recorded ancestor
nodes in path must be revalidated before folding their cached char_count into
tenant_char_count because intervening splits can change which node represents
the matched edge; update the loop that iterates path (the block that checks
node.tenant_last_access_time and updates self.tenant_char_count) to first
confirm each NodeRef still corresponds to the same matched edge/length (compare
the node's current edge length or parent/child linkage against the recorded
char_count or expected child on the path), and if any mismatch is detected abort
the replay and fall back to a fresh insertion (call self.insert_from(current,
remaining, tenant_id) or otherwise re-run insert_text path discovery) instead of
applying stale char_count or attaching the tenant to a stale node; ensure
tenant_last_access_time updates only occur for validated nodes.
| // --- 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, | ||
| } |
There was a problem hiding this comment.
match_and_insert no longer preserves the legacy access-order semantics.
On a full-match chain, the old pair does all match-side touches first and only then does the insert-side touches. This code interleaves them per node (A(match), A(insert), B(match), B(insert)…). That changes the persisted last_access_time of ancestors, so LRU/MRU/FIFO/FILO eviction order can diverge once a child is evicted and its parent becomes a leaf. The returned match counts stay the same, but the API is not actually equivalent to match_prefix_with_counts + insert_tokens for eviction behavior.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@crates/kv_index/src/token_tree.rs` around lines 907 - 939, The current
match_and_insert interleaves match-side and insert-side touches per node
(child.touch_tenant calls), changing ancestor last_access_time; instead,
traverse the full-match chain performing only match-side logic (get_any_tenant,
match_frozen, matched_tokens, child.touch_tenant for the match-side only) and
record the nodes/tenants that need insert-side touches (including duplicates
when match and insert are same tenant) into a temporary list; after the
traversal finishes (i.e., when you know the full sequence like
match_prefix_with_counts would), iterate that recorded list and call
child.touch_tenant(&tenant_id, track_lfu) in the original ancestor-to-leaf order
to preserve legacy access-order semantics; keep existing behavior around
match_frozen, last_tenant, common_len, and Step::Continue semantics.
| // 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); | ||
| } |
There was a problem hiding this comment.
The replay path can permanently overcount tenant_token_count after a concurrent split.
Phase 2 trusts the recorded (NodeRef, advance) pairs from the match walk. If one of those full-match nodes is split before Lines 1281-1299 run, the recorded ref can now be the suffix child, but the replay still adds the old advance and touches that stale node. A later request can repair reachability, but tenant_token_count stays inflated by the old prefix length, so eviction decisions are now based on bad accounting. Please revalidate the recorded path before replaying counts, or abandon replay and fall back to a fresh insert when the structure changed.
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Thorough review of the fused match+insert optimization. No issues found — approving.
What I verified:
-
insert_fromextraction (both trees): The descent loop is correctly factored out.insert_text/insert_tokensdelegate to it with(root, full_text, tenant_id). The string tree's version preserves the post-loop leaf timestamp write; the token tree's returnstokens_addedfor the caller to fold in. -
match_and_insert(token tree, interleaved single-loop): Match-side readsget_any_tenant()before insert-sidetouch_tenant()on each node, preserving the "match sees pre-insert state" invariant. Thematch_frozenflag correctly handles all-evicted nodes (match stops, insert continues). Split branches (vacant, prefix-of-child, diverge) are verbatim frominsert_tokenswith the match-side touch correctly interleaved before the split so the intermediate's clonedtenant_last_access_timeinherits it. -
match_and_insert_with(both trees, two-phase): Phase 1 records the path with read-only ops (get()/children.get()). Phase 2 replays insert's per-node work on recorded nodes and splices the remainder viainsert_from. The string tree'sresolve_tenant_readonly+finish_match_and_insertpair correctly reproduces the cache-populate and probabilistic-touch side effects ofmatch_prefix_with_counts. Theremaining.is_empty()branch correctly writes the real timestamp on the final leaf. -
cache_aware.rscall sites:record_cache_hitcorrectly usesmatch_and_insert(tenant known upfront, skips the insert when hash index isn't populated).select_worker_with_tokens/_with_textcorrectly usesmatch_and_insert_withwith closure-based tenant selection. Theselected_idxmutable capture through theFnOnceclosure is safe. -
Edge cases: empty text/tokens, sub-
PAGE_SIZEsequences (short-circuit), full-match loop-end (leaf timestamp), partial match + split, all-evicted nodes,selectreturningNone, Unicode strings. -
Concurrency: The acknowledged window between match and replay in
_withis no weaker than the original two-call sequence, which also releases all guards betweenmatchandinsert. -
Tests: Comprehensive equivalence tests covering the key scenarios. The
assert_fused_matches_pairharness is well-designed — it runs both paths step-by-step and compares match results and finaltenant_{token,char}_count.
Summary: 0 important, 0 nit, 0 pre-existing.
…th 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>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0bc14d2ef2
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| // 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); |
There was a problem hiding this comment.
When another request splits the matched path between Phase 1 and this splice, current and the recorded path can refer to nodes whose prefix changed. For example, one request can match abcdef... and pause with current at the abcdef node; a concurrent insert of abc... splits it into abc + def; resuming here inserts the suffix under the old def node and never attaches the tenant to the new abc intermediate, while the old match_prefix_with_counts + insert_text path re-walked from the root and would attach both nodes. This leaves tenant-specific lookups and eviction counts inconsistent under concurrent same-prefix inserts.
Useful? React with 👍 / 👎.
| if remaining.len() >= PAGE_SIZE { | ||
| tokens_added += | ||
| Self::insert_from(current, remaining, Arc::clone(&tenant_id), track_lfu); |
There was a problem hiding this comment.
The read-only match phase records node refs, then this insert replay uses them after all locks have been released. Under concurrent inserts that split one of those token nodes, the replay touches/counts the old suffix node and starts insert_from below it without attaching the selected tenant to the newly created intermediate prefix; the previous separate insert_tokens call re-walked from the root, so it would see the split and populate the correct prefix nodes. This can make cache routing miss prefixes and leave per-tenant token counts wrong for concurrent requests sharing a prefix.
Useful? React with 👍 / 👎.
1 participant
Description
Problem
Cache-aware routing does two full top-to-bottom descents of the same request prefix on every request —
match_prefix…(tokens)immediately followed byinsert_tokens(tokens, tenant). For long-context workloads (100K–200K-token prefixes, Kimi-class) that doubles the per-request traversal, lock acquisitions, and tenant-timestamp writes on the routing hot path.Solution
Fuse the pair into a single descent. Insert's reach is a depth-wise superset of match's, so we walk the prefix once (doing the match), then splice only the unmatched suffix — the already-matched prefix is never re-walked.
Changes
crates/kv_index/src/token_tree.rs: extractinsert_tokens' descent intoinsert_from; addmatch_and_insert(tenant known up front) andmatch_and_insert_with(tenant derived from the match result).crates/kv_index/src/string_tree.rs: same forTree.model_gateway/src/policies/cache_aware.rs:select_worker_with_tokens/_with_textand the imbalanced min-load path call the fused method instead of match-then-insert.model_gateway/benches/radix_tree_benchmark.rs:bench_match_and_insert(legacy pair vs fused, up to 131072 tokens / 524288 chars).Size: ~1.2K lines — over the 400-line target. The bulk is the second method variant for both trees + equivalence tests + the new benchmark. The two token-tree variants duplicate splice logic; collapsing them is a sensible follow-up once this is validated. Flagging per the guideline.
For reviewers:
match_and_insert_withreplays the insert onto the node chain captured during the match. If another thread splits a matched node in the microsecond window between match and replay, the tenant lands on the stale node; it self-heals on the next request and the window is shorter than today's two-call sequence. Acceptable for an approximate cache-routing tree, but worth a look. (The suffix splice usesinsert_from, which re-walks from the fall-off node and is unaffected.)Test Plan
Pre-PR gate, run locally:
cargo +nightly fmt --all→ clean.cargo clippy -p kv-index -p smg --all-targets --all-features -- -D warnings→ zero warnings.cargo test -p kv-index→ 190 passed, including 10 new equivalence tests asserting the fused path produces identicalPrefixMatchResultandtenant_{token,char}_countas separate match+insert (token: disjoint / diverge-split / prefix-of-child-split / fresh / full-match; string: basic / deep-chain / unicode; + select-and-skip on both).cargo test -p smg --lib policies::cache_aware(21) andcargo test -p smg --test routing_tests(92) → pass — exercise the new call sites end-to-end.Benchmark —
cargo bench --bench radix_tree_benchmark -- match_and_insert(legacy pair → fused):String side ~1.75× (524288 chars: 278 µs → 158 µs). The win scales with context length — largest at long context.
Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit
New Features
Performance
Tests
Benchmarks