Skip to main content

base/
hashring.rs

1// hashring -- consistent hashing: a ring of tokens with virtual nodes per member; a key belongs
2// to the member of the first token at or after its hash (going round at the end), and its
3// preference list is the next n distinct members from there (sharded stores, load balancers
4// with sticky keys).  Adding or removing a member moves only the keys of its own tokens.
5//
6//   hash64(bytes) -> u64                    FNV-1a (64 bit) with a final mix; the same on every
7//                                           target
8//   Ring::new() / add(member, key, vnodes) / remove(&member) / members() / tokens()
9//   owner(h) -> Option<&M>                  the member of the first token at or after h
10//   preference(h, n) -> Vec<M>              the first n distinct members from h on
11//   preference_by(h, n, ok)                 the same, skipping members ok() refuses
12//   shares() -> Vec<(M, u32)>               each member's part of the ring in parts per million
13//
14// The tokens of a member come from its key and the vnode number ("10.0.4.11#3"), so every
15// device that knows the members builds the same ring.
16
17use alloc::format;
18use alloc::vec::Vec;
19
20/// FNV-1a over the bytes, then the finaliser of SplitMix64 (spreads close keys over the ring)
21pub fn hash64(b: &[u8]) -> u64 {
22    let mut h: u64 = 0xCBF2_9CE4_8422_2325;
23    for &x in b {
24        h ^= x as u64;
25        h = h.wrapping_mul(0x0000_0100_0000_01B3);
26    }
27    let mut z = h;
28    z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
29    z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
30    z ^ (z >> 31)
31}
32
33/// a ring of tokens; M names a member (an address, a backend)
34#[derive(Clone, Debug, PartialEq, Eq)]
35pub struct Ring<M: Ord + Clone> {
36    /// (token, member), sorted
37    tokens: Vec<(u64, M)>,
38    members: Vec<M>,
39}
40
41impl<M: Ord + Clone> Default for Ring<M> {
42    fn default() -> Self {
43        Ring::new()
44    }
45}
46
47impl<M: Ord + Clone> Ring<M> {
48    pub fn new() -> Ring<M> {
49        Ring { tokens: Vec::new(), members: Vec::new() }
50    }
51    /// adds a member with `vnodes` tokens made from `key` (a member that is there already gets
52    /// its tokens again)
53    pub fn add(&mut self, m: M, key: &str, vnodes: u32) {
54        self.remove(&m);
55        for i in 0..vnodes.max(1) {
56            self.tokens.push((hash64(format!("{}#{}", key, i).as_bytes()), m.clone()));
57        }
58        self.tokens.sort();
59        self.members.push(m);
60        self.members.sort();
61    }
62    /// removes a member and its tokens
63    pub fn remove(&mut self, m: &M) -> bool {
64        let before = self.members.len();
65        self.members.retain(|x| x != m);
66        self.tokens.retain(|t| &t.1 != m);
67        self.members.len() != before
68    }
69    /// the members, sorted
70    pub fn members(&self) -> &[M] {
71        &self.members
72    }
73    /// the tokens with their members, sorted by token
74    pub fn tokens(&self) -> &[(u64, M)] {
75        &self.tokens
76    }
77    pub fn is_empty(&self) -> bool {
78        self.tokens.is_empty()
79    }
80    /// the index of the first token at or after h (0 after the last)
81    fn start(&self, h: u64) -> usize {
82        let i = self.tokens.partition_point(|t| t.0 < h);
83        if i == self.tokens.len() { 0 } else { i }
84    }
85    /// the member a hash belongs to
86    pub fn owner(&self, h: u64) -> Option<&M> {
87        if self.tokens.is_empty() {
88            return None;
89        }
90        Some(&self.tokens[self.start(h)].1)
91    }
92    /// the first `n` distinct members from h on, in ring order
93    pub fn preference(&self, h: u64, n: usize) -> Vec<M> {
94        self.preference_by(h, n, |_| true)
95    }
96    /// the first `n` distinct members from h on that `ok` accepts
97    pub fn preference_by(&self, h: u64, n: usize, ok: impl Fn(&M) -> bool) -> Vec<M> {
98        let mut out: Vec<M> = Vec::new();
99        if self.tokens.is_empty() || n == 0 {
100            return out;
101        }
102        let s = self.start(h);
103        for k in 0..self.tokens.len() {
104            let m = &self.tokens[(s + k) % self.tokens.len()].1;
105            if !out.contains(m) && ok(m) {
106                out.push(m.clone());
107                if out.len() == n {
108                    break;
109                }
110            }
111        }
112        out
113    }
114    /// the part of the hash space each member owns, in parts per million (sorted by member)
115    pub fn shares(&self) -> Vec<(M, u32)> {
116        let n = self.tokens.len();
117        let mut out: Vec<(M, u32)> = self.members.iter().map(|m| (m.clone(), 0)).collect();
118        if n == 0 {
119            return out;
120        }
121        let mut arcs: Vec<u128> = alloc::vec![0; self.members.len()];
122        for i in 0..n {
123            // the arc that ends at token i: from the token before it (exclusive)
124            let prev = if i == 0 { self.tokens[n - 1].0 } else { self.tokens[i - 1].0 };
125            let cur = self.tokens[i].0;
126            let len = if n == 1 { 1u128 << 64 } else { cur.wrapping_sub(prev) as u128 };
127            if let Ok(k) = self.members.binary_search(&self.tokens[i].1) {
128                arcs[k] += len;
129            }
130        }
131        for (k, a) in arcs.iter().enumerate() {
132            out[k].1 = ((a * 1_000_000) >> 64) as u32;
133        }
134        out
135    }
136}