diff --git a/src/lib.rs b/src/lib.rs index dc2c16c..3c6e5f4 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -84,6 +84,7 @@ extern crate siphasher; use siphasher::sip::SipHasher; use std::cmp::Ordering; +use std::collections::HashSet; use std::fmt::Debug; use std::hash::BuildHasher; use std::hash::Hash; @@ -236,27 +237,74 @@ impl HashRing { return None; } - let replicas = if replicas > self.ring.len() { - self.ring.len() - } else { - replicas + 1 + let replicas = std::cmp::min(replicas+1, self.ring.len()); + + let k = get_key(&self.hash_builder, key); + let n = match self.ring.binary_search_by(|node| node.key.cmp(&k)) { + Err(n) => n, + Ok(n) => n, }; + let mut replica_nodes = Vec::with_capacity(replicas); + let len = self.ring.len(); + + for i in 0..len { + let idx = (n + i) % len; + let node = &self.ring[idx]; + + replica_nodes.push(node.node.clone()); + + if replica_nodes.len() == replicas { + break; + } + } + + Some(replica_nodes) + } + + /// Get the node responsible for `key` (the primary) along with up to `replicas` + /// additional unique nodes after it, where uniqueness is determined by the + /// values produced by the `unique_key` closure. + /// Returns `None` when the ring is empty. If `replicas` is larger than the length + /// of the ring, this function will shrink to just contain each unique element of + /// the ring, up to `replicas + 1` nodes (primary plus replicas). If there are fewer + /// distinct `unique_key` values than requested, the returned vector will contain + /// fewer than `replicas + 1` nodes. + pub fn get_with_replicas_unique_by_key(&self, key: &U, replicas: usize, unique_key: F) -> Option> + where + T: Clone + Debug, + F: Fn(&T) -> K, + K: Hash + Eq, + { + if self.ring.is_empty() { + return None; + } + + let replicas = std::cmp::min(replicas + 1, self.ring.len()); let k = get_key(&self.hash_builder, key); let n = match self.ring.binary_search_by(|node| node.key.cmp(&k)) { Err(n) => n, Ok(n) => n, }; - let mut nodes = self.ring.clone(); - nodes.rotate_left(n); + let mut seen = HashSet::with_capacity(replicas); + let mut replica_nodes = Vec::with_capacity(replicas); + + let len = self.ring.len(); + for i in 0..len { + let idx = (n + i) % len; + let node = &self.ring[idx]; - let replica_nodes = nodes - .iter() - .cycle() - .take(replicas) - .map(|node| node.node.clone()) - .collect(); + let unique = unique_key(&node.node); + + if seen.insert(unique) { + replica_nodes.push(node.node.clone()); + + if replica_nodes.len() == replicas { + break; + } + } + } Some(replica_nodes) } @@ -445,6 +493,12 @@ mod tests { ring.get_with_replicas(&"foo", 4).unwrap(), vec![vnode5, vnode4, vnode3, vnode1, vnode2] ); + + assert_eq!( + ring.get_with_replicas(&"foo", 5).unwrap().len(), + 6, + "return entire ring when primary + replicas == ring.len()" + ); } #[test] @@ -475,6 +529,46 @@ mod tests { ); } + + #[test] + fn get_unique_replicas_skips_duplicates() { + let mut ring: HashRing = HashRing::new(); + + assert_eq!(ring.get_with_replicas_unique_by_key(&"foo", 1, |node| node.addr), None); + + let vnode1 = VNode::new("127.0.0.1", 1024, 1); + let vnode2 = VNode::new("127.0.0.1", 1024, 2); + let vnode3 = VNode::new("127.0.0.2", 1024, 3); + let vnode4 = VNode::new("127.0.0.2", 1024, 4); + let vnode5 = VNode::new("127.0.0.2", 1024, 5); + let vnode6 = VNode::new("127.0.0.3", 1024, 6); + + ring.add(vnode1); + ring.add(vnode2); + ring.add(vnode3); + ring.add(vnode4); + ring.add(vnode5); + ring.add(vnode6); + + assert_eq!( + ring.get_with_replicas_unique_by_key(&"bar", 20, |vnode| vnode.addr).unwrap().len(), + 3, + "replicas+1 > HashRing::len() causes the count to shrink to count of unique values" + ); + + assert_eq!( + ring.get_with_replicas_unique_by_key(&"bar", 4, |vnode| vnode.addr).unwrap().len(), + 3, + "count_of_unique_elements < replicas+1 < HashRing::len() causes the count to shrink to count of unique values" + ); + + assert_eq!( + ring.get_with_replicas_unique_by_key(&"bar", 1, |vnode| vnode.addr).unwrap().len(), + 2, + "replicas < count_of_unique_elements causes the count to match replicas + 1 (including primary)" + ); + } + #[test] fn into_iter() { let mut ring: HashRing = HashRing::new();