rust_zero_core/
hash.rs

1use std::collections::{BTreeMap, HashSet};
2
3/// A deterministic consistent-hash ring with virtual replicas.
4pub struct ConsistentHash {
5    replicas: usize,
6    nodes: HashSet<String>,
7    ring: BTreeMap<u64, Vec<String>>,
8}
9
10impl ConsistentHash {
11    pub fn new(replicas: usize) -> Self {
12        assert!(replicas > 0, "replica count must be greater than zero");
13        Self {
14            replicas,
15            nodes: HashSet::new(),
16            ring: BTreeMap::new(),
17        }
18    }
19
20    pub fn add(&mut self, node: impl Into<String>) -> bool {
21        let node = node.into();
22        if !self.nodes.insert(node.clone()) {
23            return false;
24        }
25
26        for replica in 0..self.replicas {
27            self.ring
28                .entry(hash_bytes(format!("{node}#{replica}").as_bytes()))
29                .or_default()
30                .push(node.clone());
31        }
32        true
33    }
34
35    pub fn remove(&mut self, node: &str) -> bool {
36        if !self.nodes.remove(node) {
37            return false;
38        }
39
40        for replica in 0..self.replicas {
41            let hash = hash_bytes(format!("{node}#{replica}").as_bytes());
42            let remove_entry = if let Some(nodes) = self.ring.get_mut(&hash) {
43                nodes.retain(|candidate| candidate != node);
44                nodes.is_empty()
45            } else {
46                false
47            };
48            if remove_entry {
49                self.ring.remove(&hash);
50            }
51        }
52        true
53    }
54
55    pub fn get(&self, key: impl AsRef<[u8]>) -> Option<&str> {
56        let hash = hash_bytes(key.as_ref());
57        self.ring
58            .range(hash..)
59            .next()
60            .or_else(|| self.ring.first_key_value())
61            .and_then(|(_, nodes)| nodes.first())
62            .map(String::as_str)
63    }
64
65    pub fn len(&self) -> usize {
66        self.nodes.len()
67    }
68
69    pub fn is_empty(&self) -> bool {
70        self.nodes.is_empty()
71    }
72}
73
74fn hash_bytes(bytes: &[u8]) -> u64 {
75    const OFFSET_BASIS: u64 = 0xcbf29ce484222325;
76    const PRIME: u64 = 0x100000001b3;
77
78    bytes.iter().fold(OFFSET_BASIS, |hash, byte| {
79        (hash ^ u64::from(*byte)).wrapping_mul(PRIME)
80    })
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86
87    #[test]
88    fn resolves_keys_to_added_nodes() {
89        let mut ring = ConsistentHash::new(100);
90        ring.add("a");
91        ring.add("b");
92
93        assert!(matches!(ring.get("customer-42"), Some("a" | "b")));
94    }
95
96    #[test]
97    fn removed_nodes_are_never_returned() {
98        let mut ring = ConsistentHash::new(100);
99        ring.add("a");
100        ring.add("b");
101        assert!(ring.remove("a"));
102
103        for customer in 0..100 {
104            assert_eq!(ring.get(format!("customer-{customer}")), Some("b"));
105        }
106    }
107}
rust_zero_core/hash.rs

rust_zero_core/
hash.rs
