chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[model-gateway] Optimize INSERT with leaf-only timestamp updates (#16097)

Browse Source
This commit is contained in:
Simo Lin
2025-12-29 10:43:20 -08:00
committed by GitHub
parent f44c4b3715
commit 684e148e38
1 changed files with 31 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-model-gateway/src/policies/tree.rs
+31 -44
View File
@@ -302,18 +302,14 @@ impl Tree {
// Insert text into tree with given tenant
// Use slice-based traversal to avoid Vec<char> allocation
// Get epoch for LRU ordering
let epoch = get_epoch();
// Intern the tenant ID once for reuse
let tenant_id = intern_tenant(tenant);
// Pre-compute total char count once (needed for size tracking)
let total_char_count = text.chars().count();
// Ensure tenant exists at root (don't update timestamp - root is never evicted)
self.root
.tenant_last_access_time
.insert(Arc::clone(&tenant_id), epoch);
.entry(Arc::clone(&tenant_id))
.or_insert(0);
self.tenant_char_count
.entry(Arc::clone(&tenant_id))
@@ -321,7 +317,6 @@ impl Tree {
// Track remaining text as a slice - no allocation needed
let mut remaining = text;
let mut matched_chars = 0;
let mut prev = Arc::clone(&self.root);
// Result type to carry state out of the match block
@@ -340,8 +335,11 @@ impl Tree {
// Use entry API for atomic check-and-insert semantics (required for thread safety)
let step = match prev.children.entry(first_char) {
Entry::Vacant(entry) => {
// No match - create new node with remaining text
let remaining_char_count = total_char_count - matched_chars;
// No match - create new node with remaining text (this is the leaf)
// Compute remaining char count lazily - only here when creating leaf
let remaining_char_count = remaining.chars().count();
let epoch = get_epoch();
let new_node = Arc::new(Node {
children: DashMap::with_hasher(CharHasherBuilder::default()),
text: RwLock::new(NodeText::new(remaining.to_string())),
@@ -349,7 +347,7 @@ impl Tree {
parent: RwLock::new(Some(Arc::clone(&prev))),
});
// Attach tenant to the new node and increment count
// Attach tenant to the new leaf node with timestamp
self.tenant_char_count
.entry(Arc::clone(&tenant_id))
.and_modify(|count| *count += remaining_char_count)
@@ -398,21 +396,14 @@ impl Tree {
*matched_node.text.write().unwrap() = contracted_text;
*matched_node.parent.write().unwrap() = Some(Arc::clone(&new_node));
// Atomically attach tenant to the new split node
match new_node
.tenant_last_access_time
.entry(Arc::clone(&tenant_id))
{
Entry::Vacant(v) => {
self.tenant_char_count
.entry(Arc::clone(&tenant_id))
.and_modify(|count| *count += matched_text_count)
.or_insert(matched_text_count);
v.insert(epoch);
}
Entry::Occupied(mut o) => {
o.insert(epoch);
}
// Attach tenant to the new split node (intermediate - no timestamp update)
// The cloned DashMap already has the tenant; just ensure char count is correct
if !new_node.tenant_last_access_time.contains_key(tenant_id.as_ref()) {
self.tenant_char_count
.entry(Arc::clone(&tenant_id))
.and_modify(|count| *count += matched_text_count)
.or_insert(matched_text_count);
new_node.tenant_last_access_time.insert(Arc::clone(&tenant_id), 0);
}
InsertStep::Continue {
@@ -420,24 +411,16 @@ impl Tree {
advance_chars: shared_count,
}
} else {
// Full match - move to next node
// Full match - move to next node (intermediate - no timestamp update)
drop(matched_node_text);
// Atomically attach tenant to existing node
match matched_node
.tenant_last_access_time
.entry(Arc::clone(&tenant_id))
{
Entry::Vacant(v) => {
self.tenant_char_count
.entry(Arc::clone(&tenant_id))
.and_modify(|count| *count += matched_node_text_count)
.or_insert(matched_node_text_count);
v.insert(epoch);
}
Entry::Occupied(mut o) => {
o.insert(epoch);
}
// Ensure tenant exists at this intermediate node
if !matched_node.tenant_last_access_time.contains_key(tenant_id.as_ref()) {
self.tenant_char_count
.entry(Arc::clone(&tenant_id))
.and_modify(|count| *count += matched_node_text_count)
.or_insert(matched_node_text_count);
matched_node.tenant_last_access_time.insert(Arc::clone(&tenant_id), 0);
}
InsertStep::Continue {
@@ -450,17 +433,21 @@ impl Tree {
// Entry guard is now dropped - safe to update prev
match step {
InsertStep::Done => break,
InsertStep::Done => return, // New leaf created with timestamp, we're done
InsertStep::Continue {
next_prev,
advance_chars,
} => {
prev = next_prev;
remaining = advance_by_chars(remaining, advance_chars);
matched_chars += advance_chars;
}
}
}
// Loop exited normally (remaining empty) - prev is the leaf node
// Update its timestamp for LRU ordering
let epoch = get_epoch();
prev.tenant_last_access_time.insert(Arc::clone(&tenant_id), epoch);
}
/// Performs prefix matching and returns detailed result with char counts.