fix(kv-router): release unowned radix branches after eviction

[jthomson04]

Summary

When a worker keeps a shared system prompt cached and repeatedly stores and evicts request suffixes, the single-threaded RadixTree can retain the unused branches through its parent-child links. The block lookup returns to its expected size, but process memory keeps growing.

Detach a branch when its last owner is removed. Preserve a replacement at the same parent key, and remove ownerless suffixes created by tree splits. This affects the shared KV-index implementation, including worker-side LocalKvIndexer when enabled.

Related PR: #11678 is an earlier implementation of the same fix using weak parent links and pointer checks before removal. It uses separate cleanup helpers and walks upward through empty ancestors; this PR combines cleanup at the affected node and also removes ownerless suffixes after splits. The validation below includes focused replacement-protection and split-suffix tests, plus shared-system-prompt memory measurements.

Details

• Add weak parent links, update them when a split moves children, and check pointer identity before detaching an unowned branch.
• Keep the churn test small and add focused replacement-protection and unused-suffix tests. Add a non-root debug assertion.
• Document ownership, eager cleanup, and the existing descendant-invalidation limit in the indexer README.

Validation

Base: 95cadce97dbcf3e2b0d79fdac05a2a2e03e1b1ff. Updated head: a9de76194380ffc25eb6ba89e42afe9eee1cb9c5.

• cargo test --locked --release -p dynamo-kv-router --lib --no-default-features: 952 passed, 0 failed, 2 existing ignored after restoring the final source.
• Mutation checks: removing the pointer guard fails the new replacement test (Some(3) versus Some(4)); removing the split-suffix prune fails the new suffix test ([1,1,2] versus [1,1]).
• The original regression-only control fails on the base with evicted branch is still retained.
• Package formatting, Clippy with -D warnings, and git diff --check: passed.

Shared-system-prompt memory comparison

A standalone Rust harness calls the public RadixTree API. One worker permanently caches 16 system-prompt blocks. Each cycle creates a unique eight-block request suffix, stores it under that prompt, then removes all eight blocks. Events are generated incrementally. Every cycle checks that the tracked-block count returns to 16; checkpoints also verify the prompt remains matchable and the evicted suffix is absent.

Each fresh process warms up with 5,000 cycles using two reused suffixes, then processes 1,000,000 unique suffixes. Three interleaved base/fixed pairs use identical dependencies, Rust 1.96.1 release builds (target-cpu=native, frame pointers, tokio_unstable, thin LTO), the system allocator with MALLOC_ARENA_MAX=1, and logical CPU 2 on an Intel Core Ultra 9 285K. Linux RSS is read from /proc/self/smaps_rollup after warmup and at fixed cycle counts.

Metric	Base	Fixed
RSS after warmup, median	2.118 MB	2.122 MB
RSS after 1M cycles, median	773.849 MB	2.195 MB
RSS added, median (range)	771.731 MB (same in all three runs)	0.074 MB (0.074–0.094)
Added bytes/cycle, median	771.731	0.074
Cycles/s, median (range)	1.544M (1.492–1.557M)	2.642M (2.623–2.667M)

MB uses decimal units. Median added RSS fell 99.990%. Paired cycle-rate changes were +75.78%, +71.34%, and +71.13%. All six runs completed with no operation or invariant failures. From 100,000 to 1,000,000 cycles, each fixed run added only 4 KiB of RSS. Total process execution for the six measured runs was 4.65 seconds, excluding builds.

Cycle rate includes synthetic event creation, storage, removal, assertions, and periodic checks/sampling. These are short component measurements, not end-to-end serving throughput or real-backend memory results. The local harness and raw results are excluded from this PR. Binary SHA256: base 4478cda03b36085af1429746f1bc713b2ec439be6c68c969f35479e5af91c456; fixed 93ca7ad5d2b0f072200ab63c8f9eed7d0e630a947ab6f2d9c02ff0acd95fadee.

Remaining limits

Existing mid-chain removal can leave descendant lookup entries alive; complete descendant invalidation and associated back-link cleanup are separate follow-up work. The generic cleanup refactor, duplicate parent-key field, borrow optimization, and unproven lookup-repair case are outside this patch. Worst-case high-fanout split cost and long-duration post-fix stability remain unmeasured.

Where should the reviewer start?

RadixBlock::prune_unreachable, parent-link updates in split_node, and the three cleanup regression tests in lib/kv-router/src/indexer/radix_tree.rs.

Related Issues

• Confirmed — no related issue.
• Related PR: fix(kv-router): detach dead radix nodes from their parent on removal #11678.

Summary by CodeRabbit

• Bug Fixes

  • Improved cache cleanup so inactive branches are released reliably.
  • Preserved valid replacement entries during cache updates.
  • Fixed tree state tracking after node splits and removals.
  • Prevented stale branches from accumulating during cache churn.

• Tests

  • Added regression coverage to verify evicted branches are released and cache sizes remain accurate.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 3e4b3d00-fc2f-41d5-8582-906f8cba9efe

📥 Commits

Reviewing files that changed from the base of the PR and between 95cadce and 7c7dfbd.

📒 Files selected for processing (1)

• lib/kv-router/src/indexer/radix_tree.rs

Included review availability: Your plan provides up to 12 included reviews per hour; 9 remain after this review.

---

Walkthrough

Changes

The radix tree now stores weak parent links and uses centralized pruning for unreachable nodes. Split operations repair parent links for new and moved nodes. Worker and DP-rank cleanup paths use pruning. A cache-churn regression test checks branch retention and release.

Radix tree cleanup

Layer / File(s)	Summary
Parent links and unreachable-node pruning lib/kv-router/src/indexer/radix_tree.rs	RadixBlock stores weak parent references. Constructors initialize them. prune_unreachable removes workerless nodes while preserving replacement nodes on the same edge.
Split parent-link repair lib/kv-router/src/indexer/radix_tree.rs	Child construction and split operations assign or repair parent links. Released suffix nodes pass through unreachable-node pruning.
Worker cleanup integration and regression coverage lib/kv-router/src/indexer/radix_tree.rs	Worker removal, worker-wide cleanup, and DP-rank cleanup use centralized pruning. The cache-churn test checks branch retention, eviction, edge lengths, and lookup sizes.

Priority: ➖ Normal

Estimated code review effort: 3 (Moderate) | ~25 minutes

Merge Risk: ⚪ Minimal · up to 7c7df

No concrete regression was established in the changed radix-tree cleanup paths, so the change is mergeable after normal checks.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 18.18% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 11 functions across 1 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly and concisely describes the main change: releasing unowned radix branches after cache eviction.
Description check	✅ Passed	The description provides the change summary, implementation details, validation results, reviewer starting points, and the required no-related-issue confirmation. It uses “Summary” instead of the temp…

• Fix all pre-merge checks with AI

---

_{Comment @coderabbitai help to get the list of available commands.}

[jthomson04]

Code review

Automated review of 7c7dfbd. 11 findings; 9 are inline below, 2 could not be anchored (one is in a file outside the diff, one is on an unchanged line).

The one worth acting on before merge is the first inline comment: prune_unreachable's children.clear() is itself an unfixed source of the retention this PR targets — it detaches live, worker-owned subtrees while lookup entries keep them alive, and later stores are still accepted into the detached subtree and grow it. Reproduced empirically.

Two of the three behaviors the description calls out are untested: mutation-removing the split_node prune call and mutation-removing the Rc::ptr_eq guard each leave the suite green (950 passed; 0 failed). The reparenting loop is covered — removing it fails cache_churn_releases_evicted_branches.

Two findings that could not be anchored inline

1. lib/kv-router/src/indexer/README.md L111-L118 is now stale. That block enumerates every RadixBlock field (state.edge / state.edge_index / state.worker_cutoffs / state.full_edge_workers / children / internal) and does not list the new parent: Weak<RefCell<RadixBlock>> — the single most important structural change for reasoning about ownership and drop order, and the file the crate points new contributors at. README.md#L110-L119

2. Latent: the edge_index-miss branch re-inserts an unresolvable lookup entry instead of scrubbing it, two lines above a changed hunk in the function this PR edits. radix_tree.rs#L624-L629 — lookup.remove(&block_hash) yields a node whose edge_index no longer contains the hash, and the code puts the entry straight back. The hash can never reappear in that node's edge_index, so the entry can never be removed by any future event: current_size() stays inflated for the process lifetime and the node's allocation is pinned. concurrent_radix_tree_compressed/remove.rs L93-L104 handles the identical case with a remove_lookup_hashes scrub and a comment naming the consequence. I could not construct a live trigger in the single-threaded tree today (split_node repairs lookups eagerly for every covered worker), so this is latent — but it is the same defect class and the fix is one line.

Method

Findings 1-3 are execution-backed, not inspection: prune_unreachable's branches were instrumented and all 952 lib tests run (detach path 50 hits, split-suffix prune 0 hits, ptr_eq guard 0 hits), then each of the three new behaviors was mutation-tested individually. The retention behavior in the first comment was reproduced with a scratch test. No changes were left in the working tree.

🤖 Generated with Claude Code

[jthomson04]

Follow-up for the two unanchored findings in the review:

• Indexer README: updated in a9de761. It now lists the weak parent link and explains ownership, the parent-key invariant, eager cleanup, and the deferred descendant/back-link limit.
• edge_index-miss lookup repair: deferred. No reachable trigger has been identified in the single-threaded implementation, so the existing behavior is unchanged in this focused patch.

All ten inline threads now have a disposition. The two new tests fail when their respective protections are removed; the restored source passes 952 unit tests and Clippy. The PR description includes the three-pair shared-system-prompt memory comparison. The validation harness and raw results are excluded from the PR diff.

[dynamo-review-agent]

Previously reported defects still present:

• Original discussion: Detached nodes retain weak parent links. When a detached node remains alive through a stale lookup, its Weak retains the dropped parent's Rc allocation, so the allocation-release goal remains incomplete under the known descendant-retention path.
• Original discussion: The split reparenting loop adds O(fanout) RefCell writes to every split after moving children, regressing the previously O(1) child-map move on high-fanout shared-prefix nodes.
• Original discussion: Verified still present: when a mid-edge removal leaves only cutoff ownership, prune_unreachable clears child edges even though descendant lookup entries can still retain worker-owned nodes. Those nodes become unroutable yet remain tracked and can accept stores through stale lookup entries. This predates the PR and is explicitly deferred; it requires worker-specific descendant and lookup invalidation.
• Original discussion: Verified: mid-edge removal still clears children when no worker fully covers the node, while descendant lookup entries can retain and later reuse detached worker-owned subtrees. The existing TODO in apply_removed still identifies this correctness gap; the PR documents it as deferred.
• Original discussion: prune_unreachable still clears children whenever no worker fully covers the edge, even when detached descendants remain owned through worker lookups. Those branches become unroutable and can continue growing through stale parent lookups.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[jthomson04]

/ok to test 3fbcd98