51 lines
1.6 KiB
Rust
51 lines
1.6 KiB
Rust
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
|
|
use sgl_kv_indexer::PrefixOutcome;
|
|
use sgl_router::policies::prefix_provider::RadixTreePrefixProvider;
|
|
use sgl_router::state::kv_events::{compute_block_hashes, BlockSizeOracle, HashTree, KvWorkerId};
|
|
|
|
#[test]
|
|
fn radix_tree_reports_contiguous_prefix_depth_per_worker() {
|
|
let tokens = [11_u32, 12, 13, 14];
|
|
let hashes = compute_block_hashes(&tokens, 1);
|
|
let tree = Arc::new(HashTree::new());
|
|
let oracle = BlockSizeOracle::new();
|
|
oracle.try_set(1).unwrap();
|
|
|
|
tree.insert(&KvWorkerId::new("http://deep".into(), 0), None, &hashes);
|
|
tree.insert(
|
|
&KvWorkerId::new("http://deep".into(), 1),
|
|
None,
|
|
&hashes[..3],
|
|
);
|
|
tree.insert(
|
|
&KvWorkerId::new("http://shallow".into(), 0),
|
|
None,
|
|
&hashes[..2],
|
|
);
|
|
|
|
let signal = RadixTreePrefixProvider::new(tree, oracle)
|
|
.match_request_tokens(&tokens)
|
|
.expect("established local tree must produce a prefix signal");
|
|
let PrefixOutcome::Matched {
|
|
matches,
|
|
best_prefix_blocks,
|
|
} = signal.outcome
|
|
else {
|
|
panic!("local radix-tree hit must be normalized as a match");
|
|
};
|
|
let depth_by_url: HashMap<_, _> = matches
|
|
.into_iter()
|
|
.map(|entry| (entry.address, entry.matched_prefix_blocks))
|
|
.collect();
|
|
|
|
assert_eq!(signal.query_blocks, 4);
|
|
assert_eq!(best_prefix_blocks, 4);
|
|
assert_eq!(depth_by_url.get("http://deep"), Some(&4));
|
|
assert_eq!(depth_by_url.get("http://shallow"), Some(&2));
|
|
}
|