beam: rendezvous (HRW) hashing module with tests

wizard-work/beam/rendezvous_hash.py

@@ -0,0 +1,49 @@
+"""Rendezvous (highest random weight) hashing."""
+from __future__ import annotations
+
+import hashlib
+from typing import Any, Iterable
+
+
+def _weight(key: Any, node: Any) -> int:
+    h = hashlib.sha256(repr(key).encode() + b"\x00" + repr(node).encode()).digest()
+    return int.from_bytes(h[:8], "big")
+
+
+def _nodes(nodes: Iterable[Any]) -> list[Any]:
+    seen = []
+    seen_set = set()
+    for n in nodes:
+        if n not in seen_set:
+            seen_set.add(n)
+            seen.append(n)
+    return seen
+
+
+def rendezvous(key: Any, nodes: Iterable[Any]) -> Any | None:
+    ns = _nodes(nodes)
+    if not ns:
+        return None
+    best = None
+    best_w = -1
+    for n in ns:
+        w = _weight(key, n)
+        if w > best_w:
+            best_w = w
+            best = n
+    return best
+
+
+def rendezvous_topk(key: Any, nodes: Iterable[Any], k: int) -> list[Any]:
+    ns = _nodes(nodes)
+    if k <= 0 or not ns:
+        return []
+    indexed = sorted(range(len(ns)), key=lambda i: (-_weight(key, ns[i]), i))
+    return [ns[i] for i in indexed[:k]]
+
+
+def weights(key: Any, nodes: Iterable[Any]) -> list[tuple[Any, int]]:
+    ns = _nodes(nodes)
+    out = [(n, _weight(key, n)) for n in ns]
+    out.sort(key=lambda t: (-t[1], ns.index(t[0])))
+    return out

wizard-work/beam/test_rendezvous_hash.py

@@ -0,0 +1,107 @@
+"""Rendezvous (highest random weight) hashing.
+
+Given a set of nodes and a key, pick the node with the highest weight
+hash(key, node). Deterministic, minimal disruption on node add/remove,
+no central coordinator.
+
+Public API:
+    rendezvous(key, nodes) -> node or None
+    rendezvous_topk(key, nodes, k) -> list of nodes (descending weight)
+    weights(key, nodes) -> list of (node, weight) sorted desc
+"""
+from __future__ import annotations
+
+import hashlib
+from typing import Any, Iterable
+
+
+def _weight(key: Any, node: Any) -> int:
+    h = hashlib.sha256(repr(key).encode() + b"\x00" + repr(node).encode()).digest()
+    return int.from_bytes(h[:8], "big")
+
+
+def _nodes(nodes: Iterable[Any]) -> list[Any]:
+    seen = []
+    seen_set = set()
+    for n in nodes:
+        if n not in seen_set:
+            seen_set.add(n)
+            seen.append(n)
+    return seen
+
+
+def rendezvous(key: Any, nodes: Iterable[Any]) -> Any | None:
+    ns = _nodes(nodes)
+    if not ns:
+        return None
+    best = None
+    best_w = -1
+    for n in ns:
+        w = _weight(key, n)
+        if w > best_w:
+            best_w = w
+            best = n
+    return best
+
+
+def rendezvous_topk(key: Any, nodes: Iterable[Any], k: int) -> list[Any]:
+    ns = _nodes(nodes)
+    if k <= 0 or not ns:
+        return []
+    indexed = sorted(range(len(ns)), key=lambda i: (-_weight(key, ns[i]), i))
+    return [ns[i] for i in indexed[:k]]
+
+
+def weights(key: Any, nodes: Iterable[Any]) -> list[tuple[Any, int]]:
+    ns = _nodes(nodes)
+    out = [(n, _weight(key, n)) for n in ns]
+    out.sort(key=lambda t: (-t[1], ns.index(t[0])))
+    return out
+"""Tests for rendezvous hashing."""
+import random
+import sys, os
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from rendezvous_hash import rendezvous, rendezvous_topk, weights
+
+def check(name, cond):
+    print(("PASS" if cond else "FAIL") + " " + name)
+    return cond
+
+ok = True
+nodes = ["n1","n2","n3","n4"]
+
+ok &= check("det", rendezvous("key", nodes) == rendezvous("key", nodes))
+r = rendezvous("key", nodes)
+ok &= check("inset", r in nodes)
+ok &= check("empty", rendezvous("key", []) is None)
+ok &= check("single", rendezvous("key", ["only"]) == "only")
+
+tk = rendezvous_topk("key", nodes, 2)
+ok &= check("topk len", len(tk) == 2)
+ok &= check("topk distinct", len(set(tk)) == 2)
+ok &= check("topk subset", set(tk).issubset(set(nodes)))
+ok &= check("topk 0", rendezvous_topk("key", nodes, 0) == [])
+
+ws = weights("key", nodes)
+ok &= check("weights sorted", all(ws[i][1] >= ws[i+1][1] for i in range(len(ws)-1)))
+ok &= check("weights top matches", ws[0][0] == rendezvous("key", nodes))
+
+random.seed(1)
+disrupt = 0
+trials = 2000
+for _ in range(trials):
+    ks = f"k{_}"
+    full = rendezvous(ks, nodes)
+    without = rendezvous(ks, [n for n in nodes if n != nodes[0]])
+    if full != nodes[0] and full != without:
+        disrupt += 1
+ok &= check("disruption bounded", disrupt <= trials//4 + 50)
+
+def gen():
+    for n in nodes:
+        yield n
+ok &= check("generator", rendezvous("key", gen()) in nodes)
+ok &= check("dedup", rendezvous("key", ["a","a","b"]) in ("a","b"))
+
+print("ALL" if ok else "SOME FAILED")
+sys.exit(0 if ok else 1)