diff --git a/sgl-model-gateway/src/policies/tree.rs b/sgl-model-gateway/src/policies/tree.rs index 42c09b90c..2fe544b35 100644 --- a/sgl-model-gateway/src/policies/tree.rs +++ b/sgl-model-gateway/src/policies/tree.rs @@ -302,18 +302,14 @@ impl Tree { // Insert text into tree with given tenant // Use slice-based traversal to avoid Vec 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.