feat: with replicas unique by key + some fixes & performance improvements

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<T: Hash, S: BuildHasher> HashRing<T, S> {
             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<U: Hash, F, K>(&self, key: &U, replicas: usize, unique_key: F) -> Option<Vec<T>>
+    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<VNode> = 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<VNode> = HashRing::new();