diff --git a/sgl-model-gateway/Cargo.toml b/sgl-model-gateway/Cargo.toml index 97e11becb..422d2f8c8 100644 --- a/sgl-model-gateway/Cargo.toml +++ b/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 diff --git a/sgl-model-gateway/benches/consistent_hash_bench.rs b/sgl-model-gateway/benches/consistent_hash_bench.rs new file mode 100644 index 000000000..0806bcccd --- /dev/null +++ b/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); diff --git a/sgl-model-gateway/src/mesh/consistent_hash.rs b/sgl-model-gateway/src/mesh/consistent_hash.rs index cb50dc45b..4846fcff1 100644 --- a/sgl-model-gateway/src/mesh/consistent_hash.rs +++ b/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 { + 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] diff --git a/sgl-model-gateway/src/mesh/stores.rs b/sgl-model-gateway/src/mesh/stores.rs index 3f08e9473..6429e1969 100644 --- a/sgl-model-gateway/src/mesh/stores.rs +++ b/sgl-model-gateway/src/mesh/stores.rs @@ -379,6 +379,9 @@ impl RateLimitStore { pub fn get_owners(&self, key: &str) -> Vec { 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());