From a17525d1d42af77ecf259ffae399c31274853de2 Mon Sep 17 00:00:00 2001 From: beam Date: Fri, 23 Oct 2026 13:11:38 +0000 Subject: [PATCH] beam: consistent hashing ring with virtual nodes + tests --- wizard-work/beam/consistent_hash.py | 89 ++++++++++++++++++++++ wizard-work/beam/test_consistent_hash.py | 94 ++++++++++++++++++++++++ 2 files changed, 183 insertions(+) create mode 100644 wizard-work/beam/consistent_hash.py create mode 100644 wizard-work/beam/test_consistent_hash.py diff --git a/wizard-work/beam/consistent_hash.py b/wizard-work/beam/consistent_hash.py new file mode 100644 index 0000000..31fca3a --- /dev/null +++ b/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 diff --git a/wizard-work/beam/test_consistent_hash.py b/wizard-work/beam/test_consistent_hash.py new file mode 100644 index 0000000..36485e9 --- /dev/null +++ b/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()