Repository navigation
perf(kv-index): close 4 perf gaps with Flash Indexer - #584
Merged
Merged
Conversation
Gap 1 — Per-position DashMaps: Replace DashMap<(usize, ContentHash), SeqEntry> with RwLock<Vec<DashMap<ContentHash, SeqEntry, FxBuildHasher>>>. One DashMap per block position eliminates cross-position shard contention and improves cache locality for hot prefix positions (0-3). Vec grows lazily via ensure_levels(); read lock held for the duration of queries (only blocks rare Vec growth). Gap 2 — Eliminate set cloning in query path: linear_scan_drain now accesses DashMap entries directly via Ref and checks worker membership in-place without cloning the FxHashSet<WorkerId>. Previously cloned the set at every position during drain (up to jump_size=64 clones per drain call). Gap 3 — Intern worker IDs: Add DashMap<WorkerId, (), FxBuildHasher> intern table. First call per worker does Arc::from() (heap alloc); subsequent calls return Arc::clone() (atomic increment only). Lookup by &str works via Arc<str>: Borrow<str>. Eliminates per-call allocation overhead in apply_stored, apply_removed, and remove_or_clear_worker. Gap 4 — Lazy rolling hash for Single entries: Add SeqEntry::workers_if_single() that returns the worker set directly for Single entries without requiring rolling hash computation. Since content hash collisions at 64-bit XXH3 are ~2^-64 (practically impossible), a matching content_hash at the same position is unambiguous. Applied in get_workers_lazy, count_workers_at, and linear_scan_drain — the three query-path helpers that previously always called ensure_seq_hash_computed. All query helpers (get_workers_lazy, count_workers_at, linear_scan_drain) converted to associated functions taking &[DashMap<...>] to share a single read guard acquired once in jump_search_matches. Explicit drop(levels) before acquiring worker_blocks.read() maintains lock ordering and prevents deadlock. Lock ordering: worker_blocks → index → DashMap shards → LevelIndex. Public API unchanged — all 158 kv-index tests pass. Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
Contributor
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the performance of the Highlights
Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughReplaced Arc-based WorkerId with interned u32 IDs, introduced WorkerBitset for worker sets, restructured per-worker reverse lookups and tree sizes to DashMap/Atomic types, updated SeqEntry/OverlapScores/APIs (e.g., find_matches signature) and call sites to use internal IDs. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
slin1237
changed the title
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces significant performance optimizations to the PositionalIndexer by implementing several strategies inspired by Dynamo's Flash Indexer. However, two significant security and correctness issues were identified. First, the 'Single' entry optimization in SeqEntry incorrectly assumes that a block's position and content hash uniquely identify its prefix, potentially leading to cache misrouting and incorrect load balancing decisions. Second, the new per-position indexing structure lacks bounds checks on the number of levels, creating a memory exhaustion vector that could be exploited to cause a Denial of Service (DoS). Additionally, the review highlighted a subtle race condition in the worker ID interning logic, which defeats its intended purpose of memory optimization, and a suggestion to refactor duplicated code for better maintainability. It is recommended to address these security concerns by removing the unsafe 'Single' optimization and implementing strict limits on the index depth, alongside resolving the identified race conditions and refactoring opportunities.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@kv_index/src/event_tree.rs`:
- Around line 494-503: intern_worker can race: two threads may both miss
self.intern.get and allocate separate Arc<str>s; fix by using the DashMap entry
API so only one Arc is stored and returned — inside intern_worker replace the
read-then-entry pattern with a single-entry operation (use
self.intern.entry_by_ref(worker).or_insert_with(|| Arc::from(worker)) if your
DashMap version supports entry_by_ref to avoid allocation on hits; otherwise
call self.intern.entry(Arc::from(worker)).or_insert_with(|| ()) and return the
map’s stored key; reference intern_worker, self.intern, WorkerId, entry(),
entry_by_ref and or_insert(_/with) when making the change.
…er IDs Three tightly coupled optimizations targeting the 10-115x gap with Dynamo's Flash Indexer on 128-CPU benchmarks. Dynamo deep dive confirmed their actual code uses pre-allocated index (no RwLock), per-thread worker_blocks (zero contention), and Copy WorkerWithDpRank (u64+u32). Change 1 — Remove RwLock from index: - Replace RwLock<Vec<DashMap>> with Box<[DashMap]> pre-allocated at construction via new max_num_blocks parameter (default 2048) - Eliminates atomic fetch_add cache-line bouncing on 128 cores (~50-200ns per acquisition vs ~5ns uncontended) - Remove ensure_levels(); blocks beyond capacity are truncated with warning - jump_search_matches accesses &self.index directly (no read guard) Change 2 — worker_blocks to DashMap: - Replace RwLock<FxHashMap<WorkerId, LevelIndex>> with DashMap<u32, LevelIndex, FxBuildHasher> - apply_stored reduced from 3 separate lock acquisitions to clean DashMap entry/get pattern - remove_or_clear_worker uses DashMap::remove() instead of global write lock that previously serialized ALL worker operations Change 3 — Internal u32 worker IDs: - SeqEntry uses FxHashSet<u32> internally — FxHash on u32 is 1 instruction vs 8+ for Arc<str> (20-byte string hash + pointer deref) - Eliminates Arc::clone atomic refcount bouncing in D×W inner query loop - intern_worker() returns u32 with double-checked locking pattern - Public API unchanged: &str params, OverlapScores returns Arc<str> keys - Conversion at API boundary is O(W), negligible vs inner loop savings Files changed: - kv_index/src/event_tree.rs: all three structural changes - model_gateway/src/core/kv_event_monitor.rs: pass max_num_blocks=2048 - model_gateway/src/policies/cache_aware.rs: update test constructors - model_gateway/benches/radix_tree_benchmark.rs: update bench constructors Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
github-actions
Bot
added
benchmarks
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@kv_index/src/event_tree.rs`:
- Around line 361-364: apply_removed and remove_or_clear_worker must not call
intern_worker because that creates new IDs for unknown workers and mutates
worker_to_id/id_to_worker on no-op paths; change both functions to perform a
non-mutating lookup of the worker id (e.g., use worker_to_id.get(worker) or add
a helper get_worker_id that returns Option<WorkerId>) and return early if None,
then continue using the found WorkerId to access worker_blocks; ensure no new
entries are inserted into worker_to_id or id_to_worker in these remove/clear
code paths.
In `@model_gateway/src/core/kv_event_monitor.rs`:
- Line 107: Extract the magic literal 2048 into a shared constant (e.g.,
INDEX_CAPACITY) in the kv_event_monitor.rs module (or a common constants module)
and replace the inline literal in the PositionalIndexer::new(...) call inside
the or_insert_with closure with that constant; ensure the constant is public or
re-exported if tests/reference sites need it so production and tests use the
same value and avoid drift (reference: PositionalIndexer::new and
self.jump_size).
ℹ️ Review info
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (4)
kv_index/src/event_tree.rsmodel_gateway/benches/radix_tree_benchmark.rsmodel_gateway/src/core/kv_event_monitor.rsmodel_gateway/src/policies/cache_aware.rs
github-actions
Bot
removed
benchmarks
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
kv_index/src/event_tree.rs (1)
339-341:⚠️ Potential issue | 🟠 MajorDo not intern worker IDs in remove/clear paths.
Using
intern_workerinapply_removedandremove_or_clear_workercreates IDs for unknown workers on no-op paths. That mutates state and can growworker_to_id/id_to_workerindefinitely under churn.♻️ Proposed fix (non-mutating lookup)
pub fn apply_removed(&self, worker: &str, seq_hashes: &[SequenceHash]) { - let worker_id = self.intern_worker(worker); + let Some(worker_id) = self.lookup_worker_id(worker) else { + tracing::debug!( + worker = %worker, + num_hashes = seq_hashes.len(), + "apply_removed: worker not tracked, ignoring" + ); + return; + }; @@ fn remove_or_clear_worker(&self, worker: &str, keep_worker: bool) { - let worker_id = self.intern_worker(worker); + let Some(worker_id) = self.lookup_worker_id(worker) else { + return; + }; @@ +#[inline] +fn lookup_worker_id(&self, worker: &str) -> Option<u32> { + self.worker_to_id.get(worker).map(|entry| *entry.value()) +}Also applies to: 403-405
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kv_index/src/event_tree.rs` around lines 339 - 341, The code currently calls intern_worker in apply_removed and remove_or_clear_worker which creates new IDs for unknown workers on no-op paths; change these to use a non-mutating lookup (e.g., lookup/get on worker_to_id or an existing find_worker/get_worker_id helper) instead of intern_worker, and if the lookup returns None bail out early so no new entries are inserted; update both apply_removed and remove_or_clear_worker to perform this non-mutating check before doing any removal/clear logic.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@kv_index/src/event_tree.rs`:
- Around line 200-211: The fast-path in workers_if_single currently ignores
prefix validation causing incorrect jump-search matches; change
workers_if_single to require the caller to assert the prefix is already matched
(e.g., fn workers_if_single(&self, prefix_trivial: bool) ->
Option<&FxHashSet<u32>>) and only return Some(workers) for Self::Single when
prefix_trivial is true (otherwise return None). Update all call sites that
relied on the old zero-arg workers_if_single (including the other locations
mentioned) to pass the correct prefix_trivial boolean (or perform explicit
prefix checks before using the fast-path) so Single is only used when prefix
validation is provably trivial.
---
Duplicate comments:
In `@kv_index/src/event_tree.rs`:
- Around line 339-341: The code currently calls intern_worker in apply_removed
and remove_or_clear_worker which creates new IDs for unknown workers on no-op
paths; change these to use a non-mutating lookup (e.g., lookup/get on
worker_to_id or an existing find_worker/get_worker_id helper) instead of
intern_worker, and if the lookup returns None bail out early so no new entries
are inserted; update both apply_removed and remove_or_clear_worker to perform
this non-mutating check before doing any removal/clear logic.
…apScores What changed: - kv_index/src/event_tree.rs: Remove inner RwLock from worker_blocks (DashMap shard-level locking only), add atomic tree_sizes tracking (O(1) reads vs O(n) locked iteration), add retain guard in linear_scan_drain (skip when workers.len() >= active.len()), remove id_to_worker Vec and RwLock (OverlapScores uses u32 keys directly), add early_exit parameter to find_matches, add WorkerId = u32 type alias. Convert all tests from string-key to u32 worker ID API. - model_gateway/src/policies/cache_aware.rs: Update score_overlap to use indexer.worker_id() for u32 lookups into OverlapScores, add early_exit=false to find_matches call. - model_gateway/benches/radix_tree_benchmark.rs: Add early_exit=false to find_matches calls in bench_indexer_match and bench_indexer_concurrent macros. Why: Close 5 performance gaps between our PositionalIndexer and Dynamo's Flash Indexer: (1) eliminate inner RwLock contention on worker_blocks, (2) atomic O(1) tree_sizes instead of locked iteration, (3) retain guard to skip O(active) iteration when all workers still match, (4) remove worker ID interning overhead (no Arc<str> conversion at return time), (6) early_exit support for fast existence checks. How: worker_blocks changed from DashMap<u32, RwLock<FxHashMap>> to DashMap<u32, FxHashMap> — plain HashMaps behind DashMap's shard locks. tree_sizes tracked via DashMap<u32, AtomicUsize> with fetch_add/sub on store/remove. OverlapScores now keyed by u32 internally; consumers use PositionalIndexer::worker_id() to map URLs to IDs. The workers_if_single optimization (fix 5, already present) is retained as an advantage over Dynamo. Signed-off-by: Simon Lin <simon@nvidia.com> Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
github-actions
Bot
added
benchmarks
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
kv_index/src/event_tree.rs (2)
367-369:⚠️ Potential issue | 🟠 MajorAvoid creating new worker IDs on remove/clear no-op paths.
Line 368 and Line 444 call
intern_worker(...); unknown workers then get permanently interned even though nothing is removed. This causes unboundedworker_to_idgrowth and unnecessarynext_worker_idchurn under noisy remove/clear traffic.♻️ Suggested fix
pub fn apply_removed(&self, worker: &str, seq_hashes: &[SequenceHash]) { - let worker_id = self.intern_worker(worker); + let Some(worker_id) = self.lookup_worker_id(worker) else { + tracing::debug!( + worker = %worker, + num_hashes = seq_hashes.len(), + "apply_removed: worker not tracked, ignoring" + ); + return; + }; @@ fn remove_or_clear_worker(&self, worker: &str, keep_worker: bool) { - let worker_id = self.intern_worker(worker); + let Some(worker_id) = self.lookup_worker_id(worker) else { + return; + };#[inline] fn lookup_worker_id(&self, worker: &str) -> Option<u32> { self.worker_to_id.get(worker).map(|entry| *entry.value()) }Also applies to: 443-445
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kv_index/src/event_tree.rs` around lines 367 - 369, The remove/clear paths call intern_worker(...) and permanently intern unknown worker strings even when nothing is removed; change these paths (e.g., in apply_removed and the analogous clear/remove function at the other occurrence) to first try a non-interning lookup (implement a helper like lookup_worker_id(&self, worker: &str) -> Option<u32> that checks self.worker_to_id without inserting) and early-return when it yields None, instead of calling intern_worker; this prevents unbounded growth of worker_to_id and avoids bumping next_worker_id on noop remove/clear operations.
210-221:⚠️ Potential issue | 🔴 CriticalSingle-entry fast path skips required prefix validation in jump/search path.
Line 559 and Line 599 treat any
Singleentry as a match without checkingseq_hash. That can keep/score workers that do not actually match the request prefix chain.🐛 Suggested fix
fn get_workers_lazy( @@ - if let Some(workers) = entry.value().workers_if_single() { - return Some(workers.clone()); + if position == 0 { + if let Some(workers) = entry.value().workers_if_single() { + return Some(workers.clone()); + } } @@ fn count_workers_at( @@ - if let Some(workers) = entry.value().workers_if_single() { - return workers.len(); - } Self::ensure_seq_hash_computed(seq_hashes, position, sequence); entry - .get(seq_hashes[position]) + .value() + .get(seq_hashes[position]) .map(|workers| workers.len()) .unwrap_or(0) } @@ - if let Some(workers) = entry.value().workers_if_single() { - if workers.len() < active.len() { - active.retain(|&w| { - if workers.contains(&w) { - true - } else { - internal_scores.insert(w, pos as u32); - false - } - }); - } - if early_exit && !active.is_empty() { - break; - } - continue; - } - - // Multi: need rolling hash to disambiguate. Self::ensure_seq_hash_computed(seq_hashes, pos, sequence); let seq_hash = seq_hashes[pos]; - - let Some(workers) = entry.get(seq_hash) else { + let Some(workers) = entry.value().get(seq_hash) else { for w in active.drain() { internal_scores.insert(w, pos as u32); } break; };Also applies to: 559-561, 599-617
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kv_index/src/event_tree.rs` around lines 210 - 221, The fast path method workers_if_single currently returns workers for any Self::Single without verifying the stored seq_hash, allowing false matches; change it to validate the sequence/prefix hash before returning. Modify workers_if_single (or add a new workers_if_single_matching) to accept the expected seq_hash/prefix hash parameter, pattern-match Self::Single(stored_seq_hash, workers) and only return Some(workers) when stored_seq_hash == expected_seq_hash, otherwise None; then update the jump/search call sites that used workers_if_single (the Single fast-path branches) to pass the expected seq_hash and handle None as before.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@kv_index/src/event_tree.rs`:
- Around line 339-352: The code unconditionally increments self.tree_sizes by
blocks.len(), which double-counts when seq_hash keys are overwritten; fix by
counting only newly-inserted seq_hashes and adding that delta. For each block in
the loop where you call wb_ref.insert(block.seq_hash, ...), check whether that
seq_hash already exists in wb_ref (or track with a local HashSet of seen
seq_hashes/effective pre-existence) and increment a local new_count only when
the key was absent; after drop(wb_ref) use new_count (not blocks.len()) to
fetch_add on self.tree_sizes.entry(worker_id)... or or_insert with new
AtomicUsize(new_count) so tree_sizes reflects only newly added entries.
---
Duplicate comments:
In `@kv_index/src/event_tree.rs`:
- Around line 367-369: The remove/clear paths call intern_worker(...) and
permanently intern unknown worker strings even when nothing is removed; change
these paths (e.g., in apply_removed and the analogous clear/remove function at
the other occurrence) to first try a non-interning lookup (implement a helper
like lookup_worker_id(&self, worker: &str) -> Option<u32> that checks
self.worker_to_id without inserting) and early-return when it yields None,
instead of calling intern_worker; this prevents unbounded growth of worker_to_id
and avoids bumping next_worker_id on noop remove/clear operations.
- Around line 210-221: The fast path method workers_if_single currently returns
workers for any Self::Single without verifying the stored seq_hash, allowing
false matches; change it to validate the sequence/prefix hash before returning.
Modify workers_if_single (or add a new workers_if_single_matching) to accept the
expected seq_hash/prefix hash parameter, pattern-match
Self::Single(stored_seq_hash, workers) and only return Some(workers) when
stored_seq_hash == expected_seq_hash, otherwise None; then update the
jump/search call sites that used workers_if_single (the Single fast-path
branches) to pass the expected seq_hash and handle None as before.
ℹ️ Review info
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (3)
kv_index/src/event_tree.rsmodel_gateway/benches/radix_tree_benchmark.rsmodel_gateway/src/policies/cache_aware.rs
There was a problem hiding this comment.
♻️ Duplicate comments (3)
kv_index/src/event_tree.rs (3)
691-693:⚠️ Potential issue | 🟠 MajorSame intern table pollution issue in remove/clear path.
Same concern as
apply_removed— callingintern_workerfor a worker that was never stored creates a stale entry.🐛 Proposed fix
fn remove_or_clear_worker(&self, worker: &str, keep_worker: bool) { - let worker_id = self.intern_worker(worker); + let Some(worker_id) = self.worker_to_id.get(worker).map(|e| *e.value()) else { + return; + }; if let Some((_, worker_map)) = self.worker_blocks.remove(&worker_id) {🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kv_index/src/event_tree.rs` around lines 691 - 693, remove_or_clear_worker currently calls intern_worker which inserts a new worker key if it didn't exist, polluting the intern table; change the call to use a non-inserting lookup (the same approach used to fix apply_removed) — e.g., replace self.intern_worker(worker) with the worker-id getter that does not create entries (like self.get_worker_id(worker) or look up in the intern map directly), and early-return if the worker is not found so you don't create a stale entry.
615-617:⚠️ Potential issue | 🟠 MajorIntern table grows unboundedly when removing unknown workers.
intern_workercreates entries inworker_to_ideven for unknown workers. While the function returns early ifworker_blockshas no entry, the intern table retains a stale mapping. Under worker churn, this accumulates garbage entries.A past review flagged this as addressed, but the current code still calls
intern_workerhere. Consider using a non-mutating lookup:🐛 Proposed fix to use lookup-only path
pub fn apply_removed(&self, worker: &str, seq_hashes: &[SequenceHash]) { - let worker_id = self.intern_worker(worker); - - let Some(mut wb_ref) = self.worker_blocks.get_mut(&worker_id) else { + let Some(worker_id) = self.worker_to_id.get(worker).map(|e| *e.value()) else { + tracing::debug!( + worker = %worker, + num_hashes = seq_hashes.len(), + "apply_removed: worker not interned, ignoring" + ); + return; + }; + + let Some(mut wb_ref) = self.worker_blocks.get_mut(&worker_id) else { tracing::debug!(🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kv_index/src/event_tree.rs` around lines 615 - 617, apply_removed currently calls intern_worker which inserts a mapping for unknown workers and leaks entries; change the call-site to perform a non-mutating lookup first (e.g. check self.worker_blocks.get(worker) or use a lookup-only helper like self.worker_to_id.get(worker) / self.lookup_worker(worker)), return early if no entry exists, and only call intern_worker when you actually need to create/insert a new id; reference apply_removed, intern_worker, worker_blocks and worker_to_id when making the change.
587-600:⚠️ Potential issue | 🟠 Major
tree_sizesover-counts on block re-store/overwrite.
FxHashMap::insertreturns the old value if a key is replaced, but the code ignores this and unconditionally incrementstree_sizesbyblocks.len(). Re-storing the sameseq_hashinflates the count.🐛 Proposed fix to count only new insertions
+ let mut num_new_blocks = 0usize; let mut prev_prefix = parent_prefix; for (i, block) in blocks.iter().enumerate() { let position = start_pos + i; let content_hash = block.content_hash; // Compute router prefix hash (rolling XXH3 of content hashes). // This is the SeqEntry key — consistent between store and query paths. let prefix_hash = match prev_prefix { Some(prev) => SequenceHash(Self::compute_next_seq_hash(prev.0, content_hash.0)), // Position 0: prefix_hash == content_hash (no parent to chain from). None => SequenceHash(content_hash.0), }; self.index .entry((position, content_hash)) .and_modify(|entry| entry.insert(prefix_hash, worker_id)) .or_insert_with(|| SeqEntry::new(prefix_hash, worker_id)); - wb_ref.insert(block.seq_hash, (position, content_hash, prefix_hash)); + if wb_ref.insert(block.seq_hash, (position, content_hash, prefix_hash)).is_none() { + num_new_blocks += 1; + } prev_prefix = Some(prefix_hash); } drop(wb_ref); // Atomically update tree_sizes. - let num_blocks = blocks.len(); - self.tree_sizes - .entry(worker_id) - .and_modify(|size| { - size.fetch_add(num_blocks, Ordering::Relaxed); - }) - .or_insert(AtomicUsize::new(num_blocks)); + if num_new_blocks > 0 { + self.tree_sizes + .entry(worker_id) + .and_modify(|size| { + size.fetch_add(num_new_blocks, Ordering::Relaxed); + }) + .or_insert(AtomicUsize::new(num_new_blocks)); + }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@kv_index/src/event_tree.rs` around lines 587 - 600, The tree_sizes counter is inflated because wb_ref.insert(block.seq_hash, ...) may replace existing keys and the code unconditionally adds blocks.len() to tree_sizes; change the logic to count only new insertions (use the return of wb_ref.insert or wb_ref.entry(seq_hash).or_insert(...) while tracking a new_insertions counter incremented only when there was no prior entry) and then atomically add that new_insertions value to self.tree_sizes for the given worker_id (use the same worker_id, AtomicUsize::fetch_add(new_insertions, Ordering::Relaxed) path or or_insert with AtomicUsize::new(new_insertions) if absent). Ensure you reference wb_ref.insert / wb_ref.entry, blocks.len() only for the original batch size, and self.tree_sizes so the increment reflects actual newly added blocks rather than overwritten ones.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@kv_index/src/event_tree.rs`:
- Around line 691-693: remove_or_clear_worker currently calls intern_worker
which inserts a new worker key if it didn't exist, polluting the intern table;
change the call to use a non-inserting lookup (the same approach used to fix
apply_removed) — e.g., replace self.intern_worker(worker) with the worker-id
getter that does not create entries (like self.get_worker_id(worker) or look up
in the intern map directly), and early-return if the worker is not found so you
don't create a stale entry.
- Around line 615-617: apply_removed currently calls intern_worker which inserts
a mapping for unknown workers and leaks entries; change the call-site to perform
a non-mutating lookup first (e.g. check self.worker_blocks.get(worker) or use a
lookup-only helper like self.worker_to_id.get(worker) /
self.lookup_worker(worker)), return early if no entry exists, and only call
intern_worker when you actually need to create/insert a new id; reference
apply_removed, intern_worker, worker_blocks and worker_to_id when making the
change.
- Around line 587-600: The tree_sizes counter is inflated because
wb_ref.insert(block.seq_hash, ...) may replace existing keys and the code
unconditionally adds blocks.len() to tree_sizes; change the logic to count only
new insertions (use the return of wb_ref.insert or
wb_ref.entry(seq_hash).or_insert(...) while tracking a new_insertions counter
incremented only when there was no prior entry) and then atomically add that
new_insertions value to self.tree_sizes for the given worker_id (use the same
worker_id, AtomicUsize::fetch_add(new_insertions, Ordering::Relaxed) path or
or_insert with AtomicUsize::new(new_insertions) if absent). Ensure you reference
wb_ref.insert / wb_ref.entry, blocks.len() only for the original batch size, and
self.tree_sizes so the increment reflects actual newly added blocks rather than
overwritten ones.
slin1237
force-pushed
the
slin/flash-indexer-perf
branch
from
March 3, 2026 19:57
efc5ca3 to
8a8a85b
Compare
slin1237
force-pushed
the
slin/flash-indexer-perf
branch
from
March 3, 2026 21:00
8ae11d8 to
7b57780
Compare
Phase F: Replace FxHashSet<u32> active set with Vec<u32> in jump_search_matches to avoid hash table clone overhead at position 0 and improve cache locality during retain/drain operations. - get_workers_lazy returns Option<Vec<u32>> instead of Option<FxHashSet<u32>> - linear_scan_drain takes &mut Vec<u32> with swap_remove retain pattern instead of FxHashSet::retain(), avoiding per-element hashing - Drain paths use active.iter() + active.clear() instead of active.drain() - jump_search_matches iterates with for &w in &active Phase G: Change default jump_size from 64 to 32 to match Dynamo's default. Smaller jumps reduce the linear scan window on failed jumps at the cost of more jump probes. Also fixes pre-existing clippy lint: worker_map.iter() → worker_map.values() Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
Update benchmark JUMP_SIZE constant and STORE macro to use 32 instead of 64, matching the new default in PositionalIndexer::default(). This ensures benchmarks exercise the same jump_size as production. Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
DashMap defaults to num_cpus * 4 shards. On 128-core machines this creates 512 shards per map (2048 total across 4 DashMaps), most of which are empty — wasting memory and polluting CPU caches. Constrain shard counts to match the approach used in token_tree and string_tree: - index: 32 shards (hot path, many entries, benefits from concurrency) - worker_blocks/tree_sizes/worker_to_id: 8 shards each (keyed by worker_id, at most ~500 entries, low contention) Also fix pre-existing clippy len_zero warnings in tests: assert!(indexer.index.len() > 0) → assert!(!indexer.index.is_empty()) Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
32 shards caused ~20-25% concurrent throughput regression with 32 benchmark threads — every shard was contended on every query. 64 shards gives 2:1 shard-to-thread ratio while still being 8x less than the 512 default on 128-core machines. Worker maps stay at 8 shards (low contention, few entries). Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
64 shards on the index DashMap still caused concurrent throughput regression (100w: 253K vs 387K baseline). The mixed read+write concurrent benchmark needs high shard counts to avoid contention. Revert index to DashMap default (num_cpus * 4). Only constrain the three worker-keyed maps (worker_blocks, tree_sizes, worker_to_id) to 8 shards — they hold at most ~500 entries and have low contention. Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
This was referenced Mar 4, 2026
This was referenced Jun 10, 2026
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
Closes 4 structural performance gaps between our
PositionalIndexerand Dynamo's Flash Indexer (170M block-level ops/s on 24-core Arrow Lake), identified by comparing our implementation against the Dynamo Flash Indexer blog post.All changes are internal to
PositionalIndexerinkv_index/src/event_tree.rs— the public API is unchanged.What changed
Gap 1 — Per-position DashMaps
Replace
DashMap<(usize, ContentHash), SeqEntry>withRwLock<Vec<DashMap<ContentHash, SeqEntry, FxBuildHasher>>>. One DashMap per block position eliminates cross-position shard contention and improves cache locality for hot prefix positions (0–3). Vec grows lazily viaensure_levels(); read lock held for query duration (only blocks rare Vec growth).Gap 2 — Eliminate set cloning in query path
linear_scan_drainnow accesses DashMap entries directly viaRefand checks worker membership in-place. Previously clonedFxHashSet<WorkerId>at every position during drain (up tojump_size=64 clones per drain call).Gap 3 — Intern worker IDs
Add
DashMap<WorkerId, (), FxBuildHasher>intern table. First call per worker doesArc::from()(heap alloc); subsequent calls returnArc::clone()(atomic increment only). Lookup by&strworks viaArc<str>: Borrow<str>. Eliminates per-call allocation overhead inapply_stored,apply_removed, andremove_or_clear_worker.Gap 4 — Lazy rolling hash for Single entries
Add
SeqEntry::workers_if_single()that returns the worker set directly for Single entries without requiring rolling hash computation. Since content hash collisions at 64-bit XXH3 are ~2^-64, a matchingcontent_hashat the same position is unambiguous. Applied in all three query-path helpers:get_workers_lazy,count_workers_at, andlinear_scan_drain.Structural changes
&[DashMap<...>]to share a single read guard acquired once injump_search_matchesdrop(levels)before acquiringworker_blocks.read()maintains lock orderingworker_blocks→index→ DashMap shards → LevelIndexTest plan
cargo test -p kv-index— 158/158 tests passcargo clippy -p kv-index— no warningsmake fmt— cleanbenchmark-radix-treeworkflow will run on this PR — compare before/after numbersSummary by CodeRabbit
Breaking Changes
New Features
Improvements
Behavioral Changes