basic_hash_ring/
lib.rs

1use std::collections::{HashMap, HashSet};
2use xxhash_rust::xxh3::xxh3_64;
3
4pub enum Direction {
5    Forward,
6    Backward,
7}
8
9pub struct HashRing {
10    replicas: usize,
11    sorted_keys: Vec<u64>,
12    hash_map: HashMap<u64, String>,
13}
14
15impl HashRing {
16    pub fn new_with_replicas(replicas: usize) -> Self {
17        HashRing {
18            replicas,
19            sorted_keys: Vec::new(),
20            hash_map: HashMap::new(),
21        }
22    }
23
24    pub fn is_empty(&self) -> bool {
25        self.sorted_keys.is_empty()
26    }
27
28    fn generate_hash(&self, key: &str, i: usize) -> u64 {
29      let mut s = String::with_capacity(key.len() + 10);
30      s.push_str(key);
31      s.push_str(i.to_string().as_str());
32      xxh3_64(s.as_bytes())
33    }
34
35    pub fn add(&mut self, keys: &[String]) {
36        for key in keys {
37          for i in 0..self.replicas {
38            let hash = self.generate_hash(key, i);
39            self.sorted_keys.push(hash);
40            self.hash_map.insert(hash, key.clone());
41          }
42        }
43
44        // Sort the keys
45        self.sorted_keys.sort_unstable();
46    }
47
48    pub fn remove(&mut self, key: &str) {
49        let mut i = 0;
50        while i < self.sorted_keys.len() {
51          let hash_key = &self.hash_map[&self.sorted_keys[i]];
52          if hash_key == key {
53            let hash = self.sorted_keys.remove(i);
54            self.hash_map.remove(&hash);
55          } else {
56            i += 1;
57          }
58        }
59    }
60
61    /// An optimized version of get_n that returns a single node.
62    pub fn get(&self, key: &str) -> Option<&String> {
63        if self.is_empty() {
64          return None;
65        }
66
67        let hash = xxh3_64(key.as_bytes());
68        let idx = match self.sorted_keys.binary_search(&hash) {
69          // Exact match
70          Ok(idx) => idx,
71          // Not an exact match, get the first node that has a hash greater than the key hash
72          Err(idx) => idx % self.sorted_keys.len(),
73        };
74
75        Some(&self.hash_map[&self.sorted_keys[idx]])
76    }
77
78    /// Returns the <=N unique nodes that are closest to the key in order.
79    /// Useful for choosing replica groups for distributed systems (e.g. Raft group)
80    /// By default, the nodes are returned in ascending order of the hash ring.
81    pub fn get_n(&self, key: &str, n: usize, direction: Direction) -> Vec<String> {
82      let mut result: Vec<String> = Vec::new();
83      let mut _node_map:HashSet<String> = HashSet::new();
84
85      if self.is_empty() || n == 0 {
86          return result;
87      }
88
89      let hash = xxh3_64(key.as_bytes());
90      let idx = match self.sorted_keys.binary_search(&hash) {
91        // Exact match
92        Ok(idx) => idx,
93        // Not an exact match, get the first node that has a hash greater than the key hash
94        Err(idx) => idx % self.sorted_keys.len(),
95      };
96
97      let mut current_idx = idx;
98      while result.len() < n {
99        let node = &self.hash_map[&self.sorted_keys[current_idx]];
100        if _node_map.insert(node.clone()) {
101          result.push(node.clone());
102        }
103
104        match direction {
105            Direction::Forward => {
106                current_idx = (current_idx + 1) % self.sorted_keys.len();
107            }
108            Direction::Backward => {
109                if current_idx == 0 {
110                    current_idx = self.sorted_keys.len() - 1;
111                } else {
112                    current_idx -= 1;
113                }
114            }
115        }
116
117        // Break if we hit the start point again (haven't found enough nodes)
118        if current_idx == idx {
119            break;
120        }
121      }
122
123      result
124    }
125
126}
basic_hash_ring/lib.rs

basic_hash_ring/
lib.rs
