[model-gateway]: optimize prefix_match with zero-copy tenant and deferred char count (#16099)

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()

sgl-model-gateway/src/policies/tree.rs

@@ -20,11 +20,9 @@ pub type TenantId = Arc<str>;
 /// 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<TenantId, u64>,
     /// Parent pointer for upward traversal during timestamp updates
     parent: RwLock<Option<NodeRef>>,
+    /// 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<Option<TenantId>>,
 }
 
 #[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<char> 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<TenantId> = 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)]