indexer.rs 104 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
    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;

57
use crate::kv_router::approx::{BlockEntry, PruneConfig, PruneManager};
58
use crate::kv_router::protocols::*;
59
use crate::tokens::{SequenceHash, TokenBlockSequence};
60
61
62
63
64
65
66
67
68
69
70
71

/// 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,
72
73
74

    #[error("Prune operation failed: {0}")]
    PruneFailed(String),
75
76
}

77
78
79
80
81
82
83
84
/// 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,
85
86
87

    #[error("Invalid block sequence")]
    InvalidBlockSequence,
88
89
}

90
91
92
/// A shared reference to a [`RadixBlock`].
type SharedRadixBlock = Rc<RefCell<RadixBlock>>;

93
94
95
96
pub fn compute_hash(data: &[u8]) -> u64 {
    xxh3::xxh3_64_with_seed(data, XXH3_SEED)
}

97
98
99
100
101
102
103
104
105
106
/// 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 {
107
    LocalBlockHash(compute_hash(data))
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
}

// /// 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.
129
pub fn compute_block_hash_for_seq(tokens: &[u32], kv_block_size: u32) -> Vec<LocalBlockHash> {
130
    tokens
131
        .chunks_exact(kv_block_size as usize) // Split into chunks of kv_block_size elements
132
133
134
135
136
137
138
139
140
141
142
        .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()
}

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
170
171
172
173
174
175
176
/// 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
}

177
/// A [`KvCacheEvent`] on a specific LLM worker denoted by [`WorkerId`].
178
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
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.
203
#[derive(Debug)]
204
205
206
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
207
    /// A map of workers (with dp_rank) to their external sequence block hash for this block.
208
    /// The external hash is preserved to speed up snapshotting.
Yan Ru Pei's avatar
Yan Ru Pei committed
209
    workers: HashMap<WorkerWithDpRank, ExternalSequenceBlockHash>,
210
211
212
213
214
215
216
217
218
219
220
221
222
    /// 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(),
223
            workers: HashMap::new(),
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
            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
242
    lookup: HashMap<WorkerWithDpRank, HashMap<ExternalSequenceBlockHash, SharedRadixBlock>>,
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
278
279
280
281
282
283
284
    /// 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();
285
286
287
288
289
290
291

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

        for (idx, block_hash) in sequence.iter().enumerate() {
292
293
            let next_block = {
                let current_borrow = current.borrow();
294
                current_borrow.children.get(block_hash).cloned()
295
296
            };
            if let Some(block) = next_block {
297
                scores.update_scores(block.borrow().workers.keys());
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318

                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 {
319
320
321
322
323
                tracing::trace!(
                    "RadixTree::find_matches: block not found at index {} for hash {}",
                    idx,
                    block_hash.0
                );
324
325
326
327
                break;
            }
        }

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

330
331
332
333
334
335
336
337
338
339
        // Populate tree sizes for all workers that have scores
        for worker in scores.scores.keys() {
            let tree_size = self
                .lookup
                .get(worker)
                .expect("worker in scores must exist in lookup table")
                .len();
            scores.tree_sizes.insert(*worker, tree_size);
        }

340
341
342
343
344
345
346
347
        scores
    }

    /// Apply a [`RouterEvent`] to the radix tree.
    ///
    /// ### Arguments
    ///
    /// * `event` - The `RouterEvent` to apply.
348
    pub fn apply_event(&mut self, event: RouterEvent) -> Result<(), KvCacheEventError> {
Yan Ru Pei's avatar
Yan Ru Pei committed
349
350
351
352
353
354
        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);

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

Yan Ru Pei's avatar
Yan Ru Pei committed
357
        let worker_lookup = self.lookup.entry(worker).or_default();
358
359
360
361
362
363
364

        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 {
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
                    Some(parent) => match worker_lookup.get(&parent) {
                        Some(current) => current.clone(),
                        None => {
                            tracing::warn!(
                                worker_id = worker.worker_id.to_string(),
                                dp_rank = worker.dp_rank,
                                id,
                                parent_hash = ?op.parent_hash,
                                num_blocks = op.blocks.len(),
                                "Failed to find parent block; skipping store operation"
                            );
                            return Err(KvCacheEventError::ParentBlockNotFound);
                        }
                    },
                    None => self.root.clone(),
380
381
                };

382
383
384
385
386
387
388
389
390
                fn process_blocks(
                    parent: SharedRadixBlock,
                    blocks: &[KvCacheStoredBlockData],
                    worker: WorkerWithDpRank,
                    worker_lookup: &mut HashMap<ExternalSequenceBlockHash, SharedRadixBlock>,
                    id: u64,
                ) -> Result<(), KvCacheEventError> {
                    if blocks.is_empty() {
                        return Ok(());
391
392
                    }

393
394
395
396
                    let mut parent_mut = parent.borrow_mut();
                    let block_data = &blocks[0];

                    let child = match parent_mut.children.get(&block_data.tokens_hash) {
397
398
399
400
                        Some(block) => block.clone(),
                        None => {
                            // create new block - automatically added to the lookup table
                            let new_block = worker_lookup
401
                                .get(&block_data.block_hash)
402
403
404
405
                                .cloned()
                                .unwrap_or_else(|| Rc::new(RefCell::new(RadixBlock::new())));

                            // insert into radix tree
406
                            parent_mut
407
                                .children
408
                                .insert(block_data.tokens_hash, new_block.clone());
409

410
411
412
413
                            new_block
                        }
                    };

414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
                    // Update child and check for cycles
                    {
                        // Try to borrow the child mutably - if it fails, it's already borrowed
                        // in the ancestor chain (parent_mut is alive + all ancestors in recursive stack)
                        let mut child_mut = match child.try_borrow_mut() {
                            Ok(b) => b,
                            Err(_) => {
                                tracing::warn!(
                                    worker_id = worker.worker_id.to_string(),
                                    dp_rank = worker.dp_rank,
                                    id,
                                    block_hash = ?block_data.block_hash,
                                    "Detected cycle in store event (block already in parent chain); rejecting sequence"
                                );
                                return Err(KvCacheEventError::InvalidBlockSequence);
                            }
                        };
431

432
433
434
                        // add our worker to the block with its external hash
                        child_mut.workers.insert(worker, block_data.block_hash);
                    }
435

436
437
                    // add the block to the worker_id lookup table
                    worker_lookup.insert(block_data.block_hash, child.clone());
438

439
440
                    // Recurse with the child and remaining blocks
                    process_blocks(child, &blocks[1..], worker, worker_lookup, id)
441
                }
442
443

                process_blocks(current, &op.blocks, worker, worker_lookup, id)
444
445
            }
            KvCacheEventData::Removed(remove) => {
446
                // tracing::trace!(id, "KV Remove Operation: {:?}", op);
447
448
449
450
451
452
453
454
455
456
                // 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 => {
457
                            tracing::warn!(
Yan Ru Pei's avatar
Yan Ru Pei committed
458
459
                                worker_id = worker.worker_id.to_string(),
                                dp_rank = worker.dp_rank,
460
                                id,
Yan Ru Pei's avatar
Yan Ru Pei committed
461
                                block_hash = ?block,
462
463
                                "Failed to find block to remove; skipping remove operation"
                            );
464
                            return Err(KvCacheEventError::BlockNotFound);
465
466
467
468
                        }
                    };

                    let mut guard = entry.borrow_mut();
Yan Ru Pei's avatar
Yan Ru Pei committed
469
                    guard.workers.remove(&worker);
470
                    if guard.workers.is_empty() {
471
                        // if no workers are using this block, that is true for all children
472
473
474
475
476
                        guard.children.clear();
                    }
                    // remove the block from the lookup table
                    worker_lookup.remove(&block);
                }
477
                Ok(())
478
            }
479
            KvCacheEventData::Cleared => {
Yan Ru Pei's avatar
Yan Ru Pei committed
480
                self.clear_all_blocks(worker.worker_id);
481
                Ok(())
482
            }
483
484
485
        }
    }

Yan Ru Pei's avatar
Yan Ru Pei committed
486
487
488
489
490
491
492
493
494
495
496
    /// 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();
497

Yan Ru Pei's avatar
Yan Ru Pei committed
498
499
500
501
502
503
504
505
506
        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();
                    }
                });
507

Yan Ru Pei's avatar
Yan Ru Pei committed
508
509
510
                if keep_worker {
                    // Re-insert worker with empty blocks map to keep it tracked
                    self.lookup.insert(worker_key, HashMap::new());
511
                }
512
513
514
            }
        }
    }
515

Yan Ru Pei's avatar
Yan Ru Pei committed
516
517
518
519
520
521
522
523
    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);
    }

524
    /// Get all worker IDs currently tracked in the radix tree.
Yan Ru Pei's avatar
Yan Ru Pei committed
525
    /// Returns unique worker_ids (ignoring dp_rank differences).
526
    pub fn get_workers(&self) -> Vec<WorkerId> {
Yan Ru Pei's avatar
Yan Ru Pei committed
527
528
529
530
        let mut worker_ids: Vec<WorkerId> = self.lookup.keys().map(|w| w.worker_id).collect();
        worker_ids.sort_unstable();
        worker_ids.dedup();
        worker_ids
531
532
    }

533
534
535
536
    /// 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> {
537
538
539
540
541
        tracing::debug!(
            "Dumping radix tree as events (contains information about {:?} workers)",
            self.lookup.len()
        );

542
543
544
        let mut events = Vec::new();
        let mut event_id = 0u64;

545
        // BFS queue: (current_block, parent_hashes_per_worker, tokens_hash)
Yan Ru Pei's avatar
Yan Ru Pei committed
546
        // parent_hashes_per_worker maps WorkerWithDpRank -> ExternalSequenceBlockHash
547
548
        let mut queue: VecDeque<(
            SharedRadixBlock,
Yan Ru Pei's avatar
Yan Ru Pei committed
549
            HashMap<WorkerWithDpRank, ExternalSequenceBlockHash>,
550
551
            LocalBlockHash,
        )> = VecDeque::new();
552
553
554
555

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

560
        while let Some((current_block, parent_hashes, tokens_hash)) = queue.pop_front() {
561
562
            let current_borrow = current_block.borrow();

563
564
            // Map of this block's external hashes per worker (for children to use as parent)
            let mut current_external_hashes = HashMap::new();
565
566

            // For each worker that has this block
567
568
569
570
571
572
            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
573
                    worker_id: worker_id.worker_id,
574
575
576
577
578
579
580
581
582
                    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
583
                        dp_rank: worker_id.dp_rank,
584
585
586
587
                    },
                };
                events.push(event);
                event_id += 1;
588

589
590
591
                // Track this block's external hash for this worker
                current_external_hashes.insert(*worker_id, *external_hash);
            }
592

593
            // Enqueue children with per-worker parent hashes
594
            for (child_tokens_hash, child_block) in &current_borrow.children {
595
596
597
598
599
                queue.push_back((
                    child_block.clone(),
                    current_external_hashes.clone(),
                    *child_tokens_hash,
                ));
600
601
602
603
604
            }
        }

        events
    }
605
606

    pub fn current_size(&self) -> usize {
607
        self.lookup.values().map(|m| m.len()).sum()
608
    }
609
610
}

611
612
613
614
615
616
617
618
619
620
621
/// 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";
622
pub const METRIC_STATUS_INVALID_BLOCK: &str = "invalid_block";
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641

/// 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(|| {
642
            match component.metrics().create_intcountervec(
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
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
689
690
691
692
693
694
                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,
695
                    KvCacheEventError::InvalidBlockSequence => METRIC_STATUS_INVALID_BLOCK,
696
697
698
699
700
701
702
703
704
                };
                self.kv_cache_events_applied
                    .with_label_values(&[event_type, error_label])
                    .inc_by(1);
            }
        }
    }
}

Yan Ru Pei's avatar
Yan Ru Pei committed
705
/// Scores representing the overlap of workers (with their dp_rank).
706
707
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OverlapScores {
Yan Ru Pei's avatar
Yan Ru Pei committed
708
709
    // map of worker (with dp_rank) to score
    pub scores: HashMap<WorkerWithDpRank, u32>,
710
711
    // List of frequencies that the blocks have been accessed. Entries with value 0 are omitted.
    pub frequencies: Vec<usize>,
712
713
    // Map of worker to their tree size (number of blocks in the tree for that worker)
    pub tree_sizes: HashMap<WorkerWithDpRank, usize>,
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
}

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),
732
            tree_sizes: HashMap::new(),
733
734
735
736
737
738
739
        }
    }

    /// Update the scores with a set of workers.
    ///
    /// ### Arguments
    ///
Yan Ru Pei's avatar
Yan Ru Pei committed
740
    /// * `workers` - An iterator over `WorkerWithDpRank` references.
741
742
    pub fn update_scores<'a, I>(&mut self, workers: I)
    where
Yan Ru Pei's avatar
Yan Ru Pei committed
743
        I: IntoIterator<Item = &'a WorkerWithDpRank>,
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
        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>,
}

772
773
774
775
776
777
/// A request to dump the tree as events
pub struct DumpRequest {
    /// Channel to send the dumped events
    pub resp: oneshot::Sender<Vec<RouterEvent>>,
}

778
779
780
781
782
783
/// A request to get all workers currently tracked
pub struct GetWorkersRequest {
    /// Channel to send the worker IDs
    pub resp: oneshot::Sender<Vec<WorkerId>>,
}

784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
#[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);
830
831
832
833
834
835
836

    /// 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>;
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869

    /// Process a routing decision with pre-computed hashes.
    ///
    /// ### Arguments
    ///
    /// * `worker` - The worker (with dp_rank) that was selected.
    /// * `local_hashes` - The local hashes of the tokens sent to the worker.
    /// * `sequence_hashes` - The sequence hashes of the tokens sent to the worker.
    async fn process_routing_decision(
        &self,
        worker: WorkerWithDpRank,
        local_hashes: Vec<LocalBlockHash>,
        sequence_hashes: Vec<SequenceHash>,
    ) -> Result<(), KvRouterError>;

    /// Process a routing decision for a request with tokens.
    ///
    /// ### Arguments
    ///
    /// * `tokens` - A vector of `u32` tokens.
    /// * `worker` - The worker (with dp_rank) that was selected.
    async fn process_routing_decision_for_request(
        &self,
        tokens: &[u32],
        worker: WorkerWithDpRank,
    ) -> Result<(), KvRouterError>;
}

/// A request to process a routing decision.
struct RoutingDecisionRequest {
    worker: WorkerWithDpRank,
    local_hashes: Vec<LocalBlockHash>,
    sequence_hashes: Vec<SequenceHash>,
870
871
872
873
874
875
876
877
878
879
880
881
}

/// 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>,
882
883
    /// A sender for get workers requests.
    get_workers_tx: mpsc::Sender<GetWorkersRequest>,
884
885
    /// A sender for dump requests.
    dump_tx: mpsc::Sender<DumpRequest>,
886
887
    /// A sender for routing decision requests.
    routing_tx: mpsc::Sender<RoutingDecisionRequest>,
888
889
    /// A handle to the background task managing the KV store.
    task: OnceLock<std::thread::JoinHandle<()>>,
890
    /// The size of the KV block this indexer can handle.
891
    kv_block_size: u32,
892
893
894
895
896
897
898
899
900
}

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.
901
902
    /// * `ttl` - The time-to-live for blocks before they expire.
    /// * `prune_config` - Configuration for tree-size based pruning.
903
904
905
906
907
908
909
    ///
    /// ### Returns
    ///
    /// A new `KvIndexer`.
    pub fn new_with_frequency(
        token: CancellationToken,
        expiration_duration: Option<Duration>,
910
        kv_block_size: u32,
911
        metrics: Arc<KvIndexerMetrics>,
912
        prune_config: Option<PruneConfig>,
913
914
915
916
    ) -> 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);
917
        let (get_workers_tx, get_workers_rx) = mpsc::channel::<GetWorkersRequest>(16);
918
        let (dump_tx, dump_rx) = mpsc::channel::<DumpRequest>(16);
919
920
        let (routing_tx, mut routing_rx) = mpsc::channel::<RoutingDecisionRequest>(2048);
        let (prune_tx, mut prune_rx) = mpsc::channel::<()>(1);
921
        let cancel_clone = token.clone();
922

923
        let task = std::thread::spawn(move || {
924
925
            // Create a single-threaded tokio runtime
            let runtime = tokio::runtime::Builder::new_current_thread()
926
927
928
929
                .enable_all()
                .build()
                .unwrap();

930
931
932
933
934
            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;
935
                let mut get_workers_rx = get_workers_rx;
936
937
                let mut dump_rx = dump_rx;
                let mut trie = RadixTree::new_with_frequency(expiration_duration);
938
939
940
941
942
943
944

                // Create PruneManager if prune_config is specified
                let mut prune_manager = prune_config.map(|config| {
                    PruneManager::<BlockEntry>::new(50, config)
                });
                let mut event_id_counter = 0u64;

945
                loop {
946
947
948
949
950
951
952
953
                    // Create a future that sleeps until the next expiration time
                    let expiry_fut = if let Some(ref pm) = prune_manager
                        && let Some(next_expiry) = pm.peek_next_expiry() {
                        tokio::time::sleep_until(next_expiry)
                    } else {
                        tokio::time::sleep(Duration::MAX)
                    };

954
955
956
957
958
959
960
                    tokio::select! {
                        biased;

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

962
963
964
                        Some(worker) = remove_worker_rx.recv() => {
                            trie.remove_worker(worker);
                        }
965

966
967
968
969
970
                        Some(get_workers_req) = get_workers_rx.recv() => {
                            let workers = trie.get_workers();
                            let _ = get_workers_req.resp.send(workers);
                        }

971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
                        Some(_) = prune_rx.recv() => {
                            // Tree size-based pruning triggered
                            let Some(ref mut pm) = prune_manager else { continue };
                            let Ok(pruned) = pm.prune(trie.current_size()) else { continue };

                            for p in pruned {
                                event_id_counter += 1;
                                let event = RouterEvent::new(
                                    p.worker.worker_id,
                                    KvCacheEvent {
                                        event_id: event_id_counter,
                                        data: KvCacheEventData::Removed(KvCacheRemoveData {
                                            block_hashes: vec![p.key],
                                        }),
                                        dp_rank: p.worker.dp_rank,
                                    }
                                );
                                let _ = trie.apply_event(event);
                            }
                        }

992
993
                        Some(event) = event_rx.recv() => {
                            let event_type = KvIndexerMetrics::get_event_type(&event.event.data);
994
995
                            let result = trie.apply_event(event.clone());
                            let result_is_ok = result.is_ok();
996
                            metrics.increment_event_applied(event_type, result);
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023

                            // Track blocks in PruneManager if TTL is enabled and event was stored successfully
                            let Some(ref mut pm) = prune_manager else { continue };
                            if !result_is_ok { continue };
                            let KvCacheEventData::Stored(ref store_data) = event.event.data else { continue };

                            let worker = WorkerWithDpRank::new(event.worker_id, event.event.dp_rank);
                            let block_entries: Vec<BlockEntry> = store_data.blocks.iter().enumerate().map(|(idx, block)| {
                                BlockEntry {
                                    key: block.block_hash,
                                    worker,
                                    seq_position: idx,
                                }
                            }).collect();
                            pm.insert(block_entries);

                            // Check if we need to prune due to tree size
                            let Some(ref pc) = pm.prune_config else { continue };
                            let current_size = trie.current_size();
                            if current_size > pc.max_tree_size {
                                tracing::info!(
                                    "Pruning: tree size ({}) exceeded max tree size ({}), scheduling pruning",
                                    current_size,
                                    pc.max_tree_size
                                );
                                let _ = prune_tx.try_send(());
                            }
1024
                        }
1025

1026
1027
1028
1029
                        Some(dump_req) = dump_rx.recv() => {
                            let events = trie.dump_tree_as_events();
                            let _ = dump_req.resp.send(events);
                        }
1030

1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
                        Some(routing_req) = routing_rx.recv() => {
                            // Process routing decisions when TTL/pruning is enabled
                            let Some(ref mut pm) = prune_manager else { continue };

                            event_id_counter += 1;

                            let hashes = routing_req.local_hashes.iter().zip(routing_req.sequence_hashes.iter());
                            let stored_event = KvCacheEventData::Stored(KvCacheStoreData {
                                parent_hash: None,
                                blocks: hashes.map(|(local_hash, sequence_hash)| KvCacheStoredBlockData {
                                    tokens_hash: *local_hash,
                                    block_hash: ExternalSequenceBlockHash(*sequence_hash),
                                }).collect(),
                            });

                            let event = RouterEvent::new(
                                routing_req.worker.worker_id,
                                KvCacheEvent {
                                    event_id: event_id_counter,
                                    data: stored_event,
                                    dp_rank: routing_req.worker.dp_rank,
                                }
                            );

                            if trie.apply_event(event).is_err() {
                                continue;
                            }

                            let block_entries: Vec<BlockEntry> = routing_req.sequence_hashes.iter().enumerate().map(|(idx, h)| {
                                BlockEntry {
                                    key: ExternalSequenceBlockHash(*h),
                                    worker: routing_req.worker,
                                    seq_position: idx,
                                }
                            }).collect();
                            pm.insert(block_entries);

                            // Check if we need to prune due to tree size
                            let Some(ref pc) = pm.prune_config else { continue };
                            let current_size = trie.current_size();
                            if current_size > pc.max_tree_size {
                                tracing::info!(
                                    "Pruning: tree size ({}) exceeded max tree size ({}), scheduling pruning",
                                    current_size,
                                    pc.max_tree_size
                                );
                                let _ = prune_tx.try_send(());
                            }
                        }

1081
1082
1083
                        Some(req) = match_rx.recv() => {
                            let matches = trie.find_matches(req.sequence, req.early_exit);
                            let _ = req.resp.send(matches);
1084
                        }
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105

                        _ = expiry_fut => {
                            // TTL-based expiry triggered
                            let Some(ref mut pm) = prune_manager else { continue };

                            let expired = pm.pop_expired();
                            for e in expired {
                                event_id_counter += 1;
                                let event = RouterEvent::new(
                                    e.worker.worker_id,
                                    KvCacheEvent {
                                        event_id: event_id_counter,
                                        data: KvCacheEventData::Removed(KvCacheRemoveData {
                                            block_hashes: vec![e.key],
                                        }),
                                        dp_rank: e.worker.dp_rank,
                                    }
                                );
                                let _ = trie.apply_event(event);
                            }
                        }
1106
                    }
1107
1108
                }
            });
1109

1110
            tracing::debug!("KvCacheIndexer task completed");
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
        });

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

        Self {
            cancel: token,
            event_tx,
            match_tx,
            remove_worker_tx,
1121
            get_workers_tx,
1122
            dump_tx,
1123
            routing_tx,
1124
            task: once,
1125
            kv_block_size,
1126
1127
1128
        }
    }

1129
    pub fn block_size(&self) -> u32 {
1130
1131
1132
        self.kv_block_size
    }

1133
1134
1135
1136
1137
    pub fn new(
        token: CancellationToken,
        kv_block_size: u32,
        metrics: Arc<KvIndexerMetrics>,
    ) -> Self {
1138
        Self::new_with_frequency(token, None, kv_block_size, metrics, None)
1139
1140
1141
1142
1143
1144
1145
1146
1147
    }

    /// 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()
1148
1149
1150
1151
1152
1153
1154
1155
1156
    }

    /// 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()
1157
    }
1158
1159
1160
1161
1162
1163
1164
1165
1166

    /// 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()
    }
1167
1168
1169
1170
1171
1172
1173
1174
1175

    /// 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()
    }
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
}

#[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 {
1192
            tracing::error!(
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
                "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> {
1208
        tracing::debug!(
1209
1210
1211
1212
            "Finding matches for request tokens: {:?} / len: {}",
            tokens,
            tokens.len()
        );
1213
        let sequence = compute_block_hash_for_seq(tokens, self.kv_block_size);
1214
        tracing::debug!("Computed sequence: {:?}", sequence);
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
        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");
        }
    }
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245

    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)
    }
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

    async fn process_routing_decision(
        &self,
        worker: WorkerWithDpRank,
        local_hashes: Vec<LocalBlockHash>,
        sequence_hashes: Vec<SequenceHash>,
    ) -> Result<(), KvRouterError> {
        self.routing_tx
            .send(RoutingDecisionRequest {
                worker,
                local_hashes,
                sequence_hashes,
            })
            .await
            .map_err(|_| KvRouterError::IndexerDroppedRequest)?;
        Ok(())
    }

    async fn process_routing_decision_for_request(
        &self,
        tokens: &[u32],
        worker: WorkerWithDpRank,
    ) -> Result<(), KvRouterError> {
        let local_hashes = compute_block_hash_for_seq(tokens, self.kv_block_size);
        let sequence = TokenBlockSequence::new(tokens.into(), self.kv_block_size, None);
        let sequence_hashes = sequence
            .blocks()
            .iter()
            .map(|b| b.sequence_hash())
            .collect::<Vec<_>>();

        self.process_routing_decision(worker, local_hashes, sequence_hashes)
            .await
    }
1280
1281
}

1282
1283
1284
1285
1286
1287
impl Drop for KvIndexer {
    fn drop(&mut self) {
        self.shutdown();
    }
}

1288
1289
1290
1291
1292
1293
1294
#[derive(Debug, Clone)]
pub struct ShardedMatchRequest {
    sequence: Vec<LocalBlockHash>,
    early_exit: bool,
    resp: mpsc::Sender<OverlapScores>,
}

1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
/// 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.
1309
1310
1311
pub struct KvIndexerSharded {
    /// A `CancellationToken` for managing shutdown.
    cancel: CancellationToken,
1312
    /// The size of the KV block this indexer can handle.
1313
    kv_block_size: u32,
1314
1315
1316
1317
1318
1319
    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>>,
1320
    dump_tx: Vec<mpsc::Sender<DumpRequest>>,
1321
    routing_tx: Vec<mpsc::Sender<RoutingDecisionRequest>>,
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
    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.
1333
1334
    /// * `ttl` - The time-to-live for blocks before they expire.
    /// * `prune_config` - Configuration for tree-size based pruning.
1335
1336
1337
1338
1339
1340
1341
1342
    ///
    /// ### Returns
    ///
    /// A new `KvIndexer`.
    pub fn new_with_frequency(
        token: CancellationToken,
        num_shards: usize,
        expiration_duration: Option<Duration>,
1343
        kv_block_size: u32,
1344
        metrics: Arc<KvIndexerMetrics>,
1345
        prune_config: Option<PruneConfig>,
1346
1347
1348
1349
1350
1351
    ) -> 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();
1352
        let mut get_workers_tx = Vec::new();
1353
1354
        let mut dump_tx = Vec::new();
        let mut routing_tx = Vec::new();
1355
1356
1357
1358
1359
1360
1361
1362
        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);
1363
1364
            let (shard_get_workers_tx, mut shard_get_workers_rx) =
                mpsc::channel::<GetWorkersRequest>(16);
1365
1366
1367
1368
            let (shard_dump_tx, mut shard_dump_rx) = mpsc::channel::<DumpRequest>(16);
            let (shard_routing_tx, mut shard_routing_rx) =
                mpsc::channel::<RoutingDecisionRequest>(2048);
            let (shard_prune_tx, mut shard_prune_rx) = mpsc::channel::<()>(1);
1369
1370
            let mut shard_broadcast_rx = request_broadcast_tx.subscribe();
            let cancel = token.clone();
1371
            let metrics = metrics.clone();
1372
            let prune_config_clone = prune_config.clone();
1373
1374
1375

            event_tx.push(shard_event_tx);
            remove_worker_tx.push(shard_remove_worker_tx);
1376
            get_workers_tx.push(shard_get_workers_tx);
1377
1378
            dump_tx.push(shard_dump_tx);
            routing_tx.push(shard_routing_tx);
1379

1380
            let runtime = tokio::runtime::Builder::new_current_thread()
1381
1382
1383
1384
1385
                .enable_all()
                .build()
                .unwrap();

            tasks.push(std::thread::spawn(move || {
1386
1387
                runtime.block_on(async move {
                    let mut trie = RadixTree::new_with_frequency(expiration_duration);
1388
1389
1390
1391
1392
1393
1394

                    // Create PruneManager if prune_config is specified
                    let mut prune_manager = prune_config_clone.map(|config| {
                        PruneManager::<BlockEntry>::new(50, config)
                    });
                    let mut event_id_counter = 0u64;

1395
                    loop {
1396
1397
1398
1399
1400
1401
1402
1403
                        // Create a future that sleeps until the next expiration time
                        let expiry_fut = if let Some(ref pm) = prune_manager
                            && let Some(next_expiry) = pm.peek_next_expiry() {
                            tokio::time::sleep_until(next_expiry)
                        } else {
                            tokio::time::sleep(Duration::MAX)
                        };

1404
1405
                        tokio::select! {
                            biased;
1406

1407
1408
1409
1410
                            _ = cancel.cancelled() => {
                                tracing::trace!("KvCacheIndexer progress loop shutting down");
                                return;
                            }
1411

1412
1413
1414
                            Some(worker) = shard_remove_worker_rx.recv() => {
                                trie.remove_worker(worker);
                            }
1415

1416
1417
1418
1419
1420
                            Some(get_workers_req) = shard_get_workers_rx.recv() => {
                                let workers = trie.get_workers();
                                let _ = get_workers_req.resp.send(workers);
                            }

1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
                            Some(_) = shard_prune_rx.recv() => {
                                // Tree size-based pruning triggered
                                let Some(ref mut pm) = prune_manager else { continue };
                                let Ok(pruned) = pm.prune(trie.current_size()) else { continue };

                                for p in pruned {
                                    event_id_counter += 1;
                                    let event = RouterEvent::new(
                                        p.worker.worker_id,
                                        KvCacheEvent {
                                            event_id: event_id_counter,
                                            data: KvCacheEventData::Removed(KvCacheRemoveData {
                                                block_hashes: vec![p.key],
                                            }),
                                            dp_rank: p.worker.dp_rank,
                                        }
                                    );
                                    let _ = trie.apply_event(event);
                                }
                            }

1442
1443
                            Some(event) = shard_event_rx.recv() => {
                                let event_type = KvIndexerMetrics::get_event_type(&event.event.data);
1444
1445
                                let result = trie.apply_event(event.clone());
                                let result_is_ok = result.is_ok();
1446
                                metrics.increment_event_applied(event_type, result);
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
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
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523

                                // Track blocks in PruneManager if TTL is enabled and event was stored successfully
                                let Some(ref mut pm) = prune_manager else { continue };
                                if !result_is_ok { continue };
                                let KvCacheEventData::Stored(ref store_data) = event.event.data else { continue };

                                let worker = WorkerWithDpRank::new(event.worker_id, event.event.dp_rank);
                                let block_entries: Vec<BlockEntry> = store_data.blocks.iter().enumerate().map(|(idx, block)| {
                                    BlockEntry {
                                        key: block.block_hash,
                                        worker,
                                        seq_position: idx,
                                    }
                                }).collect();
                                pm.insert(block_entries);

                                // Check if we need to prune due to tree size
                                let Some(ref pc) = pm.prune_config else { continue };
                                let current_size = trie.current_size();
                                if current_size > pc.max_tree_size {
                                    tracing::info!(
                                        "Pruning: tree size ({}) exceeded max tree size ({}), scheduling pruning",
                                        current_size,
                                        pc.max_tree_size
                                    );
                                    let _ = shard_prune_tx.try_send(());
                                }
                            }

                            Some(routing_req) = shard_routing_rx.recv() => {
                                // Process routing decisions when TTL/pruning is enabled
                                let Some(ref mut pm) = prune_manager else { continue };

                                event_id_counter += 1;

                                let hashes = routing_req.local_hashes.iter().zip(routing_req.sequence_hashes.iter());
                                let stored_event = KvCacheEventData::Stored(KvCacheStoreData {
                                    parent_hash: None,
                                    blocks: hashes.map(|(local_hash, sequence_hash)| KvCacheStoredBlockData {
                                        tokens_hash: *local_hash,
                                        block_hash: ExternalSequenceBlockHash(*sequence_hash),
                                    }).collect(),
                                });

                                let event = RouterEvent::new(
                                    routing_req.worker.worker_id,
                                    KvCacheEvent {
                                        event_id: event_id_counter,
                                        data: stored_event,
                                        dp_rank: routing_req.worker.dp_rank,
                                    }
                                );

                                if trie.apply_event(event).is_err() {
                                    continue;
                                }

                                let block_entries: Vec<BlockEntry> = routing_req.sequence_hashes.iter().enumerate().map(|(idx, h)| {
                                    BlockEntry {
                                        key: ExternalSequenceBlockHash(*h),
                                        worker: routing_req.worker,
                                        seq_position: idx,
                                    }
                                }).collect();
                                pm.insert(block_entries);

                                // Check if we need to prune due to tree size
                                let Some(ref pc) = pm.prune_config else { continue };
                                let current_size = trie.current_size();
                                if current_size > pc.max_tree_size {
                                    tracing::info!(
                                        "Pruning: tree size ({}) exceeded max tree size ({}), scheduling pruning",
                                        current_size,
                                        pc.max_tree_size
                                    );
                                    let _ = shard_prune_tx.try_send(());
                                }
1524
                            }
1525

1526
1527
1528
1529
1530
1531
1532
1533
1534
                            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);
1535
1536
                                }
                            }
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557

                            _ = expiry_fut => {
                                // TTL-based expiry triggered
                                let Some(ref mut pm) = prune_manager else { continue };

                                let expired = pm.pop_expired();
                                for e in expired {
                                    event_id_counter += 1;
                                    let event = RouterEvent::new(
                                        e.worker.worker_id,
                                        KvCacheEvent {
                                            event_id: event_id_counter,
                                            data: KvCacheEventData::Removed(KvCacheRemoveData {
                                                block_hashes: vec![e.key],
                                            }),
                                            dp_rank: e.worker.dp_rank,
                                        }
                                    );
                                    let _ = trie.apply_event(event);
                                }
                            }
1558
                        }
1559
1560
                    }
                });
1561

1562
                tracing::debug!("KvCacheIndexer task completed");
1563
1564
1565
1566
1567
            }));
        }

        Self {
            cancel: token,
1568
            kv_block_size,
1569
1570
1571
1572
1573
            worker_assignments,
            worker_counts,
            event_tx,
            request_broadcast_tx,
            remove_worker_tx,
1574
1575
            dump_tx,
            routing_tx,
1576
1577
1578
1579
            tasks,
        }
    }

1580
    pub fn block_size(&self) -> u32 {
1581
1582
1583
        self.kv_block_size
    }

1584
1585
1586
1587
1588
1589
    pub fn new(
        token: CancellationToken,
        num_shards: usize,
        kv_block_size: u32,
        metrics: Arc<KvIndexerMetrics>,
    ) -> Self {
1590
        Self::new_with_frequency(token, num_shards, None, kv_block_size, metrics, None)
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
    }
}

#[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);
1616
                        scores.tree_sizes.extend(response.tree_sizes);
1617
1618
1619
1620
1621
1622
1623
1624

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

                            if diff > 0 {
1625
                                scores.frequencies.extend(iter::repeat_n(0, diff as usize));
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
                            }

                            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> {
1648
        let sequence = compute_block_hash_for_seq(tokens, self.kv_block_size);
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
        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();
        }
    }
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717

    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)
    }
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758

    async fn process_routing_decision(
        &self,
        worker: WorkerWithDpRank,
        local_hashes: Vec<LocalBlockHash>,
        sequence_hashes: Vec<SequenceHash>,
    ) -> Result<(), KvRouterError> {
        // Route to the appropriate shard based on worker assignment
        let shard_idx = self
            .worker_assignments
            .get(&worker.worker_id)
            .copied()
            .unwrap_or(0);

        self.routing_tx[shard_idx]
            .send(RoutingDecisionRequest {
                worker,
                local_hashes,
                sequence_hashes,
            })
            .await
            .map_err(|_| KvRouterError::IndexerDroppedRequest)?;
        Ok(())
    }

    async fn process_routing_decision_for_request(
        &self,
        tokens: &[u32],
        worker: WorkerWithDpRank,
    ) -> Result<(), KvRouterError> {
        let local_hashes = compute_block_hash_for_seq(tokens, self.kv_block_size);
        let sequence = TokenBlockSequence::new(tokens.into(), self.kv_block_size, None);
        let sequence_hashes = sequence
            .blocks()
            .iter()
            .map(|b| b.sequence_hash())
            .collect::<Vec<_>>();

        self.process_routing_decision(worker, local_hashes, sequence_hashes)
            .await
    }
1759
1760
}

1761
1762
1763
1764
1765
1766
impl Drop for KvIndexerSharded {
    fn drop(&mut self) {
        self.shutdown();
    }
}

1767
1768
1769
1770
1771
#[cfg(test)]
mod tests {

    use super::*;
    use rstest::rstest;
1772
    use rstest_reuse::{self, *};
1773
1774
1775
    use tokio::time;
    use tokio_util::sync::CancellationToken;

1776
1777
1778
1779
    fn setup() {
        dynamo_runtime::logging::init();
    }

1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
1804
1805
1806
1807
1808
1809
1810
    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
1811
                dp_rank: 0,
1812
1813
1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
            },
        }
    }

    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
1827
                dp_rank: 0,
1828
1829
1830
1831
1832
1833
            },
        }
    }

    #[test]
    fn test_radix_tree() {
1834
1835
        setup();

1836
1837
        let mut trie = RadixTree::new();

1838
1839
        let worker_1 = 0;
        let worker_2 = 1;
1840

1841
1842
        trie.apply_event(create_store_event(worker_1, 1, vec![1, 2, 3], None))
            .unwrap();
1843
1844
1845
1846
1847

        let scores = trie.find_matches(
            vec![LocalBlockHash(1), LocalBlockHash(2), LocalBlockHash(3)],
            false,
        );
Yan Ru Pei's avatar
Yan Ru Pei committed
1848
1849
1850
1851
1852
1853
1854
        assert_eq!(
            scores
                .scores
                .get(&WorkerWithDpRank::from_worker_id(worker_1))
                .unwrap(),
            &3
        );
1855
1856

        assert_eq!(trie.lookup.len(), 1);
Yan Ru Pei's avatar
Yan Ru Pei committed
1857
1858
1859
1860
1861
1862
1863
        assert_eq!(
            trie.lookup
                .get(&WorkerWithDpRank::from_worker_id(worker_1))
                .unwrap()
                .len(),
            3
        );
1864
1865
1866
1867
1868
1869
1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
        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
        );

1889
1890
        trie.apply_event(create_store_event(worker_2, 1, vec![1, 4, 5], None))
            .unwrap();
1891
1892
1893
1894
1895

        let scores = trie.find_matches(
            vec![LocalBlockHash(1), LocalBlockHash(2), LocalBlockHash(3)],
            false,
        );
Yan Ru Pei's avatar
Yan Ru Pei committed
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
1907
1908
1909
        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
        );
1910
1911

        assert_eq!(trie.lookup.len(), 2);
Yan Ru Pei's avatar
Yan Ru Pei committed
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
        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
        );
1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
        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
        );

1951
1952
        trie.apply_event(create_remove_event(worker_2, 2, vec![5]))
            .unwrap();
1953
        assert_eq!(trie.lookup.len(), 2);
Yan Ru Pei's avatar
Yan Ru Pei committed
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
        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
        );
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
1980
1981
1982
1983
1984
1985
1986
1987
1988
1989
1990
1991
1992
        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
        );

1993
1994
        trie.apply_event(create_remove_event(worker_2, 3, vec![4]))
            .unwrap();
1995
1996

        assert_eq!(trie.lookup.len(), 2);
Yan Ru Pei's avatar
Yan Ru Pei committed
1997
1998
1999
2000
2001
2002
2003
2004
2005
2006
2007
2008
2009
2010
        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
        );
2011
2012
2013
2014
2015
2016
2017
2018
2019
2020
2021
2022
2023
2024
2025
2026
2027
2028
2029
2030
2031
2032
2033
2034
2035
2036
2037
2038
2039
2040
        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)),
2041
2042
        ))
        .unwrap();
2043
2044
2045
2046
2047

        let scores = trie.find_matches(
            vec![LocalBlockHash(1), LocalBlockHash(2), LocalBlockHash(3)],
            false,
        );
Yan Ru Pei's avatar
Yan Ru Pei committed
2048
2049
2050
2051
2052
2053
2054
2055
2056
2057
2058
2059
2060
2061
        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
        );
2062
2063

        assert_eq!(trie.lookup.len(), 2);
Yan Ru Pei's avatar
Yan Ru Pei committed
2064
2065
2066
2067
2068
2069
2070
2071
2072
2073
2074
2075
2076
2077
        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
        );
2078
2079
2080
2081
2082
2083
2084
2085
2086
2087
2088
2089
2090
2091
2092
2093
2094
2095
2096
2097
2098
2099
2100
2101
2102
2103
        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
2104
                .get(&WorkerWithDpRank::from_worker_id(worker_1))
2105
2106
2107
2108
2109
2110
2111
2112
2113
2114
                .unwrap()
                .get(&ExternalSequenceBlockHash(200))
                .unwrap()
                .borrow()
                .workers
                .len(),
            2
        );
        assert_eq!(
            trie.lookup
Yan Ru Pei's avatar
Yan Ru Pei committed
2115
                .get(&WorkerWithDpRank::from_worker_id(worker_2))
2116
2117
2118
2119
2120
2121
2122
2123
2124
2125
                .unwrap()
                .get(&ExternalSequenceBlockHash(200))
                .unwrap()
                .borrow()
                .workers
                .len(),
            2
        );
    }

2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
2136
2137
2138
2139
2140
2141
2142
2143
2144
2145
2146
2147
2148
2149
2150
    #[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
        ));
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
2163
2164
2165
2166
2167
2168
2169
2170
2171
2172
2173
2174
2175
2176
2177
2178

        // Parent appears in blocks: parent=1, blocks=[1, 2, 3]
        // This should be rejected as block 1 (hash 100) is the parent
        trie.apply_event(create_store_event(worker_0, 4, vec![1], None))
            .unwrap();
        let result = trie.apply_event(create_store_event(
            worker_0,
            5,
            vec![1, 2, 3],
            Some(ExternalSequenceBlockHash(100)),
        ));
        assert!(matches!(
            result.unwrap_err(),
            KvCacheEventError::InvalidBlockSequence
        ));

        // Block appears twice in sequence: parent=1, blocks=[2, 3, 2]
        // Block 2 appears at positions 0 and 2, creating a cycle
        let result = trie.apply_event(create_store_event(
            worker_0,
            6,
            vec![2, 3, 2],
            Some(ExternalSequenceBlockHash(100)),
        ));
        assert!(matches!(
            result.unwrap_err(),
            KvCacheEventError::InvalidBlockSequence
        ));
2179
2180
    }

2181
2182
    #[test]
    fn test_remove_worker() {
2183
        setup();
2184
2185
        let mut trie = RadixTree::new();

2186
2187
        let worker_0 = 0;
        let worker_1 = 1;
2188

2189
2190
2191
2192
2193
        assert!(
            trie.find_matches(vec![LocalBlockHash(0)], false)
                .scores
                .is_empty()
        );
2194

2195
2196
2197
2198
        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();
2199
2200

        let result = trie.find_matches(vec![LocalBlockHash(0)], false).scores;
Yan Ru Pei's avatar
Yan Ru Pei committed
2201
2202
2203
2204
2205
        assert!(
            result.len() == 2
                && result[&WorkerWithDpRank::from_worker_id(worker_0)] == 1
                && result[&WorkerWithDpRank::from_worker_id(worker_1)] == 1
        );
2206
2207
2208
2209

        trie.remove_worker(worker_0);

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

2213
2214
2215
2216
2217
2218
2219
    #[test]
    fn test_clear_all_blocks() {
        let mut trie = RadixTree::new();

        let worker_0 = 0;
        let worker_1 = 1;

2220
2221
2222
2223
2224
        assert!(
            trie.find_matches(vec![LocalBlockHash(0)], false)
                .scores
                .is_empty()
        );
2225
2226
2227

        // Test clearing an empty worker
        trie.clear_all_blocks(worker_0);
Yan Ru Pei's avatar
Yan Ru Pei committed
2228
2229
2230
2231
2232
        assert!(
            !trie
                .lookup
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_0))
        );
2233
2234

        // Test clearing a worker with shared blocks
2235
2236
2237
2238
        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();
2239
2240

        let result = trie.find_matches(vec![LocalBlockHash(0)], false).scores;
Yan Ru Pei's avatar
Yan Ru Pei committed
2241
2242
2243
2244
2245
        assert!(
            result.len() == 2
                && result[&WorkerWithDpRank::from_worker_id(worker_0)] == 1
                && result[&WorkerWithDpRank::from_worker_id(worker_1)] == 1
        );
2246
2247
2248

        trie.clear_all_blocks(worker_0);

Yan Ru Pei's avatar
Yan Ru Pei committed
2249
2250
2251
2252
2253
2254
2255
2256
2257
2258
        assert!(
            trie.lookup
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_0))
        );
        assert!(
            trie.lookup
                .get(&WorkerWithDpRank::from_worker_id(worker_0))
                .unwrap()
                .is_empty()
        );
2259
2260
2261
2262
        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
2263
        assert_eq!(result[&WorkerWithDpRank::from_worker_id(worker_1)], 2);
2264
2265
2266
2267
2268
2269
2270
        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
2271
        assert_eq!(result[&WorkerWithDpRank::from_worker_id(worker_1)], 1);
2272
2273

        // Test re-adding blocks after clearing worker
2274
2275
        trie.apply_event(create_store_event(worker_0, 0, vec![4, 5], None))
            .unwrap();
2276
2277
2278
2279
        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
2280
        assert_eq!(result[&WorkerWithDpRank::from_worker_id(worker_0)], 2);
2281
2282
2283
2284

        // Test multiple clears
        trie.clear_all_blocks(worker_0);
        trie.clear_all_blocks(worker_0);
Yan Ru Pei's avatar
Yan Ru Pei committed
2285
2286
2287
2288
        assert!(
            trie.lookup
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_0))
        );
2289
2290
2291
2292
2293

        // 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
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
        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()
        );
2306
2307

        // Test clearing a worker that has been removed
2308
2309
2310
2311
        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();
2312
2313
        trie.remove_worker(worker_0);
        trie.clear_all_blocks(worker_0);
Yan Ru Pei's avatar
Yan Ru Pei committed
2314
2315
2316
2317
2318
        assert!(
            !trie
                .lookup
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_0))
        );
2319
2320
        let result = trie.find_matches(vec![LocalBlockHash(6)], false).scores;
        assert_eq!(result.len(), 1);
Yan Ru Pei's avatar
Yan Ru Pei committed
2321
        assert_eq!(result[&WorkerWithDpRank::from_worker_id(worker_1)], 1);
2322
2323
2324

        // Test clearing a worker that doesn't exist
        let worker_fake = 2;
Yan Ru Pei's avatar
Yan Ru Pei committed
2325
2326
2327
2328
2329
        assert!(
            !trie
                .lookup
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_fake))
        );
2330
        trie.clear_all_blocks(worker_fake);
Yan Ru Pei's avatar
Yan Ru Pei committed
2331
2332
2333
2334
2335
2336
2337
2338
2339
        assert!(
            !trie
                .lookup
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_fake))
        );
        assert!(
            trie.lookup
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_1))
        );
2340
2341
        let result = trie.find_matches(vec![LocalBlockHash(6)], false).scores;
        assert_eq!(result.len(), 1);
Yan Ru Pei's avatar
Yan Ru Pei committed
2342
        assert_eq!(result[&WorkerWithDpRank::from_worker_id(worker_1)], 1);
2343
2344
    }

2345
2346
    #[test]
    fn test_early_stopping() {
2347
        setup();
2348
2349
        let mut trie = RadixTree::new();

2350
2351
        let worker_0 = 0;
        let worker_1 = 1;
2352

2353
2354
2355
2356
        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();
2357
2358
2359
2360
2361
2362
2363
2364

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

Yan Ru Pei's avatar
Yan Ru Pei committed
2365
2366
2367
2368
2369
        assert!(
            result.len() == 2
                && result[&WorkerWithDpRank::from_worker_id(worker_0)] == 2
                && result[&WorkerWithDpRank::from_worker_id(worker_1)] == 1
        );
2370
2371
2372
2373

        let result = trie
            .find_matches(vec![LocalBlockHash(0), LocalBlockHash(1)], true)
            .scores;
Yan Ru Pei's avatar
Yan Ru Pei committed
2374
2375
2376
2377
2378
        assert!(
            result.len() == 2
                && result[&WorkerWithDpRank::from_worker_id(worker_0)] == 2
                && result[&WorkerWithDpRank::from_worker_id(worker_1)] == 1
        );
2379
2380
    }

2381
2382
2383
2384
    #[rstest]
    #[case(11)]
    #[case(32)]
    #[case(64)]
2385
    fn test_compute_block_hash_for_seq(#[case] kv_block_size: u32) {
2386
        setup();
2387
        // create a sequence of 64 elements
2388
        let sequence = (0..kv_block_size).collect::<Vec<u32>>();
2389
        let hashes = compute_block_hash_for_seq(&sequence, kv_block_size);
2390
2391
2392
        assert_eq!(hashes.len(), 1);

        // create a sequence of 65 elements
2393
        let sequence = (0..(kv_block_size + 1)).collect::<Vec<u32>>();
2394
        let hashes = compute_block_hash_for_seq(&sequence, kv_block_size);
2395
2396
2397
        assert_eq!(hashes.len(), 1);

        // create a sequence of 129 elements
2398
        let sequence = (0..(2 * kv_block_size + 1)).collect::<Vec<u32>>();
2399
        let hashes = compute_block_hash_for_seq(&sequence, kv_block_size);
2400
2401
2402
        assert_eq!(hashes.len(), 2);
    }

2403
2404
2405
    fn make_indexer(
        token: &CancellationToken,
        num_shards: usize,
2406
        kv_block_size: u32,
2407
    ) -> Box<dyn KvIndexerInterface> {
2408
        let metrics = KvIndexerMetrics::new_unregistered();
2409
        if num_shards == 1 {
2410
            Box::new(KvIndexer::new(token.clone(), kv_block_size, metrics.into()))
2411
        } else {
2412
2413
2414
2415
            Box::new(KvIndexerSharded::new(
                token.clone(),
                num_shards,
                kv_block_size,
2416
                metrics.into(),
2417
            ))
2418
2419
2420
        }
    }

2421
    #[template]
2422
    #[rstest]
2423
2424
2425
2426
2427
2428
    fn indexer_template(
        #[values(1, 3, 8)] num_shards: usize,
        #[values(11, 32, 64)] kv_block_size: usize,
    ) {
    }

2429
    #[tokio::test]
2430
    #[apply(indexer_template)]
2431
    async fn test_kv_indexer_new(num_shards: usize, kv_block_size: u32) {
2432
        setup();
2433
2434
        let token: CancellationToken = CancellationToken::new();
        let _ = make_indexer(&token, num_shards, kv_block_size);
2435
2436
2437
    }

    #[tokio::test]
2438
    #[apply(indexer_template)]
2439
    async fn test_find_matches(num_shards: usize, kv_block_size: u32) {
2440
        setup();
2441
        let token = CancellationToken::new();
2442
        let kv_indexer = make_indexer(&token, num_shards, kv_block_size);
2443
2444
2445
2446
2447
2448
2449
2450

        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]
2451
    #[apply(indexer_template)]
2452
    async fn test_find_matches_for_request(num_shards: usize, kv_block_size: u32) {
2453
        setup();
2454
        let token = CancellationToken::new();
2455
        let kv_indexer = make_indexer(&token, num_shards, kv_block_size);
2456
2457
2458
2459
2460
2461
2462
2463

        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]
2464
    #[apply(indexer_template)]
2465
    async fn test_apply_event(num_shards: usize, kv_block_size: u32) {
2466
        setup();
2467
        let worker_id = 0;
2468
2469

        let token = CancellationToken::new();
2470
        let mut kv_indexer = make_indexer(&token, num_shards, kv_block_size);
2471
2472
2473
2474
2475
2476
2477
2478

        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]
2479
    #[apply(indexer_template)]
2480
    async fn test_shutdown(num_shards: usize, kv_block_size: u32) {
2481
        setup();
2482
        let token = CancellationToken::new();
2483
        let mut kv_indexer = make_indexer(&token, num_shards, kv_block_size);
2484
2485
2486
2487
2488

        kv_indexer.shutdown();
    }

    #[tokio::test]
2489
    #[apply(indexer_template)]
2490
    async fn test_frequency(num_shards: usize, kv_block_size: u32) {
2491
2492
2493
        const ONE_MILLIS: Duration = Duration::from_millis(1);

        setup();
2494
2495
        let mut kv_indexer: Box<dyn KvIndexerInterface>;
        let token = CancellationToken::new();
2496
        let expiration = Duration::from_millis(50);
2497
        let metrics = Arc::new(KvIndexerMetrics::new_unregistered());
2498
2499

        if num_shards == 1 {
2500
2501
            kv_indexer = Box::new(KvIndexer::new_with_frequency(
                token,
2502
                Some(expiration),
2503
                kv_block_size,
2504
                metrics,
2505
                None,
2506
            ));
2507
2508
        } else {
            kv_indexer = Box::new(KvIndexerSharded::new_with_frequency(
2509
2510
                token,
                num_shards,
2511
                Some(expiration),
2512
                kv_block_size,
2513
                metrics,
2514
                None,
2515
2516
2517
            ));
        }

2518
        // The blocks
2519
2520
2521
2522
2523
2524
2525
        let block_hashes = vec![
            LocalBlockHash(1),
            LocalBlockHash(2),
            LocalBlockHash(3),
            LocalBlockHash(4),
        ];

2526
2527
2528
2529
2530
2531
2532
2533
2534
2535
2536
        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;
2537

2538
2539
2540
2541
2542
2543
2544
2545
2546
2547
2548
2549
2550
2551
2552
2553
2554
2555
2556
2557
2558
2559
2560
2561
2562
2563
2564
2565
        // 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"
        );
2566

2567
2568
        // Let those two accesses expire
        time::sleep(expiration + Duration::from_millis(10)).await;
2569

2570
2571
2572
2573
2574
2575
2576
        // 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"
        );
2577

2578
2579
        // New second access
        let _ = kv_indexer.find_matches(block_hashes.clone()).await.unwrap();
2580

2581
2582
        // Access only the first three blocks
        let overlap = kv_indexer
2583
2584
2585
            .find_matches(block_hashes[0..3].to_vec())
            .await
            .unwrap();
2586
2587
        // We see the previous two new accesses
        assert_eq!(overlap.frequencies, vec![2, 2, 2]);
2588

2589
2590
2591
        // 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]);
2592
2593
2594
2595
    }

    #[test]
    fn test_router_event_new() {
2596
        setup();
2597
        let worker_id = 0;
2598
2599
2600
2601
2602
2603
2604
2605
2606
        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
2607
            dp_rank: 0,
2608
2609
2610
2611
2612
2613
2614
2615
2616
2617
2618
2619
2620
2621
2622
2623
2624
2625
2626
        };
        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() {
2627
        setup();
2628
2629
2630
2631
2632
2633
2634
2635
        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() {
2636
        setup();
2637
2638
2639
        let overlap_scores: OverlapScores = Default::default();
        assert!(overlap_scores.scores.is_empty());
    }
2640
2641
2642
2643
2644
2645
2646
2647

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

        // Configuration
        let kv_block_size = 32;
        let num_shards = 2;
2648
        let metrics = Arc::new(KvIndexerMetrics::new_unregistered());
2649
2650
2651

        // Build a non-trivial indexer with events
        let token1 = CancellationToken::new();
2652
2653
        let mut original_indexer =
            KvIndexerSharded::new(token1.clone(), num_shards, kv_block_size, metrics.clone());
2654
2655
2656
2657
2658
2659
2660
2661
2662
2663
2664
2665
2666
2667
2668
2669
2670
2671
2672
2673
2674
2675
2676
2677
2678
2679
2680
2681
2682
2683
2684
2685
2686
2687
2688
2689
2690
2691
2692
2693
2694
2695
2696
2697
2698

        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 =
2699
            KvIndexerSharded::new(token2.clone(), num_shards, kv_block_size, metrics);
2700
2701
2702
2703
2704
2705
2706
2707
2708
2709
2710
2711
2712
2713
2714
2715
2716
2717
2718
2719
2720
2721
2722
2723
2724
2725
2726
2727
2728
2729
2730
2731
2732
2733
2734
2735
2736
2737
2738
2739
2740
2741
2742
2743
2744
2745
2746
2747
2748
2749
2750
2751
2752
2753
2754
2755
2756
2757
2758
2759
2760
2761
2762
2763
2764
2765
2766
2767
2768
2769
2770
2771
2772
2773
2774
2775
2776
2777
2778
2779
2780
2781
2782
2783
2784
2785
2786
2787
2788
2789
2790
2791
2792
2793
2794
2795
2796
2797
2798
2799
2800
2801
2802
2803
2804
2805
2806
2807
2808
2809
2810
2811
2812
2813
2814

        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();
    }
2815
2816
2817
2818
2819
2820
2821
2822
2823
2824
2825
2826
2827
2828
2829
2830
2831
2832
2833
2834
2835
2836
2837
2838
2839
2840
2841
2842
2843
2844
2845
2846
2847
2848
2849
2850
2851
2852
2853
2854
2855
2856
2857
2858
2859

    #[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
        );
    }
2860
2861
2862
2863
2864
2865
2866
2867
2868
2869
2870
2871
2872
2873
2874
2875
2876
2877
2878

    #[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
2879
2880
2881
2882
2883
2884
2885
        assert_eq!(
            trie.lookup
                .get(&WorkerWithDpRank::from_worker_id(worker_0))
                .unwrap()
                .len(),
            3
        );
2886
2887
2888
2889

        // Verify that blocks have the correct workers
        let block_1 = trie
            .lookup
Yan Ru Pei's avatar
Yan Ru Pei committed
2890
            .get(&WorkerWithDpRank::from_worker_id(worker_0))
2891
2892
2893
2894
            .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
2895
2896
2897
2898
2899
2900
2901
2902
2903
2904
2905
2906
2907
2908
2909
2910
2911
2912
        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))
        );
2913
2914
2915
2916
2917

        // 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
2918
2919
2920
2921
2922
        assert!(
            !trie
                .lookup
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_0))
        );
2923
2924
2925
2926
2927
        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
2928
            .get(&WorkerWithDpRank::from_worker_id(worker_1))
2929
2930
2931
2932
            .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
2933
2934
2935
2936
2937
2938
2939
2940
2941
2942
2943
2944
2945
2946
2947
2948
2949
2950
        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))
        );
2951
2952
2953
2954
2955

        // 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
2956
            .get(&WorkerWithDpRank::from_worker_id(worker_1))
2957
2958
2959
2960
            .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
2961
2962
2963
2964
2965
2966
        assert!(
            block_2
                .borrow()
                .workers
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_1))
        );
2967
2968
2969
2970
2971
2972
2973
2974
2975

        // 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
2976
2977
2978
        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)));
2979
    }
2980
}