indexer.rs 81.6 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! KV RadixTree
//!
//! This module implements a key-value (KV) store using a Radix Tree structure to efficiently manage and retrieve data blocks.
//! It is designed to support LLM (Large Language Model) inference by re-using a global KV cache.
//!
//! # Overview
//!
//! The main components of this module include:
//!
//! - **Radix Tree Structure**:
//!   - The `RadixTree` struct represents the main data structure, with nodes (`RadixBlock`) containing children and associated worker IDs.
//!   - It allows efficient storage and retrieval of data blocks based on their hashes.
//!
//! - **Event Handling**:
//!   - The `RouterEvent` struct represents events emitted by LLM workers, which can be applied to the Radix Tree to update its state.
//!   - The `KvIndexer` struct manages these events and match requests asynchronously using Tokio channels.
//!
//! - **Hash Computation**:
//!   - Functions like `compute_block_hash` and `compute_block_hash_for_seq` compute hashes for data blocks and sequences of tokens, facilitating quick lookups.
//!
//! - **Concurrency and Asynchronous Operations**:
//!   - The `KvIndexer` uses a single-threaded Tokio runtime to handle events and match requests concurrently, ensuring efficient processing without blocking.
//!
//! - **Match Requests**:
//!   - The `MatchRequest` struct represents requests to find matches in the Radix Tree, returning overlap scores indicating the best matches.
//!
//! # Purpose
//!
//! This module provides a scalable and efficient way to manage and retrieve data blocks for LLM inference, leveraging a global KV cache to optimize performance.

use async_trait::async_trait;
35
36
37
use bytes::Bytes;
use dynamo_runtime::{
    component::Component,
38
    metrics::{MetricsHierarchy, prometheus_names::kvrouter},
39
40
};
use prometheus::{IntCounterVec, Opts};
41
42
43
use serde::{Deserialize, Serialize};
use std::{
    cell::RefCell,
44
    collections::{HashMap, VecDeque},
45
46
    iter,
    rc::Rc,
47
    sync::{Arc, OnceLock},
48
49
50
51
52
53
54
55
56
57
    thread::JoinHandle,
    time::{Duration, Instant},
};
use tokio::sync::{broadcast, mpsc, oneshot};
use tokio_util::sync::CancellationToken;
use xxhash_rust::xxh3;

pub const XXH3_SEED: u64 = 1337;

use crate::kv_router::protocols::*;
58
use crate::tokens::SequenceHash;
59
60
61
62
63
64
65
66
67
68
69
70
71
72

/// Errors that can occur in the KV Router.
#[derive(Debug, thiserror::Error)]
pub enum KvRouterError {
    #[error("Block not found")]
    BlockNotFound,

    #[error("Indexer is offline")]
    IndexerOffline,

    #[error("Indexer is dropped request")]
    IndexerDroppedRequest,
}

73
74
75
76
77
78
79
80
81
82
/// Errors that can occur during KV Cache Event processing.
#[derive(Debug, thiserror::Error)]
pub enum KvCacheEventError {
    #[error("Failed to find parent block")]
    ParentBlockNotFound,

    #[error("Failed to find block")]
    BlockNotFound,
}

83
84
85
/// A shared reference to a [`RadixBlock`].
type SharedRadixBlock = Rc<RefCell<RadixBlock>>;

86
87
88
89
pub fn compute_hash(data: &[u8]) -> u64 {
    xxh3::xxh3_64_with_seed(data, XXH3_SEED)
}

90
91
92
93
94
95
96
97
98
99
/// Compute the hash of a local block.
///
/// ### Arguments
///
/// * `data` - A byte slice representing the data to hash.
///
/// ### Returns
///
/// A `LocalBlockHash` representing the computed hash.
pub fn compute_block_hash(data: &[u8]) -> LocalBlockHash {
100
    LocalBlockHash(compute_hash(data))
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
}

// /// Updated version of the `compute_block_hash` function that included the lora_id
// pub fn compute_block_hash_v2(token_id: &[u32], lora_id: u64) {
//     let mut bytes = Vec::new();
//     for token in token_id {
//         bytes.extend_from_slice(&token.to_le_bytes());
//     }
//     bytes.extend_from_slice(&lora_id.to_le_bytes());
//     let hash = xxh3::xxh3_64_with_seed(&bytes, XXH3_SEED);
// }

/// Compute the hash for a sequence of tokens.
///
/// ### Arguments
///
/// * `tokens` - A vector of `u32` tokens.
///
/// ### Returns
///
/// A vector of `LocalBlockHash` representing the computed hashes for each chunk of tokens.
122
pub fn compute_block_hash_for_seq(tokens: &[u32], kv_block_size: u32) -> Vec<LocalBlockHash> {
123
    tokens
124
        .chunks_exact(kv_block_size as usize) // Split into chunks of kv_block_size elements
125
126
127
128
129
130
131
132
133
134
135
        .map(|chunk| {
            let bytes: Vec<u8> = chunk
                .iter()
                .flat_map(|&num| num.to_le_bytes()) // Convert each i32 to its little-endian bytes
                .collect();

            compute_block_hash(&Bytes::from(bytes)) // Convert the byte Vec to Bytes
        })
        .collect()
}

136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
/// Compute rolling sequence hashes for a vector of block hashes.
///
/// This mirrors the behavior in tokens.rs where:
/// - The first block's sequence hash equals its block hash
/// - Subsequent blocks' sequence hash = hash([parent_sequence_hash, current_block_hash], seed)
///
/// ### Arguments
///
/// * `block_hashes` - A vector of `LocalBlockHash` values representing the block hashes.
///
/// ### Returns
///
/// A vector of u64 values representing the sequence hashes for each block.
pub fn compute_seq_hash_for_block(block_hashes: &[LocalBlockHash]) -> Vec<SequenceHash> {
    if block_hashes.is_empty() {
        return Vec::new();
    }

    let mut sequence_hashes = Vec::with_capacity(block_hashes.len());
    sequence_hashes.push(block_hashes[0].0);

    for i in 1..block_hashes.len() {
        let parent_seq_hash = sequence_hashes[i - 1];
        let current_block_hash = block_hashes[i].0;

        let combined = [parent_seq_hash, current_block_hash];
        let bytes: Vec<u8> = combined.iter().flat_map(|&num| num.to_le_bytes()).collect();
        let seq_hash = compute_hash(&bytes);
        sequence_hashes.push(seq_hash);
    }

    sequence_hashes
}

170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
/// A [`KvCacheEvent`] on a specific LLM worker denoted by [`WorkerId`].
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RouterEvent {
    /// The ID of the worker emitting the event.
    worker_id: WorkerId,
    /// The cache event associated with the worker.
    event: KvCacheEvent,
}

impl RouterEvent {
    /// Create a new `RouterEvent`.
    ///
    /// ### Arguments
    ///
    /// * `worker_id` - The ID of the worker emitting the event.
    /// * `event` - The cache event.
    ///
    /// ### Returns
    ///
    /// A new `RouterEvent`.
    pub fn new(worker_id: WorkerId, event: KvCacheEvent) -> Self {
        Self { worker_id, event }
    }
}

/// A block in the Radix Tree.
196
#[derive(Debug)]
197
198
199
struct RadixBlock {
    /// A map of child blocks, keyed by their local block hash.
    children: HashMap<LocalBlockHash, SharedRadixBlock>,
Yan Ru Pei's avatar
Yan Ru Pei committed
200
    /// A map of workers (with dp_rank) to their external sequence block hash for this block.
201
    /// The external hash is preserved to speed up snapshotting.
Yan Ru Pei's avatar
Yan Ru Pei committed
202
    workers: HashMap<WorkerWithDpRank, ExternalSequenceBlockHash>,
203
204
205
206
207
208
209
210
211
212
213
214
215
    /// A buffer of times that this block was last traversed
    recent_uses: VecDeque<Instant>,
}

impl RadixBlock {
    /// Create a new `RadixBlock`.
    ///
    /// ### Returns
    ///
    /// A new `RadixBlock`.
    pub fn new() -> Self {
        Self {
            children: HashMap::new(),
216
            workers: HashMap::new(),
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
            recent_uses: VecDeque::new(),
        }
    }
}

pub struct RadixTree {
    /// This is the root of the radix/prefix tree
    /// This will only contain root blocks
    root: SharedRadixBlock,

    /// This is a global lookup table for all blocks which will let you jump into
    /// the radix tree at any point
    /// Lookup is best case O(1) and worst case O(N); however, even constant in-time
    /// could be expensive if N is large
    /// We should monitor the size of this table and consider using a proper radix tree.
    /// Transitioning to a radix tree only would require a change in the messaging structure
    /// as the entire prefix would need to be sent. Alternatively, we could use block_depth
    /// integers to indicate how many blocks to skip and use a radix/prefix tree at each level.
Yan Ru Pei's avatar
Yan Ru Pei committed
235
    lookup: HashMap<WorkerWithDpRank, HashMap<ExternalSequenceBlockHash, SharedRadixBlock>>,
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
    /// The time buffer the radix tree should check when considering frequence of block accesses
    expiration_duration: Option<Duration>,
}

impl Default for RadixTree {
    fn default() -> Self {
        Self::new()
    }
}

impl RadixTree {
    /// Create a new `RadixTree`.
    ///
    /// ### Returns
    ///
    /// A new `RadixTree`.
    pub fn new_with_frequency(expiration_duration: Option<Duration>) -> Self {
        Self {
            root: Rc::new(RefCell::new(RadixBlock::new())),
            lookup: HashMap::new(),
            expiration_duration,
        }
    }

    pub fn new() -> Self {
        Self::new_with_frequency(None)
    }

    /// Traverse the radix tree to find the best match for a given sequence of [`LocalBlockHash`]es.
    ///
    /// ### Arguments
    ///
    /// * `sequence` - A vector of `LocalBlockHash` representing the sequence to match.
    /// * `early_exit` - A boolean indicating whether to exit early if a single match is found.
    ///
    /// ### Returns
    ///
    /// An `OverlapScores` representing the match scores.
    pub fn find_matches(&self, sequence: Vec<LocalBlockHash>, early_exit: bool) -> OverlapScores {
        let mut scores = OverlapScores::new();
        let mut current = self.root.clone();
        let now = Instant::now();
278
279
280
281
282
283
284

        tracing::trace!(
            "RadixTree::find_matches: looking for sequence={:?}",
            sequence.iter().map(|h| h.0).collect::<Vec<_>>()
        );

        for (idx, block_hash) in sequence.iter().enumerate() {
285
286
            let next_block = {
                let current_borrow = current.borrow();
287
                current_borrow.children.get(block_hash).cloned()
288
289
            };
            if let Some(block) = next_block {
290
                scores.update_scores(block.borrow().workers.keys());
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311

                if let Some(expiration_duration) = self.expiration_duration {
                    let mut block_mut = block.borrow_mut();

                    while let Some(access_time) = block_mut.recent_uses.front() {
                        if now.duration_since(*access_time) > expiration_duration {
                            block_mut.recent_uses.pop_front();
                        } else {
                            break;
                        }
                    }
                    scores.add_frequency(block_mut.recent_uses.len());
                    block_mut.recent_uses.push_back(now);
                }

                if early_exit && block.borrow().workers.len() == 1 {
                    break;
                }

                current = block;
            } else {
312
313
314
315
316
                tracing::trace!(
                    "RadixTree::find_matches: block not found at index {} for hash {}",
                    idx,
                    block_hash.0
                );
317
318
319
320
                break;
            }
        }

321
322
        tracing::trace!("RadixTree::find_matches: final scores={:?}", scores.scores);

323
324
325
326
327
328
329
330
        scores
    }

    /// Apply a [`RouterEvent`] to the radix tree.
    ///
    /// ### Arguments
    ///
    /// * `event` - The `RouterEvent` to apply.
331
    pub fn apply_event(&mut self, event: RouterEvent) -> Result<(), KvCacheEventError> {
Yan Ru Pei's avatar
Yan Ru Pei committed
332
333
334
335
336
337
        let (worker_id, kv_event) = (event.worker_id, event.event);
        let (id, op) = (kv_event.event_id, kv_event.data);

        // Construct WorkerWithDpRank from worker_id and dp_rank from the event
        let worker = WorkerWithDpRank::new(worker_id, kv_event.dp_rank);

338
        tracing::trace!(id, "RadixTree::apply_event: Store operation: {:?}", op);
339

Yan Ru Pei's avatar
Yan Ru Pei committed
340
        let worker_lookup = self.lookup.entry(worker).or_default();
341
342
343
344
345
346
347
348
349
350
351
352
353
354

        match op {
            KvCacheEventData::Stored(op) => {
                // find the parent block - if the parent exists it must be on our worker, if not,
                // we check the radix tree's root to find it.
                // this is the single most expensive lookup
                let current = match op.parent_hash {
                    Some(parent) => worker_lookup.get(&parent),
                    None => Some(&self.root),
                };

                let mut current = match current {
                    Some(current) => current.clone(),
                    None => {
355
                        tracing::warn!(
Yan Ru Pei's avatar
Yan Ru Pei committed
356
                            worker_id = worker.worker_id.to_string(),
Yan Ru Pei's avatar
Yan Ru Pei committed
357
                            dp_rank = worker.dp_rank,
358
359
                            id,
                            parent_hash = ?op.parent_hash,
Yan Ru Pei's avatar
Yan Ru Pei committed
360
                            num_blocks = op.blocks.len(),
361
362
                            "Failed to find parent block; skipping store operation"
                        );
363
                        return Err(KvCacheEventError::ParentBlockNotFound);
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
                    }
                };

                for block_id in op.blocks {
                    let mut inner = current.borrow_mut();
                    let block = match inner.children.get(&block_id.tokens_hash) {
                        Some(block) => block.clone(),
                        None => {
                            // create new block - automatically added to the lookup table
                            let new_block = worker_lookup
                                .get(&block_id.block_hash)
                                .cloned()
                                .unwrap_or_else(|| Rc::new(RefCell::new(RadixBlock::new())));

                            // insert into radix tree
                            inner
                                .children
                                .insert(block_id.tokens_hash, new_block.clone());

                            new_block
                        }
                    };

Yan Ru Pei's avatar
Yan Ru Pei committed
387
                    // add our worker to the block with its external hash
388
389
390
                    block
                        .borrow_mut()
                        .workers
Yan Ru Pei's avatar
Yan Ru Pei committed
391
                        .insert(worker, block_id.block_hash);
392
393
394
395
396
397
398
399
400

                    // add the block to the worker_id lookup table
                    worker_lookup.insert(block_id.block_hash, block.clone());

                    // drop inner so we can shift current to this block
                    drop(inner);

                    current = block;
                }
401
                Ok(())
402
403
            }
            KvCacheEventData::Removed(remove) => {
404
                // tracing::trace!(id, "KV Remove Operation: {:?}", op);
405
406
407
408
409
410
411
412
413
414
                // let mut worker_lookup = self.lookup.get(&worker_id).expect("Worker not found");

                for block in remove.block_hashes {
                    // entry in radix tree
                    // a small optimization would be to get the next block from the reduced set of children
                    // in order to apply this optimization, we would need to know the list of blocks is always sorted
                    // by parent -> child relationship
                    let entry = match worker_lookup.get(&block) {
                        Some(entry) => entry.clone(),
                        None => {
415
                            tracing::warn!(
Yan Ru Pei's avatar
Yan Ru Pei committed
416
417
                                worker_id = worker.worker_id.to_string(),
                                dp_rank = worker.dp_rank,
418
                                id,
Yan Ru Pei's avatar
Yan Ru Pei committed
419
                                block_hash = ?block,
420
421
                                "Failed to find block to remove; skipping remove operation"
                            );
422
                            return Err(KvCacheEventError::BlockNotFound);
423
424
425
426
                        }
                    };

                    let mut guard = entry.borrow_mut();
Yan Ru Pei's avatar
Yan Ru Pei committed
427
                    guard.workers.remove(&worker);
428
                    if guard.workers.is_empty() {
429
                        // if no workers are using this block, that is true for all children
430
431
432
433
434
                        guard.children.clear();
                    }
                    // remove the block from the lookup table
                    worker_lookup.remove(&block);
                }
435
                Ok(())
436
            }
437
            KvCacheEventData::Cleared => {
Yan Ru Pei's avatar
Yan Ru Pei committed
438
                self.clear_all_blocks(worker.worker_id);
439
                Ok(())
440
            }
441
442
443
        }
    }

Yan Ru Pei's avatar
Yan Ru Pei committed
444
445
446
447
448
449
450
451
452
453
454
    /// Helper function to remove or clear blocks for a worker.
    /// If `keep_worker` is true, the worker remains in lookup with empty blocks.
    /// If `keep_worker` is false, the worker is completely removed from lookup.
    fn remove_or_clear_worker_blocks(&mut self, worker_id: WorkerId, keep_worker: bool) {
        // Collect all WorkerWithDpRank keys that match this worker_id
        let workers: Vec<WorkerWithDpRank> = self
            .lookup
            .keys()
            .filter(|w| w.worker_id == worker_id)
            .copied()
            .collect();
455

Yan Ru Pei's avatar
Yan Ru Pei committed
456
457
458
459
460
461
462
463
464
        for worker in workers {
            if let Some((worker_key, blocks)) = self.lookup.remove_entry(&worker) {
                blocks.iter().for_each(|(_, block)| {
                    block.borrow_mut().workers.remove(&worker);
                    // If no workers are using this block, that is true for all children
                    if block.borrow().workers.is_empty() {
                        block.borrow_mut().children.clear();
                    }
                });
465

Yan Ru Pei's avatar
Yan Ru Pei committed
466
467
468
                if keep_worker {
                    // Re-insert worker with empty blocks map to keep it tracked
                    self.lookup.insert(worker_key, HashMap::new());
469
                }
470
471
472
            }
        }
    }
473

Yan Ru Pei's avatar
Yan Ru Pei committed
474
475
476
477
478
479
480
481
    pub fn remove_worker(&mut self, worker_id: WorkerId) {
        self.remove_or_clear_worker_blocks(worker_id, false);
    }

    pub fn clear_all_blocks(&mut self, worker_id: WorkerId) {
        self.remove_or_clear_worker_blocks(worker_id, true);
    }

482
    /// Get all worker IDs currently tracked in the radix tree.
Yan Ru Pei's avatar
Yan Ru Pei committed
483
    /// Returns unique worker_ids (ignoring dp_rank differences).
484
    pub fn get_workers(&self) -> Vec<WorkerId> {
Yan Ru Pei's avatar
Yan Ru Pei committed
485
486
487
488
        let mut worker_ids: Vec<WorkerId> = self.lookup.keys().map(|w| w.worker_id).collect();
        worker_ids.sort_unstable();
        worker_ids.dedup();
        worker_ids
489
490
    }

491
492
493
494
    /// Dump the radix tree as a series of RouterEvents that can reconstruct the tree.
    /// Uses BFS traversal to ensure that the tree reconstruction is unique,
    /// though the exact event ordering will be lost.
    pub fn dump_tree_as_events(&self) -> Vec<RouterEvent> {
495
496
497
498
499
        tracing::debug!(
            "Dumping radix tree as events (contains information about {:?} workers)",
            self.lookup.len()
        );

500
501
502
        let mut events = Vec::new();
        let mut event_id = 0u64;

503
        // BFS queue: (current_block, parent_hashes_per_worker, tokens_hash)
Yan Ru Pei's avatar
Yan Ru Pei committed
504
        // parent_hashes_per_worker maps WorkerWithDpRank -> ExternalSequenceBlockHash
505
506
        let mut queue: VecDeque<(
            SharedRadixBlock,
Yan Ru Pei's avatar
Yan Ru Pei committed
507
            HashMap<WorkerWithDpRank, ExternalSequenceBlockHash>,
508
509
            LocalBlockHash,
        )> = VecDeque::new();
510
511
512
513

        // Process root's children first
        let root_borrow = self.root.borrow();
        for (tokens_hash, child_block) in &root_borrow.children {
514
            queue.push_back((child_block.clone(), HashMap::new(), *tokens_hash));
515
516
517
        }
        drop(root_borrow);

518
        while let Some((current_block, parent_hashes, tokens_hash)) = queue.pop_front() {
519
520
            let current_borrow = current_block.borrow();

521
522
            // Map of this block's external hashes per worker (for children to use as parent)
            let mut current_external_hashes = HashMap::new();
523
524

            // For each worker that has this block
525
526
527
528
529
530
            for (worker_id, external_hash) in &current_borrow.workers {
                // Get the correct parent hash for this worker
                let parent_hash = parent_hashes.get(worker_id).copied();

                // Create a store event for this worker
                let event = RouterEvent {
Yan Ru Pei's avatar
Yan Ru Pei committed
531
                    worker_id: worker_id.worker_id,
532
533
534
535
536
537
538
539
540
                    event: KvCacheEvent {
                        event_id,
                        data: KvCacheEventData::Stored(KvCacheStoreData {
                            parent_hash,
                            blocks: vec![KvCacheStoredBlockData {
                                block_hash: *external_hash,
                                tokens_hash,
                            }],
                        }),
Yan Ru Pei's avatar
Yan Ru Pei committed
541
                        dp_rank: worker_id.dp_rank,
542
543
544
545
                    },
                };
                events.push(event);
                event_id += 1;
546

547
548
549
                // Track this block's external hash for this worker
                current_external_hashes.insert(*worker_id, *external_hash);
            }
550

551
            // Enqueue children with per-worker parent hashes
552
            for (child_tokens_hash, child_block) in &current_borrow.children {
553
554
555
556
557
                queue.push_back((
                    child_block.clone(),
                    current_external_hashes.clone(),
                    *child_tokens_hash,
                ));
558
559
560
561
562
            }
        }

        events
    }
563
564
}

565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
/// Metrics for the KV Indexer.
#[derive(Clone)]
pub struct KvIndexerMetrics {
    /// Counter of events applied.
    pub kv_cache_events_applied: IntCounterVec,
}

/// Metric status labels.
pub const METRIC_STATUS_OK: &str = "ok";
pub const METRIC_STATUS_PARENT_NOT_FOUND: &str = "parent_block_not_found";
pub const METRIC_STATUS_BLOCK_NOT_FOUND: &str = "block_not_found";

/// Metric event labels.
pub const METRIC_EVENT_STORED: &str = "stored";
pub const METRIC_EVENT_REMOVED: &str = "removed";
pub const METRIC_EVENT_CLEARED: &str = "cleared";

static KV_INDEXER_METRICS: OnceLock<Arc<KvIndexerMetrics>> = OnceLock::new();

impl KvIndexerMetrics {
    fn new(kv_cache_events_applied: IntCounterVec) -> Self {
        Self {
            kv_cache_events_applied,
        }
    }

    /// Creates a new KvIndexerMetrics from a Component, memoizing the result in
    /// KV_INDEXER_METRICS to avoid duplicate registration issues.
    pub fn from_component(component: &Component) -> Arc<Self> {
        KV_INDEXER_METRICS.get_or_init(|| {
595
            match component.metrics().create_intcountervec(
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
                kvrouter::KV_CACHE_EVENTS_APPLIED,
                "Total number of KV cache events applied to index",
                &["event_type", "status"],
                &[],
            ) {
                Ok(kv_cache_events_applied) => Arc::new(Self::new(kv_cache_events_applied)),
                Err(e) => {
                    tracing::warn!("Failed to create kv indexer metrics from component: {}. Using unregistered metrics as fallback.", e);
                    Arc::new(Self::new_unregistered())
                }
            }
        }).clone()
    }

    /// Creates a new KvIndexerMetrics which is not registered with a MetricsRegistry.
    /// This may be used for tests or as a fallback for when a MetricsRegistry is not available / has errored.
    pub fn new_unregistered() -> Self {
        Self {
            kv_cache_events_applied: IntCounterVec::new(
                Opts::new(
                    kvrouter::KV_CACHE_EVENTS_APPLIED,
                    "Total number of KV cache events applied to index",
                ),
                &["event_type", "status"],
            )
            .unwrap(),
        }
    }

    pub fn get_event_type(event_data: &KvCacheEventData) -> &'static str {
        match event_data {
            KvCacheEventData::Stored(_) => METRIC_EVENT_STORED,
            KvCacheEventData::Removed(_) => METRIC_EVENT_REMOVED,
            KvCacheEventData::Cleared => METRIC_EVENT_CLEARED,
        }
    }

    pub fn increment_event_applied(
        &self,
        event_type: &'static str,
        result: Result<(), KvCacheEventError>,
    ) {
        match result {
            Ok(_) => {
                self.kv_cache_events_applied
                    .with_label_values(&[event_type, METRIC_STATUS_OK])
                    .inc_by(1);
            }
            Err(e) => {
                let error_label = match e {
                    KvCacheEventError::ParentBlockNotFound => METRIC_STATUS_PARENT_NOT_FOUND,
                    KvCacheEventError::BlockNotFound => METRIC_STATUS_BLOCK_NOT_FOUND,
                };
                self.kv_cache_events_applied
                    .with_label_values(&[event_type, error_label])
                    .inc_by(1);
            }
        }
    }
}

Yan Ru Pei's avatar
Yan Ru Pei committed
657
/// Scores representing the overlap of workers (with their dp_rank).
658
659
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OverlapScores {
Yan Ru Pei's avatar
Yan Ru Pei committed
660
661
    // map of worker (with dp_rank) to score
    pub scores: HashMap<WorkerWithDpRank, u32>,
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
    // List of frequencies that the blocks have been accessed. Entries with value 0 are omitted.
    pub frequencies: Vec<usize>,
}

impl Default for OverlapScores {
    fn default() -> Self {
        Self::new()
    }
}

impl OverlapScores {
    /// Create a new `OverlapScores`.
    ///
    /// ### Returns
    ///
    /// A new `OverlapScores`.
    pub fn new() -> Self {
        Self {
            scores: HashMap::new(),
            frequencies: Vec::with_capacity(32),
        }
    }

    /// Update the scores with a set of workers.
    ///
    /// ### Arguments
    ///
Yan Ru Pei's avatar
Yan Ru Pei committed
689
    /// * `workers` - An iterator over `WorkerWithDpRank` references.
690
691
    pub fn update_scores<'a, I>(&mut self, workers: I)
    where
Yan Ru Pei's avatar
Yan Ru Pei committed
692
        I: IntoIterator<Item = &'a WorkerWithDpRank>,
693
    {
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
        for worker in workers {
            let score = self.scores.entry(*worker).or_insert(0);
            *score += 1;
        }
    }

    /// Add an entry in the frequency list.
    pub fn add_frequency(&mut self, frequency: usize) {
        if frequency != 0 {
            self.frequencies
                .last()
                .inspect(|elem| debug_assert!(**elem >= frequency));
            self.frequencies.push(frequency);
        }
    }
}

/// A request to find matches in the Radix Tree.
pub struct MatchRequest {
    /// A vector of `LocalBlockHash` representing the sequence to match.
    sequence: Vec<LocalBlockHash>,
    /// A boolean indicating whether to exit early if a single match is found.
    early_exit: bool,
    /// A channel sender to send the `OverlapScores` response.
    resp: oneshot::Sender<OverlapScores>,
}

721
722
723
724
725
726
/// A request to dump the tree as events
pub struct DumpRequest {
    /// Channel to send the dumped events
    pub resp: oneshot::Sender<Vec<RouterEvent>>,
}

727
728
729
730
731
732
/// A request to get all workers currently tracked
pub struct GetWorkersRequest {
    /// Channel to send the worker IDs
    pub resp: oneshot::Sender<Vec<WorkerId>>,
}

733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
#[async_trait]
pub trait KvIndexerInterface {
    /// Find matches for a given sequence of `LocalBlockHash`es.
    ///
    /// ### Arguments
    ///
    /// * `sequence` - A vector of `LocalBlockHash` representing the sequence to match.
    ///
    /// ### Returns
    ///
    /// An `OverlapScores` representing the match scores.
    async fn find_matches(
        &self,
        sequence: Vec<LocalBlockHash>,
    ) -> Result<OverlapScores, KvRouterError>;

    /// Find matches for a given sequence of tokens.
    ///
    /// ### Arguments
    ///
    /// * `tokens` - A vector of `u32` tokens.
    ///
    /// ### Returns
    ///
    /// An `OverlapScores` representing the match scores.
    async fn find_matches_for_request(
        &self,
        tokens: &[u32],
    ) -> Result<OverlapScores, KvRouterError>;

    /// Apply a `RouterEvent` to the KV store.
    ///
    /// ### Arguments
    ///
    /// * `event` - The `RouterEvent` to apply.
    async fn apply_event(&mut self, event: RouterEvent);

    /// Remove a worker's entries from the trie.
    ///
    /// ### Arguments
    ///
    /// * `worker` - The worker to remove from the trie.
    async fn remove_worker(&mut self, worker: WorkerId);

    /// Shutdown the KV Indexer.
    fn shutdown(&mut self);
779
780
781
782
783
784
785

    /// Dump the entire tree as RouterEvents.
    ///
    /// ### Returns
    ///
    /// A vector of RouterEvents representing the current state of the tree.
    async fn dump_events(&self) -> Result<Vec<RouterEvent>, KvRouterError>;
786
787
788
789
790
791
792
793
794
795
796
797
}

/// The KV Indexer, managing the KV store and handling events and match requests.
pub struct KvIndexer {
    /// A `CancellationToken` for managing shutdown.
    cancel: CancellationToken,
    /// A sender for `RouterEvent`s.
    event_tx: mpsc::Sender<RouterEvent>,
    /// A sender for `MatchRequest`s.
    match_tx: mpsc::Sender<MatchRequest>,
    /// A sender for remove worker requests.
    remove_worker_tx: mpsc::Sender<WorkerId>,
798
799
    /// A sender for get workers requests.
    get_workers_tx: mpsc::Sender<GetWorkersRequest>,
800
801
    /// A sender for dump requests.
    dump_tx: mpsc::Sender<DumpRequest>,
802
803
    /// A handle to the background task managing the KV store.
    task: OnceLock<std::thread::JoinHandle<()>>,
804
    /// The size of the KV block this indexer can handle.
805
    kv_block_size: u32,
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
}

impl KvIndexer {
    /// Create a new `KvIndexer`.
    ///
    /// ### Arguments
    ///
    /// * `token` - A `CancellationToken` for managing shutdown.
    /// * `expiration_duration` - The amount of time that block usage should be buffered.
    ///
    /// ### Returns
    ///
    /// A new `KvIndexer`.
    pub fn new_with_frequency(
        token: CancellationToken,
        expiration_duration: Option<Duration>,
822
        kv_block_size: u32,
823
        metrics: Arc<KvIndexerMetrics>,
824
825
826
827
    ) -> Self {
        let (event_tx, event_rx) = mpsc::channel::<RouterEvent>(2048);
        let (match_tx, match_rx) = mpsc::channel::<MatchRequest>(128);
        let (remove_worker_tx, remove_worker_rx) = mpsc::channel::<WorkerId>(16);
828
        let (get_workers_tx, get_workers_rx) = mpsc::channel::<GetWorkersRequest>(16);
829
        let (dump_tx, dump_rx) = mpsc::channel::<DumpRequest>(16);
830
        let cancel_clone = token.clone();
831

832
        let task = std::thread::spawn(move || {
833
834
            // Create a single-threaded tokio runtime
            let runtime = tokio::runtime::Builder::new_current_thread()
835
836
837
838
                .enable_all()
                .build()
                .unwrap();

839
840
841
842
843
            runtime.block_on(async move {
                let cancel = cancel_clone;
                let mut match_rx = match_rx;
                let mut event_rx = event_rx;
                let mut remove_worker_rx = remove_worker_rx;
844
                let mut get_workers_rx = get_workers_rx;
845
846
847
848
849
850
851
852
853
854
                let mut dump_rx = dump_rx;
                let mut trie = RadixTree::new_with_frequency(expiration_duration);
                loop {
                    tokio::select! {
                        biased;

                        _ = cancel.cancelled() => {
                            tracing::debug!("KvCacheIndexer progress loop shutting down");
                            return;
                        }
855

856
857
858
                        Some(worker) = remove_worker_rx.recv() => {
                            trie.remove_worker(worker);
                        }
859

860
861
862
863
864
                        Some(get_workers_req) = get_workers_rx.recv() => {
                            let workers = trie.get_workers();
                            let _ = get_workers_req.resp.send(workers);
                        }

865
866
867
868
869
                        Some(event) = event_rx.recv() => {
                            let event_type = KvIndexerMetrics::get_event_type(&event.event.data);
                            let result = trie.apply_event(event);
                            metrics.increment_event_applied(event_type, result);
                        }
870

871
872
873
874
                        Some(dump_req) = dump_rx.recv() => {
                            let events = trie.dump_tree_as_events();
                            let _ = dump_req.resp.send(events);
                        }
875

876
877
878
                        Some(req) = match_rx.recv() => {
                            let matches = trie.find_matches(req.sequence, req.early_exit);
                            let _ = req.resp.send(matches);
879
880
                        }
                    }
881
882
                }
            });
883

884
            tracing::debug!("KvCacheIndexer task completed");
885
886
887
888
889
890
891
892
893
894
        });

        let once = OnceLock::new();
        once.set(task).unwrap();

        Self {
            cancel: token,
            event_tx,
            match_tx,
            remove_worker_tx,
895
            get_workers_tx,
896
            dump_tx,
897
            task: once,
898
            kv_block_size,
899
900
901
        }
    }

902
    pub fn block_size(&self) -> u32 {
903
904
905
        self.kv_block_size
    }

906
907
908
909
910
911
    pub fn new(
        token: CancellationToken,
        kv_block_size: u32,
        metrics: Arc<KvIndexerMetrics>,
    ) -> Self {
        Self::new_with_frequency(token, None, kv_block_size, metrics)
912
913
914
915
916
917
918
919
920
    }

    /// Get a sender for `RouterEvent`s.
    ///
    /// ### Returns
    ///
    /// A `mpsc::Sender` for `RouterEvent`s.
    pub fn event_sender(&self) -> mpsc::Sender<RouterEvent> {
        self.event_tx.clone()
921
922
923
924
925
926
927
928
929
    }

    /// Get a sender for dump requests (snapshot events).
    ///
    /// ### Returns
    ///
    /// A `mpsc::Sender` for `DumpRequest`s.
    pub fn snapshot_event_sender(&self) -> mpsc::Sender<DumpRequest> {
        self.dump_tx.clone()
930
    }
931
932
933
934
935
936
937
938
939

    /// Get a sender for worker removal requests.
    ///
    /// ### Returns
    ///
    /// A `mpsc::Sender` for `WorkerId`s.
    pub fn remove_worker_sender(&self) -> mpsc::Sender<WorkerId> {
        self.remove_worker_tx.clone()
    }
940
941
942
943
944
945
946
947
948

    /// Get a sender for get workers requests.
    ///
    /// ### Returns
    ///
    /// A `mpsc::Sender` for `GetWorkersRequest`s.
    pub fn get_workers_sender(&self) -> mpsc::Sender<GetWorkersRequest> {
        self.get_workers_tx.clone()
    }
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
}

#[async_trait]
impl KvIndexerInterface for KvIndexer {
    async fn find_matches(
        &self,
        sequence: Vec<LocalBlockHash>,
    ) -> Result<OverlapScores, KvRouterError> {
        let (resp_tx, resp_rx) = oneshot::channel();
        let req = MatchRequest {
            sequence,
            early_exit: false,
            resp: resp_tx,
        };

        if let Err(e) = self.match_tx.send(req).await {
965
            tracing::error!(
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
                "Failed to send match request: {:?}; the indexer maybe offline",
                e
            );
            return Err(KvRouterError::IndexerOffline);
        }

        resp_rx
            .await
            .map_err(|_| KvRouterError::IndexerDroppedRequest)
    }

    async fn find_matches_for_request(
        &self,
        tokens: &[u32],
    ) -> Result<OverlapScores, KvRouterError> {
981
        tracing::debug!(
982
983
984
985
            "Finding matches for request tokens: {:?} / len: {}",
            tokens,
            tokens.len()
        );
986
        let sequence = compute_block_hash_for_seq(tokens, self.kv_block_size);
987
        tracing::debug!("Computed sequence: {:?}", sequence);
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
        self.find_matches(sequence).await
    }

    async fn apply_event(&mut self, event: RouterEvent) {
        self.event_tx.send(event).await.unwrap();
    }

    async fn remove_worker(&mut self, worker: WorkerId) {
        self.remove_worker_tx.send(worker).await.unwrap();
    }

    fn shutdown(&mut self) {
        self.cancel.cancel();
        if let Some(task) = self.task.take() {
            task.join().expect("Failed to join kv indexer task");
        }
    }
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018

    async fn dump_events(&self) -> Result<Vec<RouterEvent>, KvRouterError> {
        let (resp_tx, resp_rx) = oneshot::channel();
        let dump_req = DumpRequest { resp: resp_tx };

        if let Err(e) = self.dump_tx.send(dump_req).await {
            tracing::error!("Failed to send dump request: {:?}", e);
            return Err(KvRouterError::IndexerOffline);
        }

        resp_rx
            .await
            .map_err(|_| KvRouterError::IndexerDroppedRequest)
    }
1019
1020
}

1021
1022
1023
1024
1025
1026
impl Drop for KvIndexer {
    fn drop(&mut self) {
        self.shutdown();
    }
}

1027
1028
1029
1030
1031
1032
1033
#[derive(Debug, Clone)]
pub struct ShardedMatchRequest {
    sequence: Vec<LocalBlockHash>,
    early_exit: bool,
    resp: mpsc::Sender<OverlapScores>,
}

1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
/// A sharded KV Indexer that partitions the RadixTree across multiple independent shards.
///
/// ## Sharding Strategy
/// - Each worker is **permanently assigned** to a single shard on first event
/// - All KV blocks from a worker exist only in that worker's assigned shard
/// - New workers are assigned to the shard with the fewest workers (load balancing)
///
/// ## Operation
/// - **Events**: Routed directly to the worker's assigned shard
/// - **Match requests**: Broadcast to all shards (scatter-gather pattern)
/// - **Threading**: Each shard runs in its own thread with a single-threaded runtime
///
/// This design ensures no cross-shard synchronization for writes while enabling
/// parallel processing and better scalability.
1048
1049
1050
pub struct KvIndexerSharded {
    /// A `CancellationToken` for managing shutdown.
    cancel: CancellationToken,
1051
    /// The size of the KV block this indexer can handle.
1052
    kv_block_size: u32,
1053
1054
1055
1056
1057
1058
    worker_assignments: HashMap<WorkerId, usize>,
    worker_counts: Vec<usize>,

    event_tx: Vec<mpsc::Sender<RouterEvent>>,
    request_broadcast_tx: broadcast::Sender<ShardedMatchRequest>,
    remove_worker_tx: Vec<mpsc::Sender<WorkerId>>,
1059
    dump_tx: Vec<mpsc::Sender<DumpRequest>>,
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
    tasks: Vec<JoinHandle<()>>,
}

impl KvIndexerSharded {
    /// Create a new `KvIndexerSharded`.
    ///
    /// ### Arguments
    ///
    /// * `token` - A `CancellationToken` for managing shutdown.
    /// * `shards` - A list of kvindexer shards.
    /// * `expiration_duration` - The amount of time that block usage should be buffered.
    ///
    /// ### Returns
    ///
    /// A new `KvIndexer`.
    pub fn new_with_frequency(
        token: CancellationToken,
        num_shards: usize,
        expiration_duration: Option<Duration>,
1079
        kv_block_size: u32,
1080
        metrics: Arc<KvIndexerMetrics>,
1081
1082
1083
1084
1085
1086
    ) -> Self {
        let worker_assignments: HashMap<WorkerId, usize> = HashMap::new();
        let worker_counts: Vec<usize> = vec![0; num_shards];

        let mut event_tx = Vec::new();
        let mut remove_worker_tx = Vec::new();
1087
        let mut get_workers_tx = Vec::new();
1088
        let mut dump_tx = Vec::new(); // Add dump channels
1089
1090
1091
1092
1093
1094
1095
1096
        let mut tasks = Vec::new();

        let (request_broadcast_tx, _) = broadcast::channel::<ShardedMatchRequest>(1048576);

        for _ in 0..num_shards {
            let (shard_event_tx, mut shard_event_rx) = mpsc::channel::<RouterEvent>(2048);
            let (shard_remove_worker_tx, mut shard_remove_worker_rx) =
                mpsc::channel::<WorkerId>(16);
1097
1098
            let (shard_get_workers_tx, mut shard_get_workers_rx) =
                mpsc::channel::<GetWorkersRequest>(16);
1099
            let (shard_dump_tx, mut shard_dump_rx) = mpsc::channel::<DumpRequest>(16); // Add dump channel
1100
1101
            let mut shard_broadcast_rx = request_broadcast_tx.subscribe();
            let cancel = token.clone();
1102
            let metrics = metrics.clone();
1103
1104
1105

            event_tx.push(shard_event_tx);
            remove_worker_tx.push(shard_remove_worker_tx);
1106
            get_workers_tx.push(shard_get_workers_tx);
1107
            dump_tx.push(shard_dump_tx); // Store dump sender
1108

1109
            let runtime = tokio::runtime::Builder::new_current_thread()
1110
1111
1112
1113
1114
                .enable_all()
                .build()
                .unwrap();

            tasks.push(std::thread::spawn(move || {
1115
1116
1117
1118
1119
                runtime.block_on(async move {
                    let mut trie = RadixTree::new_with_frequency(expiration_duration);
                    loop {
                        tokio::select! {
                            biased;
1120

1121
1122
1123
1124
                            _ = cancel.cancelled() => {
                                tracing::trace!("KvCacheIndexer progress loop shutting down");
                                return;
                            }
1125

1126
1127
1128
                            Some(worker) = shard_remove_worker_rx.recv() => {
                                trie.remove_worker(worker);
                            }
1129

1130
1131
1132
1133
1134
                            Some(get_workers_req) = shard_get_workers_rx.recv() => {
                                let workers = trie.get_workers();
                                let _ = get_workers_req.resp.send(workers);
                            }

1135
1136
1137
1138
1139
                            Some(event) = shard_event_rx.recv() => {
                                let event_type = KvIndexerMetrics::get_event_type(&event.event.data);
                                let result = trie.apply_event(event);
                                metrics.increment_event_applied(event_type, result);
                            }
1140

1141
1142
1143
1144
1145
1146
1147
1148
1149
                            Some(dump_req) = shard_dump_rx.recv() => {
                                let events = trie.dump_tree_as_events();
                                let _ = dump_req.resp.send(events);
                            }

                            Ok(req) = shard_broadcast_rx.recv() => {
                                let matches = trie.find_matches(req.sequence, req.early_exit);
                                if let Err(e) = req.resp.send(matches).await {
                                    tracing::trace!("Failed to send match response: {:?}", e);
1150
1151
1152
                                }
                            }
                        }
1153
1154
                    }
                });
1155

1156
                tracing::debug!("KvCacheIndexer task completed");
1157
1158
1159
1160
1161
            }));
        }

        Self {
            cancel: token,
1162
            kv_block_size,
1163
1164
1165
1166
1167
            worker_assignments,
            worker_counts,
            event_tx,
            request_broadcast_tx,
            remove_worker_tx,
1168
            dump_tx, // Add dump_tx field
1169
1170
1171
1172
            tasks,
        }
    }

1173
    pub fn block_size(&self) -> u32 {
1174
1175
1176
        self.kv_block_size
    }

1177
1178
1179
1180
1181
1182
1183
    pub fn new(
        token: CancellationToken,
        num_shards: usize,
        kv_block_size: u32,
        metrics: Arc<KvIndexerMetrics>,
    ) -> Self {
        Self::new_with_frequency(token, num_shards, None, kv_block_size, metrics)
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
    }
}

#[async_trait]
impl KvIndexerInterface for KvIndexerSharded {
    async fn find_matches(
        &self,
        sequence: Vec<LocalBlockHash>,
    ) -> Result<OverlapScores, KvRouterError> {
        'match_loop: loop {
            let (match_tx, mut match_rx) = mpsc::channel(self.event_tx.len());
            self.request_broadcast_tx
                .send(ShardedMatchRequest {
                    sequence: sequence.clone(),
                    early_exit: false,
                    resp: match_tx,
                })
                .map_err(|_| KvRouterError::IndexerOffline)?;

            let mut scores = OverlapScores::new();

            for response_num in 0..self.event_tx.len() {
                match match_rx.recv().await {
                    Some(response) => {
                        scores.scores.extend(response.scores);

                        if response_num == 0 {
                            scores.frequencies = response.frequencies;
                        } else {
                            let diff = (response.frequencies.len() as i64)
                                - (scores.frequencies.len() as i64);

                            if diff > 0 {
1217
                                scores.frequencies.extend(iter::repeat_n(0, diff as usize));
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
                            }

                            for i in 0..response.frequencies.len() {
                                scores.frequencies[i] += response.frequencies[i];
                            }
                        }
                    }
                    None => {
                        // This can only happen if the broadcast channel overflows.
                        // In this case, we don't want to recursively call find_matches again. Otherwise, we could overflow the stack.
                        continue 'match_loop;
                    }
                }
            }
            return Ok(scores);
        }
    }

    async fn find_matches_for_request(
        &self,
        tokens: &[u32],
    ) -> Result<OverlapScores, KvRouterError> {
1240
        let sequence = compute_block_hash_for_seq(tokens, self.kv_block_size);
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
        self.find_matches(sequence).await
    }

    async fn apply_event(&mut self, event: RouterEvent) {
        #[allow(clippy::map_entry)]
        if !self.worker_assignments.contains_key(&event.worker_id) {
            // Get the shard with the smallest amount of workers.
            let selected_shard = self
                .worker_counts
                .iter()
                .enumerate()
                .min_by_key(|&(_, value)| value)
                .unwrap()
                .0;

            self.worker_assignments
                .insert(event.worker_id, selected_shard);
            self.worker_counts[selected_shard] += 1;
        }

        self.event_tx[self.worker_assignments[&event.worker_id]]
            .send(event)
            .await
            .unwrap();
    }

    async fn remove_worker(&mut self, worker: WorkerId) {
        if let Some((_, shard)) = self.worker_assignments.remove_entry(&worker) {
            self.worker_counts[shard] -= 1;
            self.remove_worker_tx[shard].send(worker).await.unwrap();
        }
    }

    /// Shutdown the KV Indexer.
    fn shutdown(&mut self) {
        self.cancel.cancel();
        while !self.tasks.is_empty() {
            self.tasks.pop().unwrap().join().unwrap();
        }
    }
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309

    async fn dump_events(&self) -> Result<Vec<RouterEvent>, KvRouterError> {
        let mut all_events = Vec::new();

        // Create channels for each shard
        let mut receivers = Vec::new();

        for shard_dump_tx in &self.dump_tx {
            let (resp_tx, resp_rx) = oneshot::channel();
            let dump_req = DumpRequest { resp: resp_tx };

            if let Err(e) = shard_dump_tx.send(dump_req).await {
                tracing::error!("Failed to send dump request to shard: {:?}", e);
                return Err(KvRouterError::IndexerOffline);
            }

            receivers.push(resp_rx);
        }

        // Collect results from all shards
        for resp_rx in receivers {
            match resp_rx.await {
                Ok(events) => all_events.extend(events),
                Err(_) => return Err(KvRouterError::IndexerDroppedRequest),
            }
        }

        Ok(all_events)
    }
1310
1311
}

1312
1313
1314
1315
1316
1317
impl Drop for KvIndexerSharded {
    fn drop(&mut self) {
        self.shutdown();
    }
}

1318
1319
1320
1321
1322
#[cfg(test)]
mod tests {

    use super::*;
    use rstest::rstest;
1323
    use rstest_reuse::{self, *};
1324
1325
1326
    use tokio::time;
    use tokio_util::sync::CancellationToken;

1327
1328
1329
1330
    fn setup() {
        dynamo_runtime::logging::init();
    }

1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
    fn make_blocks(hashes: Vec<u64>) -> Vec<KvCacheStoredBlockData> {
        hashes
            .iter()
            .map(|i| KvCacheStoredBlockData {
                tokens_hash: LocalBlockHash(*i),
                block_hash: ExternalSequenceBlockHash(*i * 100),
            })
            .collect()
    }

    fn add_blocks(
        hashes: Vec<u64>,
        parent_hash: Option<ExternalSequenceBlockHash>,
    ) -> KvCacheEventData {
        KvCacheEventData::Stored(KvCacheStoreData {
            parent_hash,
            blocks: make_blocks(hashes),
        })
    }

    fn create_store_event(
        worker_id: WorkerId,
        event_id: u64,
        hashes: Vec<u64>,
        parent: Option<ExternalSequenceBlockHash>,
    ) -> RouterEvent {
        RouterEvent {
            worker_id,
            event: KvCacheEvent {
                event_id,
                data: add_blocks(hashes, parent),
Yan Ru Pei's avatar
Yan Ru Pei committed
1362
                dp_rank: 0,
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
            },
        }
    }

    fn create_remove_event(worker_id: WorkerId, event_id: u64, hashes: Vec<u64>) -> RouterEvent {
        RouterEvent {
            worker_id,
            event: KvCacheEvent {
                event_id,
                data: KvCacheEventData::Removed(KvCacheRemoveData {
                    block_hashes: hashes
                        .iter()
                        .map(|i| ExternalSequenceBlockHash(*i * 100))
                        .collect(),
                }),
Yan Ru Pei's avatar
Yan Ru Pei committed
1378
                dp_rank: 0,
1379
1380
1381
1382
1383
1384
            },
        }
    }

    #[test]
    fn test_radix_tree() {
1385
1386
        setup();

1387
1388
        let mut trie = RadixTree::new();

1389
1390
        let worker_1 = 0;
        let worker_2 = 1;
1391

1392
1393
        trie.apply_event(create_store_event(worker_1, 1, vec![1, 2, 3], None))
            .unwrap();
1394
1395
1396
1397
1398

        let scores = trie.find_matches(
            vec![LocalBlockHash(1), LocalBlockHash(2), LocalBlockHash(3)],
            false,
        );
Yan Ru Pei's avatar
Yan Ru Pei committed
1399
1400
1401
1402
1403
1404
1405
        assert_eq!(
            scores
                .scores
                .get(&WorkerWithDpRank::from_worker_id(worker_1))
                .unwrap(),
            &3
        );
1406
1407

        assert_eq!(trie.lookup.len(), 1);
Yan Ru Pei's avatar
Yan Ru Pei committed
1408
1409
1410
1411
1412
1413
1414
        assert_eq!(
            trie.lookup
                .get(&WorkerWithDpRank::from_worker_id(worker_1))
                .unwrap()
                .len(),
            3
        );
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
        assert_eq!(trie.root.borrow().workers.len(), 0);
        assert_eq!(trie.root.borrow().children.len(), 1);
        assert_eq!(
            trie.root
                .borrow()
                .children
                .get(&LocalBlockHash(1))
                .unwrap()
                .borrow()
                .workers
                .len(),
            1
        );
        assert_eq!(
            trie.root
                .borrow()
                .children
                .get(&LocalBlockHash(1))
                .unwrap()
                .borrow()
                .children
                .len(),
            1
        );

1440
1441
        trie.apply_event(create_store_event(worker_2, 1, vec![1, 4, 5], None))
            .unwrap();
1442
1443
1444
1445
1446

        let scores = trie.find_matches(
            vec![LocalBlockHash(1), LocalBlockHash(2), LocalBlockHash(3)],
            false,
        );
Yan Ru Pei's avatar
Yan Ru Pei committed
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
        assert_eq!(
            scores
                .scores
                .get(&WorkerWithDpRank::from_worker_id(worker_1))
                .unwrap(),
            &3
        );
        assert_eq!(
            scores
                .scores
                .get(&WorkerWithDpRank::from_worker_id(worker_2))
                .unwrap(),
            &1
        );
1461
1462

        assert_eq!(trie.lookup.len(), 2);
Yan Ru Pei's avatar
Yan Ru Pei committed
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
        assert_eq!(
            trie.lookup
                .get(&WorkerWithDpRank::from_worker_id(worker_1))
                .unwrap()
                .len(),
            3
        );
        assert_eq!(
            trie.lookup
                .get(&WorkerWithDpRank::from_worker_id(worker_2))
                .unwrap()
                .len(),
            3
        );
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
        assert_eq!(trie.root.borrow().workers.len(), 0);
        assert_eq!(trie.root.borrow().children.len(), 1);
        assert_eq!(
            trie.root
                .borrow()
                .children
                .get(&LocalBlockHash(1))
                .unwrap()
                .borrow()
                .workers
                .len(),
            2
        );
        assert_eq!(
            trie.root
                .borrow()
                .children
                .get(&LocalBlockHash(1))
                .unwrap()
                .borrow()
                .children
                .len(),
            2
        );

1502
1503
        trie.apply_event(create_remove_event(worker_2, 2, vec![5]))
            .unwrap();
1504
        assert_eq!(trie.lookup.len(), 2);
Yan Ru Pei's avatar
Yan Ru Pei committed
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
        assert_eq!(
            trie.lookup
                .get(&WorkerWithDpRank::from_worker_id(worker_1))
                .unwrap()
                .len(),
            3
        );
        assert_eq!(
            trie.lookup
                .get(&WorkerWithDpRank::from_worker_id(worker_2))
                .unwrap()
                .len(),
            2
        );
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
        assert_eq!(trie.root.borrow().workers.len(), 0);
        assert_eq!(trie.root.borrow().children.len(), 1);
        assert_eq!(
            trie.root
                .borrow()
                .children
                .get(&LocalBlockHash(1))
                .unwrap()
                .borrow()
                .workers
                .len(),
            2
        );
        assert_eq!(
            trie.root
                .borrow()
                .children
                .get(&LocalBlockHash(1))
                .unwrap()
                .borrow()
                .children
                .len(),
            2
        );

1544
1545
        trie.apply_event(create_remove_event(worker_2, 3, vec![4]))
            .unwrap();
1546
1547

        assert_eq!(trie.lookup.len(), 2);
Yan Ru Pei's avatar
Yan Ru Pei committed
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
        assert_eq!(
            trie.lookup
                .get(&WorkerWithDpRank::from_worker_id(worker_1))
                .unwrap()
                .len(),
            3
        );
        assert_eq!(
            trie.lookup
                .get(&WorkerWithDpRank::from_worker_id(worker_2))
                .unwrap()
                .len(),
            1
        );
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
        assert_eq!(trie.root.borrow().workers.len(), 0);
        assert_eq!(trie.root.borrow().children.len(), 1);
        assert_eq!(
            trie.root
                .borrow()
                .children
                .get(&LocalBlockHash(1))
                .unwrap()
                .borrow()
                .workers
                .len(),
            2
        );
        assert_eq!(
            trie.root
                .borrow()
                .children
                .get(&LocalBlockHash(1))
                .unwrap()
                .borrow()
                .children
                .len(),
            2
        );

        trie.apply_event(create_store_event(
            worker_2,
            4,
            vec![2, 6, 7],
            Some(ExternalSequenceBlockHash(100)),
1592
1593
        ))
        .unwrap();
1594
1595
1596
1597
1598

        let scores = trie.find_matches(
            vec![LocalBlockHash(1), LocalBlockHash(2), LocalBlockHash(3)],
            false,
        );
Yan Ru Pei's avatar
Yan Ru Pei committed
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
        assert_eq!(
            scores
                .scores
                .get(&WorkerWithDpRank::from_worker_id(worker_1))
                .unwrap(),
            &3
        );
        assert_eq!(
            scores
                .scores
                .get(&WorkerWithDpRank::from_worker_id(worker_2))
                .unwrap(),
            &2
        );
1613
1614

        assert_eq!(trie.lookup.len(), 2);
Yan Ru Pei's avatar
Yan Ru Pei committed
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
        assert_eq!(
            trie.lookup
                .get(&WorkerWithDpRank::from_worker_id(worker_1))
                .unwrap()
                .len(),
            3
        );
        assert_eq!(
            trie.lookup
                .get(&WorkerWithDpRank::from_worker_id(worker_2))
                .unwrap()
                .len(),
            4
        );
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
        assert_eq!(trie.root.borrow().workers.len(), 0);
        assert_eq!(trie.root.borrow().children.len(), 1);
        assert_eq!(
            trie.root
                .borrow()
                .children
                .get(&LocalBlockHash(1))
                .unwrap()
                .borrow()
                .workers
                .len(),
            2
        );
        assert_eq!(
            trie.root
                .borrow()
                .children
                .get(&LocalBlockHash(1))
                .unwrap()
                .borrow()
                .children
                .len(),
            2
        );
        assert_eq!(
            trie.lookup
Yan Ru Pei's avatar
Yan Ru Pei committed
1655
                .get(&WorkerWithDpRank::from_worker_id(worker_1))
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
                .unwrap()
                .get(&ExternalSequenceBlockHash(200))
                .unwrap()
                .borrow()
                .workers
                .len(),
            2
        );
        assert_eq!(
            trie.lookup
Yan Ru Pei's avatar
Yan Ru Pei committed
1666
                .get(&WorkerWithDpRank::from_worker_id(worker_2))
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
                .unwrap()
                .get(&ExternalSequenceBlockHash(200))
                .unwrap()
                .borrow()
                .workers
                .len(),
            2
        );
    }

1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
    #[test]
    fn test_radix_tree_apply_event_errors() {
        let mut trie = RadixTree::new();
        let worker_0 = 0;

        // Parent block not found
        let result = trie.apply_event(create_store_event(
            worker_0,
            0,
            vec![1, 2, 3],
            Some(ExternalSequenceBlockHash(12345)),
        ));
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            KvCacheEventError::ParentBlockNotFound
        ));

        // Block not found for remove event.
        let result = trie.apply_event(create_remove_event(worker_0, 0, vec![1, 2, 3]));
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            KvCacheEventError::BlockNotFound
        ));
    }

1704
1705
    #[test]
    fn test_remove_worker() {
1706
        setup();
1707
1708
        let mut trie = RadixTree::new();

1709
1710
        let worker_0 = 0;
        let worker_1 = 1;
1711

1712
1713
1714
1715
1716
        assert!(
            trie.find_matches(vec![LocalBlockHash(0)], false)
                .scores
                .is_empty()
        );
1717

1718
1719
1720
1721
        trie.apply_event(create_store_event(worker_0, 0, vec![0], None))
            .unwrap();
        trie.apply_event(create_store_event(worker_1, 0, vec![0], None))
            .unwrap();
1722
1723

        let result = trie.find_matches(vec![LocalBlockHash(0)], false).scores;
Yan Ru Pei's avatar
Yan Ru Pei committed
1724
1725
1726
1727
1728
        assert!(
            result.len() == 2
                && result[&WorkerWithDpRank::from_worker_id(worker_0)] == 1
                && result[&WorkerWithDpRank::from_worker_id(worker_1)] == 1
        );
1729
1730
1731
1732

        trie.remove_worker(worker_0);

        let result = trie.find_matches(vec![LocalBlockHash(0)], false).scores;
Yan Ru Pei's avatar
Yan Ru Pei committed
1733
        assert!(result.len() == 1 && result[&WorkerWithDpRank::from_worker_id(worker_1)] == 1);
1734
1735
    }

1736
1737
1738
1739
1740
1741
1742
    #[test]
    fn test_clear_all_blocks() {
        let mut trie = RadixTree::new();

        let worker_0 = 0;
        let worker_1 = 1;

1743
1744
1745
1746
1747
        assert!(
            trie.find_matches(vec![LocalBlockHash(0)], false)
                .scores
                .is_empty()
        );
1748
1749
1750

        // Test clearing an empty worker
        trie.clear_all_blocks(worker_0);
Yan Ru Pei's avatar
Yan Ru Pei committed
1751
1752
1753
1754
1755
        assert!(
            !trie
                .lookup
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_0))
        );
1756
1757

        // Test clearing a worker with shared blocks
1758
1759
1760
1761
        trie.apply_event(create_store_event(worker_0, 0, vec![0, 1, 3], None))
            .unwrap();
        trie.apply_event(create_store_event(worker_1, 0, vec![0, 2, 3], None))
            .unwrap();
1762
1763

        let result = trie.find_matches(vec![LocalBlockHash(0)], false).scores;
Yan Ru Pei's avatar
Yan Ru Pei committed
1764
1765
1766
1767
1768
        assert!(
            result.len() == 2
                && result[&WorkerWithDpRank::from_worker_id(worker_0)] == 1
                && result[&WorkerWithDpRank::from_worker_id(worker_1)] == 1
        );
1769
1770
1771

        trie.clear_all_blocks(worker_0);

Yan Ru Pei's avatar
Yan Ru Pei committed
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
        assert!(
            trie.lookup
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_0))
        );
        assert!(
            trie.lookup
                .get(&WorkerWithDpRank::from_worker_id(worker_0))
                .unwrap()
                .is_empty()
        );
1782
1783
1784
1785
        let result = trie
            .find_matches(vec![LocalBlockHash(0), LocalBlockHash(2)], false)
            .scores;
        assert_eq!(result.len(), 1);
Yan Ru Pei's avatar
Yan Ru Pei committed
1786
        assert_eq!(result[&WorkerWithDpRank::from_worker_id(worker_1)], 2);
1787
1788
1789
1790
1791
1792
1793
        let result = trie
            .find_matches(
                vec![LocalBlockHash(0), LocalBlockHash(1), LocalBlockHash(3)],
                false,
            )
            .scores;
        assert_eq!(result.len(), 1);
Yan Ru Pei's avatar
Yan Ru Pei committed
1794
        assert_eq!(result[&WorkerWithDpRank::from_worker_id(worker_1)], 1);
1795
1796

        // Test re-adding blocks after clearing worker
1797
1798
        trie.apply_event(create_store_event(worker_0, 0, vec![4, 5], None))
            .unwrap();
1799
1800
1801
1802
        let result = trie
            .find_matches(vec![LocalBlockHash(4), LocalBlockHash(5)], false)
            .scores;
        assert_eq!(result.len(), 1);
Yan Ru Pei's avatar
Yan Ru Pei committed
1803
        assert_eq!(result[&WorkerWithDpRank::from_worker_id(worker_0)], 2);
1804
1805
1806
1807

        // Test multiple clears
        trie.clear_all_blocks(worker_0);
        trie.clear_all_blocks(worker_0);
Yan Ru Pei's avatar
Yan Ru Pei committed
1808
1809
1810
1811
        assert!(
            trie.lookup
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_0))
        );
1812
1813
1814
1815
1816

        // Test clearing all workers
        trie.clear_all_blocks(worker_0);
        trie.clear_all_blocks(worker_1);
        assert!(!trie.lookup.is_empty());
Yan Ru Pei's avatar
Yan Ru Pei committed
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
        assert!(
            trie.lookup
                .get(&WorkerWithDpRank::from_worker_id(worker_0))
                .unwrap()
                .is_empty()
        );
        assert!(
            trie.lookup
                .get(&WorkerWithDpRank::from_worker_id(worker_1))
                .unwrap()
                .is_empty()
        );
1829
1830

        // Test clearing a worker that has been removed
1831
1832
1833
1834
        trie.apply_event(create_store_event(worker_0, 0, vec![6], None))
            .unwrap();
        trie.apply_event(create_store_event(worker_1, 0, vec![6], None))
            .unwrap();
1835
1836
        trie.remove_worker(worker_0);
        trie.clear_all_blocks(worker_0);
Yan Ru Pei's avatar
Yan Ru Pei committed
1837
1838
1839
1840
1841
        assert!(
            !trie
                .lookup
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_0))
        );
1842
1843
        let result = trie.find_matches(vec![LocalBlockHash(6)], false).scores;
        assert_eq!(result.len(), 1);
Yan Ru Pei's avatar
Yan Ru Pei committed
1844
        assert_eq!(result[&WorkerWithDpRank::from_worker_id(worker_1)], 1);
1845
1846
1847

        // Test clearing a worker that doesn't exist
        let worker_fake = 2;
Yan Ru Pei's avatar
Yan Ru Pei committed
1848
1849
1850
1851
1852
        assert!(
            !trie
                .lookup
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_fake))
        );
1853
        trie.clear_all_blocks(worker_fake);
Yan Ru Pei's avatar
Yan Ru Pei committed
1854
1855
1856
1857
1858
1859
1860
1861
1862
        assert!(
            !trie
                .lookup
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_fake))
        );
        assert!(
            trie.lookup
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_1))
        );
1863
1864
        let result = trie.find_matches(vec![LocalBlockHash(6)], false).scores;
        assert_eq!(result.len(), 1);
Yan Ru Pei's avatar
Yan Ru Pei committed
1865
        assert_eq!(result[&WorkerWithDpRank::from_worker_id(worker_1)], 1);
1866
1867
    }

1868
1869
    #[test]
    fn test_early_stopping() {
1870
        setup();
1871
1872
        let mut trie = RadixTree::new();

1873
1874
        let worker_0 = 0;
        let worker_1 = 1;
1875

1876
1877
1878
1879
        trie.apply_event(create_store_event(worker_0, 0, vec![0, 1, 2], None))
            .unwrap();
        trie.apply_event(create_store_event(worker_1, 0, vec![0], None))
            .unwrap();
1880
1881
1882
1883
1884
1885
1886
1887

        let result = trie
            .find_matches(
                vec![LocalBlockHash(0), LocalBlockHash(1), LocalBlockHash(2)],
                true,
            )
            .scores;

Yan Ru Pei's avatar
Yan Ru Pei committed
1888
1889
1890
1891
1892
        assert!(
            result.len() == 2
                && result[&WorkerWithDpRank::from_worker_id(worker_0)] == 2
                && result[&WorkerWithDpRank::from_worker_id(worker_1)] == 1
        );
1893
1894
1895
1896

        let result = trie
            .find_matches(vec![LocalBlockHash(0), LocalBlockHash(1)], true)
            .scores;
Yan Ru Pei's avatar
Yan Ru Pei committed
1897
1898
1899
1900
1901
        assert!(
            result.len() == 2
                && result[&WorkerWithDpRank::from_worker_id(worker_0)] == 2
                && result[&WorkerWithDpRank::from_worker_id(worker_1)] == 1
        );
1902
1903
    }

1904
1905
1906
1907
    #[rstest]
    #[case(11)]
    #[case(32)]
    #[case(64)]
1908
    fn test_compute_block_hash_for_seq(#[case] kv_block_size: u32) {
1909
        setup();
1910
        // create a sequence of 64 elements
1911
        let sequence = (0..kv_block_size).collect::<Vec<u32>>();
1912
        let hashes = compute_block_hash_for_seq(&sequence, kv_block_size);
1913
1914
1915
        assert_eq!(hashes.len(), 1);

        // create a sequence of 65 elements
1916
        let sequence = (0..(kv_block_size + 1)).collect::<Vec<u32>>();
1917
        let hashes = compute_block_hash_for_seq(&sequence, kv_block_size);
1918
1919
1920
        assert_eq!(hashes.len(), 1);

        // create a sequence of 129 elements
1921
        let sequence = (0..(2 * kv_block_size + 1)).collect::<Vec<u32>>();
1922
        let hashes = compute_block_hash_for_seq(&sequence, kv_block_size);
1923
1924
1925
        assert_eq!(hashes.len(), 2);
    }

1926
1927
1928
    fn make_indexer(
        token: &CancellationToken,
        num_shards: usize,
1929
        kv_block_size: u32,
1930
    ) -> Box<dyn KvIndexerInterface> {
1931
        let metrics = KvIndexerMetrics::new_unregistered();
1932
        if num_shards == 1 {
1933
            Box::new(KvIndexer::new(token.clone(), kv_block_size, metrics.into()))
1934
        } else {
1935
1936
1937
1938
            Box::new(KvIndexerSharded::new(
                token.clone(),
                num_shards,
                kv_block_size,
1939
                metrics.into(),
1940
            ))
1941
1942
1943
        }
    }

1944
    #[template]
1945
    #[rstest]
1946
1947
1948
1949
1950
1951
    fn indexer_template(
        #[values(1, 3, 8)] num_shards: usize,
        #[values(11, 32, 64)] kv_block_size: usize,
    ) {
    }

1952
    #[tokio::test]
1953
    #[apply(indexer_template)]
1954
    async fn test_kv_indexer_new(num_shards: usize, kv_block_size: u32) {
1955
        setup();
1956
1957
        let token: CancellationToken = CancellationToken::new();
        let _ = make_indexer(&token, num_shards, kv_block_size);
1958
1959
1960
    }

    #[tokio::test]
1961
    #[apply(indexer_template)]
1962
    async fn test_find_matches(num_shards: usize, kv_block_size: u32) {
1963
        setup();
1964
        let token = CancellationToken::new();
1965
        let kv_indexer = make_indexer(&token, num_shards, kv_block_size);
1966
1967
1968
1969
1970
1971
1972
1973

        let sequence = vec![compute_block_hash(b"test data")];
        let scores = kv_indexer.find_matches(sequence).await;

        assert!(scores.unwrap().scores.is_empty());
    }

    #[tokio::test]
1974
    #[apply(indexer_template)]
1975
    async fn test_find_matches_for_request(num_shards: usize, kv_block_size: u32) {
1976
        setup();
1977
        let token = CancellationToken::new();
1978
        let kv_indexer = make_indexer(&token, num_shards, kv_block_size);
1979
1980
1981
1982
1983
1984
1985
1986

        let tokens = vec![1, 2, 3, 4];
        let scores = kv_indexer.find_matches_for_request(&tokens).await;

        assert!(scores.unwrap().scores.is_empty());
    }

    #[tokio::test]
1987
    #[apply(indexer_template)]
1988
    async fn test_apply_event(num_shards: usize, kv_block_size: u32) {
1989
        setup();
1990
        let worker_id = 0;
1991
1992

        let token = CancellationToken::new();
1993
        let mut kv_indexer = make_indexer(&token, num_shards, kv_block_size);
1994
1995
1996
1997
1998
1999
2000
2001

        let event = create_store_event(worker_id, 1, vec![1, 2, 3], None);
        kv_indexer.apply_event(event).await;

        // No assertion here, just ensuring it runs without panic
    }

    #[tokio::test]
2002
    #[apply(indexer_template)]
2003
    async fn test_shutdown(num_shards: usize, kv_block_size: u32) {
2004
        setup();
2005
        let token = CancellationToken::new();
2006
        let mut kv_indexer = make_indexer(&token, num_shards, kv_block_size);
2007
2008
2009
2010
2011

        kv_indexer.shutdown();
    }

    #[tokio::test]
2012
    #[apply(indexer_template)]
2013
    async fn test_frequency(num_shards: usize, kv_block_size: u32) {
2014
2015
2016
        const ONE_MILLIS: Duration = Duration::from_millis(1);

        setup();
2017
2018
        let mut kv_indexer: Box<dyn KvIndexerInterface>;
        let token = CancellationToken::new();
2019
        let expiration = Duration::from_millis(50);
2020
        let metrics = Arc::new(KvIndexerMetrics::new_unregistered());
2021
2022

        if num_shards == 1 {
2023
2024
            kv_indexer = Box::new(KvIndexer::new_with_frequency(
                token,
2025
                Some(expiration),
2026
                kv_block_size,
2027
                metrics,
2028
            ));
2029
2030
        } else {
            kv_indexer = Box::new(KvIndexerSharded::new_with_frequency(
2031
2032
                token,
                num_shards,
2033
                Some(expiration),
2034
                kv_block_size,
2035
                metrics,
2036
2037
2038
            ));
        }

2039
        // The blocks
2040
2041
2042
2043
2044
2045
2046
        let block_hashes = vec![
            LocalBlockHash(1),
            LocalBlockHash(2),
            LocalBlockHash(3),
            LocalBlockHash(4),
        ];

2047
2048
2049
2050
2051
2052
2053
2054
2055
2056
2057
        let overlap = kv_indexer.find_matches(block_hashes.clone()).await.unwrap();
        assert_eq!(
            overlap.frequencies.len(),
            0,
            "Should be no cached blocks yet"
        );

        // Blocks go in cache
        let worker_id = 0;
        let event = create_store_event(worker_id, 0, vec![1, 2, 3, 4], None);
        kv_indexer.apply_event(event).await;
2058

2059
2060
2061
2062
2063
2064
2065
2066
2067
2068
2069
2070
2071
2072
2073
2074
2075
2076
2077
2078
2079
2080
2081
2082
2083
2084
2085
2086
        // First access
        // The store event is applied async so poll briefly
        let mut overlap = OverlapScores::default();
        let timeout = Duration::from_millis(10);
        let start = Instant::now();
        while overlap.scores.is_empty() && Instant::now().duration_since(start) < timeout {
            time::sleep(ONE_MILLIS).await;
            overlap = kv_indexer.find_matches(block_hashes.clone()).await.unwrap();
        }
        assert_eq!(
            overlap.scores.len(),
            1,
            "One worker has these blocks cached"
        );
        assert_eq!(
            overlap.frequencies.len(),
            0,
            "Blocks have not previously been accessed"
        );

        // Second access
        let overlap = kv_indexer.find_matches(block_hashes.clone()).await.unwrap();
        assert_eq!(overlap.scores.len(), 1, "Still one worker matches");
        assert_eq!(
            overlap.frequencies,
            vec![1, 1, 1, 1],
            "We should see the first access now"
        );
2087

2088
2089
        // Let those two accesses expire
        time::sleep(expiration + Duration::from_millis(10)).await;
2090

2091
2092
2093
2094
2095
2096
2097
        // New first access
        let overlap = kv_indexer.find_matches(block_hashes.clone()).await.unwrap();
        assert_eq!(
            overlap.frequencies.len(),
            0,
            "Blocks were accessed too long ago"
        );
2098

2099
2100
        // New second access
        let _ = kv_indexer.find_matches(block_hashes.clone()).await.unwrap();
2101

2102
2103
        // Access only the first three blocks
        let overlap = kv_indexer
2104
2105
2106
            .find_matches(block_hashes[0..3].to_vec())
            .await
            .unwrap();
2107
2108
        // We see the previous two new accesses
        assert_eq!(overlap.frequencies, vec![2, 2, 2]);
2109

2110
2111
2112
        // The third access did not touch the last block
        let overlap = kv_indexer.find_matches(block_hashes.clone()).await.unwrap();
        assert_eq!(overlap.frequencies, vec![3, 3, 3, 2]);
2113
2114
2115
2116
    }

    #[test]
    fn test_router_event_new() {
2117
        setup();
2118
        let worker_id = 0;
2119
2120
2121
2122
2123
2124
2125
2126
2127
        let kv_cache_event = KvCacheEvent {
            event_id: 1,
            data: KvCacheEventData::Stored(KvCacheStoreData {
                parent_hash: None,
                blocks: vec![KvCacheStoredBlockData {
                    block_hash: ExternalSequenceBlockHash(0),
                    tokens_hash: LocalBlockHash(13226331709069118873),
                }],
            }),
Yan Ru Pei's avatar
Yan Ru Pei committed
2128
            dp_rank: 0,
2129
2130
2131
2132
2133
2134
2135
2136
2137
2138
2139
2140
2141
2142
2143
2144
2145
2146
2147
        };
        let router_event = RouterEvent::new(worker_id, kv_cache_event);

        assert_eq!(router_event.worker_id, worker_id);
        assert_eq!(router_event.event.event_id, 1);
        if let KvCacheEventData::Stored(store_op) = &router_event.event.data {
            assert_eq!(store_op.blocks.len(), 1);
            assert_eq!(
                store_op.blocks[0].tokens_hash,
                compute_block_hash(b"test data")
            );
            assert_eq!(store_op.blocks[0].block_hash, ExternalSequenceBlockHash(0));
        } else {
            panic!("Expected KvCacheEventData::Stored");
        }
    }

    #[test]
    fn test_radix_tree_default() {
2148
        setup();
2149
2150
2151
2152
2153
2154
2155
2156
        let radix_tree: RadixTree = Default::default();
        assert!(radix_tree.root.borrow().children.is_empty());
        assert!(radix_tree.root.borrow().workers.is_empty());
        assert!(radix_tree.lookup.is_empty());
    }

    #[test]
    fn test_overlap_scores_default() {
2157
        setup();
2158
2159
2160
        let overlap_scores: OverlapScores = Default::default();
        assert!(overlap_scores.scores.is_empty());
    }
2161
2162
2163
2164
2165
2166
2167
2168

    #[tokio::test]
    async fn test_dump_tree_as_events_round_trip() {
        setup();

        // Configuration
        let kv_block_size = 32;
        let num_shards = 2;
2169
        let metrics = Arc::new(KvIndexerMetrics::new_unregistered());
2170
2171
2172

        // Build a non-trivial indexer with events
        let token1 = CancellationToken::new();
2173
2174
        let mut original_indexer =
            KvIndexerSharded::new(token1.clone(), num_shards, kv_block_size, metrics.clone());
2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
2206
2207
2208
2209
2210
2211
2212
2213
2214
2215
2216
2217
2218
2219

        let worker_0 = 0;
        let worker_1 = 1;
        let worker_2 = 2;

        // Apply events to the original indexer
        original_indexer
            .apply_event(create_store_event(worker_0, 0, vec![1, 2, 3], None))
            .await;

        original_indexer
            .apply_event(create_store_event(worker_1, 1, vec![1, 2, 3], None))
            .await;
        original_indexer
            .apply_event(create_store_event(
                worker_1,
                2,
                vec![4, 5],
                Some(ExternalSequenceBlockHash(100)),
            ))
            .await;

        original_indexer
            .apply_event(create_store_event(worker_2, 3, vec![6, 7], None))
            .await;

        original_indexer
            .apply_event(create_store_event(
                worker_0,
                4,
                vec![4],
                Some(ExternalSequenceBlockHash(100)),
            ))
            .await;

        // Allow some time for events to be processed
        tokio::time::sleep(Duration::from_millis(50)).await;

        // Dump the original indexer
        let dump1 = original_indexer.dump_events().await.unwrap();
        println!("Dumped {} events", dump1.len());

        // Create a new indexer and apply all dumped events
        let token2 = CancellationToken::new();
        let mut reconstructed_indexer =
2220
            KvIndexerSharded::new(token2.clone(), num_shards, kv_block_size, metrics);
2221
2222
2223
2224
2225
2226
2227
2228
2229
2230
2231
2232
2233
2234
2235
2236
2237
2238
2239
2240
2241
2242
2243
2244
2245
2246
2247
2248
2249
2250
2251
2252
2253
2254
2255
2256
2257
2258
2259
2260
2261
2262
2263
2264
2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
2308
2309
2310
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
2322
2323
2324
2325
2326
2327
2328
2329
2330
2331
2332
2333
2334
2335

        for event in &dump1 {
            reconstructed_indexer.apply_event(event.clone()).await;
        }

        // Allow some time for events to be processed
        tokio::time::sleep(Duration::from_millis(50)).await;

        // Dump the reconstructed indexer
        let dump2 = reconstructed_indexer.dump_events().await.unwrap();

        // Sort both dumps for comparison (order might differ due to HashMap iteration and sharding)
        let mut sorted_dump1 = dump1.clone();
        let mut sorted_dump2 = dump2.clone();

        // Sort by (worker_id, tokens_hash, parent_hash)
        let sort_key = |event: &RouterEvent| {
            if let KvCacheEventData::Stored(ref data) = event.event.data {
                (
                    event.worker_id,
                    data.blocks.first().map(|b| b.tokens_hash.0).unwrap_or(0),
                    data.parent_hash.map(|h| h.0).unwrap_or(0),
                )
            } else {
                (event.worker_id, 0, 0)
            }
        };

        sorted_dump1.sort_by_key(sort_key);
        sorted_dump2.sort_by_key(sort_key);

        // Verify the dumps have the same length
        assert_eq!(
            sorted_dump1.len(),
            sorted_dump2.len(),
            "Dumps have different lengths: {} vs {}",
            sorted_dump1.len(),
            sorted_dump2.len()
        );

        // Verify each event matches
        for (i, (event1, event2)) in sorted_dump1.iter().zip(sorted_dump2.iter()).enumerate() {
            assert_eq!(
                event1.worker_id, event2.worker_id,
                "Event {} worker_id mismatch",
                i
            );

            if let (KvCacheEventData::Stored(data1), KvCacheEventData::Stored(data2)) =
                (&event1.event.data, &event2.event.data)
            {
                assert_eq!(
                    data1.parent_hash, data2.parent_hash,
                    "Event {} parent_hash mismatch",
                    i
                );
                assert_eq!(
                    data1.blocks.len(),
                    data2.blocks.len(),
                    "Event {} blocks length mismatch",
                    i
                );

                for (j, (block1, block2)) in
                    data1.blocks.iter().zip(data2.blocks.iter()).enumerate()
                {
                    assert_eq!(
                        block1.tokens_hash, block2.tokens_hash,
                        "Event {} block {} tokens_hash mismatch",
                        i, j
                    );
                    assert_eq!(
                        block1.block_hash, block2.block_hash,
                        "Event {} block {} block_hash mismatch",
                        i, j
                    );
                }
            } else {
                panic!("Expected Stored events in both dumps");
            }
        }

        // Also verify that both indexers produce the same match results
        for test_seq in [
            vec![LocalBlockHash(1), LocalBlockHash(2), LocalBlockHash(3)],
            vec![LocalBlockHash(1), LocalBlockHash(4), LocalBlockHash(5)],
            vec![LocalBlockHash(6), LocalBlockHash(7)],
            vec![LocalBlockHash(1)],
        ] {
            let scores1 = original_indexer
                .find_matches(test_seq.clone())
                .await
                .unwrap();
            let scores2 = reconstructed_indexer
                .find_matches(test_seq.clone())
                .await
                .unwrap();

            // Sort the scores to compare
            let mut scores1_sorted: Vec<_> = scores1.scores.iter().collect();
            let mut scores2_sorted: Vec<_> = scores2.scores.iter().collect();
            scores1_sorted.sort_by_key(|(k, _)| *k);
            scores2_sorted.sort_by_key(|(k, _)| *k);

            assert_eq!(
                scores1_sorted, scores2_sorted,
                "Match scores differ for sequence {:?}",
                test_seq
            );
        }

        // Clean up
        original_indexer.shutdown();
        reconstructed_indexer.shutdown();
    }
2336
2337
2338
2339
2340
2341
2342
2343
2344
2345
2346
2347
2348
2349
2350
2351
2352
2353
2354
2355
2356
2357
2358
2359
2360
2361
2362
2363
2364
2365
2366
2367
2368
2369
2370
2371
2372
2373
2374
2375
2376
2377
2378
2379
2380

    #[test]
    fn test_increment_event_applied() {
        let metrics = KvIndexerMetrics::new_unregistered();

        metrics.increment_event_applied(METRIC_EVENT_STORED, Ok(()));
        assert_eq!(
            metrics
                .kv_cache_events_applied
                .get_metric_with_label_values(&[METRIC_EVENT_STORED, METRIC_STATUS_OK])
                .unwrap()
                .get(),
            1
        );

        metrics.increment_event_applied(
            METRIC_EVENT_STORED,
            Err(KvCacheEventError::ParentBlockNotFound),
        );
        assert_eq!(
            metrics
                .kv_cache_events_applied
                .get_metric_with_label_values(&[
                    METRIC_EVENT_STORED,
                    METRIC_STATUS_PARENT_NOT_FOUND
                ])
                .unwrap()
                .get(),
            1
        );

        metrics
            .increment_event_applied(METRIC_EVENT_REMOVED, Err(KvCacheEventError::BlockNotFound));
        assert_eq!(
            metrics
                .kv_cache_events_applied
                .get_metric_with_label_values(&[
                    METRIC_EVENT_REMOVED,
                    METRIC_STATUS_BLOCK_NOT_FOUND
                ])
                .unwrap()
                .get(),
            1
        );
    }
2381
2382
2383
2384
2385
2386
2387
2388
2389
2390
2391
2392
2393
2394
2395
2396
2397
2398
2399

    #[test]
    fn test_remove_worker_verifies_hash_removal() {
        setup();
        let mut trie = RadixTree::new();

        let worker_0 = 0;
        let worker_1 = 1;
        let worker_2 = 2;

        // Add blocks for multiple workers
        trie.apply_event(create_store_event(worker_0, 0, vec![1, 2, 3], None))
            .unwrap();
        trie.apply_event(create_store_event(worker_1, 0, vec![1, 2, 3], None))
            .unwrap();
        trie.apply_event(create_store_event(worker_2, 0, vec![1, 4, 5], None))
            .unwrap();

        // Verify worker_0 has 3 blocks in lookup
Yan Ru Pei's avatar
Yan Ru Pei committed
2400
2401
2402
2403
2404
2405
2406
        assert_eq!(
            trie.lookup
                .get(&WorkerWithDpRank::from_worker_id(worker_0))
                .unwrap()
                .len(),
            3
        );
2407
2408
2409
2410

        // Verify that blocks have the correct workers
        let block_1 = trie
            .lookup
Yan Ru Pei's avatar
Yan Ru Pei committed
2411
            .get(&WorkerWithDpRank::from_worker_id(worker_0))
2412
2413
2414
2415
            .unwrap()
            .get(&ExternalSequenceBlockHash(100))
            .unwrap();
        assert_eq!(block_1.borrow().workers.len(), 3); // worker_0, worker_1, and worker_2 (all have hash 1)
Yan Ru Pei's avatar
Yan Ru Pei committed
2416
2417
2418
2419
2420
2421
2422
2423
2424
2425
2426
2427
2428
2429
2430
2431
2432
2433
        assert!(
            block_1
                .borrow()
                .workers
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_0))
        );
        assert!(
            block_1
                .borrow()
                .workers
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_1))
        );
        assert!(
            block_1
                .borrow()
                .workers
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_2))
        );
2434
2435
2436
2437
2438

        // Remove worker_0
        trie.remove_worker(worker_0);

        // Verify worker_0 is completely removed from lookup table
Yan Ru Pei's avatar
Yan Ru Pei committed
2439
2440
2441
2442
2443
        assert!(
            !trie
                .lookup
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_0))
        );
2444
2445
2446
2447
2448
        assert_eq!(trie.lookup.len(), 2);

        // Verify that worker_0's hash is removed from the workers set
        let block_1 = trie
            .lookup
Yan Ru Pei's avatar
Yan Ru Pei committed
2449
            .get(&WorkerWithDpRank::from_worker_id(worker_1))
2450
2451
2452
2453
            .unwrap()
            .get(&ExternalSequenceBlockHash(100))
            .unwrap();
        assert_eq!(block_1.borrow().workers.len(), 2); // worker_1 and worker_2 remain
Yan Ru Pei's avatar
Yan Ru Pei committed
2454
2455
2456
2457
2458
2459
2460
2461
2462
2463
2464
2465
2466
2467
2468
2469
2470
2471
        assert!(
            !block_1
                .borrow()
                .workers
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_0))
        );
        assert!(
            block_1
                .borrow()
                .workers
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_1))
        );
        assert!(
            block_1
                .borrow()
                .workers
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_2))
        );
2472
2473
2474
2475
2476

        // Verify that blocks with no remaining workers have their children cleared
        // This tests the optimization where empty blocks clear their children
        let block_2 = trie
            .lookup
Yan Ru Pei's avatar
Yan Ru Pei committed
2477
            .get(&WorkerWithDpRank::from_worker_id(worker_1))
2478
2479
2480
2481
            .unwrap()
            .get(&ExternalSequenceBlockHash(200))
            .unwrap();
        assert_eq!(block_2.borrow().workers.len(), 1); // only worker_1
Yan Ru Pei's avatar
Yan Ru Pei committed
2482
2483
2484
2485
2486
2487
        assert!(
            block_2
                .borrow()
                .workers
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_1))
        );
2488
2489
2490
2491
2492
2493
2494
2495
2496

        // Verify match results no longer include worker_0
        let result = trie
            .find_matches(
                vec![LocalBlockHash(1), LocalBlockHash(2), LocalBlockHash(3)],
                false,
            )
            .scores;
        assert_eq!(result.len(), 2);
Yan Ru Pei's avatar
Yan Ru Pei committed
2497
2498
2499
        assert!(!result.contains_key(&WorkerWithDpRank::from_worker_id(worker_0)));
        assert!(result.contains_key(&WorkerWithDpRank::from_worker_id(worker_1)));
        assert!(result.contains_key(&WorkerWithDpRank::from_worker_id(worker_2)));
2500
    }
2501
}