From ea009e3b4d39afbd3b4bb1f34dde14ed0d40cffb Mon Sep 17 00:00:00 2001 From: bojiang-li <327132355+bojiang-li@users.noreply.github.com> Date: Thu, 1 Oct 2026 03:05:03 -0700 Subject: [PATCH 1/2] fix(kv-router): fence stale TTL removals after refresh Signed-off-by: bojiang-li <327132355+bojiang-li@users.noreply.github.com> --- .../sync_impl.rs | 23 + lib/kv-router/src/indexer/kv_indexer.rs | 63 ++- lib/kv-router/src/indexer/lower_tier.rs | 31 +- lib/kv-router/src/indexer/positional.rs | 31 +- lib/kv-router/src/indexer/pruning.rs | 185 +++++++- lib/kv-router/src/indexer/thread_pool.rs | 435 ++++++++++++++++-- lib/kv-router/src/indexer/types.rs | 5 + 7 files changed, 728 insertions(+), 45 deletions(-) diff --git a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/sync_impl.rs b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/sync_impl.rs index 74c7b6d31d2f..6088b85e4287 100644 --- a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/sync_impl.rs +++ b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/sync_impl.rs @@ -47,6 +47,21 @@ impl SyncIndexer for ConcurrentRadixTreeCompressed { } let _ = resp.send(applied); } + WorkerTask::ApproximateTtl(task) => task.apply(|event| { + let kind = EventKind::of(&event.event.data); + let result = self.apply_event(&mut lookup, event, counters.as_ref()); + let applied = result.is_ok(); + if result.is_err() { + tracing::warn!( + "Failed to apply approximate TTL event: {:?}", + result.as_ref().err() + ); + } + if let Some(ref counters) = counters { + counters.inc(kind, result); + } + applied + }), WorkerTask::ApproximateLru(task) => { approximate_lru.observe_task(&task); let ApproximateLruTask { @@ -116,9 +131,13 @@ impl SyncIndexer for ConcurrentRadixTreeCompressed { } WorkerTask::RemoveWorker { worker_id, + prune_manager, sweep_tree, resp, } => { + if let Some(manager) = prune_manager { + manager.remove_worker(worker_id); + } approximate_lru.forget_worker(worker_id); self.erase_worker_coverage( &mut lookup, @@ -130,8 +149,12 @@ impl SyncIndexer for ConcurrentRadixTreeCompressed { WorkerTask::RemoveWorkerDpRank { worker_id, dp_rank, + prune_manager, sweep_tree, } => { + if let Some(manager) = prune_manager { + manager.remove_worker_dp_rank(WorkerWithDpRank::new(worker_id, dp_rank)); + } approximate_lru.forget_rank(WorkerWithDpRank::new(worker_id, dp_rank)); self.erase_worker_coverage( &mut lookup, diff --git a/lib/kv-router/src/indexer/kv_indexer.rs b/lib/kv-router/src/indexer/kv_indexer.rs index ecf97d1f25fa..f353303f7973 100644 --- a/lib/kv-router/src/indexer/kv_indexer.rs +++ b/lib/kv-router/src/indexer/kv_indexer.rs @@ -101,7 +101,12 @@ fn apply_routing_decision_with_prune_tracking( pm.insert_block_entries(block_entries); } -fn apply_prune_removes(trie: &mut RadixTree, entries: Vec, event_id_counter: &mut u64) { +fn apply_prune_removes( + trie: &mut RadixTree, + entries: Vec, + prune_manager: &WorkerPruneManager, + event_id_counter: &mut u64, +) { let mut entries_by_worker = BTreeMap::>::new(); for entry in entries { entries_by_worker @@ -110,7 +115,11 @@ fn apply_prune_removes(trie: &mut RadixTree, entries: Vec, event_id_ .push(entry); } - for (worker, mut entries) in entries_by_worker { + for (worker, entries) in entries_by_worker { + let mut entries = prune_manager.take_pending_expirations(worker, entries); + if entries.is_empty() { + continue; + } entries.sort_unstable_by_key(|entry| entry.seq_position); *event_id_counter += 1; let event = RouterEvent::new( @@ -288,7 +297,7 @@ fn drain_pending_mutations( if let Some(pm) = prune_manager { let entries = pm.drain_due_and_pending(tokio::time::Instant::now()); - apply_prune_removes(trie, entries, event_id_counter); + apply_prune_removes(trie, entries, pm, event_id_counter); } } @@ -623,6 +632,7 @@ impl KvIndexer { apply_prune_removes( &mut trie, entries, + pm, &mut event_id_counter, ); } @@ -1082,3 +1092,50 @@ mod fifo_reset_tests { assert_eq!(new_scores.scores.get(&worker), Some(&1)); } } + +#[cfg(test)] +mod prune_ordering_tests { + use super::*; + + #[tokio::test(start_paused = true)] + async fn direct_routing_refresh_invalidates_a_dequeued_expiration() { + let ttl = std::time::Duration::from_secs(10); + let manager = WorkerPruneManager::new(PruneConfig { ttl }); + let worker = WorkerWithDpRank::new(7, 0); + let hashes = vec![LocalBlockHash(1)]; + let sequence_hashes = crate::protocols::compute_seq_hash_for_block(&hashes); + let mut trie = RadixTree::new(); + let mut event_id = 0; + let route = || RoutingDecisionRequest { + worker, + local_hashes: hashes.clone(), + sequence_hashes: sequence_hashes.clone(), + }; + apply_routing_decision_with_prune_tracking( + &mut trie, + route(), + &Some(manager.clone()), + &mut event_id, + ); + tokio::time::advance(ttl).await; + let expired = manager.drain_due_and_pending(tokio::time::Instant::now()); + assert_eq!(expired.len(), 1); + apply_routing_decision_with_prune_tracking( + &mut trie, + route(), + &Some(manager.clone()), + &mut event_id, + ); + apply_prune_removes(&mut trie, expired, &manager, &mut event_id); + assert_eq!( + trie.find_matches(hashes.clone(), false).scores.get(&worker), + Some(&1) + ); + + tokio::time::advance(ttl).await; + let expired = manager.drain_due_and_pending(tokio::time::Instant::now()); + apply_prune_removes(&mut trie, expired, &manager, &mut event_id); + assert!(trie.find_matches(hashes, false).scores.is_empty()); + manager.shutdown(); + } +} diff --git a/lib/kv-router/src/indexer/lower_tier.rs b/lib/kv-router/src/indexer/lower_tier.rs index b485fcdad66a..50ddba18bfb6 100644 --- a/lib/kv-router/src/indexer/lower_tier.rs +++ b/lib/kv-router/src/indexer/lower_tier.rs @@ -1064,6 +1064,21 @@ impl SyncIndexer for LowerTierIndexer { } let _ = resp.send(applied); } + WorkerTask::ApproximateTtl(task) => task.apply(|event| { + let kind = EventKind::of(&event.event.data); + let result = self.apply_event(&mut worker_blocks, event); + let applied = result.is_ok(); + if result.is_err() { + tracing::warn!( + "Failed to apply approximate TTL event: {:?}", + result.as_ref().err() + ); + } + if let Some(ref counters) = counters { + counters.inc(kind, result); + } + applied + }), WorkerTask::ApproximateLru(task) => task.complete(Err(KvRouterError::Unsupported( "approximate LRU is not supported for lower-tier indexers".to_string(), ))), @@ -1096,14 +1111,26 @@ impl SyncIndexer for LowerTierIndexer { } } WorkerTask::RemoveWorker { - worker_id, resp, .. + worker_id, + prune_manager, + resp, + .. } => { + if let Some(manager) = prune_manager { + manager.remove_worker(worker_id); + } self.remove_worker(&mut worker_blocks, worker_id); let _ = resp.send(()); } WorkerTask::RemoveWorkerDpRank { - worker_id, dp_rank, .. + worker_id, + dp_rank, + prune_manager, + .. } => { + if let Some(manager) = prune_manager { + manager.remove_worker_dp_rank(WorkerWithDpRank::new(worker_id, dp_rank)); + } self.remove_worker_dp_rank(&mut worker_blocks, worker_id, dp_rank); } WorkerTask::DumpEvents(sender) => { diff --git a/lib/kv-router/src/indexer/positional.rs b/lib/kv-router/src/indexer/positional.rs index 274230c7d1c9..5b2d2095420b 100644 --- a/lib/kv-router/src/indexer/positional.rs +++ b/lib/kv-router/src/indexer/positional.rs @@ -253,6 +253,21 @@ impl SyncIndexer for PositionalIndexer { } let _ = resp.send(applied); } + WorkerTask::ApproximateTtl(task) => task.apply(|event| { + let kind = EventKind::of(&event.event.data); + let result = self.apply_event(&mut worker_blocks, event, counters.as_ref()); + let applied = result.is_ok(); + if result.is_err() { + tracing::warn!( + "Failed to apply approximate TTL event: {:?}", + result.as_ref().err() + ); + } + if let Some(ref counters) = counters { + counters.inc(kind, result); + } + applied + }), WorkerTask::ApproximateLru(task) => task.complete(Err(KvRouterError::Unsupported( "approximate LRU requires ConcurrentRadixTreeCompressed".to_string(), ))), @@ -285,14 +300,26 @@ impl SyncIndexer for PositionalIndexer { } } WorkerTask::RemoveWorker { - worker_id, resp, .. + worker_id, + prune_manager, + resp, + .. } => { + if let Some(manager) = prune_manager { + manager.remove_worker(worker_id); + } self.remove_worker_blocks_impl(&mut worker_blocks, worker_id); let _ = resp.send(()); } WorkerTask::RemoveWorkerDpRank { - worker_id, dp_rank, .. + worker_id, + dp_rank, + prune_manager, + .. } => { + if let Some(manager) = prune_manager { + manager.remove_worker_dp_rank(WorkerWithDpRank::new(worker_id, dp_rank)); + } self.remove_worker_dp_rank_impl(&mut worker_blocks, worker_id, dp_rank); } WorkerTask::CleanupStaleChildren => { diff --git a/lib/kv-router/src/indexer/pruning.rs b/lib/kv-router/src/indexer/pruning.rs index 29df35a88818..9f031a080e50 100644 --- a/lib/kv-router/src/indexer/pruning.rs +++ b/lib/kv-router/src/indexer/pruning.rs @@ -13,11 +13,14 @@ use std::sync::{Arc, Mutex}; use dashmap::DashMap; use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet}; -use tokio::sync::watch; +use tokio::sync::{oneshot, watch}; use tokio::time::{Duration, Instant}; use tokio_util::sync::CancellationToken; -use crate::protocols::{ExternalSequenceBlockHash, WorkerId, WorkerWithDpRank}; +use crate::protocols::{ + ExternalSequenceBlockHash, KvCacheEvent, KvCacheEventData, KvCacheRemoveData, RouterEvent, + WorkerId, WorkerWithDpRank, +}; const WORKER_EXPIRY_HEAP_REBUILD_THRESHOLD: usize = 10; /// Approximate TTL expirations are rounded up to this interval. A non-zero TTL @@ -222,25 +225,36 @@ fn round_up_duration(duration: Duration, bucket: Duration) -> Duration { #[derive(Debug)] struct WorkerPruneState { timers: PruneManager, + // Expiration authorizes a removal until a refresh or explicit reset revokes + // it. Keep this separate from the delivery queue: a removal may already be + // queued on an indexer worker when the same block is refreshed. + pending_expirations: FxHashSet, } impl WorkerPruneState { fn new(config: PruneConfig) -> Self { Self { timers: PruneManager::new(config), + pending_expirations: FxHashSet::default(), } } fn insert_block_entries(&mut self, entries: Vec, now: Instant) { + for entry in &entries { + self.pending_expirations.remove(entry); + } self.timers.insert_at(entries, now); } fn remove_block_entry(&mut self, entry: &BlockEntry) { self.timers.remove(entry); + self.pending_expirations.remove(entry); } fn pop_expired(&mut self, now: Instant) -> Vec { - self.timers.pop_expired(now) + let expired = self.timers.pop_expired(now); + self.pending_expirations.extend(expired.iter().copied()); + expired } fn peek_next_valid_expiry(&mut self) -> Option { @@ -474,6 +488,24 @@ impl WorkerPruneManager { .collect() } + /// Consume still-current expiration ownership immediately before applying + /// a remove. The caller must serialize this check and backend application + /// with successful stores and retention changes for the same worker. An old + /// delivery may consume a newer authorization only after the current copy + /// has independently expired again; that copy is then eligible for removal. + pub(crate) fn take_pending_expirations( + &self, + worker: WorkerWithDpRank, + mut entries: Vec, + ) -> Vec { + let Some(state) = self.inner.workers.get(&worker) else { + return Vec::new(); + }; + let mut state = state.lock().expect("worker prune state mutex poisoned"); + entries.retain(|entry| state.pending_expirations.remove(entry)); + entries + } + pub fn subscribe_ready(&self) -> watch::Receiver { self.inner.ready_tx.subscribe() } @@ -483,6 +515,90 @@ impl WorkerPruneManager { } } +/// TTL mutations applied on the same FIFO as the worker's backend events. +/// Registration follows successful application even if the caller drops its +/// acknowledgement; expiration is revalidated only when removal reaches the lane. +pub struct ApproximateTtlTask { + manager: WorkerPruneManager, + worker: WorkerWithDpRank, + entries: Vec, + operation: ApproximateTtlOperation, +} + +enum ApproximateTtlOperation { + Store { + event: RouterEvent, + response: oneshot::Sender, + }, + Remove { + event_id: u64, + }, +} + +impl ApproximateTtlTask { + pub(crate) fn store( + manager: WorkerPruneManager, + event: RouterEvent, + entries: Vec, + response: oneshot::Sender, + ) -> Self { + Self { + manager, + worker: WorkerWithDpRank::new(event.worker_id, event.event.dp_rank), + entries, + operation: ApproximateTtlOperation::Store { event, response }, + } + } + + pub(crate) fn remove( + manager: WorkerPruneManager, + worker: WorkerWithDpRank, + entries: Vec, + event_id: u64, + ) -> Self { + Self { + manager, + worker, + entries, + operation: ApproximateTtlOperation::Remove { event_id }, + } + } + + pub(crate) fn apply(self, mut apply_event: impl FnMut(RouterEvent) -> bool) { + match self.operation { + ApproximateTtlOperation::Store { event, response } => { + let applied = apply_event(event); + if applied { + self.manager + .insert_worker_block_entries(self.worker, self.entries); + } + let _ = response.send(applied); + } + ApproximateTtlOperation::Remove { event_id } => { + let mut entries = self + .manager + .take_pending_expirations(self.worker, self.entries); + if entries.is_empty() { + return; + } + entries.sort_unstable_by_key(|entry| entry.key); + entries.dedup_by_key(|entry| entry.key); + let event = RouterEvent::new( + self.worker.worker_id, + KvCacheEvent { + event_id, + data: KvCacheEventData::Removed(KvCacheRemoveData { + block_hashes: entries.into_iter().map(|entry| entry.key).collect(), + }), + dp_rank: self.worker.dp_rank, + }, + ); + apply_event(event); + } + } + } +} + impl WorkerPruneManagerInner { fn bump_sender(sender: &watch::Sender) { let next = sender.borrow().wrapping_add(1); @@ -658,6 +774,69 @@ mod tests { } } + #[rstest::rstest] + #[case::refresh("refresh")] + #[case::forget("forget")] + #[case::rank_reset("rank")] + #[case::worker_removal("worker")] + #[tokio::test(start_paused = true)] + async fn pending_expiration_is_revoked_after_dequeue(#[case] change: &str) { + let ttl = Duration::from_secs(10); + let manager = WorkerPruneManager::new(PruneConfig { ttl }); + let worker = WorkerWithDpRank::new(7, 0); + let block = test_block(worker, 42, 0); + manager.insert_block_entries(vec![block]); + time::advance(ttl).await; + let expired = manager.drain_due_and_pending(Instant::now()); + assert_eq!(expired, vec![block]); + + match change { + "refresh" => manager.insert_block_entries(vec![block]), + "forget" => manager.remove_block_entries(&[block]), + "rank" => manager.remove_worker_dp_rank(worker), + "worker" => manager.remove_worker(worker.worker_id), + _ => unreachable!(), + } + assert!(manager.take_pending_expirations(worker, expired).is_empty()); + + // Revocation must not prevent a later, independently authorized expiry. + manager.insert_block_entries(vec![block]); + time::advance(ttl).await; + let expired = manager.drain_due_and_pending(Instant::now()); + assert_eq!( + manager.take_pending_expirations(worker, expired.clone()), + vec![block] + ); + assert!(manager.take_pending_expirations(worker, expired).is_empty()); + manager.shutdown(); + } + + #[tokio::test(start_paused = true)] + async fn delayed_delivery_can_consume_a_later_legitimate_expiration_once() { + let ttl = Duration::from_secs(10); + let manager = WorkerPruneManager::new(PruneConfig { ttl }); + let worker = WorkerWithDpRank::new(7, 0); + let block = test_block(worker, 42, 0); + manager.insert_block_entries(vec![block]); + time::advance(ttl).await; + let old_delivery = manager.drain_due_and_pending(Instant::now()); + manager.insert_block_entries(vec![block]); + time::advance(ttl).await; + let new_delivery = manager.drain_due_and_pending(Instant::now()); + assert_eq!(old_delivery, vec![block]); + assert_eq!(new_delivery, vec![block]); + assert_eq!( + manager.take_pending_expirations(worker, old_delivery), + vec![block] + ); + assert!( + manager + .take_pending_expirations(worker, new_delivery) + .is_empty() + ); + manager.shutdown(); + } + /// Validate basic insert / expiry behaviour of [`PruneManager`]. #[test] fn test_prune_manager_expiry() { diff --git a/lib/kv-router/src/indexer/thread_pool.rs b/lib/kv-router/src/indexer/thread_pool.rs index 135115cd0876..d7d33551b4b3 100644 --- a/lib/kv-router/src/indexer/thread_pool.rs +++ b/lib/kv-router/src/indexer/thread_pool.rs @@ -27,7 +27,7 @@ use super::{ EventCompletionBuffer, EventCompletionWriter, ObservationError, ObservationSeal, ObservedEnqueueReceipt, ThreadPoolObservationPlan, ThreadPoolObservationSnapshot, }; -use crate::indexer::pruning::{BlockEntry, PruneConfig, WorkerPruneManager}; +use crate::indexer::pruning::{ApproximateTtlTask, BlockEntry, PruneConfig, WorkerPruneManager}; use crate::protocols::*; use crate::scheduling::AttemptId; use dynamo_tokens::SequenceHash; @@ -831,6 +831,7 @@ impl ThreadPoolIndexer { channel .send(WorkerTask::RemoveWorker { worker_id, + prune_manager: None, sweep_tree: false, resp: resp_tx, }) @@ -850,6 +851,7 @@ impl ThreadPoolIndexer { self.worker_event_channels[0] .send(WorkerTask::RemoveWorker { worker_id, + prune_manager: self.prune_manager.clone(), sweep_tree: true, resp: resp_tx, }) @@ -916,33 +918,26 @@ impl ThreadPoolIndexer { worker_assignment_count: &AtomicUsize, num_workers: usize, synthetic_event_id: &AtomicU64, + prune_manager: &WorkerPruneManager, entries: Vec, ) { - let mut by_worker: BTreeMap> = - BTreeMap::new(); + let mut by_worker: BTreeMap> = BTreeMap::new(); for entry in entries { - by_worker.entry(entry.worker).or_default().insert(entry.key); + by_worker.entry(entry.worker).or_default().push(entry); } - for (worker, hashes) in by_worker { + for (worker, entries) in by_worker { let event_id = Self::next_synthetic_event_id(synthetic_event_id); - let event = RouterEvent::new( - worker.worker_id, - KvCacheEvent { - event_id, - data: KvCacheEventData::Removed(KvCacheRemoveData { - block_hashes: hashes.into_iter().collect(), - }), - dp_rank: worker.dp_rank, - }, - ); + let task = ApproximateTtlTask::remove(prune_manager.clone(), worker, entries, event_id); let thread_idx = Self::get_or_assign_thread_idx( worker_assignments, worker_assignment_count, worker, num_workers, ); - if let Err(error) = worker_event_channels[thread_idx].send(WorkerTask::Event(event)) { + if let Err(error) = + worker_event_channels[thread_idx].send(WorkerTask::ApproximateTtl(task)) + { tracing::warn!( thread_idx, ?error, @@ -982,6 +977,7 @@ impl ThreadPoolIndexer { &worker_assignment_count, num_workers, &synthetic_event_id, + &prune_manager, entries, ); } @@ -1057,18 +1053,17 @@ impl ThreadPoolIndexer { let (resp_tx, resp_rx) = oneshot::channel(); self.worker_event_channels[thread_idx] - .send(WorkerTask::EventWithAck { + .send(WorkerTask::ApproximateTtl(ApproximateTtlTask::store( + prune_manager.clone(), event, - resp: resp_tx, - }) + prune_entries, + resp_tx, + ))) .map_err(|_| KvRouterError::IndexerOffline)?; - let applied = resp_rx + resp_rx .await .map_err(|_| KvRouterError::IndexerDroppedRequest)?; - if applied { - prune_manager.insert_worker_block_entries(worker, prune_entries); - } Ok(()) } @@ -1283,10 +1278,6 @@ impl KvIndexerInterface for ThreadPoolIndexer { } async fn remove_worker(&self, worker_id: WorkerId) { - if let Some(prune_manager) = &self.prune_manager { - prune_manager.remove_worker(worker_id); - } - if let Err(error) = self.remove_worker_across_lanes_and_wait(worker_id).await { tracing::error!(worker_id, %error, "Failed to remove worker across mutation lanes"); } @@ -1305,16 +1296,13 @@ impl KvIndexerInterface for ThreadPoolIndexer { return; } - if let Some(prune_manager) = &self.prune_manager { - prune_manager.remove_worker_dp_rank(WorkerWithDpRank::new(worker_id, dp_rank)); - } - let worker = WorkerWithDpRank::new(worker_id, dp_rank); self.enqueue_rank_removal( worker, WorkerTask::RemoveWorkerDpRank { worker_id, dp_rank, + prune_manager: self.prune_manager.clone(), sweep_tree: true, }, ); @@ -1333,10 +1321,6 @@ impl KvIndexerInterface for ThreadPoolIndexer { .await; } - if let Some(prune_manager) = &self.prune_manager { - prune_manager.remove_worker_dp_rank(WorkerWithDpRank::new(worker_id, dp_rank)); - } - let worker = WorkerWithDpRank::new(worker_id, dp_rank); let thread_idx = Self::get_or_assign_thread_idx( &self.worker_assignments, @@ -1348,6 +1332,7 @@ impl KvIndexerInterface for ThreadPoolIndexer { .send(WorkerTask::RemoveWorkerDpRank { worker_id, dp_rank, + prune_manager: self.prune_manager.clone(), sweep_tree: true, }) .map_err(|_| KvRouterError::IndexerOffline)?; @@ -1460,6 +1445,7 @@ impl KvIndexerInterface for ThreadPoolIndexer { &self.worker_assignment_count, self.num_workers, &self.synthetic_event_id, + prune_manager, entries, ); self.flush().await; @@ -1472,6 +1458,7 @@ impl KvIndexerInterface for ThreadPoolIndexer { &self.worker_assignment_count, self.num_workers, &self.synthetic_event_id, + prune_manager, entries, ); if has_entries { @@ -1517,6 +1504,384 @@ mod tests { indexer.worker_assignments.get(&worker).map(|entry| *entry) } + fn positional_with_delegate( + delegate: Arc, + ) -> crate::indexer::positional::PositionalIndexer { + crate::indexer::positional::PositionalIndexer::new_with_delegate( + 1, + crate::indexer::positional::SearchMode::Strided, + delegate, + ) + } + + // Pause one real mutation worker at a delegate callback, so a test can + // deterministically enqueue a store and a removal before its ack is read. + struct MutationGate { + hash: ExternalSequenceBlockHash, + entered: flume::Sender<()>, + release: flume::Receiver<()>, + } + + impl crate::indexer::KvIndexerDelegate for MutationGate { + fn on_create(&self, hash: ExternalSequenceBlockHash) { + if hash == self.hash { + let _ = self.entered.send(()); + // Bound the gate even if the test panics before releasing it. + let _ = self.release.recv_timeout(std::time::Duration::from_secs(5)); + } + } + + fn on_remove(&self, _hash: ExternalSequenceBlockHash) {} + } + + fn gated_ttl_indexer( + ttl: std::time::Duration, + make_backend: impl FnOnce(Arc) -> T, + ) -> (ThreadPoolIndexer, flume::Receiver<()>, flume::Sender<()>) { + let (entered_tx, entered_rx) = flume::bounded(1); + let (release_tx, release_rx) = flume::bounded(1); + let hash = crate::protocols::compute_seq_hash_for_block(&[LocalBlockHash(999)])[0]; + let backend = make_backend(Arc::new(MutationGate { + hash: ExternalSequenceBlockHash(hash), + entered: entered_tx, + release: release_rx, + })); + let indexer = ThreadPoolIndexer::new_with_pruning(backend, 1, 16, PruneConfig { ttl }); + // Tests deliver expired batches explicitly to control their FIFO position. + indexer.prune_pump_cancel.as_ref().unwrap().cancel(); + (indexer, entered_rx, release_tx) + } + + fn enqueue_expired(indexer: &ThreadPoolIndexer, entries: Vec) { + ThreadPoolIndexer::::enqueue_prune_removes( + &indexer.worker_event_channels, + &indexer.worker_assignments, + &indexer.worker_assignment_count, + indexer.num_workers, + &indexer.synthetic_event_id, + indexer.prune_manager.as_ref().unwrap(), + entries, + ); + } + + #[rstest::rstest] + #[case::compressed(ConcurrentRadixTreeCompressed::new_with_delegate)] + #[case::positional(positional_with_delegate)] + #[case::lower_tier(crate::indexer::LowerTierIndexer::new_with_delegate)] + #[tokio::test] + async fn ttl_refresh_fences_a_remove_queued_before_its_store_ack_is_read( + #[case] make_backend: fn(Arc) -> T, + ) { + tokio::time::timeout(std::time::Duration::from_secs(10), async { + let ttl = std::time::Duration::from_millis(100); + let (indexer, entered, release) = gated_ttl_indexer(ttl, make_backend); + let manager = indexer.prune_manager.as_ref().unwrap(); + let mut ready = manager.subscribe_ready(); + let worker = WorkerWithDpRank::new(7, 0); + let hashes = [LocalBlockHash(1)]; + let sequence_hashes = crate::protocols::compute_seq_hash_for_block(&hashes); + indexer + .record_ttl_fallback_hashes(worker, &hashes, &sequence_hashes) + .await + .unwrap(); + ready.changed().await.unwrap(); + // Wait for the real TTL notification, then control subsequent expiry + // delivery explicitly; no timing-sensitive wall-clock sleep is needed. + manager.shutdown(); + let expired = manager.drain_pending_removes(); + assert_eq!(expired.len(), 1); + + indexer.apply_event(make_store_event(99, &[999])).await; + entered.recv_async().await.unwrap(); + let mut refresh = + Box::pin(indexer.record_ttl_fallback_hashes(worker, &hashes, &sequence_hashes)); + assert!(futures_util::poll!(refresh.as_mut()).is_pending()); + enqueue_expired(&indexer, expired); + release.send(()).unwrap(); + // Store and old remove both run before the caller can register TTL. + indexer.flush().await; + refresh.await.unwrap(); + assert_score(&indexer, &[1], worker, 1).await; + + // The refresh must still expire normally under its own timer. + let expired = manager.drain_due_and_pending(tokio::time::Instant::now() + ttl * 2); + assert_eq!(expired.len(), 1); + enqueue_expired(&indexer, expired); + indexer.flush().await; + assert_no_scores(&indexer, &[1]).await; + }) + .await + .expect("TTL worker regression timed out"); + } + + #[rstest::rstest] + #[case::compressed(ConcurrentRadixTreeCompressed::new_with_delegate)] + #[case::positional(positional_with_delegate)] + #[case::lower_tier(crate::indexer::LowerTierIndexer::new_with_delegate)] + #[tokio::test] + async fn ttl_store_keeps_expiry_when_its_ack_receiver_is_cancelled( + #[case] make_backend: fn(Arc) -> T, + ) { + tokio::time::timeout(std::time::Duration::from_secs(10), async { + let ttl = std::time::Duration::from_secs(60); + let (indexer, entered, release) = gated_ttl_indexer(ttl, make_backend); + let manager = indexer.prune_manager.as_ref().unwrap(); + manager.shutdown(); + let worker = WorkerWithDpRank::new(7, 0); + let hashes = [LocalBlockHash(1)]; + let sequence_hashes = crate::protocols::compute_seq_hash_for_block(&hashes); + indexer.apply_event(make_store_event(99, &[999])).await; + entered.recv_async().await.unwrap(); + let mut store = + Box::pin(indexer.record_ttl_fallback_hashes(worker, &hashes, &sequence_hashes)); + assert!(futures_util::poll!(store.as_mut()).is_pending()); + drop(store); + release.send(()).unwrap(); + indexer.flush().await; + assert_score(&indexer, &[1], worker, 1).await; + + let expired = manager.drain_due_and_pending(tokio::time::Instant::now() + ttl * 2); + assert_eq!( + expired.len(), + 1, + "applied store must retain its expiry after caller cancellation" + ); + enqueue_expired(&indexer, expired); + indexer.flush().await; + assert_no_scores(&indexer, &[1]).await; + }) + .await + .expect("TTL worker regression timed out"); + } + + #[tokio::test] + async fn ttl_failed_refresh_does_not_revoke_the_previous_expiration() { + tokio::time::timeout(std::time::Duration::from_secs(10), async { + let ttl = std::time::Duration::from_millis(100); + let (indexer, _, _) = + gated_ttl_indexer(ttl, ConcurrentRadixTreeCompressed::new_with_delegate); + let manager = indexer.prune_manager.as_ref().unwrap(); + let mut ready = manager.subscribe_ready(); + let worker = WorkerWithDpRank::new(7, 0); + let hashes = [LocalBlockHash(1)]; + let sequence_hashes = crate::protocols::compute_seq_hash_for_block(&hashes); + indexer + .record_ttl_fallback_hashes(worker, &hashes, &sequence_hashes) + .await + .unwrap(); + ready.changed().await.unwrap(); + manager.shutdown(); + let expired = manager.drain_pending_removes(); + assert_eq!(expired.len(), 1); + + let mut event = + ThreadPoolIndexer::::stored_event_for_hashes( + worker, + &hashes, + &sequence_hashes, + 10, + ); + let KvCacheEventData::Stored(ref mut stored) = event.event.data else { + unreachable!() + }; + stored.parent_hash = Some(ExternalSequenceBlockHash(999)); + let (response, result) = oneshot::channel(); + indexer.worker_event_channels[0] + .send(WorkerTask::ApproximateTtl(ApproximateTtlTask::store( + manager.clone(), + event, + expired.clone(), + response, + ))) + .unwrap(); + assert!( + !result.await.unwrap(), + "missing parent must reject the store" + ); + + enqueue_expired(&indexer, expired); + indexer.flush().await; + assert_no_scores(&indexer, &[1]).await; + assert!( + manager + .drain_due_and_pending(tokio::time::Instant::now() + ttl * 2) + .is_empty() + ); + }) + .await + .expect("TTL worker regression timed out"); + } + + #[tokio::test] + async fn ttl_remove_cannot_erase_a_new_lru_incarnation() { + tokio::time::timeout(std::time::Duration::from_secs(10), async { + let ttl = std::time::Duration::from_millis(100); + let indexer = ThreadPoolIndexer::new_with_metrics_and_approximate_retention( + ConcurrentRadixTreeCompressed::new(), + 1, + 16, + None, + Some(ApproximateRetentionConfig::Lru { + fallback_ttl: PruneConfig { ttl }, + }), + ); + indexer.prune_pump_cancel.as_ref().unwrap().cancel(); + let manager = indexer.prune_manager.as_ref().unwrap(); + let mut ready = manager.subscribe_ready(); + let worker = WorkerWithDpRank::new(7, 0); + let hashes = [LocalBlockHash(1)]; + let sequence_hashes = crate::protocols::compute_seq_hash_for_block(&hashes); + let blocks = vec![crate::indexer::ApproximateLruBlock { + local_hash: hashes[0], + sequence_hash: sequence_hashes[0], + }]; + let first = 1; + indexer + .set_approximate_lru_capacity(worker, first, None) + .await + .unwrap(); + let lease = indexer + .begin_approximate_lru_request(worker, first, AttemptId::new(1)) + .unwrap(); + lease.acquire(blocks.clone(), 0).await.unwrap(); + ready.changed().await.unwrap(); + manager.shutdown(); + let expired = manager.drain_pending_removes(); + assert_eq!(expired.len(), 1); + + let second = 2; + indexer + .set_approximate_lru_capacity(worker, second, Some(2)) + .await + .unwrap(); + let lease = indexer + .begin_approximate_lru_request(worker, second, AttemptId::new(2)) + .unwrap(); + lease.acquire(blocks, 0).await.unwrap(); + enqueue_expired(&indexer, expired); + indexer.flush().await; + assert_score(&indexer, &[1], worker, 1).await; + assert_eq!( + indexer.approximate_lru_stats().await.unwrap().active_blocks, + 1 + ); + }) + .await + .expect("TTL worker regression timed out"); + } + + #[rstest::rstest] + #[case::whole_worker(true)] + #[case::rank_reset(false)] + #[tokio::test] + async fn ttl_remove_cannot_erase_content_restored_after_reset(#[case] whole_worker: bool) { + tokio::time::timeout(std::time::Duration::from_secs(10), async { + let ttl = std::time::Duration::from_millis(100); + let (indexer, _, _) = + gated_ttl_indexer(ttl, ConcurrentRadixTreeCompressed::new_with_delegate); + let manager = indexer.prune_manager.as_ref().unwrap(); + let mut ready = manager.subscribe_ready(); + let worker = WorkerWithDpRank::new(7, 0); + let hashes = [LocalBlockHash(1)]; + let sequence_hashes = crate::protocols::compute_seq_hash_for_block(&hashes); + indexer + .record_ttl_fallback_hashes(worker, &hashes, &sequence_hashes) + .await + .unwrap(); + ready.changed().await.unwrap(); + manager.shutdown(); + let expired = manager.drain_pending_removes(); + assert_eq!(expired.len(), 1); + + if whole_worker { + indexer.remove_worker(worker.worker_id).await; + } else { + indexer + .reset_worker_dp_rank_and_wait(worker.worker_id, worker.dp_rank) + .await + .unwrap(); + } + // Event-driven content restored after reset has no TTL ownership. + indexer + .apply_event(make_store_event(worker.worker_id, &[1])) + .await; + indexer.flush().await; + enqueue_expired(&indexer, expired); + indexer.flush().await; + assert_score(&indexer, &[1], worker, 1).await; + assert!( + manager + .drain_due_and_pending(tokio::time::Instant::now() + ttl * 2) + .is_empty() + ); + }) + .await + .expect("TTL worker regression timed out"); + } + + #[rstest::rstest] + #[case::compressed(ConcurrentRadixTreeCompressed::new_with_delegate)] + #[case::positional(positional_with_delegate)] + #[case::lower_tier(crate::indexer::LowerTierIndexer::new_with_delegate)] + #[tokio::test] + async fn ttl_store_queued_before_reset_cannot_leave_a_timer( + #[case] make_backend: fn(Arc) -> T, + #[values("worker", "rank", "rank_reset")] reset: &str, + ) { + tokio::time::timeout(std::time::Duration::from_secs(10), async { + let ttl = std::time::Duration::from_secs(60); + let (indexer, entered, release) = gated_ttl_indexer(ttl, make_backend); + let manager = indexer.prune_manager.as_ref().unwrap(); + manager.shutdown(); + let worker = WorkerWithDpRank::new(7, 0); + let hashes = [LocalBlockHash(1)]; + let sequence_hashes = crate::protocols::compute_seq_hash_for_block(&hashes); + indexer.apply_event(make_store_event(99, &[999])).await; + entered.recv_async().await.unwrap(); + let mut store = + Box::pin(indexer.record_ttl_fallback_hashes(worker, &hashes, &sequence_hashes)); + assert!(futures_util::poll!(store.as_mut()).is_pending()); + let mut removal = Box::pin(async { + match reset { + "worker" => indexer.remove_worker(worker.worker_id).await, + "rank" => { + indexer + .remove_worker_dp_rank(worker.worker_id, worker.dp_rank) + .await; + } + "rank_reset" => { + indexer + .reset_worker_dp_rank_and_wait(worker.worker_id, worker.dp_rank) + .await + .unwrap(); + } + _ => unreachable!(), + } + }); + let pending = futures_util::poll!(removal.as_mut()).is_pending(); + release.send(()).unwrap(); + store.await.unwrap(); + if pending { + removal.await; + } + indexer.flush().await; + assert_no_scores(&indexer, &[1]).await; + + // The retired store must not later expire unrelated event-driven + // content restored under the same worker and block identity. + indexer + .apply_event(make_store_event(worker.worker_id, &[1])) + .await; + indexer.flush().await; + let expired = manager.drain_due_and_pending(tokio::time::Instant::now() + ttl * 2); + enqueue_expired(&indexer, expired); + indexer.flush().await; + assert_score(&indexer, &[1], worker, 1).await; + }) + .await + .expect("TTL worker reset regression timed out"); + } + #[derive(Clone, Copy)] enum ColdRemoval { Worker, diff --git a/lib/kv-router/src/indexer/types.rs b/lib/kv-router/src/indexer/types.rs index edf72d2777f8..1c4db25be6d5 100644 --- a/lib/kv-router/src/indexer/types.rs +++ b/lib/kv-router/src/indexer/types.rs @@ -612,6 +612,7 @@ pub enum WorkerTask { resp: oneshot::Sender, }, ApproximateLru(super::ApproximateLruTask), + ApproximateTtl(super::pruning::ApproximateTtlTask), #[cfg(feature = "bench")] InstallObservation { writer: EventCompletionWriter, @@ -633,6 +634,8 @@ pub enum WorkerTask { /// Permanently remove a worker from tracking. RemoveWorker { worker_id: WorkerId, + /// Cleared on the final lane barrier after all retired stores have applied. + prune_manager: Option, /// True for the one shared-state backend task that owns structural cleanup. sweep_tree: bool, /// Acknowledges completion of this lane's cold-path removal phase. @@ -642,6 +645,8 @@ pub enum WorkerTask { RemoveWorkerDpRank { worker_id: WorkerId, dp_rank: DpRank, + /// Cleared on the same FIFO lane as successful TTL stores. + prune_manager: Option, /// True for the one shared-state backend task that owns structural cleanup. sweep_tree: bool, }, From 7086119ccc7ba18a231bd65c4084f3c12fb5a3cf Mon Sep 17 00:00:00 2001 From: bojiang-li <327132355+bojiang-li@users.noreply.github.com> Date: Thu, 1 Oct 2026 04:42:54 -0700 Subject: [PATCH 2/2] fix(kv-router): preserve worker API while ordering TTL mutations Signed-off-by: bojiang-li <327132355+bojiang-li@users.noreply.github.com> --- .../sync_impl.rs | 32 +- lib/kv-router/src/indexer/lower_tier.rs | 39 +- lib/kv-router/src/indexer/mod.rs | 2 + lib/kv-router/src/indexer/positional.rs | 39 +- lib/kv-router/src/indexer/pruning.rs | 99 +++- lib/kv-router/src/indexer/thread_pool.rs | 550 +++++++++++++++++- lib/kv-router/src/indexer/traits.rs | 14 + lib/kv-router/src/indexer/types.rs | 5 - lib/kv-router/src/indexer/worker_receiver.rs | 469 +++++++++++++++ 9 files changed, 1101 insertions(+), 148 deletions(-) create mode 100644 lib/kv-router/src/indexer/worker_receiver.rs diff --git a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/sync_impl.rs b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/sync_impl.rs index 6088b85e4287..60e9137f025e 100644 --- a/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/sync_impl.rs +++ b/lib/kv-router/src/indexer/concurrent_radix_tree_compressed/sync_impl.rs @@ -16,6 +16,15 @@ impl SyncIndexer for ConcurrentRadixTreeCompressed { &self, event_receiver: flume::Receiver, metrics: Option>, + ) -> anyhow::Result<()> { + self.worker_with_retention(event_receiver.into(), metrics) + } + + #[cfg_attr(feature = "profile", inline(never))] + fn worker_with_retention( + &self, + mut event_receiver: crate::indexer::WorkerTaskReceiver, + metrics: Option>, ) -> anyhow::Result<()> { let mut lookup = FxHashMap::default(); let counters = metrics.as_ref().map(|m| m.prebind()); @@ -47,21 +56,6 @@ impl SyncIndexer for ConcurrentRadixTreeCompressed { } let _ = resp.send(applied); } - WorkerTask::ApproximateTtl(task) => task.apply(|event| { - let kind = EventKind::of(&event.event.data); - let result = self.apply_event(&mut lookup, event, counters.as_ref()); - let applied = result.is_ok(); - if result.is_err() { - tracing::warn!( - "Failed to apply approximate TTL event: {:?}", - result.as_ref().err() - ); - } - if let Some(ref counters) = counters { - counters.inc(kind, result); - } - applied - }), WorkerTask::ApproximateLru(task) => { approximate_lru.observe_task(&task); let ApproximateLruTask { @@ -131,13 +125,9 @@ impl SyncIndexer for ConcurrentRadixTreeCompressed { } WorkerTask::RemoveWorker { worker_id, - prune_manager, sweep_tree, resp, } => { - if let Some(manager) = prune_manager { - manager.remove_worker(worker_id); - } approximate_lru.forget_worker(worker_id); self.erase_worker_coverage( &mut lookup, @@ -149,12 +139,8 @@ impl SyncIndexer for ConcurrentRadixTreeCompressed { WorkerTask::RemoveWorkerDpRank { worker_id, dp_rank, - prune_manager, sweep_tree, } => { - if let Some(manager) = prune_manager { - manager.remove_worker_dp_rank(WorkerWithDpRank::new(worker_id, dp_rank)); - } approximate_lru.forget_rank(WorkerWithDpRank::new(worker_id, dp_rank)); self.erase_worker_coverage( &mut lookup, diff --git a/lib/kv-router/src/indexer/lower_tier.rs b/lib/kv-router/src/indexer/lower_tier.rs index 50ddba18bfb6..509c3a4e618c 100644 --- a/lib/kv-router/src/indexer/lower_tier.rs +++ b/lib/kv-router/src/indexer/lower_tier.rs @@ -1034,6 +1034,14 @@ impl SyncIndexer for LowerTierIndexer { &self, event_receiver: flume::Receiver, metrics: Option>, + ) -> anyhow::Result<()> { + self.worker_with_retention(event_receiver.into(), metrics) + } + + fn worker_with_retention( + &self, + mut event_receiver: crate::indexer::WorkerTaskReceiver, + metrics: Option>, ) -> anyhow::Result<()> { let mut worker_blocks = WorkerBlockIndex::default(); let counters = metrics.as_ref().map(|m| m.prebind()); @@ -1064,21 +1072,6 @@ impl SyncIndexer for LowerTierIndexer { } let _ = resp.send(applied); } - WorkerTask::ApproximateTtl(task) => task.apply(|event| { - let kind = EventKind::of(&event.event.data); - let result = self.apply_event(&mut worker_blocks, event); - let applied = result.is_ok(); - if result.is_err() { - tracing::warn!( - "Failed to apply approximate TTL event: {:?}", - result.as_ref().err() - ); - } - if let Some(ref counters) = counters { - counters.inc(kind, result); - } - applied - }), WorkerTask::ApproximateLru(task) => task.complete(Err(KvRouterError::Unsupported( "approximate LRU is not supported for lower-tier indexers".to_string(), ))), @@ -1111,26 +1104,14 @@ impl SyncIndexer for LowerTierIndexer { } } WorkerTask::RemoveWorker { - worker_id, - prune_manager, - resp, - .. + worker_id, resp, .. } => { - if let Some(manager) = prune_manager { - manager.remove_worker(worker_id); - } self.remove_worker(&mut worker_blocks, worker_id); let _ = resp.send(()); } WorkerTask::RemoveWorkerDpRank { - worker_id, - dp_rank, - prune_manager, - .. + worker_id, dp_rank, .. } => { - if let Some(manager) = prune_manager { - manager.remove_worker_dp_rank(WorkerWithDpRank::new(worker_id, dp_rank)); - } self.remove_worker_dp_rank(&mut worker_blocks, worker_id, dp_rank); } WorkerTask::DumpEvents(sender) => { diff --git a/lib/kv-router/src/indexer/mod.rs b/lib/kv-router/src/indexer/mod.rs index 5a9e98205605..5029046ad0b5 100644 --- a/lib/kv-router/src/indexer/mod.rs +++ b/lib/kv-router/src/indexer/mod.rs @@ -62,6 +62,8 @@ fn warn_on_unit_block_size(indexer_type: &'static str, kv_block_size: u32) { ); } } +mod worker_receiver; +pub use worker_receiver::WorkerTaskReceiver; mod kv_indexer; mod local; mod lower_tier; diff --git a/lib/kv-router/src/indexer/positional.rs b/lib/kv-router/src/indexer/positional.rs index 5b2d2095420b..37865c610fd3 100644 --- a/lib/kv-router/src/indexer/positional.rs +++ b/lib/kv-router/src/indexer/positional.rs @@ -223,6 +223,14 @@ impl SyncIndexer for PositionalIndexer { &self, event_receiver: flume::Receiver, metrics: Option>, + ) -> anyhow::Result<()> { + self.worker_with_retention(event_receiver.into(), metrics) + } + + fn worker_with_retention( + &self, + mut event_receiver: crate::indexer::WorkerTaskReceiver, + metrics: Option>, ) -> anyhow::Result<()> { let mut worker_blocks = FxHashMap::default(); let counters = metrics.as_ref().map(|m| m.prebind()); @@ -253,21 +261,6 @@ impl SyncIndexer for PositionalIndexer { } let _ = resp.send(applied); } - WorkerTask::ApproximateTtl(task) => task.apply(|event| { - let kind = EventKind::of(&event.event.data); - let result = self.apply_event(&mut worker_blocks, event, counters.as_ref()); - let applied = result.is_ok(); - if result.is_err() { - tracing::warn!( - "Failed to apply approximate TTL event: {:?}", - result.as_ref().err() - ); - } - if let Some(ref counters) = counters { - counters.inc(kind, result); - } - applied - }), WorkerTask::ApproximateLru(task) => task.complete(Err(KvRouterError::Unsupported( "approximate LRU requires ConcurrentRadixTreeCompressed".to_string(), ))), @@ -300,26 +293,14 @@ impl SyncIndexer for PositionalIndexer { } } WorkerTask::RemoveWorker { - worker_id, - prune_manager, - resp, - .. + worker_id, resp, .. } => { - if let Some(manager) = prune_manager { - manager.remove_worker(worker_id); - } self.remove_worker_blocks_impl(&mut worker_blocks, worker_id); let _ = resp.send(()); } WorkerTask::RemoveWorkerDpRank { - worker_id, - dp_rank, - prune_manager, - .. + worker_id, dp_rank, .. } => { - if let Some(manager) = prune_manager { - manager.remove_worker_dp_rank(WorkerWithDpRank::new(worker_id, dp_rank)); - } self.remove_worker_dp_rank_impl(&mut worker_blocks, worker_id, dp_rank); } WorkerTask::CleanupStaleChildren => { diff --git a/lib/kv-router/src/indexer/pruning.rs b/lib/kv-router/src/indexer/pruning.rs index 9f031a080e50..03da75bf4891 100644 --- a/lib/kv-router/src/indexer/pruning.rs +++ b/lib/kv-router/src/indexer/pruning.rs @@ -518,7 +518,7 @@ impl WorkerPruneManager { /// TTL mutations applied on the same FIFO as the worker's backend events. /// Registration follows successful application even if the caller drops its /// acknowledgement; expiration is revalidated only when removal reaches the lane. -pub struct ApproximateTtlTask { +pub(crate) struct ApproximateTtlTask { manager: WorkerPruneManager, worker: WorkerWithDpRank, entries: Vec, @@ -535,6 +535,33 @@ enum ApproximateTtlOperation { }, } +pub(crate) struct PendingTtlStore { + manager: WorkerPruneManager, + worker: WorkerWithDpRank, + entries: Vec, + response: oneshot::Sender, +} + +pub(crate) enum PreparedTtlTask { + Store(RouterEvent, PendingTtlStore), + Remove(ApproximateTtlTask), +} + +impl PendingTtlStore { + pub(crate) fn complete( + self, + result: Result, + ) -> Result<(), super::KvRouterError> { + let applied = result?; + if applied { + self.manager + .insert_worker_block_entries(self.worker, self.entries); + } + let _ = self.response.send(applied); + Ok(()) + } +} + impl ApproximateTtlTask { pub(crate) fn store( manager: WorkerPruneManager, @@ -564,38 +591,46 @@ impl ApproximateTtlTask { } } - pub(crate) fn apply(self, mut apply_event: impl FnMut(RouterEvent) -> bool) { - match self.operation { - ApproximateTtlOperation::Store { event, response } => { - let applied = apply_event(event); - if applied { - self.manager - .insert_worker_block_entries(self.worker, self.entries); - } - let _ = response.send(applied); - } - ApproximateTtlOperation::Remove { event_id } => { - let mut entries = self - .manager - .take_pending_expirations(self.worker, self.entries); - if entries.is_empty() { - return; - } - entries.sort_unstable_by_key(|entry| entry.key); - entries.dedup_by_key(|entry| entry.key); - let event = RouterEvent::new( - self.worker.worker_id, - KvCacheEvent { - event_id, - data: KvCacheEventData::Removed(KvCacheRemoveData { - block_hashes: entries.into_iter().map(|entry| entry.key).collect(), - }), - dp_rank: self.worker.dp_rank, - }, - ); - apply_event(event); - } + pub(crate) fn prepare(self) -> PreparedTtlTask { + if !matches!(self.operation, ApproximateTtlOperation::Store { .. }) { + return PreparedTtlTask::Remove(self); + } + let ApproximateTtlOperation::Store { event, response } = self.operation else { + unreachable!() + }; + PreparedTtlTask::Store( + event, + PendingTtlStore { + manager: self.manager, + worker: self.worker, + entries: self.entries, + response, + }, + ) + } + + pub(crate) fn into_remove_event(self) -> Option { + let ApproximateTtlOperation::Remove { event_id } = self.operation else { + unreachable!("store must be extracted before removal") + }; + let mut entries = self + .manager + .take_pending_expirations(self.worker, self.entries); + if entries.is_empty() { + return None; } + entries.sort_unstable_by_key(|entry| entry.key); + entries.dedup_by_key(|entry| entry.key); + Some(RouterEvent::new( + self.worker.worker_id, + KvCacheEvent { + event_id, + data: KvCacheEventData::Removed(KvCacheRemoveData { + block_hashes: entries.into_iter().map(|entry| entry.key).collect(), + }), + dp_rank: self.worker.dp_rank, + }, + )) } } diff --git a/lib/kv-router/src/indexer/thread_pool.rs b/lib/kv-router/src/indexer/thread_pool.rs index d7d33551b4b3..b6b77439bde9 100644 --- a/lib/kv-router/src/indexer/thread_pool.rs +++ b/lib/kv-router/src/indexer/thread_pool.rs @@ -4,7 +4,7 @@ use std::{ collections::{BTreeMap, BTreeSet}, sync::{ - Arc, Mutex, + Arc, Mutex, OnceLock, atomic::{AtomicU64, AtomicUsize}, }, thread::JoinHandle, @@ -16,6 +16,7 @@ use rustc_hash::FxBuildHasher; use tokio::sync::oneshot; use super::concurrent_radix_tree_compressed::ConcurrentRadixTreeCompressed; +use super::worker_receiver::{OrderedTask, WorkerTaskReceiver}; use super::{ ApproximateLruClient, ApproximateLruCommandSink, ApproximateLruIncarnation, ApproximateLruLease, ApproximateLruStats, ApproximateLruTask, ApproximateRetentionConfig, @@ -73,7 +74,7 @@ pub struct ThreadPoolIndexer { /// Channels to send tasks to worker threads (one per thread). /// Sending `WorkerTask::Terminate` signals the thread to shut down. - worker_event_channels: Vec>, + worker_event_channels: Vec, /// Number of worker threads. num_workers: usize, @@ -99,8 +100,78 @@ pub struct ThreadPoolIndexer { observation_active: AtomicBool, } +/// Keep private retention envelopes out of the public worker protocol. +#[derive(Clone)] +enum WorkerTaskSender { + Direct(flume::Sender), + Retained { + sender: flume::Sender, + legacy_queue: Arc>>, + }, +} + +enum WorkerInput { + Direct(flume::Receiver), + Retained(WorkerTaskReceiver), +} + +impl WorkerTaskSender { + fn send(&self, task: WorkerTask) -> Result<(), KvRouterError> { + match self { + Self::Direct(sender) => sender.send(task).map_err(|_| KvRouterError::IndexerOffline), + Self::Retained { sender, .. } => sender + .send(OrderedTask::Task(task)) + .map_err(|_| KvRouterError::IndexerOffline), + } + } + + fn send_ttl(&self, task: ApproximateTtlTask) -> Result<(), KvRouterError> { + let Self::Retained { sender, .. } = self else { + unreachable!("TTL requires a retained lane") + }; + sender + .send(OrderedTask::Ttl(task)) + .map_err(|_| KvRouterError::IndexerOffline) + } + + fn send_removal( + &self, + task: WorkerTask, + manager: Option, + ) -> Result<(), KvRouterError> { + if let Some(manager) = manager { + let Self::Retained { sender, .. } = self else { + unreachable!("TTL requires a retained lane") + }; + sender + .send(OrderedTask::Removal { task, manager }) + .map_err(|_| KvRouterError::IndexerOffline) + } else { + self.send(task) + } + } + + fn len(&self) -> usize { + match self { + Self::Direct(sender) => sender.len(), + Self::Retained { + sender, + legacy_queue, + } => sender.len() + legacy_queue.get().map_or(0, flume::Sender::len), + } + } + + #[cfg(test)] + fn is_disconnected(&self) -> bool { + match self { + Self::Direct(sender) => sender.is_disconnected(), + Self::Retained { sender, .. } => sender.is_disconnected(), + } + } +} + struct ThreadPoolLruSink { - sender: flume::Sender, + sender: WorkerTaskSender, fallback_prune_manager: WorkerPruneManager, } @@ -271,8 +342,24 @@ impl ThreadPoolIndexer { let worker_assignment_count = Arc::new(AtomicUsize::new(0)); let synthetic_event_id = Arc::new(AtomicU64::new(0)); for worker_idx in 0..num_workers { - let (event_sender, event_receiver) = flume::unbounded::(); - worker_event_senders.push(event_sender); + let (sender, input) = if prune_config.is_some() { + let (sender, receiver) = flume::unbounded(); + let receiver = WorkerTaskReceiver::retained(receiver); + ( + WorkerTaskSender::Retained { + sender, + legacy_queue: receiver.legacy_queue(), + }, + WorkerInput::Retained(receiver), + ) + } else { + let (sender, receiver) = flume::unbounded(); + ( + WorkerTaskSender::Direct(sender), + WorkerInput::Direct(receiver), + ) + }; + worker_event_senders.push(sender); let backend = Arc::clone(&backend); let metrics = metrics.clone(); @@ -281,7 +368,13 @@ impl ThreadPoolIndexer { // This is observability, not recovery: if the worker panics, log // through tracing and then preserve the panic for join(). let panic_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - if let Err(error) = backend.worker(event_receiver, metrics) { + let result = match input { + WorkerInput::Direct(receiver) => backend.worker(receiver, metrics), + WorkerInput::Retained(receiver) => { + backend.worker_with_retention(receiver, metrics) + } + }; + if let Err(error) = result { tracing::error!( worker_thread_index = worker_idx, ?error, @@ -807,7 +900,9 @@ impl ThreadPoolIndexer { self.num_workers, ); - if let Err(error) = self.worker_event_channels[thread_idx].send(task) { + if let Err(error) = + self.worker_event_channels[thread_idx].send_removal(task, self.prune_manager.clone()) + { tracing::error!( worker_id = worker.worker_id, dp_rank = worker.dp_rank, @@ -831,7 +926,6 @@ impl ThreadPoolIndexer { channel .send(WorkerTask::RemoveWorker { worker_id, - prune_manager: None, sweep_tree: false, resp: resp_tx, }) @@ -849,12 +943,14 @@ impl ThreadPoolIndexer { // before this barrier completes belong to the retiring source and may be removed. let (resp_tx, resp_rx) = oneshot::channel(); self.worker_event_channels[0] - .send(WorkerTask::RemoveWorker { - worker_id, - prune_manager: self.prune_manager.clone(), - sweep_tree: true, - resp: resp_tx, - }) + .send_removal( + WorkerTask::RemoveWorker { + worker_id, + sweep_tree: true, + resp: resp_tx, + }, + self.prune_manager.clone(), + ) .map_err(|_| KvRouterError::IndexerOffline)?; resp_rx .await @@ -913,7 +1009,7 @@ impl ThreadPoolIndexer { } fn enqueue_prune_removes( - worker_event_channels: &[flume::Sender], + worker_event_channels: &[WorkerTaskSender], worker_assignments: &DashMap, worker_assignment_count: &AtomicUsize, num_workers: usize, @@ -935,9 +1031,7 @@ impl ThreadPoolIndexer { worker, num_workers, ); - if let Err(error) = - worker_event_channels[thread_idx].send(WorkerTask::ApproximateTtl(task)) - { + if let Err(error) = worker_event_channels[thread_idx].send_ttl(task) { tracing::warn!( thread_idx, ?error, @@ -949,7 +1043,7 @@ impl ThreadPoolIndexer { fn spawn_prune_pump( prune_manager: WorkerPruneManager, - worker_event_channels: Vec>, + worker_event_channels: Vec, worker_assignments: Arc>, worker_assignment_count: Arc, num_workers: usize, @@ -1053,12 +1147,12 @@ impl ThreadPoolIndexer { let (resp_tx, resp_rx) = oneshot::channel(); self.worker_event_channels[thread_idx] - .send(WorkerTask::ApproximateTtl(ApproximateTtlTask::store( + .send_ttl(ApproximateTtlTask::store( prune_manager.clone(), event, prune_entries, resp_tx, - ))) + )) .map_err(|_| KvRouterError::IndexerOffline)?; resp_rx @@ -1302,7 +1396,6 @@ impl KvIndexerInterface for ThreadPoolIndexer { WorkerTask::RemoveWorkerDpRank { worker_id, dp_rank, - prune_manager: self.prune_manager.clone(), sweep_tree: true, }, ); @@ -1329,12 +1422,14 @@ impl KvIndexerInterface for ThreadPoolIndexer { self.num_workers, ); self.worker_event_channels[thread_idx] - .send(WorkerTask::RemoveWorkerDpRank { - worker_id, - dp_rank, - prune_manager: self.prune_manager.clone(), - sweep_tree: true, - }) + .send_removal( + WorkerTask::RemoveWorkerDpRank { + worker_id, + dp_rank, + sweep_tree: true, + }, + self.prune_manager.clone(), + ) .map_err(|_| KvRouterError::IndexerOffline)?; self.flush_worker_lane_and_wait(worker).await @@ -1514,6 +1609,368 @@ mod tests { ) } + /// A downstream-style backend using only the original public worker API. + /// Its exhaustive match and removal constructors intentionally reject API + /// additions. It supports TTL events without implementing private LRU tasks. + struct LegacyIndexer { + blocks: Mutex< + rustc_hash::FxHashMap< + WorkerWithDpRank, + rustc_hash::FxHashMap, + >, + >, + delegate: Arc, + } + + impl LegacyIndexer { + fn new_with_delegate(delegate: Arc) -> Self { + Self { + blocks: Mutex::default(), + delegate, + } + } + + fn apply(&self, event: RouterEvent) { + let worker = WorkerWithDpRank::new(event.worker_id, event.event.dp_rank); + let mut blocks = self.blocks.lock().unwrap(); + match event.event.data { + KvCacheEventData::Stored(stored) => { + for block in stored.blocks { + self.delegate.on_create(block.block_hash); + blocks + .entry(worker) + .or_default() + .insert(block.block_hash, block.tokens_hash); + } + } + KvCacheEventData::Removed(removed) => { + if let Some(blocks) = blocks.get_mut(&worker) { + for hash in removed.block_hashes { + blocks.remove(&hash); + } + } + } + KvCacheEventData::Cleared => { + blocks.remove(&worker); + } + } + } + } + + impl SyncIndexer for LegacyIndexer { + fn worker( + &self, + events: flume::Receiver, + _metrics: Option>, + ) -> anyhow::Result<()> { + while let Ok(task) = events.recv() { + match task { + WorkerTask::Event(event) => self.apply(event), + WorkerTask::EventWithAck { event, resp } => { + self.apply(event); + let _ = resp.send(true); + } + WorkerTask::ApproximateLru(_) => anyhow::bail!("LRU unsupported"), + #[cfg(feature = "bench")] + WorkerTask::InstallObservation { .. } + | WorkerTask::ObservedEvent { .. } + | WorkerTask::SealObservation(_) + | WorkerTask::HarvestObservation(_) => anyhow::bail!("observation unsupported"), + WorkerTask::Anchor { .. } => anyhow::bail!("anchors unsupported"), + WorkerTask::RemoveWorker { + worker_id, + sweep_tree, + resp, + } => { + // Keep construction source-compatible as well as matching. + let WorkerTask::RemoveWorker { resp, .. } = (WorkerTask::RemoveWorker { + worker_id, + sweep_tree, + resp, + }) else { + unreachable!() + }; + self.blocks + .lock() + .unwrap() + .retain(|worker, _| worker.worker_id != worker_id); + let _ = resp.send(()); + } + WorkerTask::RemoveWorkerDpRank { + worker_id, + dp_rank, + sweep_tree, + } => { + let _ = WorkerTask::RemoveWorkerDpRank { + worker_id, + dp_rank, + sweep_tree, + }; + self.blocks + .lock() + .unwrap() + .remove(&WorkerWithDpRank::new(worker_id, dp_rank)); + } + WorkerTask::CleanupStaleChildren => {} + WorkerTask::DumpEvents(resp) => { + let _ = resp.send(Ok(Vec::new())); + } + WorkerTask::Stats(resp) => { + let _ = resp.send(WorkerLookupStats::from_worker_block_counts( + self.blocks + .lock() + .unwrap() + .iter() + .map(|(worker, blocks)| (*worker, blocks.len())), + )); + } + WorkerTask::ContainsWorkerBlock { + worker, + block_hash, + resp, + } => { + let _ = resp.send( + self.blocks + .lock() + .unwrap() + .get(&worker) + .is_some_and(|blocks| blocks.contains_key(&block_hash)), + ); + } + WorkerTask::Flush(resp) => { + let _ = resp.send(()); + } + WorkerTask::Terminate => break, + } + } + Ok(()) + } + + fn find_matches(&self, sequence: &[LocalBlockHash], _early_exit: bool) -> OverlapScores { + let mut result = OverlapScores::default(); + for (worker, blocks) in self.blocks.lock().unwrap().iter() { + let matched = sequence + .iter() + .take_while(|hash| blocks.values().any(|stored| stored == *hash)) + .count(); + if matched > 0 { + result.scores.insert(*worker, matched as u32); + } + } + result + } + } + + struct FailingWorker(&'static str); + + struct ExitingRetentionWorker(Option); + + impl SyncIndexer for ExitingRetentionWorker { + fn worker( + &self, + _events: flume::Receiver, + _metrics: Option>, + ) -> anyhow::Result<()> { + anyhow::bail!("retention must use the inline receiver") + } + + fn worker_with_retention( + &self, + mut events: WorkerTaskReceiver, + _metrics: Option>, + ) -> anyhow::Result<()> { + let WorkerTask::EventWithAck { resp, .. } = events.recv()? else { + anyhow::bail!("expected a store"); + }; + if let Some(applied) = self.0 { + let _ = resp.send(applied); + } + // Return before the next recv: the receiver must settle a ready ack + // from Drop, including when its original caller was cancelled. + Ok(()) + } + + fn find_matches(&self, _sequence: &[LocalBlockHash], _early_exit: bool) -> OverlapScores { + OverlapScores::default() + } + } + + #[rstest::rstest] + #[case::applied(Some(true))] + #[case::rejected(Some(false))] + #[case::unacknowledged(None)] + #[tokio::test] + async fn ttl_inline_receiver_settles_an_ack_on_early_worker_return( + #[case] applied: Option, + #[values(false, true)] cancelled: bool, + ) { + tokio::time::timeout(std::time::Duration::from_secs(5), async { + let indexer = ThreadPoolIndexer::new_with_pruning( + ExitingRetentionWorker(applied), + 1, + 16, + PruneConfig::default(), + ); + let hashes = [LocalBlockHash(1)]; + let sequence = crate::protocols::compute_seq_hash_for_block(&hashes); + let mut request = Box::pin(indexer.record_ttl_fallback_hashes( + WorkerWithDpRank::new(7, 0), + &hashes, + &sequence, + )); + let result = futures_util::poll!(request.as_mut()); + if cancelled { + drop(request); + } else { + let result = match result { + std::task::Poll::Ready(result) => result, + std::task::Poll::Pending => request.await, + }; + assert_eq!(result.is_ok(), applied.is_some()); + } + while !indexer.worker_event_channels[0].is_disconnected() { + tokio::task::yield_now().await; + } + assert_eq!( + indexer + .prune_manager + .as_ref() + .unwrap() + .drain_due_and_pending( + tokio::time::Instant::now() + std::time::Duration::from_secs(300) + ) + .len(), + usize::from(applied == Some(true)) + ); + }) + .await + .expect("inline receiver early return timed out"); + } + + impl SyncIndexer for FailingWorker { + fn worker( + &self, + events: flume::Receiver, + _metrics: Option>, + ) -> anyhow::Result<()> { + while let Ok(task) = events.recv() { + match task { + WorkerTask::Event(_) if self.0 == "idle_exit" => return Ok(()), + WorkerTask::EventWithAck { resp, .. } => match self.0 { + "panic" => panic!("injected backend failure"), + "error" => anyhow::bail!("injected backend error"), + "dropped_ack" => drop(resp), + _ => unreachable!(), + }, + WorkerTask::Terminate => break, + _ => anyhow::bail!("unexpected task"), + } + } + Ok(()) + } + + fn find_matches(&self, _sequence: &[LocalBlockHash], _early_exit: bool) -> OverlapScores { + OverlapScores::default() + } + } + + #[rstest::rstest] + #[case::panic("panic")] + #[case::error("error")] + #[case::dropped_ack("dropped_ack")] + #[case::idle_exit("idle_exit")] + fn ttl_legacy_adapter_releases_waiters_and_joins_after_backend_failure( + #[case] failure: &'static str, + ) { + // Run the single-thread runtime behind an OS-thread deadline: a broken + // synchronous Drop cannot be interrupted by a timer on the blocked runtime. + let (done, result) = flume::bounded(1); + let thread = std::thread::spawn(move || { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(async { + let indexer = ThreadPoolIndexer::new_with_pruning( + FailingWorker(failure), + 1, + 16, + PruneConfig::default(), + ); + let worker = WorkerWithDpRank::new(7, 0); + if failure == "idle_exit" { + indexer.apply_event(make_store_event(7, &[1])).await; + tokio::time::timeout(std::time::Duration::from_secs(2), async { + while !indexer.worker_event_channels[0].is_disconnected() { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + } + let result = indexer + .process_routing_decision_with_hashes( + worker, + vec![LocalBlockHash(1)], + vec![1], + ) + .await; + if failure == "idle_exit" { + assert!(matches!(result, Err(KvRouterError::IndexerOffline))); + } else { + assert!(matches!(result, Err(KvRouterError::IndexerDroppedRequest))); + } + assert!( + indexer + .prune_manager + .as_ref() + .unwrap() + .drain_due_and_pending( + tokio::time::Instant::now() + std::time::Duration::from_secs(300) + ) + .is_empty() + ); + drop(indexer); + }); + done.send(()).unwrap(); + }); + result + .recv_timeout(std::time::Duration::from_secs(5)) + .unwrap(); + thread.join().unwrap(); + } + + #[rstest::rstest] + #[case::drop(false)] + #[case::shutdown(true)] + fn ttl_workers_join_on_a_single_thread_runtime(#[case] explicit_shutdown: bool) { + let (done, result) = flume::bounded(1); + let thread = std::thread::spawn(move || { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(async { + let indexer = ThreadPoolIndexer::new_with_pruning( + ConcurrentRadixTreeCompressed::new(), + 2, + 16, + PruneConfig::default(), + ); + indexer.apply_event(make_store_event(7, &[1])).await; + if explicit_shutdown { + indexer.shutdown(); + } + drop(indexer); + }); + done.send(()).unwrap(); + }); + result + .recv_timeout(std::time::Duration::from_secs(5)) + .unwrap(); + thread.join().unwrap(); + } + // Pause one real mutation worker at a delegate callback, so a test can // deterministically enqueue a store and a removal before its ack is read. struct MutationGate { @@ -1564,10 +2021,41 @@ mod tests { ); } + #[tokio::test] + async fn ttl_legacy_queue_depth_includes_forwarded_backend_tasks() { + tokio::time::timeout(std::time::Duration::from_secs(5), async { + let (indexer, entered, release) = gated_ttl_indexer( + std::time::Duration::from_secs(60), + LegacyIndexer::new_with_delegate, + ); + indexer.apply_event(make_store_event(99, &[999])).await; + entered.recv_async().await.unwrap(); + for key in 0..16 { + indexer.apply_event(make_store_event(7, &[key])).await; + } + let WorkerTaskSender::Retained { + sender, + legacy_queue, + } = &indexer.worker_event_channels[0] + else { + unreachable!() + }; + while !sender.is_empty() || legacy_queue.get().unwrap().len() != 16 { + tokio::task::yield_now().await; + } + assert_eq!(indexer.worker_event_channels[0].len(), 16); + release.send(()).unwrap(); + indexer.flush().await; + }) + .await + .expect("legacy queue observation timed out"); + } + #[rstest::rstest] #[case::compressed(ConcurrentRadixTreeCompressed::new_with_delegate)] #[case::positional(positional_with_delegate)] #[case::lower_tier(crate::indexer::LowerTierIndexer::new_with_delegate)] + #[case::legacy(LegacyIndexer::new_with_delegate)] #[tokio::test] async fn ttl_refresh_fences_a_remove_queued_before_its_store_ack_is_read( #[case] make_backend: fn(Arc) -> T, @@ -1618,6 +2106,7 @@ mod tests { #[case::compressed(ConcurrentRadixTreeCompressed::new_with_delegate)] #[case::positional(positional_with_delegate)] #[case::lower_tier(crate::indexer::LowerTierIndexer::new_with_delegate)] + #[case::legacy(LegacyIndexer::new_with_delegate)] #[tokio::test] async fn ttl_store_keeps_expiry_when_its_ack_receiver_is_cancelled( #[case] make_backend: fn(Arc) -> T, @@ -1687,12 +2176,12 @@ mod tests { stored.parent_hash = Some(ExternalSequenceBlockHash(999)); let (response, result) = oneshot::channel(); indexer.worker_event_channels[0] - .send(WorkerTask::ApproximateTtl(ApproximateTtlTask::store( + .send_ttl(ApproximateTtlTask::store( manager.clone(), event, expired.clone(), response, - ))) + )) .unwrap(); assert!( !result.await.unwrap(), @@ -1823,6 +2312,7 @@ mod tests { #[case::compressed(ConcurrentRadixTreeCompressed::new_with_delegate)] #[case::positional(positional_with_delegate)] #[case::lower_tier(crate::indexer::LowerTierIndexer::new_with_delegate)] + #[case::legacy(LegacyIndexer::new_with_delegate)] #[tokio::test] async fn ttl_store_queued_before_reset_cannot_leave_a_timer( #[case] make_backend: fn(Arc) -> T, diff --git a/lib/kv-router/src/indexer/traits.rs b/lib/kv-router/src/indexer/traits.rs index b193c50605c4..2cc46a481fb1 100644 --- a/lib/kv-router/src/indexer/traits.rs +++ b/lib/kv-router/src/indexer/traits.rs @@ -216,6 +216,20 @@ pub trait SyncIndexer: Send + Sync + 'static { metrics: Option>, ) -> anyhow::Result<()>; + /// Process a retention-aware task stream without changing the public worker protocol. + /// + /// The default keeps existing implementations source-compatible by forwarding + /// original `WorkerTask` values through one additional dispatch thread. Backends + /// can override this and use `receiver.recv()` directly in their blocking worker + /// loop to perform retention bookkeeping inline, without the extra thread/hop. + fn worker_with_retention( + &self, + receiver: super::WorkerTaskReceiver, + metrics: Option>, + ) -> anyhow::Result<()> { + receiver.run_legacy(|receiver| self.worker(receiver, metrics)) + } + /// Find matches for a sequence of block hashes. fn find_matches(&self, sequence: &[LocalBlockHash], early_exit: bool) -> OverlapScores; diff --git a/lib/kv-router/src/indexer/types.rs b/lib/kv-router/src/indexer/types.rs index 1c4db25be6d5..edf72d2777f8 100644 --- a/lib/kv-router/src/indexer/types.rs +++ b/lib/kv-router/src/indexer/types.rs @@ -612,7 +612,6 @@ pub enum WorkerTask { resp: oneshot::Sender, }, ApproximateLru(super::ApproximateLruTask), - ApproximateTtl(super::pruning::ApproximateTtlTask), #[cfg(feature = "bench")] InstallObservation { writer: EventCompletionWriter, @@ -634,8 +633,6 @@ pub enum WorkerTask { /// Permanently remove a worker from tracking. RemoveWorker { worker_id: WorkerId, - /// Cleared on the final lane barrier after all retired stores have applied. - prune_manager: Option, /// True for the one shared-state backend task that owns structural cleanup. sweep_tree: bool, /// Acknowledges completion of this lane's cold-path removal phase. @@ -645,8 +642,6 @@ pub enum WorkerTask { RemoveWorkerDpRank { worker_id: WorkerId, dp_rank: DpRank, - /// Cleared on the same FIFO lane as successful TTL stores. - prune_manager: Option, /// True for the one shared-state backend task that owns structural cleanup. sweep_tree: bool, }, diff --git a/lib/kv-router/src/indexer/worker_receiver.rs b/lib/kv-router/src/indexer/worker_receiver.rs new file mode 100644 index 000000000000..3001622d6405 --- /dev/null +++ b/lib/kv-router/src/indexer/worker_receiver.rs @@ -0,0 +1,469 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::{ + sync::{Arc, OnceLock}, + thread::JoinHandle, +}; + +use tokio::sync::oneshot; + +use super::pruning::{ApproximateTtlTask, PendingTtlStore, PreparedTtlTask, WorkerPruneManager}; +use super::{KvRouterError, WorkerTask}; +use crate::protocols::{WorkerId, WorkerWithDpRank}; + +pub(super) enum OrderedTask { + Task(WorkerTask), + Ttl(ApproximateTtlTask), + Removal { + task: WorkerTask, + manager: WorkerPruneManager, + }, +} + +enum Source { + Plain(flume::Receiver), + Retained(flume::Receiver), +} + +enum Pending { + Store(PendingTtlStore, oneshot::Receiver), + WorkerRemoval { + manager: WorkerPruneManager, + worker_id: WorkerId, + result: oneshot::Receiver<()>, + response: oneshot::Sender<()>, + }, + RankRemoval { + manager: WorkerPruneManager, + worker: WorkerWithDpRank, + }, +} + +/// A blocking worker receiver that preserves the public [`WorkerTask`] protocol. +/// +/// Retention bookkeeping runs between tasks on the mutation worker: successful +/// stores register TTL before their callers are acknowledged, and expiry checks +/// run immediately before removal. A backend must finish each returned task +/// before calling [`recv`](Self::recv) again. Use this receiver in +/// [`SyncIndexer::worker_with_retention`](super::SyncIndexer::worker_with_retention) +/// to avoid the compatibility dispatch thread used by that method's default. +/// Normal event-driven receivers can be converted with [`From`]. +pub struct WorkerTaskReceiver { + source: Source, + pending: Option, + // Only the default adapter needs backend barriers and an independent stop + // signal. Builtin workers execute the returned task themselves. + legacy_backend: Option>, + legacy_queue: Arc>>, + stop: Option>, +} + +impl From> for WorkerTaskReceiver { + fn from(receiver: flume::Receiver) -> Self { + Self { + source: Source::Plain(receiver), + pending: None, + legacy_backend: None, + legacy_queue: Arc::default(), + stop: None, + } + } +} + +impl WorkerTaskReceiver { + pub(super) fn retained(receiver: flume::Receiver) -> Self { + Self { + source: Source::Retained(receiver), + pending: None, + legacy_backend: None, + legacy_queue: Arc::default(), + stop: None, + } + } + + pub(super) fn legacy_queue(&self) -> Arc>> { + Arc::clone(&self.legacy_queue) + } + + fn receive_ack(&self, mut result: oneshot::Receiver) -> Result { + if self.legacy_backend.is_some() { + result + .blocking_recv() + .map_err(|_| KvRouterError::IndexerDroppedRequest) + } else { + // The inline worker has already finished the previous task. An + // absent acknowledgement is a broken contract, not a reason to wait + // for the same worker to make progress while it is inside recv. + result + .try_recv() + .map_err(|_| KvRouterError::IndexerDroppedRequest) + } + } + + fn legacy_barrier(&self) -> Result<(), KvRouterError> { + if let Some(backend) = &self.legacy_backend { + let (response, result) = oneshot::channel(); + backend + .send(WorkerTask::Flush(response)) + .map_err(|_| KvRouterError::IndexerOffline)?; + result + .blocking_recv() + .map_err(|_| KvRouterError::IndexerDroppedRequest)?; + } + Ok(()) + } + + fn complete_pending(&mut self) -> Result<(), KvRouterError> { + match self.pending.take() { + Some(Pending::Store(completion, result)) => { + completion.complete(self.receive_ack(result))? + } + Some(Pending::WorkerRemoval { + manager, + worker_id, + result, + response, + }) => { + self.receive_ack(result)?; + manager.remove_worker(worker_id); + let _ = response.send(()); + } + Some(Pending::RankRemoval { manager, worker }) => { + self.legacy_barrier()?; + manager.remove_worker_dp_rank(worker); + } + None => {} + } + Ok(()) + } + + /// Wait for the next task, completing retention bookkeeping for the previous + /// task first. Call this only from the backend's blocking mutation loop. + /// Disconnection or a missing required acknowledgement stops the receiver. + pub fn recv(&mut self) -> Result { + self.complete_pending()?; + loop { + let message = match &self.source { + Source::Plain(receiver) => { + return receiver.recv().map_err(|_| KvRouterError::IndexerOffline); + } + Source::Retained(receiver) => { + if let Some(stop) = &self.stop { + flume::Selector::new() + .recv(receiver, |result| result.ok()) + .recv(stop, |_| None) + .wait() + .ok_or(KvRouterError::IndexerOffline)? + } else { + receiver.recv().map_err(|_| KvRouterError::IndexerOffline)? + } + } + }; + match message { + OrderedTask::Task(task) => return Ok(task), + OrderedTask::Ttl(task) => match task.prepare() { + PreparedTtlTask::Store(event, completion) => { + let (response, result) = oneshot::channel(); + self.pending = Some(Pending::Store(completion, result)); + return Ok(WorkerTask::EventWithAck { + event, + resp: response, + }); + } + PreparedTtlTask::Remove(task) => { + self.legacy_barrier()?; + if let Some(event) = task.into_remove_event() { + return Ok(WorkerTask::Event(event)); + } + } + }, + OrderedTask::Removal { task, manager } => match task { + WorkerTask::RemoveWorker { + worker_id, + sweep_tree, + resp, + } => { + let (response, result) = oneshot::channel(); + self.pending = Some(Pending::WorkerRemoval { + manager, + worker_id, + result, + response: resp, + }); + return Ok(WorkerTask::RemoveWorker { + worker_id, + sweep_tree, + resp: response, + }); + } + WorkerTask::RemoveWorkerDpRank { + worker_id, + dp_rank, + sweep_tree, + } => { + self.pending = Some(Pending::RankRemoval { + manager, + worker: WorkerWithDpRank::new(worker_id, dp_rank), + }); + return Ok(WorkerTask::RemoveWorkerDpRank { + worker_id, + dp_rank, + sweep_tree, + }); + } + _ => unreachable!("only removals carry TTL cleanup"), + }, + } + } + } + + pub(super) fn run_legacy( + mut self, + worker: impl FnOnce(flume::Receiver) -> anyhow::Result<()>, + ) -> anyhow::Result<()> { + if let Source::Plain(receiver) = &self.source { + // A caller can explicitly pass a converted ordinary receiver to the + // provided trait method. No retention coordination is needed, and + // an idle plain recv has no adapter stop channel to wake on return. + return worker(receiver.clone()); + } + let (backend, receiver) = flume::unbounded(); + let (stop, stopped) = flume::bounded(1); + self.legacy_backend = Some(backend.clone()); + let _ = self.legacy_queue.set(backend.clone()); + self.stop = Some(stopped); + let thread = std::thread::spawn(move || { + let _terminate = TerminateBackendOnDrop(backend.clone()); + while let Ok(task) = self.recv() { + let terminate = matches!(task, WorkerTask::Terminate); + if backend.send(task).is_err() || terminate { + break; + } + } + }); + // Joining from Drop also covers a backend panic. Its receiver and any + // in-flight ack sender unwind before this guard stops the adapter. + let _bridge = LegacyBridge { + stop, + thread: Some(thread), + }; + worker(receiver) + } +} + +impl Drop for WorkerTaskReceiver { + fn drop(&mut self) { + // A backend may acknowledge its last mutation and return without another + // recv. Settle ready acknowledgements, but never wait in Drop: the sender + // might still be held by the exiting worker or another task. + match self.pending.take() { + Some(Pending::Store(completion, mut result)) => { + let _ = completion.complete( + result + .try_recv() + .map_err(|_| KvRouterError::IndexerDroppedRequest), + ); + } + Some(Pending::WorkerRemoval { + manager, + worker_id, + mut result, + response, + }) => { + if result.try_recv().is_ok() { + manager.remove_worker(worker_id); + let _ = response.send(()); + } + } + // Rank removal has no public acknowledgement. If a worker exits + // before the next recv, completion is unproven; leave its metadata + // intact. The disconnected lane cannot accept restored coverage. + Some(Pending::RankRemoval { .. }) | None => {} + } + } +} + +struct LegacyBridge { + stop: flume::Sender<()>, + thread: Option>, +} + +impl Drop for LegacyBridge { + fn drop(&mut self) { + let _ = self.stop.send(()); + if let Some(thread) = self.thread.take() { + let _ = thread.join(); + } + } +} + +struct TerminateBackendOnDrop(flume::Sender); + +impl Drop for TerminateBackendOnDrop { + fn drop(&mut self) { + let _ = self.0.send(WorkerTask::Terminate); + } +} + +#[cfg(test)] +mod tests { + use super::super::pruning::{BlockEntry, PruneConfig}; + use super::*; + use crate::protocols::ExternalSequenceBlockHash; + use crate::test_utils::make_store_event; + + fn block(worker: WorkerWithDpRank) -> BlockEntry { + BlockEntry { + key: ExternalSequenceBlockHash(1), + worker, + seq_position: 0, + } + } + + #[test] + fn legacy_adapter_returns_with_an_open_plain_input_channel() { + let (sender, tasks) = flume::unbounded(); + let (done, result) = flume::bounded(1); + let thread = std::thread::spawn(move || { + WorkerTaskReceiver::from(tasks) + .run_legacy(|_| Ok(())) + .unwrap(); + let _ = done.send(()); + }); + result + .recv_timeout(std::time::Duration::from_secs(2)) + .unwrap(); + thread.join().unwrap(); + assert!(sender.is_disconnected()); + } + + #[tokio::test] + async fn dropping_ttl_receiver_does_not_wait_for_an_outstanding_store_ack() { + let manager = WorkerPruneManager::new(PruneConfig::default()); + manager.shutdown(); + let worker = WorkerWithDpRank::new(7, 0); + let (sender, tasks) = flume::unbounded(); + let (response, caller) = oneshot::channel(); + sender + .send(OrderedTask::Ttl(ApproximateTtlTask::store( + manager.clone(), + make_store_event(7, &[1]), + vec![block(worker)], + response, + ))) + .unwrap(); + let mut receiver = WorkerTaskReceiver::retained(tasks); + let WorkerTask::EventWithAck { resp, .. } = receiver.recv().unwrap() else { + panic!("expected store"); + }; + let (done, dropped) = oneshot::channel(); + let thread = std::thread::spawn(move || { + drop(receiver); + let _ = done.send(()); + }); + tokio::time::timeout(std::time::Duration::from_secs(2), dropped) + .await + .unwrap() + .unwrap(); + thread.join().unwrap(); + assert!(caller.await.is_err()); + assert!( + resp.send(true).is_err(), + "late ack must not outlive receiver ownership" + ); + assert!( + manager + .drain_due_and_pending( + tokio::time::Instant::now() + std::time::Duration::from_secs(300) + ) + .is_empty() + ); + } + + #[rstest::rstest] + #[case::acknowledged(true)] + #[case::unacknowledged(false)] + #[tokio::test] + async fn dropping_ttl_receiver_settles_only_acknowledged_worker_removal( + #[case] acknowledged: bool, + ) { + let manager = WorkerPruneManager::new(PruneConfig::default()); + manager.shutdown(); + let worker = WorkerWithDpRank::new(7, 0); + manager.insert_block_entries(vec![block(worker)]); + let (sender, tasks) = flume::unbounded(); + let (response, caller) = oneshot::channel(); + sender + .send(OrderedTask::Removal { + task: WorkerTask::RemoveWorker { + worker_id: 7, + sweep_tree: true, + resp: response, + }, + manager: manager.clone(), + }) + .unwrap(); + let mut receiver = WorkerTaskReceiver::retained(tasks); + let WorkerTask::RemoveWorker { resp, .. } = receiver.recv().unwrap() else { + panic!("expected worker removal"); + }; + if acknowledged { + resp.send(()).unwrap(); + } else { + drop(resp); + } + drop(receiver); + assert_eq!(caller.await.is_ok(), acknowledged); + assert_eq!( + manager + .drain_due_and_pending( + tokio::time::Instant::now() + std::time::Duration::from_secs(300) + ) + .len(), + usize::from(!acknowledged) + ); + } + + #[tokio::test] + async fn dropping_ttl_receiver_does_not_assume_rank_removal_completed() { + let manager = WorkerPruneManager::new(PruneConfig::default()); + manager.shutdown(); + let worker = WorkerWithDpRank::new(7, 0); + manager.insert_block_entries(vec![block(worker)]); + let (sender, tasks) = flume::unbounded(); + sender + .send(OrderedTask::Removal { + task: WorkerTask::RemoveWorkerDpRank { + worker_id: 7, + dp_rank: 0, + sweep_tree: true, + }, + manager: manager.clone(), + }) + .unwrap(); + let mut receiver = WorkerTaskReceiver::retained(tasks); + assert!(matches!( + receiver.recv().unwrap(), + WorkerTask::RemoveWorkerDpRank { .. } + )); + drop(receiver); + assert_eq!( + manager + .drain_due_and_pending( + tokio::time::Instant::now() + std::time::Duration::from_secs(300) + ) + .len(), + 1 + ); + assert!( + sender + .send(OrderedTask::Task(WorkerTask::Event(make_store_event( + 7, + &[1] + )))) + .is_err(), + "retired lane cannot accept restored coverage" + ); + } +}