[model-gateway]: ASCII byte comparison and probabilistic timestamp updates (#16181)

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

@@ -13,6 +13,28 @@ use tracing::debug;
 
 type NodeRef = Arc<Node>;
 
+/// Shard counts for DashMaps to balance concurrency vs allocation overhead.
+/// Default DashMap uses num_cpus * 4 shards (e.g., 256 on 64-core machines).
+///
+/// Root node uses higher shard count since ALL requests pass through it.
+/// Other nodes use lower count as traffic diverges through the tree.
+///
+/// This reduces memory by ~90% vs default while maintaining good concurrency.
+const ROOT_SHARD_COUNT: usize = 32;
+const NODE_SHARD_COUNT: usize = 8;
+
+/// Create a children DashMap for non-root nodes
+#[inline]
+fn new_children_map() -> DashMap<char, NodeRef, CharHasherBuilder> {
+    DashMap::with_hasher_and_shard_amount(CharHasherBuilder::default(), NODE_SHARD_COUNT)
+}
+
+/// Create a tenant access time DashMap for non-root nodes
+#[inline]
+fn new_tenant_map() -> DashMap<TenantId, u64> {
+    DashMap::with_shard_amount(NODE_SHARD_COUNT)
+}
+
 /// Interned tenant ID to avoid repeated string allocations.
 /// Using Arc<str> allows cheap cloning and comparison.
 pub type TenantId = Arc<str>;
@@ -65,11 +87,22 @@ type CharHasherBuilder = BuildHasherDefault<CharHasher>;
 
 /// Advance a string slice by N characters, returning the remaining slice.
 /// Returns empty string if n >= char count.
+/// Optimized: uses direct byte slicing for ASCII, falls back to char_indices for UTF-8.
 #[inline]
 fn advance_by_chars(s: &str, n: usize) -> &str {
     if n == 0 {
         return s;
     }
+    if n >= s.len() {
+        return "";
+    }
+    // Fast path: if first N bytes are all ASCII, we can slice directly
+    let bytes = s.as_bytes();
+    if bytes[..n].is_ascii() {
+        // Safe: we verified all bytes in [0..n] are ASCII (valid UTF-8 boundary)
+        return &s[n..];
+    }
+    // Slow path: UTF-8 requires char-by-char traversal
     s.char_indices()
         .nth(n)
         .map(|(idx, _)| &s[idx..])
@@ -251,9 +284,31 @@ impl PartialEq for EvictionEntry {
 
 /// Count matching prefix characters between two strings.
 /// Returns the number of characters that match from the start.
-/// Uses iterator-based comparison - no allocation required.
+/// Optimized: uses fast byte comparison for ASCII, falls back to char iteration for UTF-8.
 #[inline]
 fn shared_prefix_count(a: &str, b: &str) -> usize {
+    let a_bytes = a.as_bytes();
+    let b_bytes = b.as_bytes();
+
+    // Find common byte prefix length using iterator (potentially SIMD-optimized)
+    let common_byte_len = a_bytes
+        .iter()
+        .zip(b_bytes)
+        .position(|(&a_byte, &b_byte)| a_byte != b_byte)
+        .unwrap_or_else(|| a_bytes.len().min(b_bytes.len()));
+
+    // If the common byte prefix is all ASCII, byte count == char count
+    // Otherwise, fall back to char-by-char comparison for UTF-8 safety
+    if a_bytes[..common_byte_len].is_ascii() {
+        common_byte_len
+    } else {
+        shared_prefix_count_chars(a, b)
+    }
+}
+
+/// Fallback char-by-char comparison for strings with non-ASCII characters.
+#[inline]
+fn shared_prefix_count_chars(a: &str, b: &str) -> usize {
     a.chars()
         .zip(b.chars())
         .take_while(|(a_char, b_char)| a_char == b_char)
@@ -289,14 +344,18 @@ impl Tree {
 
     pub fn new() -> Self {
         Tree {
+            // Root uses higher shard count since ALL requests pass through it
             root: Arc::new(Node {
-                children: DashMap::with_hasher(CharHasherBuilder::default()),
+                children: DashMap::with_hasher_and_shard_amount(
+                    CharHasherBuilder::default(),
+                    ROOT_SHARD_COUNT,
+                ),
                 text: RwLock::new(NodeText::empty()),
-                tenant_last_access_time: DashMap::new(),
+                tenant_last_access_time: DashMap::with_shard_amount(ROOT_SHARD_COUNT),
                 parent: RwLock::new(None),
                 last_tenant: parking_lot::RwLock::new(None),
             }),
-            tenant_char_count: DashMap::new(),
+            tenant_char_count: DashMap::with_shard_amount(ROOT_SHARD_COUNT),
         }
     }
 
@@ -343,9 +402,9 @@ impl Tree {
                     let epoch = get_epoch();
 
                     let new_node = Arc::new(Node {
-                        children: DashMap::with_hasher(CharHasherBuilder::default()),
+                        children: new_children_map(),
                         text: RwLock::new(NodeText::new(remaining.to_string())),
-                        tenant_last_access_time: DashMap::new(),
+                        tenant_last_access_time: new_tenant_map(),
                         parent: RwLock::new(Some(Arc::clone(&prev))),
                         last_tenant: parking_lot::RwLock::new(Some(Arc::clone(&tenant_id))),
                     });
@@ -384,7 +443,7 @@ impl Tree {
 
                         let new_node = Arc::new(Node {
                             text: RwLock::new(matched_text),
-                            children: DashMap::with_hasher(CharHasherBuilder::default()),
+                            children: new_children_map(),
                             parent: RwLock::new(Some(Arc::clone(&prev))),
                             tenant_last_access_time: matched_node.tenant_last_access_time.clone(),
                             last_tenant: parking_lot::RwLock::new(
@@ -539,10 +598,13 @@ impl Tree {
             }
         };
 
-        // Update timestamp on the matched node only (O(1)).
+        // Update timestamp probabilistically (1 in 8 matches) to reduce DashMap contention.
+        // LRU eviction doesn't need perfect accuracy - approximate timestamps suffice.
         let epoch = get_epoch();
+        if epoch & 0x7 == 0 {
             curr.tenant_last_access_time
                 .insert(Arc::clone(&tenant), epoch);
+        }
 
         // Compute input char count from matched + remaining (deferred from start)
         let input_char_count = matched_chars + remaining.chars().count();