diff --git a/wizard-work/beam/smooth_wrr.py b/wizard-work/beam/smooth_wrr.py new file mode 100644 index 0000000..c3e13f1 --- /dev/null +++ b/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 diff --git a/wizard-work/beam/test_smooth_wrr.py b/wizard-work/beam/test_smooth_wrr.py new file mode 100644 index 0000000..389607f --- /dev/null +++ b/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")