kv_index/
lib.rs

1//! Radix tree implementations for prefix matching and cache-aware routing.
2//!
3//! This module provides radix tree data structures that mirror SGLang's
4//! scheduler memory management patterns. Two implementations are available:
5//!
6//! - [`StringTree`]: Character-based tree for HTTP router (text input)
7//! - [`TokenTree`]: Token-based tree for gRPC router (pre-tokenized input)
8//!
9//! Both implementations support:
10//! - Multi-tenant prefix tracking with LRU eviction
11//! - Concurrent access via DashMap and RwLock
12//! - Efficient prefix matching with match counts
13
14mod common;
15mod event_tree;
16mod path_hash;
17pub mod snapshot;
18mod string_tree;
19mod token_tree;
20
21pub use common::{MatchResult, TenantId};
22pub use event_tree::{
23    compute_content_hash, compute_request_content_hashes, ApplyError, ContentHash, OverlapScores,
24    PositionalIndexer, SequenceHash, StoredBlock, WorkerBlockMap, WorkerId,
25};
26pub use path_hash::{hash_node_path, hash_token_path, GLOBAL_EVICTION_HASH};
27// Re-export under names matching old tree.rs API for easier migration
28pub use string_tree::Tree;
29pub use string_tree::{
30    PrefixMatchResult as StringMatchResult, PrefixMatchResult, Tree as StringTree,
31};
32pub use token_tree::{PrefixMatchResult as TokenMatchResult, TokenTree};
33
34/// Trait for radix tree implementations.
35///
36/// This trait provides a unified interface for both string-based and token-based
37/// radix trees used in cache-aware routing.
38pub trait RadixTree: Send + Sync {
39    /// The key type for this tree (e.g., &str or &[u32])
40    type Key: ?Sized;
41
42    /// The result type returned by prefix matching
43    type MatchResult: MatchResult;
44
45    /// Insert a key with associated tenant.
46    ///
47    /// If the key already exists or shares a prefix with existing keys,
48    /// the tree structure is updated to track the tenant association.
49    fn insert(&self, key: &Self::Key, tenant: &str);
50
51    /// Find the longest matching prefix and return the associated tenant.
52    ///
53    /// Returns `None` if no prefix matches.
54    fn prefix_match(&self, key: &Self::Key) -> Option<TenantId>;
55
56    /// Find the longest matching prefix with detailed match counts.
57    ///
58    /// Returns match result with:
59    /// - `tenant`: The tenant that owns the matched prefix
60    /// - `matched_count`: Number of units (chars/tokens) matched
61    /// - `input_count`: Total units in the input key
62    fn prefix_match_with_counts(&self, key: &Self::Key) -> Self::MatchResult;
63
64    /// Evict cached entries for a tenant to reduce memory usage.
65    ///
66    /// # Arguments
67    /// * `tenant` - The tenant whose entries should be evicted
68    /// * `max_units` - Maximum units (chars/tokens) to retain for this tenant
69    fn evict(&self, tenant: &TenantId, max_units: usize);
70
71    // TODO: Add remove_tenant(&self, tenant: &str) with efficient implementation.
72    // Current naive approach requires O(n) full tree traversal.
73    // Efficient implementation options:
74    // 1. Maintain reverse index: DashMap<TenantId, Vec<Weak<Node>>> for O(k) removal
75    // 2. Lazy tombstone marking with background cleanup
76    // For now, rely on LRU eviction to clean up stale entries.
77
78    /// Get the current size (in units) for a tenant.
79    fn tenant_size(&self, tenant: &TenantId) -> usize;
80
81    /// Reset the tree to empty state.
82    fn reset(&self);
83}
kv_index/lib.rs

kv_index/
lib.rs
