oximedia_distributed/
shard_map.rs

1#![allow(dead_code)]
2//! Shard mapping for distributing data across nodes using consistent hashing.
3//!
4//! This module provides a virtual-node-based consistent hashing ring that maps
5//! arbitrary keys to physical node assignments. It supports:
6//! - Adding/removing nodes with automatic rebalancing
7//! - Configurable virtual node count per physical node
8//! - Key-to-node lookups with O(log n) performance
9//! - Shard statistics and load factor computation
10
11use std::collections::{BTreeMap, HashMap};
12use std::fmt;
13
14/// Identifier for a physical node in the cluster.
15#[derive(Debug, Clone, PartialEq, Eq, Hash)]
16pub struct NodeId {
17    /// Unique string identifier for the node.
18    pub id: String,
19    /// Human-readable label.
20    pub label: String,
21    /// Weight factor for virtual node count (1.0 = normal).
22    pub weight: u32,
23}
24
25impl NodeId {
26    /// Create a new node identifier.
27    #[must_use]
28    pub fn new(id: &str, label: &str) -> Self {
29        Self {
30            id: id.to_string(),
31            label: label.to_string(),
32            weight: 100,
33        }
34    }
35
36    /// Create a node with a custom weight (percentage; 100 = normal).
37    #[must_use]
38    pub fn with_weight(mut self, weight: u32) -> Self {
39        self.weight = weight;
40        self
41    }
42}
43
44impl fmt::Display for NodeId {
45    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
46        write!(f, "{}({})", self.label, self.id)
47    }
48}
49
50/// A virtual node on the hash ring.
51#[derive(Debug, Clone)]
52struct VNode {
53    /// Hash position on the ring.
54    hash: u64,
55    /// The physical node this virtual node belongs to.
56    node_id: String,
57    /// Virtual node index.
58    vnode_index: u32,
59}
60
61/// Configuration for the shard map.
62#[derive(Debug, Clone)]
63pub struct ShardMapConfig {
64    /// Base number of virtual nodes per physical node.
65    pub vnodes_per_node: u32,
66    /// Whether to use node weight for virtual node count scaling.
67    pub use_weights: bool,
68}
69
70impl Default for ShardMapConfig {
71    fn default() -> Self {
72        Self {
73            vnodes_per_node: 150,
74            use_weights: true,
75        }
76    }
77}
78
79/// Consistent hash ring for shard mapping.
80#[derive(Debug, Clone)]
81pub struct ShardMap {
82    /// The hash ring: hash -> `node_id`.
83    ring: BTreeMap<u64, String>,
84    /// Registered physical nodes.
85    nodes: HashMap<String, NodeId>,
86    /// Configuration.
87    config: ShardMapConfig,
88}
89
90impl ShardMap {
91    /// Create a new empty shard map with default configuration.
92    #[must_use]
93    pub fn new() -> Self {
94        Self {
95            ring: BTreeMap::new(),
96            nodes: HashMap::new(),
97            config: ShardMapConfig::default(),
98        }
99    }
100
101    /// Create a shard map with custom configuration.
102    #[must_use]
103    pub fn with_config(config: ShardMapConfig) -> Self {
104        Self {
105            ring: BTreeMap::new(),
106            nodes: HashMap::new(),
107            config,
108        }
109    }
110
111    /// Add a node to the ring.
112    pub fn add_node(&mut self, node: NodeId) {
113        let vnode_count = self.effective_vnodes(&node);
114        for i in 0..vnode_count {
115            let key = format!("{}:vnode:{}", node.id, i);
116            let hash = Self::hash_key(&key);
117            self.ring.insert(hash, node.id.clone());
118        }
119        self.nodes.insert(node.id.clone(), node);
120    }
121
122    /// Remove a node from the ring.
123    pub fn remove_node(&mut self, node_id: &str) -> bool {
124        if let Some(node) = self.nodes.remove(node_id) {
125            let vnode_count = self.effective_vnodes(&node);
126            for i in 0..vnode_count {
127                let key = format!("{}:vnode:{}", node.id, i);
128                let hash = Self::hash_key(&key);
129                self.ring.remove(&hash);
130            }
131            true
132        } else {
133            false
134        }
135    }
136
137    /// Look up which node a key maps to.
138    #[must_use]
139    pub fn lookup(&self, key: &str) -> Option<&str> {
140        if self.ring.is_empty() {
141            return None;
142        }
143        let hash = Self::hash_key(key);
144        // Find the first node at or after the hash
145        if let Some((_h, node_id)) = self.ring.range(hash..).next() {
146            return Some(node_id.as_str());
147        }
148        // Wrap around to the first node in the ring
149        self.ring.values().next().map(std::string::String::as_str)
150    }
151
152    /// Get the number of physical nodes.
153    #[must_use]
154    pub fn node_count(&self) -> usize {
155        self.nodes.len()
156    }
157
158    /// Get the total number of virtual nodes on the ring.
159    #[must_use]
160    pub fn vnode_count(&self) -> usize {
161        self.ring.len()
162    }
163
164    /// Check if the ring contains a specific node.
165    #[must_use]
166    pub fn has_node(&self, node_id: &str) -> bool {
167        self.nodes.contains_key(node_id)
168    }
169
170    /// Get all registered node IDs.
171    #[must_use]
172    pub fn node_ids(&self) -> Vec<&str> {
173        self.nodes.keys().map(std::string::String::as_str).collect()
174    }
175
176    /// Compute load distribution: how many virtual nodes each physical node owns.
177    #[must_use]
178    pub fn load_distribution(&self) -> HashMap<String, usize> {
179        let mut dist: HashMap<String, usize> = HashMap::new();
180        for node_id in self.ring.values() {
181            *dist.entry(node_id.clone()).or_insert(0) += 1;
182        }
183        dist
184    }
185
186    /// Compute the load factor (std dev / mean of vnode counts).
187    #[allow(clippy::cast_precision_loss)]
188    #[must_use]
189    pub fn load_factor(&self) -> f64 {
190        let dist = self.load_distribution();
191        if dist.is_empty() {
192            return 0.0;
193        }
194        let counts: Vec<f64> = dist.values().map(|&c| c as f64).collect();
195        let mean = counts.iter().sum::<f64>() / counts.len() as f64;
196        if mean == 0.0 {
197            return 0.0;
198        }
199        let variance = counts.iter().map(|c| (c - mean).powi(2)).sum::<f64>() / counts.len() as f64;
200        variance.sqrt() / mean
201    }
202
203    /// Compute effective virtual node count for a given physical node.
204    #[allow(
205        clippy::cast_precision_loss,
206        clippy::cast_possible_truncation,
207        clippy::cast_sign_loss
208    )]
209    fn effective_vnodes(&self, node: &NodeId) -> u32 {
210        if self.config.use_weights {
211            let scaled = f64::from(self.config.vnodes_per_node) * (f64::from(node.weight) / 100.0);
212            scaled.round() as u32
213        } else {
214            self.config.vnodes_per_node
215        }
216    }
217
218    /// Simple FNV-1a-style hash for deterministic key hashing.
219    fn hash_key(key: &str) -> u64 {
220        let mut hash: u64 = 0xcbf2_9ce4_8422_2325;
221        for byte in key.as_bytes() {
222            hash ^= u64::from(*byte);
223            hash = hash.wrapping_mul(0x0100_0000_01b3);
224        }
225        hash
226    }
227}
228
229impl Default for ShardMap {
230    fn default() -> Self {
231        Self::new()
232    }
233}
234
235/// Shard assignment result for a batch of keys.
236#[derive(Debug, Clone)]
237pub struct BatchAssignment {
238    /// Map of key -> assigned `node_id`.
239    pub assignments: HashMap<String, String>,
240    /// Keys that could not be assigned (empty ring).
241    pub unassigned: Vec<String>,
242}
243
244impl BatchAssignment {
245    /// Create a new empty batch assignment.
246    #[must_use]
247    pub fn new() -> Self {
248        Self {
249            assignments: HashMap::new(),
250            unassigned: Vec::new(),
251        }
252    }
253
254    /// Number of successfully assigned keys.
255    #[must_use]
256    pub fn assigned_count(&self) -> usize {
257        self.assignments.len()
258    }
259
260    /// Number of unassigned keys.
261    #[must_use]
262    pub fn unassigned_count(&self) -> usize {
263        self.unassigned.len()
264    }
265}
266
267impl Default for BatchAssignment {
268    fn default() -> Self {
269        Self::new()
270    }
271}
272
273/// Assign a batch of keys to nodes using the given shard map.
274#[must_use]
275pub fn batch_assign(shard_map: &ShardMap, keys: &[&str]) -> BatchAssignment {
276    let mut result = BatchAssignment::new();
277    for &key in keys {
278        if let Some(node_id) = shard_map.lookup(key) {
279            result
280                .assignments
281                .insert(key.to_string(), node_id.to_string());
282        } else {
283            result.unassigned.push(key.to_string());
284        }
285    }
286    result
287}
288
289#[cfg(test)]
290mod tests {
291    use super::*;
292
293    #[test]
294    fn test_node_id_creation() {
295        let node = NodeId::new("node-1", "Worker 1");
296        assert_eq!(node.id, "node-1");
297        assert_eq!(node.label, "Worker 1");
298        assert_eq!(node.weight, 100);
299    }
300
301    #[test]
302    fn test_node_id_with_weight() {
303        let node = NodeId::new("n1", "N1").with_weight(200);
304        assert_eq!(node.weight, 200);
305    }
306
307    #[test]
308    fn test_node_id_display() {
309        let node = NodeId::new("n1", "Worker");
310        assert_eq!(node.to_string(), "Worker(n1)");
311    }
312
313    #[test]
314    fn test_shard_map_empty() {
315        let sm = ShardMap::new();
316        assert_eq!(sm.node_count(), 0);
317        assert_eq!(sm.vnode_count(), 0);
318        assert_eq!(sm.lookup("any-key"), None);
319    }
320
321    #[test]
322    fn test_shard_map_add_node() {
323        let mut sm = ShardMap::new();
324        sm.add_node(NodeId::new("n1", "Node 1"));
325        assert_eq!(sm.node_count(), 1);
326        assert!(sm.has_node("n1"));
327        assert!(sm.vnode_count() > 0);
328    }
329
330    #[test]
331    fn test_shard_map_remove_node() {
332        let mut sm = ShardMap::new();
333        sm.add_node(NodeId::new("n1", "Node 1"));
334        assert!(sm.remove_node("n1"));
335        assert_eq!(sm.node_count(), 0);
336        assert_eq!(sm.vnode_count(), 0);
337        assert!(!sm.has_node("n1"));
338    }
339
340    #[test]
341    fn test_shard_map_remove_nonexistent() {
342        let mut sm = ShardMap::new();
343        assert!(!sm.remove_node("nonexistent"));
344    }
345
346    #[test]
347    fn test_shard_map_lookup_single_node() {
348        let mut sm = ShardMap::new();
349        sm.add_node(NodeId::new("n1", "Node 1"));
350        // With a single node, all keys must map to it
351        assert_eq!(sm.lookup("key-a"), Some("n1"));
352        assert_eq!(sm.lookup("key-b"), Some("n1"));
353        assert_eq!(sm.lookup("key-c"), Some("n1"));
354    }
355
356    #[test]
357    fn test_shard_map_lookup_deterministic() {
358        let mut sm = ShardMap::new();
359        sm.add_node(NodeId::new("n1", "Node 1"));
360        sm.add_node(NodeId::new("n2", "Node 2"));
361        let result1 = sm
362            .lookup("my-key")
363            .expect("lookup should succeed")
364            .to_string();
365        let result2 = sm
366            .lookup("my-key")
367            .expect("lookup should succeed")
368            .to_string();
369        assert_eq!(result1, result2);
370    }
371
372    #[test]
373    fn test_shard_map_distribution() {
374        let mut sm = ShardMap::new();
375        sm.add_node(NodeId::new("n1", "N1"));
376        sm.add_node(NodeId::new("n2", "N2"));
377        sm.add_node(NodeId::new("n3", "N3"));
378
379        let dist = sm.load_distribution();
380        assert_eq!(dist.len(), 3);
381        // Each node should have roughly vnodes_per_node vnodes
382        for count in dist.values() {
383            assert!(*count > 0);
384        }
385    }
386
387    #[test]
388    fn test_shard_map_load_factor() {
389        let mut sm = ShardMap::new();
390        sm.add_node(NodeId::new("n1", "N1"));
391        sm.add_node(NodeId::new("n2", "N2"));
392        sm.add_node(NodeId::new("n3", "N3"));
393        let lf = sm.load_factor();
394        // Load factor should be small for equal-weight nodes
395        assert!(lf < 0.5, "load factor too high: {}", lf);
396    }
397
398    #[test]
399    fn test_shard_map_load_factor_empty() {
400        let sm = ShardMap::new();
401        assert_eq!(sm.load_factor(), 0.0);
402    }
403
404    #[test]
405    fn test_shard_map_node_ids() {
406        let mut sm = ShardMap::new();
407        sm.add_node(NodeId::new("a", "A"));
408        sm.add_node(NodeId::new("b", "B"));
409        let mut ids = sm.node_ids();
410        ids.sort();
411        assert_eq!(ids, vec!["a", "b"]);
412    }
413
414    #[test]
415    fn test_weighted_nodes() {
416        let mut sm = ShardMap::with_config(ShardMapConfig {
417            vnodes_per_node: 100,
418            use_weights: true,
419        });
420        sm.add_node(NodeId::new("n1", "N1").with_weight(100));
421        sm.add_node(NodeId::new("n2", "N2").with_weight(200));
422
423        let dist = sm.load_distribution();
424        let n1_count = dist.get("n1").copied().unwrap_or(0);
425        let n2_count = dist.get("n2").copied().unwrap_or(0);
426        // n2 should have roughly twice as many vnodes as n1
427        assert!(
428            n2_count > n1_count,
429            "n2={} should be > n1={}",
430            n2_count,
431            n1_count
432        );
433    }
434
435    #[test]
436    fn test_batch_assign() {
437        let mut sm = ShardMap::new();
438        sm.add_node(NodeId::new("n1", "N1"));
439        sm.add_node(NodeId::new("n2", "N2"));
440
441        let keys = vec!["key1", "key2", "key3"];
442        let result = batch_assign(&sm, &keys);
443        assert_eq!(result.assigned_count(), 3);
444        assert_eq!(result.unassigned_count(), 0);
445    }
446
447    #[test]
448    fn test_batch_assign_empty_ring() {
449        let sm = ShardMap::new();
450        let keys = vec!["key1", "key2"];
451        let result = batch_assign(&sm, &keys);
452        assert_eq!(result.assigned_count(), 0);
453        assert_eq!(result.unassigned_count(), 2);
454    }
455
456    #[test]
457    fn test_default_config() {
458        let config = ShardMapConfig::default();
459        assert_eq!(config.vnodes_per_node, 150);
460        assert!(config.use_weights);
461    }
462
463    #[test]
464    fn test_shard_map_default_trait() {
465        let sm = ShardMap::default();
466        assert_eq!(sm.node_count(), 0);
467    }
468
469    /// Verify that the BTreeMap-based O(log n) lookup produces identical results
470    /// to a naive linear scan over all virtual nodes.
471    #[test]
472    fn test_consistent_hash_lookup_binary_matches_linear() {
473        let mut sm = ShardMap::new();
474        for i in 0..5_u32 {
475            sm.add_node(NodeId::new(&format!("node-{i}"), &format!("Node {i}")));
476        }
477
478        // Collect ring entries as a sorted Vec for a reference linear scan
479        let ring_vec: Vec<(u64, String)> = sm.ring.iter().map(|(&h, v)| (h, v.clone())).collect();
480
481        let linear_lookup = |key: &str| -> Option<&str> {
482            if ring_vec.is_empty() {
483                return None;
484            }
485            let hash = ShardMap::hash_key(key);
486            // Linear scan for first entry >= hash (ring is sorted)
487            ring_vec
488                .iter()
489                .find(|(h, _)| *h >= hash)
490                .or_else(|| ring_vec.first())
491                .map(|(_, id)| id.as_str())
492        };
493
494        // Run 1000 queries and compare
495        let queries: Vec<String> = (0..1000_u32).map(|i| format!("key-{i}")).collect();
496        let mut mismatches = 0_u32;
497        for q in &queries {
498            let btree_result = sm.lookup(q);
499            let linear_result = linear_lookup(q);
500            if btree_result != linear_result {
501                mismatches += 1;
502            }
503        }
504        assert_eq!(
505            mismatches, 0,
506            "BTreeMap lookup and linear scan disagree on {mismatches} of 1000 queries"
507        );
508    }
509}
oximedia_distributed/shard_map.rs

oximedia_distributed/
shard_map.rs
