From 8c6f865afb01da4947710b31508861adffc1da3b Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Mon, 29 Dec 2025 11:32:57 -0800 Subject: [PATCH] [model-gateway]: optimize prefix_match with zero-copy tenant and deferred char count (#16099) --- sgl-model-gateway/src/policies/cache_aware.rs | 2 +- sgl-model-gateway/src/policies/tree.rs | 83 ++++++++++++------- 2 files changed, 54 insertions(+), 31 deletions(-) diff --git a/sgl-model-gateway/src/policies/cache_aware.rs b/sgl-model-gateway/src/policies/cache_aware.rs index e3c5239f3..044e3a35a 100644 --- a/sgl-model-gateway/src/policies/cache_aware.rs +++ b/sgl-model-gateway/src/policies/cache_aware.rs @@ -302,7 +302,7 @@ impl LoadBalancingPolicy for CacheAwarePolicy { }; let selected_url = if match_rate > self.config.cache_threshold { - result.tenant + result.tenant.to_string() } else { let min_load_idx = *healthy_indices .iter() diff --git a/sgl-model-gateway/src/policies/tree.rs b/sgl-model-gateway/src/policies/tree.rs index 99ee75540..fd178b0d1 100644 --- a/sgl-model-gateway/src/policies/tree.rs +++ b/sgl-model-gateway/src/policies/tree.rs @@ -20,11 +20,9 @@ pub type TenantId = Arc; /// Result of a prefix match operation, including char counts to avoid recomputation. #[derive(Debug, Clone)] pub struct PrefixMatchResult { - /// The matched prefix text - pub matched_text: String, - /// The tenant that owns the matched prefix - pub tenant: String, - /// Number of characters matched (avoids chars().count()) + /// The tenant that owns the matched prefix (zero-copy) + pub tenant: TenantId, + /// Number of characters matched pub matched_char_count: usize, /// Total number of characters in the input text pub input_char_count: usize, @@ -206,6 +204,9 @@ struct Node { tenant_last_access_time: DashMap, /// Parent pointer for upward traversal during timestamp updates parent: RwLock>, + /// Cached last-accessed tenant for O(1) lookup during prefix match. + /// Avoids O(shards) DashMap iteration in the common case. + last_tenant: parking_lot::RwLock>, } #[derive(Debug)] @@ -293,6 +294,7 @@ impl Tree { text: RwLock::new(NodeText::empty()), tenant_last_access_time: DashMap::new(), parent: RwLock::new(None), + last_tenant: parking_lot::RwLock::new(None), }), tenant_char_count: DashMap::new(), } @@ -345,6 +347,7 @@ impl Tree { text: RwLock::new(NodeText::new(remaining.to_string())), tenant_last_access_time: DashMap::new(), parent: RwLock::new(Some(Arc::clone(&prev))), + last_tenant: parking_lot::RwLock::new(Some(Arc::clone(&tenant_id))), }); // Attach tenant to the new leaf node with timestamp @@ -384,6 +387,9 @@ impl Tree { children: DashMap::with_hasher(CharHasherBuilder::default()), parent: RwLock::new(Some(Arc::clone(&prev))), tenant_last_access_time: matched_node.tenant_last_access_time.clone(), + last_tenant: parking_lot::RwLock::new( + matched_node.last_tenant.read().clone(), + ), }); let first_new_char = contracted_text.first_char().unwrap(); @@ -462,11 +468,8 @@ impl Tree { } /// Performs prefix matching and returns detailed result with char counts. - /// This is the optimized version that avoids redundant chars().count() calls. + /// Optimized: no string allocations, deferred char counting. pub fn prefix_match_with_counts(&self, text: &str) -> PrefixMatchResult { - // Use slice-based traversal - no Vec allocation - let input_char_count = text.chars().count(); - let mut remaining = text; let mut matched_chars = 0; let mut prev = Arc::clone(&self.root); @@ -503,40 +506,60 @@ impl Tree { let curr = prev; - // Select the first tenant (key in the map) - let tenant: Option = curr - .tenant_last_access_time - .iter() - .next() - .map(|kv| Arc::clone(kv.key())); + // Try cached tenant first (O(1)) before falling back to O(shards) DashMap iteration. + // The cache is valid if the tenant still exists in tenant_last_access_time. + let tenant: TenantId = { + let cached = curr.last_tenant.read(); + if let Some(ref t) = *cached { + if curr.tenant_last_access_time.contains_key(t.as_ref()) { + Arc::clone(t) + } else { + drop(cached); + // Cache stale, fall back to iteration and update cache + let t = curr + .tenant_last_access_time + .iter() + .next() + .map(|kv| Arc::clone(kv.key())) + .unwrap_or_else(|| Arc::from("empty")); + *curr.last_tenant.write() = Some(Arc::clone(&t)); + t + } + } else { + drop(cached); + // No cache, iterate and populate cache + let t = curr + .tenant_last_access_time + .iter() + .next() + .map(|kv| Arc::clone(kv.key())) + .unwrap_or_else(|| Arc::from("empty")); + *curr.last_tenant.write() = Some(Arc::clone(&t)); + t + } + }; // Update timestamp on the matched node only (O(1)). - // Ancestor propagation is unnecessary - eviction only uses leaf timestamps. - if let Some(ref tenant_id) = tenant { - let epoch = get_epoch(); - curr.tenant_last_access_time - .insert(Arc::clone(tenant_id), epoch); - } + let epoch = get_epoch(); + curr.tenant_last_access_time + .insert(Arc::clone(&tenant), epoch); - // Build matched text from original input using char count - let matched_text = take_chars(text, matched_chars); - let tenant_str = tenant - .map(|t| t.to_string()) - .unwrap_or_else(|| "empty".to_string()); + // Compute input char count from matched + remaining (deferred from start) + let input_char_count = matched_chars + remaining.chars().count(); PrefixMatchResult { - matched_text, - tenant: tenant_str, + tenant, matched_char_count: matched_chars, input_char_count, } } /// Legacy prefix_match API for backward compatibility. - /// Prefer prefix_match_with_counts() for better performance. + /// Note: This computes matched_text which has allocation overhead. pub fn prefix_match(&self, text: &str) -> (String, String) { let result = self.prefix_match_with_counts(text); - (result.matched_text, result.tenant) + let matched_text = take_chars(text, result.matched_char_count); + (matched_text, result.tenant.to_string()) } #[allow(dead_code)]