beam: add smooth weighted round-robin scheduler with tests

wizard-work/beam/smooth_wrr.py

@@ -0,0 +1,73 @@
+"""Smooth weighted round-robin scheduler (nginx algorithm).
+
+Each backend has a static weight. Selection is deterministic and smooth:
+no backend is starved; over a full cycle each backend is picked `weight`
+times, and consecutive picks of the same backend are spaced as evenly as
+possible.
+
+Correctness notes:
+- Zero/negative weights are rejected (use weight 0 to disable? we reject <=0
+  to keep invariants; callers filter disabled backends).
+- Empty pool raises.
+- Thread-safe via a lock; the algorithm mutates current_weight in place.
+- `peak`/`reset` helpers for introspection and reuse.
+"""
+from __future__ import annotations
+
+import threading
+from dataclasses import dataclass, field
+from typing import Generic, Iterable, List, Optional, TypeVar
+
+T = TypeVar("T")
+
+
+@dataclass
+class _Node(Generic[T]):
+    item: T
+    weight: int
+    current: int = field(default=0, init=False)
+
+
+class SmoothWeightedRoundRobin(Generic[T]):
+    def __init__(self, items: Iterable[tuple[T, int]]):
+        nodes: List[_Node[T]] = []
+        total = 0
+        for item, w in items:
+            if not isinstance(w, int) or isinstance(w, bool):
+                raise TypeError("weight must be int")
+            if w <= 0:
+                raise ValueError("weight must be > 0")
+            nodes.append(_Node(item=item, weight=w))
+            total += w
+        if not nodes:
+            raise ValueError("pool must be non-empty")
+        self._nodes = nodes
+        self._total = total
+        self._lock = threading.Lock()
+
+    @property
+    def total_weight(self) -> int:
+        return self._total
+
+    def size(self) -> int:
+        return len(self._nodes)
+
+    def next(self) -> T:
+        with self._lock:
+            best = None
+            for n in self._nodes:
+                n.current += n.weight
+                if best is None or n.current > best.current:
+                    best = n
+            assert best is not None
+            best.current -= self._total
+            return best.item
+
+    def peek_weights(self) -> list[int]:
+        with self._lock:
+            return [n.current for n in self._nodes]
+
+    def reset(self) -> None:
+        with self._lock:
+            for n in self._nodes:
+                n.current = 0

wizard-work/beam/test_smooth_wrr.py

@@ -0,0 +1,69 @@
+import threading
+from collections import Counter
+from smooth_wrr import SmoothWeightedRoundRobin as S
+
+def test_basic_counts():
+    s = S([("a", 3), ("b", 2), ("c", 1)])
+    seq = [s.next() for _ in range(12)]
+    c = Counter(seq)
+    assert c == {"a": 6, "b": 4, "c": 2}, c
+    # full cycle repeats
+    seq2 = [s.next() for _ in range(12)]
+    assert seq2 == seq
+
+def test_single():
+    s = S([("x", 5)])
+    assert [s.next() for _ in range(4)] == ["x"]*4
+
+def test_smoothness_no_adjacent_same_when_possible():
+    s = S([("a", 2), ("b", 1)])
+    seq = [s.next() for _ in range(6)]
+    # a appears twice as often; never three a in a row
+    for i in range(len(seq)-2):
+        assert not (seq[i]==seq[i+1]==seq[i+2]=="a"), seq
+
+def test_validation():
+    try:
+        S([])
+        raise AssertionError("empty should raise")
+    except ValueError: pass
+    for bad in [0, -1]:
+        try:
+            S([("a", bad)])
+            raise AssertionError("bad weight should raise")
+        except ValueError: pass
+    try:
+        S([("a", 1.5)])
+        raise AssertionError("float weight should raise")
+    except TypeError: pass
+    try:
+        S([("a", True)])
+        raise AssertionError("bool weight should raise")
+    except TypeError: pass
+
+def test_thread_safety_counts():
+    s = S([("a", 3), ("b", 2)])
+    N = 10000
+    out = []
+    lock = threading.Lock()
+    def worker(k):
+        local = [s.next() for _ in range(k)]
+        with lock:
+            out.extend(local)
+    ts = [threading.Thread(target=worker, args=(N//4,)) for _ in range(4)]
+    for t in ts: t.start()
+    for t in ts: t.join()
+    c = Counter(out)
+    assert c["a"] == 3*N//5, c
+    assert c["b"] == 2*N//5, c
+
+def test_reset():
+    s = S([("a",1),("b",1)])
+    s.next(); s.next()
+    s.reset()
+    assert s.peek_weights() == [0,0]
+
+if __name__ == "__main__":
+    test_basic_counts(); test_single(); test_smoothness_no_adjacent_same_when_possible()
+    test_validation(); test_thread_safety_counts(); test_reset()
+    print("ALL OK")