Skip to content

feat(kv-router): warn on duplicate store replays #8473

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
PeaBrane merged 2 commits into from
Apr 22, 2026
Merged

feat(kv-router): warn on duplicate store replays #8473

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b4440cf
teach router to side-eye duplicate stores
PeaBrane Apr 21, 2026
4302eac
Merge branch 'main' into rupei/router-duplicate-store-warning
PeaBrane Apr 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 25 additions & 12 deletions lib/kv-router/src/indexer/concurrent_radix_tree.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,9 @@ use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet};
use std::collections::VecDeque;
use std::sync::atomic::{AtomicUsize, Ordering};

use super::{EventKind, KvIndexerMetrics, SyncIndexer, WorkerTask};
use super::{
EventKind, EventWarningKind, KvIndexerMetrics, PreBoundEventCounters, SyncIndexer, WorkerTask,
};
use crate::active_set::reconcile_active_workers;
use crate::cleanup::{self, CleanableNode, CleanupGuard, CleanupState};
use crate::protocols::*;
Expand Down Expand Up @@ -309,6 +311,7 @@ impl ConcurrentRadixTree {
&self,
lookup: &mut FxHashMap<WorkerWithDpRank, WorkerLookup>,
event: RouterEvent,
counters: Option<&PreBoundEventCounters>,
) -> Result<(), KvCacheEventError> {
let (worker_id, kv_event) = (event.worker_id, event.event);
let (id, op) = (kv_event.event_id, kv_event.data);
Expand All @@ -317,7 +320,7 @@ impl ConcurrentRadixTree {
let worker = WorkerWithDpRank::new(worker_id, kv_event.dp_rank);

match op {
KvCacheEventData::Stored(op) => self.apply_stored(lookup, worker, op, id),
KvCacheEventData::Stored(op) => self.apply_stored(lookup, worker, op, id, counters),
KvCacheEventData::Removed(op) => self.apply_removed(lookup, worker, op, id),
KvCacheEventData::Cleared => {
// Ensure the worker is tracked in lookup before clearing,
Expand All @@ -340,6 +343,7 @@ impl ConcurrentRadixTree {
worker: WorkerWithDpRank,
op: KvCacheStoreData,
id: u64,
counters: Option<&PreBoundEventCounters>,
) -> Result<(), KvCacheEventError> {
// Ensure this worker has an entry in the outer map.
let worker_lookup = lookup.entry(worker).or_default();
Expand All @@ -364,6 +368,7 @@ impl ConcurrentRadixTree {
};

let mut needs_worker_insert = false;
let mut duplicate_store = !op.blocks.is_empty();

let mut num_blocks_added = 0;

Expand All @@ -379,8 +384,8 @@ impl ConcurrentRadixTree {
// Insert worker into this node if it was the child from the
// previous iteration (skip for the initial parent, which is
// not one of the blocks being stored).
if needs_worker_insert {
parent_guard.workers.insert(worker);
if needs_worker_insert && parent_guard.workers.insert(worker) {
duplicate_store = false;
}
needs_worker_insert = true;

Expand All @@ -390,6 +395,7 @@ impl ConcurrentRadixTree {
{
let existing_guard = existing.read();
if existing_guard.block_hash != Some(block_data.block_hash) {
duplicate_store = false;
tracing::warn!(
expected = ?block_data.block_hash,
actual = ?existing_guard.block_hash,
Expand All @@ -400,6 +406,7 @@ impl ConcurrentRadixTree {
existing.clone()
}
None => {
duplicate_store = false;
// Reuse from lookup or create new
let new_block = worker_lookup
.get(&block_data.block_hash)
Expand All @@ -417,20 +424,22 @@ impl ConcurrentRadixTree {
};

// Update lookup
if worker_lookup
.insert(block_data.block_hash, child.clone())
.is_none()
{
num_blocks_added += 1;
match worker_lookup.insert(block_data.block_hash, child.clone()) {
Some(existing) if Arc::ptr_eq(&existing, &child) => {}
Some(_) => duplicate_store = false,
None => {
num_blocks_added += 1;
duplicate_store = false;
}
}

current = child;
}

// Insert worker into the last child (not yet handled since there is
// no subsequent iteration to pick it up).
if needs_worker_insert {
current.write().workers.insert(worker);
if needs_worker_insert && current.write().workers.insert(worker) {
duplicate_store = false;
}

match self.tree_sizes.get(&worker) {
Expand All @@ -443,6 +452,10 @@ impl ConcurrentRadixTree {
}
}

if duplicate_store && let Some(counters) = counters {
counters.inc_warning(EventWarningKind::DuplicateStore);
}

Ok(())
}

Expand Down Expand Up @@ -649,7 +662,7 @@ impl SyncIndexer for ConcurrentRadixTree {
match task {
WorkerTask::Event(event) => {
let kind = EventKind::of(&event.event.data);
let result = self.apply_event(&mut lookup, event);
let result = self.apply_event(&mut lookup, event, counters.as_ref());
if result.is_err() {
tracing::warn!("Failed to apply event: {:?}", result.as_ref().err());
}
Expand Down
97 changes: 74 additions & 23 deletions lib/kv-router/src/indexer/concurrent_radix_tree_compressed.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,9 @@ use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet};
use std::collections::VecDeque;
use std::sync::atomic::{AtomicUsize, Ordering};

use super::{EventKind, KvIndexerMetrics, SyncIndexer, WorkerTask};
use super::{
EventKind, EventWarningKind, KvIndexerMetrics, PreBoundEventCounters, SyncIndexer, WorkerTask,
};
use crate::cleanup::{self, CleanableNode, CleanupGuard, CleanupState};
use crate::protocols::*;

Expand Down Expand Up @@ -156,10 +158,13 @@ impl Node {
}

#[inline]
fn promote_to_full(&mut self, worker: WorkerWithDpRank) {
fn promote_to_full(&mut self, worker: WorkerWithDpRank) -> bool {
if !self.full_edge_workers.contains(&worker) {
self.worker_cutoffs.remove(&worker);
self.full_edge_workers.insert(worker);
true
} else {
false
}
}

Expand Down Expand Up @@ -250,6 +255,11 @@ struct RemoveOutcome {
stale_hashes: Vec<ExternalSequenceBlockHash>,
}

struct StoreInsertOutcome {
num_blocks_added: usize,
duplicate_store: bool,
}

/// Thread-safe radix tree (compressed trie) for concurrent KV cache lookups.
pub struct ConcurrentRadixTreeCompressed {
/// The root of the radix tree. Has an empty edge and only contains children.
Expand Down Expand Up @@ -614,13 +624,14 @@ impl ConcurrentRadixTreeCompressed {
&self,
lookup: &mut FxHashMap<WorkerWithDpRank, WorkerLookup>,
event: RouterEvent,
counters: Option<&PreBoundEventCounters>,
) -> Result<(), KvCacheEventError> {
let (worker_id, kv_event) = (event.worker_id, event.event);
let (id, op) = (kv_event.event_id, kv_event.data);
let worker = WorkerWithDpRank::new(worker_id, kv_event.dp_rank);

match op {
KvCacheEventData::Stored(op) => self.apply_stored(lookup, worker, op, id),
KvCacheEventData::Stored(op) => self.apply_stored(lookup, worker, op, id, counters),
KvCacheEventData::Removed(op) => self.apply_removed(lookup, worker, op, id),
KvCacheEventData::Cleared => {
lookup.entry(worker).or_default();
Expand All @@ -643,6 +654,7 @@ impl ConcurrentRadixTreeCompressed {
worker: WorkerWithDpRank,
op: KvCacheStoreData,
id: u64,
counters: Option<&PreBoundEventCounters>,
) -> Result<(), KvCacheEventError> {
lookup.entry(worker).or_default();

Expand Down Expand Up @@ -723,19 +735,24 @@ impl ConcurrentRadixTreeCompressed {
None => self.root.clone(),
};

let num_blocks_added =
self.insert_blocks_from(lookup, worker, &parent, op.parent_hash, &op.blocks);
let outcome = self.insert_blocks_from(lookup, worker, &parent, op.parent_hash, &op.blocks);

match self.tree_sizes.get(&worker) {
Some(size) => {
size.fetch_add(num_blocks_added, Ordering::Relaxed);
size.fetch_add(outcome.num_blocks_added, Ordering::Relaxed);
}
None => {
self.tree_sizes
.insert(worker, AtomicUsize::new(num_blocks_added));
.insert(worker, AtomicUsize::new(outcome.num_blocks_added));
}
}

if outcome.duplicate_store
&& let Some(counters) = counters
{
counters.inc_warning(EventWarningKind::DuplicateStore);
}

Ok(())
}

Expand All @@ -746,10 +763,11 @@ impl ConcurrentRadixTreeCompressed {
parent: &SharedNode,
seed_hash: Option<ExternalSequenceBlockHash>,
blocks: &[KvCacheStoredBlockData],
) -> usize {
) -> StoreInsertOutcome {
let mut current_parent = parent.clone();
let mut remaining = blocks;
let mut num_blocks_added = 0usize;
let mut duplicate_store = !blocks.is_empty();
// Track the last ExternalSequenceBlockHash we matched to detect if
// `current_parent` was split by a concurrent thread between iterations.
// A split shortens `current_parent`'s edge and moves our last-matched
Expand Down Expand Up @@ -811,11 +829,18 @@ impl ConcurrentRadixTreeCompressed {

let wl = lookup.get_mut(&worker).unwrap();
for b in remaining {
if wl.insert(b.block_hash, new_node.clone()).is_none() {
num_blocks_added += 1;
match wl.insert(b.block_hash, new_node.clone()) {
Some(existing) if Arc::ptr_eq(&existing, &new_node) => {}
Some(_) => {}
None => {
num_blocks_added += 1;
}
}
}
return num_blocks_added;
return StoreInsertOutcome {
num_blocks_added,
duplicate_store: false,
};
}
}
};
Expand All @@ -830,6 +855,7 @@ impl ConcurrentRadixTreeCompressed {
break;
}
if edge_elem.1 != rem_elem.block_hash {
duplicate_store = false;
tracing::warn!(
expected = ?rem_elem.block_hash,
actual = ?edge_elem.1,
Expand Down Expand Up @@ -882,13 +908,21 @@ impl ConcurrentRadixTreeCompressed {

let wl = lookup.get_mut(&worker).unwrap();
for b in &remaining[..match_len] {
if wl.insert(b.block_hash, child.clone()).is_none() {
num_blocks_added += 1;
match wl.insert(b.block_hash, child.clone()) {
Some(existing) if Arc::ptr_eq(&existing, &child) => {}
Some(_) => {}
None => {
num_blocks_added += 1;
}
}
}
for b in tail {
if wl.insert(b.block_hash, new_node.clone()).is_none() {
num_blocks_added += 1;
match wl.insert(b.block_hash, new_node.clone()) {
Some(existing) if Arc::ptr_eq(&existing, &new_node) => {}
Some(_) => {}
None => {
num_blocks_added += 1;
}
}
}
} else {
Expand All @@ -897,22 +931,36 @@ impl ConcurrentRadixTreeCompressed {

let wl = lookup.get_mut(&worker).unwrap();
for b in &remaining[..match_len] {
if wl.insert(b.block_hash, child.clone()).is_none() {
num_blocks_added += 1;
match wl.insert(b.block_hash, child.clone()) {
Some(existing) if Arc::ptr_eq(&existing, &child) => {}
Some(_) => {}
None => {
num_blocks_added += 1;
}
}
}
}
return num_blocks_added;
return StoreInsertOutcome {
num_blocks_added,
duplicate_store: false,
};
}

// Full edge match: upgrade worker to full coverage if necessary.
child_guard.promote_to_full(worker);
if child_guard.promote_to_full(worker) {
duplicate_store = false;
}
drop(child_guard);

let wl = lookup.get_mut(&worker).unwrap();
for b in &remaining[..edge_len] {
if wl.insert(b.block_hash, child.clone()).is_none() {
num_blocks_added += 1;
match wl.insert(b.block_hash, child.clone()) {
Some(existing) if Arc::ptr_eq(&existing, &child) => {}
Some(_) => duplicate_store = false,
None => {
num_blocks_added += 1;
duplicate_store = false;
}
}
}

Expand All @@ -922,7 +970,10 @@ impl ConcurrentRadixTreeCompressed {
}
}

num_blocks_added
StoreInsertOutcome {
num_blocks_added,
duplicate_store,
}
}

// ------------------------------------------------------------------
Expand Down Expand Up @@ -1245,7 +1296,7 @@ impl SyncIndexer for ConcurrentRadixTreeCompressed {
match task {
WorkerTask::Event(event) => {
let kind = EventKind::of(&event.event.data);
let result = self.apply_event(&mut lookup, event);
let result = self.apply_event(&mut lookup, event, counters.as_ref());
if result.is_err() {
tracing::warn!("Failed to apply event: {:?}", result.as_ref().err());
}
Expand Down
2 changes: 1 addition & 1 deletion lib/kv-router/src/indexer/kv_indexer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ fn apply_event_with_prune_tracking(
let event_id = event.event.event_id;
let worker_id = event.worker_id;
let event_for_prune = prune_manager.is_some().then(|| event.clone());
let result = trie.apply_event(event);
let result = trie.apply_event_with_counters(event, Some(counters));
Comment thread
PeaBrane marked this conversation as resolved.
Resolved
let result_is_ok = result.is_ok();
let tree_size = trie.current_size();
tracing::trace!(
Expand Down
Loading
Loading