[model-gateway] Optimize consistent hashing hot path to eliminate allocations (#17467)

sgl-model-gateway/Cargo.toml

@@ -141,6 +141,10 @@ tonic-v12 = { version = "0.12.3", package = "tonic" }
 serial_test = "3.0"
 rsa = { version = "0.9", features = ["sha2"] }
 
+[[bench]]
+name = "consistent_hash_bench"
+harness = false
+path = "benches/consistent_hash_bench.rs"
 [[bench]]
 name = "wasm_middleware_latency"
 harness = false

sgl-model-gateway/benches/consistent_hash_bench.rs

@@ -0,0 +1,36 @@
+use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};
+use smg::mesh::consistent_hash::ConsistentHashRing;
+
+fn setup_ring(node_count: usize) -> ConsistentHashRing {
+    let mut ring = ConsistentHashRing::new();
+    for i in 0..node_count {
+        ring.add_node(&format!("node-{}", i));
+    }
+    ring
+}
+
+fn bench_consistent_hash(c: &mut Criterion) {
+    let mut group = c.benchmark_group("ConsistentHashRing");
+
+    for size in [10, 100, 500].iter() {
+        let ring = setup_ring(*size);
+        let key = "test-request-key-for-rate-limiting";
+        let node_name = "node-5";
+
+        group.bench_with_input(BenchmarkId::new("get_owners", size), size, |b, _| {
+            b.iter(|| {
+                black_box(ring.get_owners(key));
+            });
+        });
+        group.bench_with_input(BenchmarkId::new("is_owner", size), size, |b, _| {
+            b.iter(|| {
+                black_box(ring.is_owner(key, node_name));
+            });
+        });
+    }
+
+    group.finish();
+}
+
+criterion_group!(benches, bench_consistent_hash);
+criterion_main!(benches);

sgl-model-gateway/src/mesh/consistent_hash.rs

@@ -57,27 +57,24 @@ impl ConsistentHashRing {
         }
     }
 
-    /// Get K owners for a key
-    pub fn get_owners(&self, key: &str) -> Vec<String> {
+    pub fn get_owners(&self, key: &str) -> Vec<&str> {
         if self.ring.is_empty() {
             return Vec::new();
         }
 
         let key_hash = Self::hash(key);
-        let mut owners = Vec::new();
-        let mut seen_nodes = HashSet::new();
+
+        let mut owners = Vec::with_capacity(NUM_OWNERS);
         let total_unique_nodes = self.node_hashes.len();
 
-        // Find the first node >= key_hash (clockwise)
         let mut iter = self.ring.range(key_hash..);
-        while owners.len() < NUM_OWNERS && seen_nodes.len() < total_unique_nodes {
+
+        while owners.len() < NUM_OWNERS && owners.len() < total_unique_nodes {
             if let Some((_, node)) = iter.next() {
-                if !seen_nodes.contains(node) {
-                    owners.push(node.clone());
-                    seen_nodes.insert(node.clone());
+                if !owners.contains(&node.as_str()) {
+                    owners.push(node.as_str());
                 }
             } else {
-                // Wrap around to the beginning
                 iter = self.ring.range(..);
             }
         }
@@ -87,7 +84,7 @@ impl ConsistentHashRing {
 
     /// Check if a node is an owner of a key
     pub fn is_owner(&self, key: &str, node_name: &str) -> bool {
-        self.get_owners(key).contains(&node_name.to_string())
+        self.get_owners(key).contains(&node_name)
     }
 
     /// Get all nodes in the ring
@@ -197,8 +194,8 @@ mod tests {
         let owners = ring.get_owners("test_key");
         // Should return all available nodes (2) without infinite loop
         assert_eq!(owners.len(), 2);
-        assert!(owners.contains(&"node1".to_string()));
-        assert!(owners.contains(&"node2".to_string()));
+        assert!(owners.contains(&"node1"));
+        assert!(owners.contains(&"node2"));
     }
 
     #[test]

sgl-model-gateway/src/mesh/stores.rs

@@ -379,6 +379,9 @@ impl RateLimitStore {
     pub fn get_owners(&self, key: &str) -> Vec<String> {
         let ring = self.hash_ring.read();
         ring.get_owners(key)
+            .into_iter()
+            .map(|s| s.to_string())
+            .collect()
     }
 
     /// Get or create counter (only if this node is an owner)
@@ -441,8 +444,11 @@ impl RateLimitStore {
 
         for key in counters.keys() {
             let owners = ring.get_owners(key);
-            // Check if any owner has failed
-            if owners.iter().any(|owner| failed_nodes.contains(owner)) {
+
+            if owners
+                .iter()
+                .any(|&owner| failed_nodes.iter().any(|f| f == owner))
+            {
                 // Check if we are now an owner
                 if ring.is_owner(key, &self.self_name) {
                     affected_keys.push(key.clone());