beam: consistent hashing ring with virtual nodes + tests

wizard-work/beam/consistent_hash.py

@@ -0,0 +1,89 @@
+"""Consistent hashing ring with virtual nodes.
+
+A minimal, dependency-free implementation suitable for distributing keys
+across a dynamic set of nodes with minimal rebalancing.
+"""
+import bisect
+import hashlib
+from typing import Any, Iterable, List, Optional
+
+
+class ConsistentHashRing:
+    def __init__(self, num_vnodes: int = 100):
+        if not isinstance(num_vnodes, int) or num_vnodes < 1:
+            raise ValueError("num_vnodes must be a positive int")
+        self.num_vnodes = num_vnodes
+        self._ring: dict[int, Any] = {}          # position -> node
+        self._sorted: List[int] = []             # sorted positions
+        self._nodes: set = set()
+
+    @staticmethod
+    def _hash(key: Any) -> int:
+        # 128-bit digest collapsed to int; deterministic across processes.
+        return int(hashlib.md5(str(key).encode("utf-8")).hexdigest(), 16)
+
+    def add_node(self, node: Any) -> None:
+        if node in self._nodes:
+            return
+        self._nodes.add(node)
+        for i in range(self.num_vnodes):
+            pos = self._hash(f"{node}#{i}")
+            # handle (extremely unlikely) collision by nudging
+            while pos in self._ring:
+                pos = (pos + 1) & ((1 << 128) - 1)
+            self._ring[pos] = node
+            bisect.insort(self._sorted, pos)
+
+    def remove_node(self, node: Any) -> None:
+        if node not in self._nodes:
+            return
+        self._nodes.discard(node)
+        # rebuild positions belonging to node
+        to_delete = [p for p, n in self._ring.items() if n == node]
+        for p in to_delete:
+            del self._ring[p]
+            idx = bisect.bisect_left(self._sorted, p)
+            if idx < len(self._sorted) and self._sorted[idx] == p:
+                del self._sorted[idx]
+
+    @property
+    def nodes(self) -> List[Any]:
+        return sorted(self._nodes, key=lambda x: str(x))
+
+    def get_node(self, key: Any) -> Optional[Any]:
+        if not self._sorted:
+            return None
+        pos = self._hash(key)
+        idx = bisect.bisect_left(self._sorted, pos)
+        if idx == len(self._sorted):
+            idx = 0
+        return self._ring[self._sorted[idx]]
+
+    def get_nodes(self, key: Any, n: int) -> List[Any]:
+        if n < 1:
+            raise ValueError("n must be >= 1")
+        if not self._sorted:
+            return []
+        result: List[Any] = []
+        seen: set = set()
+        pos = self._hash(key)
+        idx = bisect.bisect_left(self._sorted, pos)
+        if idx == len(self._sorted):
+            idx = 0
+        start = idx
+        # walk ring collecting distinct nodes
+        i = 0
+        total = len(self._sorted)
+        while len(result) < n and i < total:
+            node = self._ring[self._sorted[(start + i) % total]]
+            if node not in seen:
+                seen.add(node)
+                result.append(node)
+            i += 1
+        return result
+
+    def __len__(self) -> int:
+        return len(self._nodes)
+
+    def __contains__(self, node: Any) -> bool:
+        return node in self._nodes

wizard-work/beam/test_consistent_hash.py

@@ -0,0 +1,94 @@
+import unittest
+from consistent_hash import ConsistentHashRing
+
+
+class TestConsistentHash(unittest.TestCase):
+    def test_empty_ring_returns_none(self):
+        r = ConsistentHashRing()
+        self.assertIsNone(r.get_node("k"))
+
+    def test_single_node_always(self):
+        r = ConsistentHashRing()
+        r.add_node("A")
+        for k in ["x", "y", "z", 1, 2]:
+            self.assertEqual(r.get_node(k), "A")
+
+    def test_add_remove(self):
+        r = ConsistentHashRing()
+        r.add_node("A"); r.add_node("B")
+        self.assertEqual(len(r), 2)
+        r.remove_node("A")
+        self.assertEqual(len(r), 1)
+        self.assertNotIn("A", r)
+        for k in ["k1", "k2", "k3"]:
+            self.assertEqual(r.get_node(k), "B")
+
+    def test_duplicate_add_idempotent(self):
+        r = ConsistentHashRing(num_vnodes=5)
+        r.add_node("A"); r.add_node("A")
+        self.assertEqual(len(r), 1)
+        self.assertEqual(len(r._ring), 5)
+
+    def test_remove_nonexistent_noop(self):
+        r = ConsistentHashRing()
+        r.add_node("A")
+        r.remove_node("Z")
+        self.assertEqual(len(r), 1)
+
+    def test_determinism(self):
+        r1 = ConsistentHashRing(num_vnodes=10)
+        r2 = ConsistentHashRing(num_vnodes=10)
+        for n in ["n1", "n2", "n3"]:
+            r1.add_node(n); r2.add_node(n)
+        for k in ["a", "b", "c", "d"]:
+            self.assertEqual(r1.get_node(k), r2.get_node(k))
+
+    def test_minimal_rebalancing(self):
+        # adding a node should move only a fraction of keys
+        r = ConsistentHashRing(num_vnodes=50)
+        for n in ["n1", "n2"]:
+            r.add_node(n)
+        before = {k: r.get_node(k) for k in [f"k{i}" for i in range(200)]}
+        r.add_node("n3")
+        moved = sum(1 for k in before if r.get_node(k) != before[k])
+        # expect roughly 1/3 moved, definitely not all
+        self.assertLess(moved, 200)
+        self.assertGreater(moved, 0)
+
+    def test_get_nodes_replication(self):
+        r = ConsistentHashRing(num_vnodes=20)
+        for n in ["a", "b", "c"]:
+            r.add_node(n)
+        ns = r.get_nodes("key", 2)
+        self.assertEqual(len(ns), 2)
+        self.assertEqual(len(set(ns)), 2)
+        # n larger than node count
+        ns2 = r.get_nodes("key", 10)
+        self.assertEqual(len(ns2), 3)
+        self.assertEqual(len(set(ns2)), 3)
+
+    def test_get_nodes_empty(self):
+        r = ConsistentHashRing()
+        self.assertEqual(r.get_nodes("k", 3), [])
+
+    def test_invalid_vnodes(self):
+        with self.assertRaises(ValueError):
+            ConsistentHashRing(0)
+        with self.assertRaises(ValueError):
+            ConsistentHashRing(-5)
+
+    def test_invalid_n(self):
+        r = ConsistentHashRing()
+        r.add_node("A")
+        with self.assertRaises(ValueError):
+            r.get_nodes("k", 0)
+
+    def test_nodes_sorted(self):
+        r = ConsistentHashRing()
+        for n in ["z", "a", "m"]:
+            r.add_node(n)
+        self.assertEqual(r.nodes, ["a", "m", "z"])
+
+
+if __name__ == "__main__":
+    unittest.main()