fix(prompt_cache): make radix trie iterative to prevent stack overflow#63
Merged
Conversation
The deep-branching regression test pop_prefixes_deep_branching_chain_does_not_overflow aborted the entire `cargo test --lib` run with a stack overflow. The trie's pop_prefixes and Drop were already iterative, but insert_into, remove_from, and dfs (reached via for_each_candidate) still recursed one stack frame per trie node, so the trie was only partially iterative. On an adversarially deep prompt trie — one node per token, which defeats path compression because every sibling branch is unique — the recursive insert overflowed even the test's 32 MiB build thread in the full debug build and aborted the process. That is why sibling tests appeared to fail under the full run yet passed when rerun individually. The same recursion also executes on ~2 MiB Tokio worker stacks in production through PromptCacheStore insert and eviction and the per-request longest-prefix candidate walk, so a crafted prompt sequence is a real stack-overflow denial-of-service vector, not just a test artifact. Rewrite all three remaining recursive operations to run in O(1) stack: - insert_into: descend while recording the path, mutate at the terminal node, then repair subtree_count bottom-up along the recorded path. - remove_from: descend recording (node, descend-key) ancestors, then prune empty children, merge a lone empty non-root node, and repair counts bottom-up; the now-unused is_root parameter is dropped. - dfs: walk an explicit heap-allocated stack instead of the call stack (traversal order remains unspecified, as the for_each_candidate contract already allows). The regression test now builds and exercises the whole trie — insert, longest-prefix lookup plus candidate walk, remove, pop_prefixes, and the final Drop — directly on a 2 MiB stack that matches a Tokio worker, with N=4096 (about 4x the depth that overflows a 2 MiB stack). This replaces the previous 32 MiB build-thread workaround: if any path is still recursive the spawned thread overflows and aborts, its join() returns Err, and the test fails.
inureyes
added
type:bug
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The radix trie backing the prompt cache was only partially iterative:
pop_prefixesandDropalready avoided recursion, butinsert_into,remove_from, anddfs(reached throughfor_each_candidate) still recursed one stack frame per trie node. On an adversarially deep prompt trie — one node per token, which defeats path compression because every sibling branch is unique — that recursion overflows the stack.This is a real denial-of-service vector, not just a test artifact. The same recursive paths run on ~2 MiB Tokio worker stacks in production via
PromptCacheStoreinsert/eviction and the per-request longest-prefix candidate walk, so a crafted prompt sequence can crash a worker. It also aborted the entirecargo test --librun (the recursive insert overflowed even the test's 32 MiB build thread), which is why unrelated sibling tests appeared to fail under the full run yet passed when rerun individually.Changes
Rewrites the three remaining recursive operations to run in O(1) stack:
insert_into— descend while recording the path, mutate at the terminal node, then repairsubtree_countbottom-up along the recorded path.remove_from— descend recording(node, descend-key)ancestors, then prune empty children, merge a lone empty non-root node, and repair counts bottom-up; the now-unusedis_rootparameter is dropped.dfs— walks an explicit heap-allocated stack instead of the call stack (traversal order stays unspecified, as thefor_each_candidatecontract already permits).Testing
pop_prefixes_deep_branching_chain_does_not_overflownow builds and exercises the whole trie (insert, longest-prefix lookup + candidate walk, remove,pop_prefixes, and Drop) directly on a 2 MiB Tokio-worker-equivalent stack withN=4096(~4× the depth that overflows a 2 MiB stack), replacing the prior 32 MiB build-thread workaround.cargo test --features metal,accelerate --lib prompt_cache::trie→ 21 passed, 0 failed.cargo clippy --features metal,accelerate --lib --tests -- -D warnings→ clean.cargo fmt --all -- --check→ clean.