indexer.rs 128 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
use dynamo_runtime::{
    component::Component,
37
    metrics::{MetricsHierarchy, prometheus_names::kvrouter},
38
39
};
use prometheus::{IntCounterVec, Opts};
40
41
42
use serde::{Deserialize, Serialize};
use std::{
    cell::RefCell,
43
    collections::{HashMap, VecDeque},
44
45
    iter,
    rc::Rc,
46
    sync::{Arc, Mutex, OnceLock},
47
48
49
50
51
52
53
54
55
    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;

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

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

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

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

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

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

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

96
97
98
99
100
101
102
103
104
105
/// 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 {
106
    LocalBlockHash(compute_hash(data))
107
108
109
110
111
112
113
114
115
116
117
118
}

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

119
120
121
122
123
/// Compute the hash for a sequence of tokens, optionally including multimodal metadata.
///
/// When multimodal extra info is provided, the mm_hashes are included in the hash computation
/// to ensure that blocks with identical tokens but different multimodal objects produce
/// different hashes.
124
125
126
127
///
/// ### Arguments
///
/// * `tokens` - A vector of `u32` tokens.
128
129
/// * `kv_block_size` - The size of each block in tokens.
/// * `block_mm_infos` - Optional per-block multimodal metadata.
130
131
132
133
///
/// ### Returns
///
/// A vector of `LocalBlockHash` representing the computed hashes for each chunk of tokens.
134
135
136
137
138
pub fn compute_block_hash_for_seq(
    tokens: &[u32],
    kv_block_size: u32,
    block_mm_infos: Option<&[Option<BlockExtraInfo>]>,
) -> Vec<LocalBlockHash> {
139
    tokens
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
        .chunks_exact(kv_block_size as usize)
        .enumerate()
        .map(|(block_idx, chunk)| {
            let mut bytes: Vec<u8> = chunk.iter().flat_map(|&num| num.to_le_bytes()).collect();

            // Include MM hashes in the block hash computation if present
            if let Some(mm_infos) = block_mm_infos
                && let Some(Some(block_mm_info)) = mm_infos.get(block_idx)
            {
                // The order of different multimodal hashes does not matter.
                // Only which multimodal infos are present in a block is important.
                // The order may differ in different code paths, so the hashes are sorted
                // to keep the block hash stable.
                let mut mm_hashes: Vec<u64> = block_mm_info
                    .mm_objects
                    .iter()
                    .map(|obj| obj.mm_hash)
                    .collect();
                mm_hashes.sort_unstable();

                // Append sorted mm_hashes to the byte array
                for mm_hash in mm_hashes {
                    bytes.extend_from_slice(&mm_hash.to_le_bytes());
                }
            }
165

166
            compute_block_hash(&bytes)
167
168
169
170
        })
        .collect()
}

171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
/// 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
}

205
/// A [`KvCacheEvent`] on a specific LLM worker denoted by [`WorkerId`].
206
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
207
208
pub struct RouterEvent {
    /// The ID of the worker emitting the event.
209
    pub worker_id: WorkerId,
210
    /// The cache event associated with the worker.
211
    pub event: KvCacheEvent,
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
}

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

230
231
232
233
234
235
236
237
238
239
// -------
// Distributed router - Worker KV Query types
// -------

/// Request to query a worker's local KV indexer.
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct WorkerKvQueryRequest {
    /// The worker ID of the worker to query.
    pub worker_id: WorkerId,

240
    /// Start event ID (inclusive). If `None`, dumps entire tree.
241
    pub start_event_id: Option<u64>,
242
    /// End event ID (inclusive). If `None`, returns up to newest available.
243
244
245
246
247
    pub end_event_id: Option<u64>,
}

/// Response from a worker's local KV indexer.
#[derive(Serialize, Deserialize, Debug, Clone)]
248
249
250
251
252
253
254
255
256
257
258
259
260
pub enum WorkerKvQueryResponse {
    /// Events served from the circular buffer (with original event IDs)
    Events(Vec<RouterEvent>),
    /// Full tree dump (with synthetic 0-indexed event IDs)
    TreeDump(Vec<RouterEvent>),
    /// Requested range is newer than available data
    TooNew {
        requested_start: Option<u64>,
        requested_end: Option<u64>,
        newest_available: u64,
    },
    /// Invalid range: end_id < start_id
    InvalidRange { start_id: u64, end_id: u64 },
261
262
}

263
/// A block in the Radix Tree.
264
#[derive(Debug)]
265
266
267
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
268
    /// A map of workers (with dp_rank) to their external sequence block hash for this block.
269
    /// The external hash is preserved to speed up snapshotting.
Yan Ru Pei's avatar
Yan Ru Pei committed
270
    workers: HashMap<WorkerWithDpRank, ExternalSequenceBlockHash>,
271
272
273
274
275
276
277
278
279
280
281
282
283
    /// 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(),
284
            workers: HashMap::new(),
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
            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
303
    lookup: HashMap<WorkerWithDpRank, HashMap<ExternalSequenceBlockHash, SharedRadixBlock>>,
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
    /// 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();
346
347
348
349
350
351
352

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

        for (idx, block_hash) in sequence.iter().enumerate() {
353
354
            let next_block = {
                let current_borrow = current.borrow();
355
                current_borrow.children.get(block_hash).cloned()
356
357
            };
            if let Some(block) = next_block {
358
                scores.update_scores(block.borrow().workers.keys());
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379

                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 {
380
381
382
383
384
                tracing::trace!(
                    "RadixTree::find_matches: block not found at index {} for hash {}",
                    idx,
                    block_hash.0
                );
385
386
387
388
                break;
            }
        }

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

391
392
393
394
395
396
397
398
399
400
        // 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);
        }

401
402
403
404
405
406
407
408
        scores
    }

    /// Apply a [`RouterEvent`] to the radix tree.
    ///
    /// ### Arguments
    ///
    /// * `event` - The `RouterEvent` to apply.
409
    pub fn apply_event(&mut self, event: RouterEvent) -> Result<(), KvCacheEventError> {
Yan Ru Pei's avatar
Yan Ru Pei committed
410
411
412
413
414
415
        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);

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

Yan Ru Pei's avatar
Yan Ru Pei committed
418
        let worker_lookup = self.lookup.entry(worker).or_default();
419
420
421
422
423
424
425

        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 {
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
                    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(),
441
442
                };

443
444
445
446
447
448
449
450
451
                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(());
452
453
                    }

454
455
456
457
                    let mut parent_mut = parent.borrow_mut();
                    let block_data = &blocks[0];

                    let child = match parent_mut.children.get(&block_data.tokens_hash) {
458
459
460
461
                        Some(block) => block.clone(),
                        None => {
                            // create new block - automatically added to the lookup table
                            let new_block = worker_lookup
462
                                .get(&block_data.block_hash)
463
464
465
466
                                .cloned()
                                .unwrap_or_else(|| Rc::new(RefCell::new(RadixBlock::new())));

                            // insert into radix tree
467
                            parent_mut
468
                                .children
469
                                .insert(block_data.tokens_hash, new_block.clone());
470

471
472
473
474
                            new_block
                        }
                    };

475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
                    // 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);
                            }
                        };
492

493
494
495
                        // add our worker to the block with its external hash
                        child_mut.workers.insert(worker, block_data.block_hash);
                    }
496

497
498
                    // add the block to the worker_id lookup table
                    worker_lookup.insert(block_data.block_hash, child.clone());
499

500
501
                    // Recurse with the child and remaining blocks
                    process_blocks(child, &blocks[1..], worker, worker_lookup, id)
502
                }
503
504

                process_blocks(current, &op.blocks, worker, worker_lookup, id)
505
506
            }
            KvCacheEventData::Removed(remove) => {
507
                // tracing::trace!(id, "KV Remove Operation: {:?}", op);
508
509
510
511
512
513
514
515
516
517
                // 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 => {
518
                            tracing::warn!(
Yan Ru Pei's avatar
Yan Ru Pei committed
519
520
                                worker_id = worker.worker_id.to_string(),
                                dp_rank = worker.dp_rank,
521
                                id,
Yan Ru Pei's avatar
Yan Ru Pei committed
522
                                block_hash = ?block,
523
524
                                "Failed to find block to remove; skipping remove operation"
                            );
525
                            return Err(KvCacheEventError::BlockNotFound);
526
527
528
529
                        }
                    };

                    let mut guard = entry.borrow_mut();
Yan Ru Pei's avatar
Yan Ru Pei committed
530
                    guard.workers.remove(&worker);
531
                    if guard.workers.is_empty() {
532
                        // if no workers are using this block, that is true for all children
533
534
535
536
537
                        guard.children.clear();
                    }
                    // remove the block from the lookup table
                    worker_lookup.remove(&block);
                }
538
                Ok(())
539
            }
540
            KvCacheEventData::Cleared => {
Yan Ru Pei's avatar
Yan Ru Pei committed
541
                self.clear_all_blocks(worker.worker_id);
542
                Ok(())
543
            }
544
545
546
        }
    }

Yan Ru Pei's avatar
Yan Ru Pei committed
547
548
549
550
551
552
553
554
555
556
557
    /// 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();
558

Yan Ru Pei's avatar
Yan Ru Pei committed
559
560
561
562
563
564
565
566
567
        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();
                    }
                });
568

Yan Ru Pei's avatar
Yan Ru Pei committed
569
570
571
                if keep_worker {
                    // Re-insert worker with empty blocks map to keep it tracked
                    self.lookup.insert(worker_key, HashMap::new());
572
                }
573
574
575
            }
        }
    }
576

Yan Ru Pei's avatar
Yan Ru Pei committed
577
578
579
580
581
582
583
584
    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);
    }

585
    /// Get all worker IDs currently tracked in the radix tree.
Yan Ru Pei's avatar
Yan Ru Pei committed
586
    /// Returns unique worker_ids (ignoring dp_rank differences).
587
    pub fn get_workers(&self) -> Vec<WorkerId> {
Yan Ru Pei's avatar
Yan Ru Pei committed
588
589
590
591
        let mut worker_ids: Vec<WorkerId> = self.lookup.keys().map(|w| w.worker_id).collect();
        worker_ids.sort_unstable();
        worker_ids.dedup();
        worker_ids
592
593
    }

594
595
596
597
    /// 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> {
598
599
600
601
602
        tracing::debug!(
            "Dumping radix tree as events (contains information about {:?} workers)",
            self.lookup.len()
        );

603
604
605
        let mut events = Vec::new();
        let mut event_id = 0u64;

606
        // BFS queue: (current_block, parent_hashes_per_worker, tokens_hash)
Yan Ru Pei's avatar
Yan Ru Pei committed
607
        // parent_hashes_per_worker maps WorkerWithDpRank -> ExternalSequenceBlockHash
608
609
        let mut queue: VecDeque<(
            SharedRadixBlock,
Yan Ru Pei's avatar
Yan Ru Pei committed
610
            HashMap<WorkerWithDpRank, ExternalSequenceBlockHash>,
611
612
            LocalBlockHash,
        )> = VecDeque::new();
613
614
615
616

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

621
        while let Some((current_block, parent_hashes, tokens_hash)) = queue.pop_front() {
622
623
            let current_borrow = current_block.borrow();

624
625
            // Map of this block's external hashes per worker (for children to use as parent)
            let mut current_external_hashes = HashMap::new();
626
627

            // For each worker that has this block
628
629
630
631
632
633
            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
634
                    worker_id: worker_id.worker_id,
635
636
637
638
639
640
                    event: KvCacheEvent {
                        event_id,
                        data: KvCacheEventData::Stored(KvCacheStoreData {
                            parent_hash,
                            blocks: vec![KvCacheStoredBlockData {
                                block_hash: *external_hash,
641
                                mm_extra_info: None,
642
643
644
                                tokens_hash,
                            }],
                        }),
Yan Ru Pei's avatar
Yan Ru Pei committed
645
                        dp_rank: worker_id.dp_rank,
646
647
648
649
                    },
                };
                events.push(event);
                event_id += 1;
650

651
652
653
                // Track this block's external hash for this worker
                current_external_hashes.insert(*worker_id, *external_hash);
            }
654

655
            // Enqueue children with per-worker parent hashes
656
            for (child_tokens_hash, child_block) in &current_borrow.children {
657
658
659
660
661
                queue.push_back((
                    child_block.clone(),
                    current_external_hashes.clone(),
                    *child_tokens_hash,
                ));
662
663
664
665
666
            }
        }

        events
    }
667
668

    pub fn current_size(&self) -> usize {
669
        self.lookup.values().map(|m| m.len()).sum()
670
    }
671
672
}

673
674
675
676
677
678
679
680
681
682
683
/// 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";
684
pub const METRIC_STATUS_INVALID_BLOCK: &str = "invalid_block";
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703

/// 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(|| {
704
            match component.metrics().create_intcountervec(
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
                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,
757
                    KvCacheEventError::InvalidBlockSequence => METRIC_STATUS_INVALID_BLOCK,
758
759
760
761
762
763
764
765
766
                };
                self.kv_cache_events_applied
                    .with_label_values(&[event_type, error_label])
                    .inc_by(1);
            }
        }
    }
}

Yan Ru Pei's avatar
Yan Ru Pei committed
767
/// Scores representing the overlap of workers (with their dp_rank).
768
769
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OverlapScores {
Yan Ru Pei's avatar
Yan Ru Pei committed
770
771
    // map of worker (with dp_rank) to score
    pub scores: HashMap<WorkerWithDpRank, u32>,
772
773
    // List of frequencies that the blocks have been accessed. Entries with value 0 are omitted.
    pub frequencies: Vec<usize>,
774
775
    // Map of worker to their tree size (number of blocks in the tree for that worker)
    pub tree_sizes: HashMap<WorkerWithDpRank, usize>,
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
}

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),
794
            tree_sizes: HashMap::new(),
795
796
797
798
799
800
801
        }
    }

    /// Update the scores with a set of workers.
    ///
    /// ### Arguments
    ///
Yan Ru Pei's avatar
Yan Ru Pei committed
802
    /// * `workers` - An iterator over `WorkerWithDpRank` references.
803
804
    pub fn update_scores<'a, I>(&mut self, workers: I)
    where
Yan Ru Pei's avatar
Yan Ru Pei committed
805
        I: IntoIterator<Item = &'a WorkerWithDpRank>,
806
    {
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
        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>,
}

834
835
836
837
838
839
/// A request to dump the tree as events
pub struct DumpRequest {
    /// Channel to send the dumped events
    pub resp: oneshot::Sender<Vec<RouterEvent>>,
}

840
841
842
843
844
845
/// A request to get all workers currently tracked
pub struct GetWorkersRequest {
    /// Channel to send the worker IDs
    pub resp: oneshot::Sender<Vec<WorkerId>>,
}

846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
#[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);
892
893
894
895
896
897
898

    /// 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>;
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931

    /// 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>,
932
933
934
}

/// The KV Indexer, managing the KV store and handling events and match requests.
935
#[derive(Clone)]
936
937
938
939
940
941
942
943
944
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>,
945
946
    /// A sender for get workers requests.
    get_workers_tx: mpsc::Sender<GetWorkersRequest>,
947
948
    /// A sender for dump requests.
    dump_tx: mpsc::Sender<DumpRequest>,
949
950
    /// A sender for routing decision requests.
    routing_tx: mpsc::Sender<RoutingDecisionRequest>,
951
    /// The size of the KV block this indexer can handle.
952
    kv_block_size: u32,
953
954
955
    /// Reference counter for Clone-aware Drop.
    /// Only the last clone should cancel the token on drop.
    _ref_count: Arc<()>,
956
957
958
959
960
961
962
963
964
}

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.
965
966
    /// * `ttl` - The time-to-live for blocks before they expire.
    /// * `prune_config` - Configuration for tree-size based pruning.
967
968
969
970
971
972
973
    ///
    /// ### Returns
    ///
    /// A new `KvIndexer`.
    pub fn new_with_frequency(
        token: CancellationToken,
        expiration_duration: Option<Duration>,
974
        kv_block_size: u32,
975
        metrics: Arc<KvIndexerMetrics>,
976
        prune_config: Option<PruneConfig>,
977
978
979
980
    ) -> 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);
981
        let (get_workers_tx, get_workers_rx) = mpsc::channel::<GetWorkersRequest>(16);
982
        let (dump_tx, dump_rx) = mpsc::channel::<DumpRequest>(16);
983
984
        let (routing_tx, mut routing_rx) = mpsc::channel::<RoutingDecisionRequest>(2048);
        let (prune_tx, mut prune_rx) = mpsc::channel::<()>(1);
985

986
        let cancel_clone = token.clone();
987

988
        std::thread::spawn(move || {
989
990
            // Create a single-threaded tokio runtime
            let runtime = tokio::runtime::Builder::new_current_thread()
991
992
993
994
                .enable_all()
                .build()
                .unwrap();

995
996
997
998
999
            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;
1000
                let mut get_workers_rx = get_workers_rx;
1001
1002
                let mut dump_rx = dump_rx;
                let mut trie = RadixTree::new_with_frequency(expiration_duration);
1003
1004
1005
1006
1007
1008
1009

                // 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;

1010
                loop {
1011
1012
1013
1014
1015
1016
1017
1018
                    // 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)
                    };

1019
1020
1021
1022
1023
1024
1025
                    tokio::select! {
                        biased;

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

1027
1028
1029
                        Some(worker) = remove_worker_rx.recv() => {
                            trie.remove_worker(worker);
                        }
1030

1031
1032
1033
1034
1035
                        Some(get_workers_req) = get_workers_rx.recv() => {
                            let workers = trie.get_workers();
                            let _ = get_workers_req.resp.send(workers);
                        }

1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
                        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);
                            }
                        }

1057
1058
                        Some(event) = event_rx.recv() => {
                            let event_type = KvIndexerMetrics::get_event_type(&event.event.data);
1059
1060
                            let result = trie.apply_event(event.clone());
                            let result_is_ok = result.is_ok();
1061
                            metrics.increment_event_applied(event_type, result);
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088

                            // 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(());
                            }
1089
                        }
1090

1091
1092
1093
1094
                        Some(dump_req) = dump_rx.recv() => {
                            let events = trie.dump_tree_as_events();
                            let _ = dump_req.resp.send(events);
                        }
1095

1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
                        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),
1108
                                mm_extra_info: None,
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
                                }).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(());
                            }
                        }

1147
1148
1149
                        Some(req) = match_rx.recv() => {
                            let matches = trie.find_matches(req.sequence, req.early_exit);
                            let _ = req.resp.send(matches);
1150
                        }
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171

                        _ = 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);
                            }
                        }
1172
                    }
1173
1174
                }
            });
1175

1176
            tracing::debug!("KvCacheIndexer task completed");
1177
1178
1179
1180
1181
1182
1183
        });

        Self {
            cancel: token,
            event_tx,
            match_tx,
            remove_worker_tx,
1184
            get_workers_tx,
1185
            dump_tx,
1186
            routing_tx,
1187
            kv_block_size,
1188
            _ref_count: Arc::new(()),
1189
1190
1191
        }
    }

1192
    pub fn block_size(&self) -> u32 {
1193
1194
1195
        self.kv_block_size
    }

1196
1197
1198
1199
1200
    pub fn new(
        token: CancellationToken,
        kv_block_size: u32,
        metrics: Arc<KvIndexerMetrics>,
    ) -> Self {
1201
        Self::new_with_frequency(token, None, kv_block_size, metrics, None)
1202
1203
1204
1205
1206
1207
1208
1209
1210
    }

    /// 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()
1211
1212
1213
1214
1215
1216
1217
1218
1219
    }

    /// 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()
1220
    }
1221
1222
1223
1224
1225
1226
1227
1228
1229

    /// 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()
    }
1230
1231
1232
1233
1234
1235
1236
1237
1238

    /// 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()
    }
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
}

#[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 {
1255
            tracing::error!(
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
                "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> {
1271
        tracing::debug!(
1272
1273
1274
1275
            "Finding matches for request tokens: {:?} / len: {}",
            tokens,
            tokens.len()
        );
1276
        let sequence = compute_block_hash_for_seq(tokens, self.kv_block_size, None);
1277
        tracing::debug!("Computed sequence: {:?}", sequence);
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
        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();
    }
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305

    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)
    }
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328

    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> {
1329
        let local_hashes = compute_block_hash_for_seq(tokens, self.kv_block_size, None);
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
        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
    }
1340
1341
}

1342
1343
impl Drop for KvIndexer {
    fn drop(&mut self) {
1344
1345
1346
1347
1348
        // Only cancel the token if we're the last reference.
        // This allows clones to be dropped without killing the background task.
        if Arc::strong_count(&self._ref_count) == 1 {
            self.shutdown();
        }
1349
1350
1351
    }
}

1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
// -------------------------------------------------
// Decentralized router: LocalKvIndexer for workers
// -------------------------------------------------

/// A thin wrapper around KvIndexer that buffers recent events
/// (e.g. which may be queued by router upon startup)
///
pub struct LocalKvIndexer {
    /// The underlying indexer
    indexer: KvIndexer,
    /// Circular buffer of recent events
    event_buffer: Mutex<VecDeque<RouterEvent>>,
    /// Maximum number of events to keep in buffer
    max_buffer_size: usize, // Router sets this to WORKER_KV_INDEXER_BUFFER_SIZE
}

impl LocalKvIndexer {
    /// create a new LocalKvIndexer pointing to a KvIndexer.
    pub fn new(
        token: CancellationToken,
        kv_block_size: u32,
        metrics: Arc<KvIndexerMetrics>,
        max_buffer_size: usize,
    ) -> Self {
        Self {
            indexer: KvIndexer::new(token, kv_block_size, metrics),
            event_buffer: Mutex::new(VecDeque::with_capacity(max_buffer_size)),
            max_buffer_size,
        }
    }

    /// Get all buffered events (oldest first).
    pub fn get_all_events_in_buffer(&self) -> Vec<RouterEvent> {
        let buffer = self.event_buffer.lock().unwrap();
        buffer.iter().cloned().collect()
    }

    /// Query events by ID range, returning events in `[start_id, end_id]` (both inclusive).
    ///
    /// ### Arguments
    ///
1393
    /// * `start_id` - Starting event ID (inclusive). If `None`, dumps entire tree.
1394
1395
1396
1397
    /// * `end_id` - Ending event ID (inclusive). If `None`, returns up to newest available.
    ///
    /// ### Returns
    ///
1398
1399
1400
1401
    /// - `Events`: Buffered events with original IDs (when range is within buffer)
    /// - `TreeDump`: Full tree dump with synthetic IDs (when range is too old or unspecified)
    /// - `TooNew`: Error when requested range is newer than available data
    /// - `InvalidRange`: Error when end_id < start_id
1402
1403
1404
1405
    pub async fn get_events_in_id_range(
        &self,
        start_id: Option<u64>,
        end_id: Option<u64>,
1406
    ) -> WorkerKvQueryResponse {
1407
1408
        // Validate range if both specified
        if let (Some(s), Some(e)) = (start_id, end_id)
1409
            && e < s
1410
        {
1411
1412
1413
1414
1415
            tracing::warn!(start_id = s, end_id = e, "Invalid range: end_id < start_id");
            return WorkerKvQueryResponse::InvalidRange {
                start_id: s,
                end_id: e,
            };
1416
1417
        }

1418
1419
        // Get buffer state
        let (first_id, last_id) = {
1420
1421
            let buffer = self.event_buffer.lock().unwrap();
            if buffer.is_empty() {
1422
                (None, None)
1423
            } else {
1424
1425
1426
1427
                (
                    Some(buffer.front().unwrap().event.event_id),
                    Some(buffer.back().unwrap().event.event_id),
                )
1428
1429
1430
            }
        };

1431
1432
1433
1434
1435
        // If no start_id specified, dump entire tree
        if start_id.is_none() {
            tracing::debug!("No start_id specified, dumping entire tree");
            let events = self.dump_events().await.unwrap_or_default();
            return WorkerKvQueryResponse::TreeDump(events);
1436
1437
        }

1438
1439
        let start_id = start_id.unwrap();
        let end_id = end_id.unwrap_or_else(|| last_id.unwrap_or(start_id));
1440

1441
1442
1443
1444
1445
1446
1447
        // Check for empty buffer
        let Some(first_buffered) = first_id else {
            tracing::debug!("Buffer empty, dumping entire tree");
            let events = self.dump_events().await.unwrap_or_default();
            return WorkerKvQueryResponse::TreeDump(events);
        };
        let last_buffered = last_id.unwrap();
1448

1449
1450
        // Check if request is too new
        if start_id > last_buffered {
1451
1452
            tracing::warn!(
                start_id,
1453
1454
                last_buffered,
                "Requested start_id is newer than buffer"
1455
            );
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
            return WorkerKvQueryResponse::TooNew {
                requested_start: Some(start_id),
                requested_end: Some(end_id),
                newest_available: last_buffered,
            };
        }

        // Check if start_id is too old (before buffer) -> tree dump
        if start_id < first_buffered {
            tracing::info!(
                start_id,
                first_buffered,
                "Requested start_id is older than buffer, dumping entire tree"
            );
            let events = self.dump_events().await.unwrap_or_default();
            return WorkerKvQueryResponse::TreeDump(events);
1472
1473
        }

1474
1475
1476
        // Serve from buffer
        let buffer = self.event_buffer.lock().unwrap();

1477
1478
1479
1480
1481
        let start_idx = match buffer.binary_search_by_key(&start_id, |e| e.event.event_id) {
            Ok(idx) => idx,
            Err(insertion_point) => insertion_point,
        };

1482
1483
1484
        // Clamp end_id to buffer bounds
        let clamped_end_id = end_id.min(last_buffered);
        let end_idx = match buffer.binary_search_by_key(&clamped_end_id, |e| e.event.event_id) {
1485
1486
1487
1488
            Ok(idx) => idx + 1, // Include the matched element
            Err(insertion_point) => insertion_point,
        };

1489
        let events: Vec<RouterEvent> = buffer
1490
1491
1492
1493
            .iter()
            .skip(start_idx)
            .take(end_idx.saturating_sub(start_idx))
            .cloned()
1494
1495
1496
            .collect();

        WorkerKvQueryResponse::Events(events)
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
    }

    /// Record an event in the buffer
    fn record_event(&self, event: RouterEvent) {
        let mut buffer = self.event_buffer.lock().unwrap();

        // Check that event id is consecutive to last one
        if let Some(last_event) = buffer.back()
            && event.event.event_id != last_event.event.event_id + 1
        {
            let expected = last_event.event.event_id + 1;
            tracing::error!(
                worker_id = event.worker_id,
                expected,
                got = event.event.event_id,
                "Non-consecutive KV event id; buffer may have gaps"
            );
        }
1515
        tracing::debug!(
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
            "Recorded event {:?} in buffer, now size is {}",
            event,
            buffer.len()
        );

        // Add to back
        buffer.push_back(event);

        // Remove from front if over capacity (circular buffer behavior)
        while buffer.len() > self.max_buffer_size {
            buffer.pop_front();
        }
    }

    /// Apply event with buffering.
    ///
    /// This records the event in the buffer and forwards it to the underlying indexer.
    pub async fn apply_event_with_buffer(&self, event: RouterEvent) -> Result<(), KvRouterError> {
        // Record in buffer
        self.record_event(event.clone());

        // Forward to underlying indexer
        self.indexer
            .event_sender()
            .send(event)
            .await
            .map_err(|_| KvRouterError::IndexerOffline)
    }

    /// Clear the event buffer.
    pub fn clear_buffer(&self) {
        let mut buffer = self.event_buffer.lock().unwrap();
        buffer.clear();
    }

    /// Get the current buffer size.
    pub fn buffer_len(&self) -> usize {
        let buffer = self.event_buffer.lock().unwrap();
        buffer.len()
    }

    // Delegation methods to underlying KvIndexer
    /// Get a sender for `RouterEvent`s.
    pub fn event_sender(&self) -> mpsc::Sender<RouterEvent> {
        self.indexer.event_sender()
    }

    /// Get a sender for dump requests (snapshot events).
    pub fn snapshot_event_sender(&self) -> mpsc::Sender<DumpRequest> {
        self.indexer.snapshot_event_sender()
    }

    /// Get a sender for worker removal requests.
    pub fn remove_worker_sender(&self) -> mpsc::Sender<WorkerId> {
        self.indexer.remove_worker_sender()
    }

    /// Get a sender for get workers requests.
    pub fn get_workers_sender(&self) -> mpsc::Sender<GetWorkersRequest> {
        self.indexer.get_workers_sender()
    }

    /// Get the KV block size.
    pub fn block_size(&self) -> u32 {
        self.indexer.block_size()
    }
}

1584
1585
1586
1587
1588
1589
1590
1591
// Implement KvIndexerInterface by delegating to the underlying indexer
#[async_trait]
impl KvIndexerInterface for LocalKvIndexer {
    async fn find_matches(
        &self,
        sequence: Vec<LocalBlockHash>,
    ) -> Result<OverlapScores, KvRouterError> {
        self.indexer.find_matches(sequence).await
1592
1593
    }

1594
1595
1596
1597
1598
1599
    async fn find_matches_for_request(
        &self,
        tokens: &[u32],
    ) -> Result<OverlapScores, KvRouterError> {
        self.indexer.find_matches_for_request(tokens).await
    }
1600

1601
1602
1603
1604
    async fn apply_event(&mut self, event: RouterEvent) {
        // Use the buffering version
        let _ = self.apply_event_with_buffer(event).await;
    }
1605

1606
1607
1608
    async fn remove_worker(&mut self, worker: WorkerId) {
        let _ = self.indexer.remove_worker_sender().send(worker).await;
    }
1609

1610
1611
1612
1613
1614
    fn shutdown(&mut self) {
        // Note: Since indexer is Arc<KvIndexer>, we can't call mutable methods directly.
        // The indexer will be shut down when the CancellationToken is cancelled
        // or when the last Arc reference is dropped.
    }
1615

1616
1617
1618
    async fn dump_events(&self) -> Result<Vec<RouterEvent>, KvRouterError> {
        self.indexer.dump_events().await
    }
1619

1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
    async fn process_routing_decision(
        &self,
        worker: WorkerWithDpRank,
        local_hashes: Vec<LocalBlockHash>,
        sequence_hashes: Vec<SequenceHash>,
    ) -> Result<(), KvRouterError> {
        // TODO I guess the local kvindexers have little use for this method?
        // Keeping it here now to implement the trait fully
        self.indexer
            .process_routing_decision(worker, local_hashes, sequence_hashes)
            .await
    }
1632

1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
    async fn process_routing_decision_for_request(
        &self,
        tokens: &[u32],
        worker: WorkerWithDpRank,
    ) -> Result<(), KvRouterError> {
        // TODO I guess the local kvindexers have little use for this method?
        // Keeping it here now to implement the trait fully
        self.indexer
            .process_routing_decision_for_request(tokens, worker)
            .await
1643
    }
1644
}
1645

1646
1647
1648
1649
1650
1651
#[derive(Debug, Clone)]
pub struct ShardedMatchRequest {
    sequence: Vec<LocalBlockHash>,
    early_exit: bool,
    resp: mpsc::Sender<OverlapScores>,
}
1652

1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
/// 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.
pub struct KvIndexerSharded {
    /// A `CancellationToken` for managing shutdown.
    cancel: CancellationToken,
    /// The size of the KV block this indexer can handle.
    kv_block_size: u32,
    worker_assignments: HashMap<WorkerId, usize>,
    worker_counts: Vec<usize>,
1674

1675
1676
1677
1678
1679
1680
1681
    event_tx: Vec<mpsc::Sender<RouterEvent>>,
    request_broadcast_tx: broadcast::Sender<ShardedMatchRequest>,
    remove_worker_tx: Vec<mpsc::Sender<WorkerId>>,
    dump_tx: Vec<mpsc::Sender<DumpRequest>>,
    routing_tx: Vec<mpsc::Sender<RoutingDecisionRequest>>,
    tasks: Vec<JoinHandle<()>>,
}
1682

1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
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.
    /// * `ttl` - The time-to-live for blocks before they expire.
    /// * `prune_config` - Configuration for tree-size based pruning.
    ///
    /// ### Returns
    ///
    /// A new `KvIndexer`.
    pub fn new_with_frequency(
        token: CancellationToken,
        num_shards: usize,
        expiration_duration: Option<Duration>,
        kv_block_size: u32,
        metrics: Arc<KvIndexerMetrics>,
        prune_config: Option<PruneConfig>,
    ) -> Self {
        let worker_assignments: HashMap<WorkerId, usize> = HashMap::new();
        let worker_counts: Vec<usize> = vec![0; num_shards];
1707

1708
1709
1710
1711
1712
1713
        let mut event_tx = Vec::new();
        let mut remove_worker_tx = Vec::new();
        let mut get_workers_tx = Vec::new();
        let mut dump_tx = Vec::new();
        let mut routing_tx = Vec::new();
        let mut tasks = Vec::new();
1714

1715
        let (request_broadcast_tx, _) = broadcast::channel::<ShardedMatchRequest>(1048576);
1716
1717
1718
1719
1720

        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);
1721
1722
            let (shard_get_workers_tx, mut shard_get_workers_rx) =
                mpsc::channel::<GetWorkersRequest>(16);
1723
1724
1725
1726
            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);
1727
1728
            let mut shard_broadcast_rx = request_broadcast_tx.subscribe();
            let cancel = token.clone();
1729
            let metrics = metrics.clone();
1730
            let prune_config_clone = prune_config.clone();
1731
1732
1733

            event_tx.push(shard_event_tx);
            remove_worker_tx.push(shard_remove_worker_tx);
1734
            get_workers_tx.push(shard_get_workers_tx);
1735
1736
            dump_tx.push(shard_dump_tx);
            routing_tx.push(shard_routing_tx);
1737

1738
            let runtime = tokio::runtime::Builder::new_current_thread()
1739
1740
1741
1742
1743
                .enable_all()
                .build()
                .unwrap();

            tasks.push(std::thread::spawn(move || {
1744
1745
                runtime.block_on(async move {
                    let mut trie = RadixTree::new_with_frequency(expiration_duration);
1746
1747
1748
1749
1750
1751
1752

                    // 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;

1753
                    loop {
1754
1755
1756
1757
1758
1759
1760
1761
                        // 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)
                        };

1762
1763
                        tokio::select! {
                            biased;
1764

1765
1766
1767
1768
                            _ = cancel.cancelled() => {
                                tracing::trace!("KvCacheIndexer progress loop shutting down");
                                return;
                            }
1769

1770
1771
1772
                            Some(worker) = shard_remove_worker_rx.recv() => {
                                trie.remove_worker(worker);
                            }
1773

1774
1775
1776
1777
1778
                            Some(get_workers_req) = shard_get_workers_rx.recv() => {
                                let workers = trie.get_workers();
                                let _ = get_workers_req.resp.send(workers);
                            }

1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
1797
1798
1799
                            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);
                                }
                            }

1800
1801
                            Some(event) = shard_event_rx.recv() => {
                                let event_type = KvIndexerMetrics::get_event_type(&event.event.data);
1802
1803
                                let result = trie.apply_event(event.clone());
                                let result_is_ok = result.is_ok();
1804
                                metrics.increment_event_applied(event_type, result);
1805
1806
1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845

                                // 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),
1846
                                mm_extra_info: None,
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
                                    }).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(());
                                }
1883
                            }
1884

1885
1886
1887
1888
1889
1890
1891
1892
1893
                            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);
1894
1895
                                }
                            }
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916

                            _ = 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);
                                }
                            }
1917
                        }
1918
1919
                    }
                });
1920

1921
                tracing::debug!("KvCacheIndexer task completed");
1922
1923
1924
1925
1926
            }));
        }

        Self {
            cancel: token,
1927
            kv_block_size,
1928
1929
1930
1931
1932
            worker_assignments,
            worker_counts,
            event_tx,
            request_broadcast_tx,
            remove_worker_tx,
1933
1934
            dump_tx,
            routing_tx,
1935
1936
1937
1938
            tasks,
        }
    }

1939
    pub fn block_size(&self) -> u32 {
1940
1941
1942
        self.kv_block_size
    }

1943
1944
1945
1946
1947
1948
    pub fn new(
        token: CancellationToken,
        num_shards: usize,
        kv_block_size: u32,
        metrics: Arc<KvIndexerMetrics>,
    ) -> Self {
1949
        Self::new_with_frequency(token, num_shards, None, kv_block_size, metrics, None)
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
    }
}

#[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);
1975
                        scores.tree_sizes.extend(response.tree_sizes);
1976
1977
1978
1979
1980
1981
1982
1983

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

                            if diff > 0 {
1984
                                scores.frequencies.extend(iter::repeat_n(0, diff as usize));
1985
1986
1987
1988
1989
1990
1991
1992
1993
1994
1995
1996
1997
1998
1999
2000
2001
2002
2003
2004
2005
2006
                            }

                            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> {
2007
        let sequence = compute_block_hash_for_seq(tokens, self.kv_block_size, None);
2008
2009
2010
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
2041
2042
2043
2044
2045
2046
2047
        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();
        }
    }
2048
2049
2050
2051
2052
2053
2054
2055
2056
2057
2058
2059
2060
2061
2062
2063
2064
2065
2066
2067
2068
2069
2070
2071
2072
2073
2074
2075
2076

    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)
    }
2077
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
2104
2105
2106

    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> {
2107
        let local_hashes = compute_block_hash_for_seq(tokens, self.kv_block_size, None);
2108
2109
2110
2111
2112
2113
2114
2115
2116
2117
        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
    }
2118
2119
}

2120
2121
2122
2123
2124
2125
impl Drop for KvIndexerSharded {
    fn drop(&mut self) {
        self.shutdown();
    }
}

2126
2127
2128
#[cfg(test)]
mod tests {
    use super::*;
2129
    use crate::kv_router::protocols::{ExternalSequenceBlockHash, LocalBlockHash};
2130
    use rstest::rstest;
2131
    use rstest_reuse::{self, *};
2132
2133
2134
    use tokio::time;
    use tokio_util::sync::CancellationToken;

2135
2136
2137
2138
    fn setup() {
        dynamo_runtime::logging::init();
    }

2139
2140
2141
2142
2143
2144
    fn make_blocks(hashes: Vec<u64>) -> Vec<KvCacheStoredBlockData> {
        hashes
            .iter()
            .map(|i| KvCacheStoredBlockData {
                tokens_hash: LocalBlockHash(*i),
                block_hash: ExternalSequenceBlockHash(*i * 100),
2145
                mm_extra_info: None,
2146
2147
2148
2149
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
2163
2164
2165
2166
2167
2168
2169
2170
            })
            .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
2171
                dp_rank: 0,
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186
            },
        }
    }

    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
2187
                dp_rank: 0,
2188
2189
2190
2191
2192
2193
            },
        }
    }

    #[test]
    fn test_radix_tree() {
2194
2195
        setup();

2196
2197
        let mut trie = RadixTree::new();

2198
2199
        let worker_1 = 0;
        let worker_2 = 1;
2200

2201
2202
        trie.apply_event(create_store_event(worker_1, 1, vec![1, 2, 3], None))
            .unwrap();
2203
2204
2205
2206
2207

        let scores = trie.find_matches(
            vec![LocalBlockHash(1), LocalBlockHash(2), LocalBlockHash(3)],
            false,
        );
Yan Ru Pei's avatar
Yan Ru Pei committed
2208
2209
2210
2211
2212
2213
2214
        assert_eq!(
            scores
                .scores
                .get(&WorkerWithDpRank::from_worker_id(worker_1))
                .unwrap(),
            &3
        );
2215
2216

        assert_eq!(trie.lookup.len(), 1);
Yan Ru Pei's avatar
Yan Ru Pei committed
2217
2218
2219
2220
2221
2222
2223
        assert_eq!(
            trie.lookup
                .get(&WorkerWithDpRank::from_worker_id(worker_1))
                .unwrap()
                .len(),
            3
        );
2224
2225
2226
2227
2228
2229
2230
2231
2232
2233
2234
2235
2236
2237
2238
2239
2240
2241
2242
2243
2244
2245
2246
2247
2248
        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
        );

2249
2250
        trie.apply_event(create_store_event(worker_2, 1, vec![1, 4, 5], None))
            .unwrap();
2251
2252
2253
2254
2255

        let scores = trie.find_matches(
            vec![LocalBlockHash(1), LocalBlockHash(2), LocalBlockHash(3)],
            false,
        );
Yan Ru Pei's avatar
Yan Ru Pei committed
2256
2257
2258
2259
2260
2261
2262
2263
2264
2265
2266
2267
2268
2269
        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
        );
2270
2271

        assert_eq!(trie.lookup.len(), 2);
Yan Ru Pei's avatar
Yan Ru Pei committed
2272
2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
2284
2285
        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
        );
2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
2308
2309
2310
        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
        );

2311
2312
        trie.apply_event(create_remove_event(worker_2, 2, vec![5]))
            .unwrap();
2313
        assert_eq!(trie.lookup.len(), 2);
Yan Ru Pei's avatar
Yan Ru Pei committed
2314
2315
2316
2317
2318
2319
2320
2321
2322
2323
2324
2325
2326
2327
        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
        );
2328
2329
2330
2331
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341
2342
2343
2344
2345
2346
2347
2348
2349
2350
2351
2352
        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
        );

2353
2354
        trie.apply_event(create_remove_event(worker_2, 3, vec![4]))
            .unwrap();
2355
2356

        assert_eq!(trie.lookup.len(), 2);
Yan Ru Pei's avatar
Yan Ru Pei committed
2357
2358
2359
2360
2361
2362
2363
2364
2365
2366
2367
2368
2369
2370
        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
        );
2371
2372
2373
2374
2375
2376
2377
2378
2379
2380
2381
2382
2383
2384
2385
2386
2387
2388
2389
2390
2391
2392
2393
2394
2395
2396
2397
2398
2399
2400
        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)),
2401
2402
        ))
        .unwrap();
2403
2404
2405
2406
2407

        let scores = trie.find_matches(
            vec![LocalBlockHash(1), LocalBlockHash(2), LocalBlockHash(3)],
            false,
        );
Yan Ru Pei's avatar
Yan Ru Pei committed
2408
2409
2410
2411
2412
2413
2414
2415
2416
2417
2418
2419
2420
2421
        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
        );
2422
2423

        assert_eq!(trie.lookup.len(), 2);
Yan Ru Pei's avatar
Yan Ru Pei committed
2424
2425
2426
2427
2428
2429
2430
2431
2432
2433
2434
2435
2436
2437
        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
        );
2438
2439
2440
2441
2442
2443
2444
2445
2446
2447
2448
2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
2461
2462
2463
        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
2464
                .get(&WorkerWithDpRank::from_worker_id(worker_1))
2465
2466
2467
2468
2469
2470
2471
2472
2473
2474
                .unwrap()
                .get(&ExternalSequenceBlockHash(200))
                .unwrap()
                .borrow()
                .workers
                .len(),
            2
        );
        assert_eq!(
            trie.lookup
Yan Ru Pei's avatar
Yan Ru Pei committed
2475
                .get(&WorkerWithDpRank::from_worker_id(worker_2))
2476
2477
2478
2479
2480
2481
2482
2483
2484
2485
                .unwrap()
                .get(&ExternalSequenceBlockHash(200))
                .unwrap()
                .borrow()
                .workers
                .len(),
            2
        );
    }

2486
2487
2488
2489
2490
2491
2492
2493
2494
2495
2496
2497
2498
2499
2500
2501
2502
2503
2504
2505
2506
2507
2508
2509
2510
    #[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
        ));
2511
2512
2513
2514
2515
2516
2517
2518
2519
2520
2521
2522
2523
2524
2525
2526
2527
2528
2529
2530
2531
2532
2533
2534
2535
2536
2537
2538

        // 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
        ));
2539
2540
    }

2541
2542
    #[test]
    fn test_remove_worker() {
2543
        setup();
2544
2545
        let mut trie = RadixTree::new();

2546
2547
        let worker_0 = 0;
        let worker_1 = 1;
2548

2549
2550
2551
2552
2553
        assert!(
            trie.find_matches(vec![LocalBlockHash(0)], false)
                .scores
                .is_empty()
        );
2554

2555
2556
2557
2558
        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();
2559
2560

        let result = trie.find_matches(vec![LocalBlockHash(0)], false).scores;
Yan Ru Pei's avatar
Yan Ru Pei committed
2561
2562
2563
2564
2565
        assert!(
            result.len() == 2
                && result[&WorkerWithDpRank::from_worker_id(worker_0)] == 1
                && result[&WorkerWithDpRank::from_worker_id(worker_1)] == 1
        );
2566
2567
2568
2569

        trie.remove_worker(worker_0);

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

2573
2574
2575
2576
2577
2578
2579
    #[test]
    fn test_clear_all_blocks() {
        let mut trie = RadixTree::new();

        let worker_0 = 0;
        let worker_1 = 1;

2580
2581
2582
2583
2584
        assert!(
            trie.find_matches(vec![LocalBlockHash(0)], false)
                .scores
                .is_empty()
        );
2585
2586
2587

        // Test clearing an empty worker
        trie.clear_all_blocks(worker_0);
Yan Ru Pei's avatar
Yan Ru Pei committed
2588
2589
2590
2591
2592
        assert!(
            !trie
                .lookup
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_0))
        );
2593
2594

        // Test clearing a worker with shared blocks
2595
2596
2597
2598
        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();
2599
2600

        let result = trie.find_matches(vec![LocalBlockHash(0)], false).scores;
Yan Ru Pei's avatar
Yan Ru Pei committed
2601
2602
2603
2604
2605
        assert!(
            result.len() == 2
                && result[&WorkerWithDpRank::from_worker_id(worker_0)] == 1
                && result[&WorkerWithDpRank::from_worker_id(worker_1)] == 1
        );
2606
2607
2608

        trie.clear_all_blocks(worker_0);

Yan Ru Pei's avatar
Yan Ru Pei committed
2609
2610
2611
2612
2613
2614
2615
2616
2617
2618
        assert!(
            trie.lookup
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_0))
        );
        assert!(
            trie.lookup
                .get(&WorkerWithDpRank::from_worker_id(worker_0))
                .unwrap()
                .is_empty()
        );
2619
2620
2621
2622
        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
2623
        assert_eq!(result[&WorkerWithDpRank::from_worker_id(worker_1)], 2);
2624
2625
2626
2627
2628
2629
2630
        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
2631
        assert_eq!(result[&WorkerWithDpRank::from_worker_id(worker_1)], 1);
2632
2633

        // Test re-adding blocks after clearing worker
2634
2635
        trie.apply_event(create_store_event(worker_0, 0, vec![4, 5], None))
            .unwrap();
2636
2637
2638
2639
        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
2640
        assert_eq!(result[&WorkerWithDpRank::from_worker_id(worker_0)], 2);
2641
2642
2643
2644

        // Test multiple clears
        trie.clear_all_blocks(worker_0);
        trie.clear_all_blocks(worker_0);
Yan Ru Pei's avatar
Yan Ru Pei committed
2645
2646
2647
2648
        assert!(
            trie.lookup
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_0))
        );
2649
2650
2651
2652
2653

        // 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
2654
2655
2656
2657
2658
2659
2660
2661
2662
2663
2664
2665
        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()
        );
2666
2667

        // Test clearing a worker that has been removed
2668
2669
2670
2671
        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();
2672
2673
        trie.remove_worker(worker_0);
        trie.clear_all_blocks(worker_0);
Yan Ru Pei's avatar
Yan Ru Pei committed
2674
2675
2676
2677
2678
        assert!(
            !trie
                .lookup
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_0))
        );
2679
2680
        let result = trie.find_matches(vec![LocalBlockHash(6)], false).scores;
        assert_eq!(result.len(), 1);
Yan Ru Pei's avatar
Yan Ru Pei committed
2681
        assert_eq!(result[&WorkerWithDpRank::from_worker_id(worker_1)], 1);
2682
2683
2684

        // Test clearing a worker that doesn't exist
        let worker_fake = 2;
Yan Ru Pei's avatar
Yan Ru Pei committed
2685
2686
2687
2688
2689
        assert!(
            !trie
                .lookup
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_fake))
        );
2690
        trie.clear_all_blocks(worker_fake);
Yan Ru Pei's avatar
Yan Ru Pei committed
2691
2692
2693
2694
2695
2696
2697
2698
2699
        assert!(
            !trie
                .lookup
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_fake))
        );
        assert!(
            trie.lookup
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_1))
        );
2700
2701
        let result = trie.find_matches(vec![LocalBlockHash(6)], false).scores;
        assert_eq!(result.len(), 1);
Yan Ru Pei's avatar
Yan Ru Pei committed
2702
        assert_eq!(result[&WorkerWithDpRank::from_worker_id(worker_1)], 1);
2703
2704
    }

2705
2706
    #[test]
    fn test_early_stopping() {
2707
        setup();
2708
2709
        let mut trie = RadixTree::new();

2710
2711
        let worker_0 = 0;
        let worker_1 = 1;
2712

2713
2714
2715
2716
        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();
2717
2718
2719
2720
2721
2722
2723
2724

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

Yan Ru Pei's avatar
Yan Ru Pei committed
2725
2726
2727
2728
2729
        assert!(
            result.len() == 2
                && result[&WorkerWithDpRank::from_worker_id(worker_0)] == 2
                && result[&WorkerWithDpRank::from_worker_id(worker_1)] == 1
        );
2730
2731
2732
2733

        let result = trie
            .find_matches(vec![LocalBlockHash(0), LocalBlockHash(1)], true)
            .scores;
Yan Ru Pei's avatar
Yan Ru Pei committed
2734
2735
2736
2737
2738
        assert!(
            result.len() == 2
                && result[&WorkerWithDpRank::from_worker_id(worker_0)] == 2
                && result[&WorkerWithDpRank::from_worker_id(worker_1)] == 1
        );
2739
2740
    }

2741
2742
2743
2744
    #[rstest]
    #[case(11)]
    #[case(32)]
    #[case(64)]
2745
    fn test_compute_block_hash_for_seq(#[case] kv_block_size: u32) {
2746
        setup();
2747
        // create a sequence of 64 elements
2748
        let sequence = (0..kv_block_size).collect::<Vec<u32>>();
2749
        let hashes = compute_block_hash_for_seq(&sequence, kv_block_size, None);
2750
2751
2752
        assert_eq!(hashes.len(), 1);

        // create a sequence of 65 elements
2753
        let sequence = (0..(kv_block_size + 1)).collect::<Vec<u32>>();
2754
        let hashes = compute_block_hash_for_seq(&sequence, kv_block_size, None);
2755
2756
2757
        assert_eq!(hashes.len(), 1);

        // create a sequence of 129 elements
2758
        let sequence = (0..(2 * kv_block_size + 1)).collect::<Vec<u32>>();
2759
        let hashes = compute_block_hash_for_seq(&sequence, kv_block_size, None);
2760
2761
2762
        assert_eq!(hashes.len(), 2);
    }

2763
2764
2765
    fn make_indexer(
        token: &CancellationToken,
        num_shards: usize,
2766
        kv_block_size: u32,
2767
    ) -> Box<dyn KvIndexerInterface> {
2768
        let metrics = KvIndexerMetrics::new_unregistered();
2769
        if num_shards == 1 {
2770
            Box::new(KvIndexer::new(token.clone(), kv_block_size, metrics.into()))
2771
        } else {
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
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
2860
2861
2862
2863
2864
2865
2866
2867
2868
2869
2870
2871
2872
2873
2874
2875
2876
2877
2878
2879
2880
2881
2882
2883
2884
2885
2886
2887
2888
2889
2890
2891
2892
2893
2894
2895
2896
2897
2898
2899
2900
2901
2902
2903
2904
2905
2906
2907
2908
2909
2910
2911
2912
2913
2914
2915
2916
2917
2918
2919
2920
2921
2922
2923
2924
2925
2926
2927
2928
2929
2930
2931
2932
2933
2934
2935
2936
2937
2938
2939
2940
2941
2942
2943
2944
2945
2946
2947
2948
2949
2950
2951
2952
2953
2954
2955
2956
2957
2958
2959
2960
2961
2962
2963
            Box::new(KvIndexerSharded::new(
                token.clone(),
                num_shards,
                kv_block_size,
                metrics.into(),
            ))
        }
    }

    #[template]
    #[rstest]
    fn indexer_template(
        #[values(1, 3, 8)] num_shards: usize,
        #[values(11, 32, 64)] kv_block_size: usize,
    ) {
    }

    #[tokio::test]
    #[apply(indexer_template)]
    async fn test_kv_indexer_new(num_shards: usize, kv_block_size: u32) {
        setup();
        let token: CancellationToken = CancellationToken::new();
        let _ = make_indexer(&token, num_shards, kv_block_size);
    }

    #[tokio::test]
    #[apply(indexer_template)]
    async fn test_find_matches(num_shards: usize, kv_block_size: u32) {
        setup();
        let token = CancellationToken::new();
        let kv_indexer = make_indexer(&token, num_shards, kv_block_size);

        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]
    #[apply(indexer_template)]
    async fn test_find_matches_for_request(num_shards: usize, kv_block_size: u32) {
        setup();
        let token = CancellationToken::new();
        let kv_indexer = make_indexer(&token, num_shards, kv_block_size);

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

        let token = CancellationToken::new();
        let mut kv_indexer = make_indexer(&token, num_shards, kv_block_size);

        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]
    #[apply(indexer_template)]
    async fn test_shutdown(num_shards: usize, kv_block_size: u32) {
        setup();
        let token = CancellationToken::new();
        let mut kv_indexer = make_indexer(&token, num_shards, kv_block_size);

        kv_indexer.shutdown();
    }

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

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

        if num_shards == 1 {
            kv_indexer = Box::new(KvIndexer::new_with_frequency(
                token,
                Some(expiration),
                kv_block_size,
                metrics,
                None,
            ));
        } else {
            kv_indexer = Box::new(KvIndexerSharded::new_with_frequency(
                token,
                num_shards,
                Some(expiration),
                kv_block_size,
                metrics,
                None,
            ));
        }

        // The blocks
        let block_hashes = vec![
            LocalBlockHash(1),
            LocalBlockHash(2),
            LocalBlockHash(3),
            LocalBlockHash(4),
        ];

        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;

        // 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"
        );

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

        // 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"
        );

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

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

        // 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]);
    }

    #[test]
    fn test_router_event_new() {
        setup();
        let worker_id = 0;
        let kv_cache_event = KvCacheEvent {
            event_id: 1,
            data: KvCacheEventData::Stored(KvCacheStoreData {
                parent_hash: None,
                blocks: vec![KvCacheStoredBlockData {
                    block_hash: ExternalSequenceBlockHash(0),
2964
                    mm_extra_info: None,
2965
2966
2967
2968
2969
2970
2971
2972
2973
2974
2975
2976
2977
2978
2979
2980
2981
2982
                    tokens_hash: LocalBlockHash(13226331709069118873),
                }],
            }),
            dp_rank: 0,
        };
        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");
2983
2984
2985
        }
    }

2986
2987
2988
2989
2990
2991
2992
    #[test]
    fn test_radix_tree_default() {
        setup();
        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());
2993
2994
    }

2995
2996
    #[test]
    fn test_overlap_scores_default() {
2997
        setup();
2998
2999
        let overlap_scores: OverlapScores = Default::default();
        assert!(overlap_scores.scores.is_empty());
3000
3001
3002
    }

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

3006
3007
3008
3009
        // Configuration
        let kv_block_size = 32;
        let num_shards = 2;
        let metrics = Arc::new(KvIndexerMetrics::new_unregistered());
3010

3011
3012
3013
3014
        // Build a non-trivial indexer with events
        let token1 = CancellationToken::new();
        let mut original_indexer =
            KvIndexerSharded::new(token1.clone(), num_shards, kv_block_size, metrics.clone());
3015

3016
3017
3018
        let worker_0 = 0;
        let worker_1 = 1;
        let worker_2 = 2;
3019

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

3025
3026
3027
3028
3029
3030
3031
3032
3033
3034
3035
        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;
3036

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

3041
3042
3043
3044
3045
3046
3047
3048
        original_indexer
            .apply_event(create_store_event(
                worker_0,
                4,
                vec![4],
                Some(ExternalSequenceBlockHash(100)),
            ))
            .await;
3049

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

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

3057
3058
3059
3060
        // Create a new indexer and apply all dumped events
        let token2 = CancellationToken::new();
        let mut reconstructed_indexer =
            KvIndexerSharded::new(token2.clone(), num_shards, kv_block_size, metrics);
3061

3062
3063
3064
        for event in &dump1 {
            reconstructed_indexer.apply_event(event.clone()).await;
        }
3065

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

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

3072
3073
3074
3075
3076
3077
3078
3079
3080
3081
3082
3083
3084
3085
3086
3087
3088
3089
3090
3091
3092
3093
3094
3095
3096
3097
3098
3099
3100
3101
3102
3103
3104
3105
3106
3107
3108
3109
3110
3111
3112
3113
3114
3115
3116
3117
3118
3119
3120
3121
3122
3123
3124
3125
3126
3127
3128
3129
3130
3131
3132
3133
3134
3135
3136
3137
3138
3139
3140
        // 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");
            }
3141
3142
        }

3143
3144
3145
3146
3147
3148
3149
3150
3151
3152
3153
3154
3155
3156
3157
3158
3159
3160
3161
3162
3163
3164
3165
3166
3167
3168
3169
3170
3171
3172
3173
3174
3175
        // 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();
    }
3176

3177
3178
3179
3180
3181
    #[test]
    fn test_increment_event_applied() {
        let metrics = KvIndexerMetrics::new_unregistered();

        metrics.increment_event_applied(METRIC_EVENT_STORED, Ok(()));
3182
        assert_eq!(
3183
3184
3185
3186
3187
3188
            metrics
                .kv_cache_events_applied
                .get_metric_with_label_values(&[METRIC_EVENT_STORED, METRIC_STATUS_OK])
                .unwrap()
                .get(),
            1
3189
3190
        );

3191
3192
3193
        metrics.increment_event_applied(
            METRIC_EVENT_STORED,
            Err(KvCacheEventError::ParentBlockNotFound),
3194
3195
        );
        assert_eq!(
3196
3197
3198
3199
3200
3201
3202
3203
3204
            metrics
                .kv_cache_events_applied
                .get_metric_with_label_values(&[
                    METRIC_EVENT_STORED,
                    METRIC_STATUS_PARENT_NOT_FOUND
                ])
                .unwrap()
                .get(),
            1
3205
3206
        );

3207
3208
        metrics
            .increment_event_applied(METRIC_EVENT_REMOVED, Err(KvCacheEventError::BlockNotFound));
3209
        assert_eq!(
3210
3211
3212
3213
3214
3215
3216
3217
3218
            metrics
                .kv_cache_events_applied
                .get_metric_with_label_values(&[
                    METRIC_EVENT_REMOVED,
                    METRIC_STATUS_BLOCK_NOT_FOUND
                ])
                .unwrap()
                .get(),
            1
3219
        );
3220
    }
3221

3222
3223
3224
3225
    #[test]
    fn test_remove_worker_verifies_hash_removal() {
        setup();
        let mut trie = RadixTree::new();
3226

3227
3228
3229
3230
3231
3232
3233
3234
3235
3236
3237
3238
3239
        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
3240
        assert_eq!(
3241
3242
3243
3244
3245
            trie.lookup
                .get(&WorkerWithDpRank::from_worker_id(worker_0))
                .unwrap()
                .len(),
            3
3246
        );
3247

3248
3249
3250
3251
3252
3253
        // Verify that blocks have the correct workers
        let block_1 = trie
            .lookup
            .get(&WorkerWithDpRank::from_worker_id(worker_0))
            .unwrap()
            .get(&ExternalSequenceBlockHash(100))
3254
            .unwrap();
3255
3256
3257
3258
3259
3260
3261
3262
3263
3264
3265
3266
3267
3268
3269
3270
3271
3272
3273
        assert_eq!(block_1.borrow().workers.len(), 3); // worker_0, worker_1, and worker_2 (all have hash 1)
        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))
        );
3274

3275
3276
        // Remove worker_0
        trie.remove_worker(worker_0);
3277

3278
3279
3280
3281
3282
3283
3284
        // Verify worker_0 is completely removed from lookup table
        assert!(
            !trie
                .lookup
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_0))
        );
        assert_eq!(trie.lookup.len(), 2);
3285

3286
3287
3288
3289
3290
3291
3292
3293
3294
3295
3296
3297
3298
3299
3300
3301
3302
3303
3304
3305
3306
3307
3308
3309
3310
3311
        // Verify that worker_0's hash is removed from the workers set
        let block_1 = trie
            .lookup
            .get(&WorkerWithDpRank::from_worker_id(worker_1))
            .unwrap()
            .get(&ExternalSequenceBlockHash(100))
            .unwrap();
        assert_eq!(block_1.borrow().workers.len(), 2); // worker_1 and worker_2 remain
        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))
        );
3312

3313
3314
3315
3316
3317
3318
3319
3320
3321
3322
3323
3324
3325
3326
3327
3328
3329
3330
3331
3332
3333
3334
3335
3336
3337
3338
3339
        // 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
            .get(&WorkerWithDpRank::from_worker_id(worker_1))
            .unwrap()
            .get(&ExternalSequenceBlockHash(200))
            .unwrap();
        assert_eq!(block_2.borrow().workers.len(), 1); // only worker_1
        assert!(
            block_2
                .borrow()
                .workers
                .contains_key(&WorkerWithDpRank::from_worker_id(worker_1))
        );

        // 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);
        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)));
3340
3341
    }

3342
3343
3344
3345
3346
3347
3348
3349
3350
3351
3352
3353
3354
3355
3356
3357
3358
3359
3360
3361
3362
3363
    // LocalKvIndexer tests
    fn make_indexer_with_events(ids: &[u64]) -> LocalKvIndexer {
        let indexer = LocalKvIndexer::new(
            CancellationToken::new(),
            4,
            Arc::new(KvIndexerMetrics::new_unregistered()),
            32,
        );
        {
            let mut buffer = indexer.event_buffer.lock().unwrap();
            for &id in ids {
                buffer.push_back(RouterEvent::new(
                    0,
                    KvCacheEvent {
                        event_id: id,
                        data: KvCacheEventData::Cleared,
                        dp_rank: 0,
                    },
                ));
            }
        }
        indexer
3364
    }
3365
3366

    #[tokio::test]
3367
3368
    async fn returns_slice_within_range() {
        let indexer = make_indexer_with_events(&[1, 2, 3, 4, 5]);
3369

3370
3371
3372
3373
3374
3375
3376
3377
        // Helper to extract events from response
        let extract_events = |resp: WorkerKvQueryResponse| -> Vec<RouterEvent> {
            match resp {
                WorkerKvQueryResponse::Events(e) => e,
                WorkerKvQueryResponse::TreeDump(e) => e,
                _ => panic!("Unexpected response type"),
            }
        };
3378

3379
3380
3381
        let get_ids = |events: Vec<RouterEvent>| -> Vec<u64> {
            events.iter().map(|e| e.event.event_id).collect()
        };
3382

3383
3384
3385
3386
3387
        // Test get_events_in_id_range (buffer queries)
        // Range is [start, end] inclusive
        let result = indexer.get_events_in_id_range(Some(2), Some(4)).await;
        let ids = get_ids(extract_events(result));
        assert_eq!(ids, vec![2, 3, 4]); // inclusive range [2, 4]
3388

3389
3390
3391
        let result = indexer.get_events_in_id_range(Some(2), Some(6)).await;
        let ids = get_ids(extract_events(result));
        assert_eq!(ids, vec![2, 3, 4, 5]); // clamp end to buffer max
3392

3393
3394
3395
        // start_id=0 is before buffer (first is 1), so should trigger tree dump
        let result = indexer.get_events_in_id_range(Some(0), Some(4)).await;
        assert!(matches!(result, WorkerKvQueryResponse::TreeDump(_)));
3396

3397
3398
3399
        let result = indexer.get_events_in_id_range(Some(3), Some(3)).await;
        let ids = get_ids(extract_events(result));
        assert_eq!(ids, vec![3]); // single element when start == end
3400

3401
3402
3403
3404
        // Invalid range: end < start
        let result = indexer.get_events_in_id_range(Some(5), Some(2)).await;
        assert!(matches!(result, WorkerKvQueryResponse::InvalidRange { .. }));
    }
3405

3406
3407
3408
3409
3410
3411
3412
3413
3414
3415
    #[tokio::test]
    async fn test_get_events_in_id_range_all_cases() {
        // Create indexer with small buffer (5 events max)
        // This way older events will only be in the tree, not the buffer
        let indexer = LocalKvIndexer::new(
            CancellationToken::new(),
            4, // block_size
            Arc::new(KvIndexerMetrics::new_unregistered()),
            5, // max_buffer_size - only keeps 5 most recent events
        );
3416

3417
3418
3419
3420
3421
3422
3423
3424
3425
3426
3427
        // Helper to create a test event
        let make_event = |id: u64| {
            RouterEvent::new(
                0, // worker_id
                KvCacheEvent {
                    event_id: id,
                    data: KvCacheEventData::Stored(KvCacheStoreData {
                        parent_hash: None,
                        blocks: vec![KvCacheStoredBlockData {
                            block_hash: ExternalSequenceBlockHash(id * 100),
                            tokens_hash: LocalBlockHash(id * 200),
3428
                            mm_extra_info: None,
3429
3430
3431
3432
3433
3434
                        }],
                    }),
                    dp_rank: 0,
                },
            )
        };
3435

3436
3437
3438
3439
3440
3441
3442
3443
        // Add 10 events (IDs 5-14)
        // Buffer will only keep the last 5: events 10-14
        // Tree will have all blocks
        for id in 5..15 {
            indexer
                .apply_event_with_buffer(make_event(id))
                .await
                .unwrap();
3444
3445
        }

3446
3447
        // Wait for events to be processed by the tree
        tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
3448

3449
3450
3451
3452
3453
3454
        // Helper to extract events from response
        let extract_events = |resp: WorkerKvQueryResponse| -> Vec<RouterEvent> {
            match resp {
                WorkerKvQueryResponse::Events(e) => e,
                WorkerKvQueryResponse::TreeDump(e) => e,
                _ => panic!("Unexpected response type: {:?}", resp),
3455
3456
3457
            }
        };

3458
3459
3460
3461
        // Helper to extract event IDs from result
        let get_ids = |events: Vec<RouterEvent>| -> Vec<u64> {
            events.iter().map(|e| e.event.event_id).collect()
        };
3462

3463
3464
        // Verify buffer state: should have events 10-14 (last 5)
        let buffer_events = indexer.get_all_events_in_buffer();
3465
        assert_eq!(
3466
3467
3468
            get_ids(buffer_events),
            vec![10, 11, 12, 13, 14],
            "Buffer should have events 10-14"
3469
3470
        );

3471
3472
        // ========== BUFFER PATH TESTS (start_id >= first_buffered) ==========
        // Range is [start, end] inclusive
3473

3474
3475
3476
        // Test: start_id within buffer, no end
        let result = indexer.get_events_in_id_range(Some(11), None).await;
        assert!(matches!(result, WorkerKvQueryResponse::Events(_)));
3477
        assert_eq!(
3478
3479
3480
            get_ids(extract_events(result)),
            vec![11, 12, 13, 14],
            "start_id=11 (in buffer) should return [11, 14]"
3481
3482
        );

3483
3484
3485
        // Test: start_id at buffer boundary
        let result = indexer.get_events_in_id_range(Some(10), None).await;
        assert!(matches!(result, WorkerKvQueryResponse::Events(_)));
3486
        assert_eq!(
3487
3488
3489
            get_ids(extract_events(result)),
            vec![10, 11, 12, 13, 14],
            "start_id=10 (buffer start) should return [10, 14]"
3490
3491
        );

3492
3493
3494
        // Test: both start and end within buffer (inclusive)
        let result = indexer.get_events_in_id_range(Some(11), Some(13)).await;
        assert!(matches!(result, WorkerKvQueryResponse::Events(_)));
3495
        assert_eq!(
3496
3497
3498
            get_ids(extract_events(result)),
            vec![11, 12, 13],
            "range [11, 13] inclusive should return 3 events"
3499
        );
3500

3501
3502
3503
3504
3505
3506
3507
        let result = indexer.get_events_in_id_range(Some(10), Some(14)).await;
        assert!(matches!(result, WorkerKvQueryResponse::Events(_)));
        assert_eq!(
            get_ids(extract_events(result)),
            vec![10, 11, 12, 13, 14],
            "range [10, 14] should return all buffer events"
        );
3508

3509
3510
3511
        // ========== TREE DUMP PATH TESTS (range extends before buffer) ==========
        // Note: Tree dumps return synthetic 0-indexed event IDs, so we just check
        // that we get events back (the IDs won't match original IDs)
3512

3513
3514
3515
        // Test: (None, None) dumps entire tree
        let result = indexer.get_events_in_id_range(None, None).await;
        assert!(matches!(result, WorkerKvQueryResponse::TreeDump(_)));
Yan Ru Pei's avatar
Yan Ru Pei committed
3516
        assert_eq!(
3517
3518
3519
            extract_events(result).len(),
            10,
            "(None, None) should dump entire tree (10 events)"
Yan Ru Pei's avatar
Yan Ru Pei committed
3520
        );
3521

3522
3523
3524
3525
3526
3527
3528
        // Test: (None, Some(_)) dumps entire tree
        let result = indexer.get_events_in_id_range(None, Some(8)).await;
        assert!(matches!(result, WorkerKvQueryResponse::TreeDump(_)));
        assert_eq!(
            extract_events(result).len(),
            10,
            "(None, Some(_)) dumps entire tree - end_id is ignored for tree dumps"
Yan Ru Pei's avatar
Yan Ru Pei committed
3529
        );
3530
3531
3532
3533
3534
3535
3536
3537

        // Test: start_id before buffer triggers tree dump
        let result = indexer.get_events_in_id_range(Some(7), None).await;
        assert!(matches!(result, WorkerKvQueryResponse::TreeDump(_)));
        assert_eq!(
            extract_events(result).len(),
            10,
            "start_id=7 (before buffer) should dump entire tree"
Yan Ru Pei's avatar
Yan Ru Pei committed
3538
        );
3539
3540
3541
3542
3543
3544
3545

        let result = indexer.get_events_in_id_range(Some(5), Some(12)).await;
        assert!(matches!(result, WorkerKvQueryResponse::TreeDump(_)));
        assert_eq!(
            extract_events(result).len(),
            10,
            "range [5, 12] extending before buffer should dump entire tree"
Yan Ru Pei's avatar
Yan Ru Pei committed
3546
        );
3547

3548
        // ========== EDGE CASES ==========
3549

3550
3551
3552
3553
3554
3555
3556
        // Single element when start == end (inclusive range)
        let result = indexer.get_events_in_id_range(Some(12), Some(12)).await;
        assert!(matches!(result, WorkerKvQueryResponse::Events(_)));
        assert_eq!(
            get_ids(extract_events(result)),
            vec![12],
            "start == end should return single event"
Yan Ru Pei's avatar
Yan Ru Pei committed
3557
        );
3558

3559
3560
        // InvalidRange when start > end
        let result = indexer.get_events_in_id_range(Some(15), Some(10)).await;
Yan Ru Pei's avatar
Yan Ru Pei committed
3561
        assert!(
3562
3563
            matches!(result, WorkerKvQueryResponse::InvalidRange { .. }),
            "start > end should return InvalidRange"
Yan Ru Pei's avatar
Yan Ru Pei committed
3564
        );
3565

3566
3567
        // TooNew when start_id is beyond buffer
        let result = indexer.get_events_in_id_range(Some(100), Some(200)).await;
Yan Ru Pei's avatar
Yan Ru Pei committed
3568
        assert!(
3569
3570
            matches!(result, WorkerKvQueryResponse::TooNew { .. }),
            "start_id beyond buffer should return TooNew"
Yan Ru Pei's avatar
Yan Ru Pei committed
3571
        );
3572

3573
3574
3575
3576
3577
3578
3579
3580
        // Request with end beyond buffer but valid start -> buffer returns what it has
        let result = indexer.get_events_in_id_range(Some(12), Some(100)).await;
        assert!(matches!(result, WorkerKvQueryResponse::Events(_)));
        assert_eq!(
            get_ids(extract_events(result)),
            vec![12, 13, 14],
            "range with end beyond buffer should return available buffer events"
        );
3581
3582
3583
3584
3585
3586
3587
3588
3589
3590
3591
3592
3593
3594
3595
3596
3597
3598
3599
3600
3601
3602
3603
    }

    #[tokio::test]
    async fn test_local_indexer_buffer_and_serialization() {
        // Tests components of the LocalKvIndexer query without using nats

        let worker_id = 42u64;

        // Create a local indexer
        let token = CancellationToken::new();
        let metrics = Arc::new(KvIndexerMetrics::new_unregistered());
        let local_indexer = Arc::new(LocalKvIndexer::new(token.clone(), 4, metrics, 100));

        // Add events to local indexer's buffer
        let test_event_1 = RouterEvent::new(
            worker_id,
            KvCacheEvent {
                event_id: 1,
                data: KvCacheEventData::Stored(KvCacheStoreData {
                    parent_hash: None,
                    blocks: vec![KvCacheStoredBlockData {
                        block_hash: ExternalSequenceBlockHash(100),
                        tokens_hash: LocalBlockHash(200),
3604
                        mm_extra_info: None,
3605
3606
3607
3608
3609
3610
3611
3612
3613
3614
3615
3616
3617
3618
3619
3620
3621
3622
3623
3624
3625
3626
3627
                    }],
                }),
                dp_rank: 0,
            },
        );

        // Apply events with buffer
        local_indexer
            .apply_event_with_buffer(test_event_1)
            .await
            .unwrap();

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

        // Get buffered events (what the query service would return)
        let buffered_events = local_indexer.get_all_events_in_buffer();

        // Verify buffer contents
        assert_eq!(buffered_events.len(), 1, "Buffer should have 1 event");
        assert_eq!(buffered_events[0].worker_id, worker_id);
        assert_eq!(buffered_events[0].event.event_id, 1);

3628
3629
        // Build the response that would be sent (Events variant)
        let response = WorkerKvQueryResponse::Events(buffered_events.clone());
3630
3631
3632
3633
3634
3635

        // Test serialization/deserialization (simulating NATS round-trip)
        let serialized = serde_json::to_vec(&response).unwrap();
        let deserialized: WorkerKvQueryResponse = serde_json::from_slice(&serialized).unwrap();

        // Verify response correctness
3636
3637
3638
3639
3640
3641
3642
        let events = match deserialized {
            WorkerKvQueryResponse::Events(e) => e,
            _ => panic!("Expected Events variant"),
        };
        assert_eq!(events.len(), 1);
        assert_eq!(events[0].worker_id, worker_id);
        assert_eq!(events[0].event.event_id, 1);
3643
3644

        // Verify event data
3645
        match &events[0].event.data {
3646
3647
3648
3649
3650
3651
3652
3653
3654
            KvCacheEventData::Stored(store_data) => {
                assert_eq!(store_data.blocks.len(), 1);
                assert_eq!(store_data.blocks[0].block_hash.0, 100);
                assert_eq!(store_data.blocks[0].tokens_hash.0, 200);
            }
            _ => panic!("Expected Stored event"),
        }
    }
}