Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
8 changes: 6 additions & 2 deletions lib/kv-router/src/indexer/positional.rs
Original file line number Diff line number Diff line change
Expand Up @@ -477,6 +477,9 @@ impl PositionalIndexer {
KvCacheEventError::BlockNotFound
})?;

// A missing hash skips only itself, as in the other indexers: returning early
// would leave every later block of the batch indexed for this worker.
let mut result = Ok(());
for seq_hash in seq_hashes {
let Some((position, local_hash)) = worker_map.remove(seq_hash) else {
tracing::warn!(
Expand All @@ -487,7 +490,8 @@ impl PositionalIndexer {
"Failed to find block to remove; skipping remove operation"
);

return Err(KvCacheEventError::BlockNotFound);
result = Err(KvCacheEventError::BlockNotFound);
continue;
};

if let Some(mut entry) = self.index.get_mut(&(position, local_hash)) {
Expand All @@ -496,7 +500,7 @@ impl PositionalIndexer {
self.lifecycle.remove(worker, *seq_hash);
}

Ok(())
result
}

fn remove_worker_dp_rank_impl(
Expand Down
34 changes: 34 additions & 0 deletions lib/kv-router/src/indexer/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1472,6 +1472,40 @@ mod interface_tests {
.await;
}

/// A hash the index does not hold, or already dropped earlier in the same batch,
/// must not stop the batch: the blocks after it are still removed.
#[tokio::test]
#[apply(matching_indexer_template)]
async fn test_remove_batch_continues_past_missing_blocks(variant: &str) {
let worker = WorkerWithDpRank::new(0, 0);
let workers = [worker];
let full_hashes: Vec<LocalBlockHash> = (1..=3).map(LocalBlockHash).collect();
let seq_hashes = compute_seq_hash_for_block(&full_hashes);
let block_2 = ExternalSequenceBlockHash(seq_hashes[1]);
let block_3 = ExternalSequenceBlockHash(seq_hashes[2]);

for batch in [
vec![ExternalSequenceBlockHash(999), block_3, block_2],
vec![block_3, block_3, block_2],
] {
let index = make_matching_indexer(variant, &workers);
index.apply_event(make_store_event(0, &[1, 2, 3])).await;
index.apply_event(remove_event(0, 0, 0, batch)).await;

flush_and_settle(index.as_ref()).await;

assert_query_scores_with_semantics(
variant,
index.as_ref(),
&[1, 2, 3],
&workers,
&[(worker, 1)],
MatchSemantics::Exact,
)
.await;
}
}

#[tokio::test]
#[apply(matching_indexer_template)]
async fn test_clear_then_reuse(variant: &str) {
Expand Down
Loading