Skip to main content

atomr_core/routing/
consistent_hash.rs

1//! Consistent-hash router.
2
3use std::collections::BTreeMap;
4use std::hash::{Hash, Hasher};
5
6use crate::actor::ActorRef;
7
8pub struct ConsistentHashRouter<M: Send + Clone + 'static> {
9    ring: BTreeMap<u64, usize>,
10    routees: Vec<ActorRef<M>>,
11    vnodes: u32,
12}
13
14impl<M: Send + Clone + 'static> ConsistentHashRouter<M> {
15    pub fn new(routees: Vec<ActorRef<M>>, virtual_nodes_factor: u32) -> Self {
16        let mut ring = BTreeMap::new();
17        for (i, r) in routees.iter().enumerate() {
18            for v in 0..virtual_nodes_factor {
19                let mut h = std::collections::hash_map::DefaultHasher::new();
20                r.path().to_string().hash(&mut h);
21                v.hash(&mut h);
22                ring.insert(h.finish(), i);
23            }
24        }
25        Self { ring, routees, vnodes: virtual_nodes_factor }
26    }
27
28    pub fn route_by_key<K: Hash>(&self, key: K, msg: M) {
29        if self.routees.is_empty() {
30            return;
31        }
32        let mut h = std::collections::hash_map::DefaultHasher::new();
33        key.hash(&mut h);
34        let k = h.finish();
35        let idx = self.ring.range(k..).next().or_else(|| self.ring.iter().next()).map(|(_, i)| *i);
36        if let Some(i) = idx {
37            self.routees[i].tell(msg);
38        }
39    }
40
41    pub fn virtual_nodes(&self) -> u32 {
42        self.vnodes
43    }
44}