diff --git a/wizard-work/beam/singleflight.py b/wizard-work/beam/singleflight.py new file mode 100644 index 0000000..42cbe59 --- /dev/null +++ b/wizard-work/beam/singleflight.py @@ -0,0 +1,84 @@ +"""Single-flight (request coalescing): dedupe concurrent calls for one key. + +Only one caller (the leader) actually runs ``fn`` for a given key; every other +caller that arrives while it is in flight waits and receives the same result or +exception. Once the in-flight call completes, the key is freed so the next call +runs ``fn`` again. + +Correctness notes: +- A global lock protects the in-flight registry; each in-flight call has its + own condition variable for waiters. +- The leader runs ``fn`` outside any lock so other keys (and registry + operations) are not blocked. +- Result or exception is broadcast to all waiters; waiters re-raise the same + exception object the leader raised. +- Registry entry is removed in a ``finally`` so a crashing leader cannot leak + a stuck key. +- ``fn`` is called exactly once per in-flight window. +""" +from __future__ import annotations + +import threading +from typing import Any, Callable, Dict, Optional + + +class _InFlight: + __slots__ = ("key", "fn", "cond", "result", "exc", "done", "waiters") + + def __init__(self, key: str, fn: Callable[[], Any]) -> None: + self.key = key + self.fn = fn + self.cond = threading.Condition(threading.Lock()) + self.result: Any = None + self.exc: Optional[BaseException] = None + self.done: bool = False + self.waiters: int = 0 + + +class SingleFlight: + def __init__(self) -> None: + self._mu = threading.Lock() + self._inflight: Dict[Any, _InFlight] = {} + + @property + def inflight_count(self) -> int: + with self._mu: + return len(self._inflight) + + def do(self, key: Any, fn: Callable[[], Any]) -> Any: + # Register (or join) an in-flight call. + with self._mu: + c = self._inflight.get(key) + leader = c is None + if leader: + c = _InFlight(key, fn) + self._inflight[key] = c + c.waiters += 1 + + if leader: + try: + result = fn() + except BaseException as e: # noqa: BLE001 - must propagate anything + with c.cond: + c.exc = e + c.done = True + c.cond.notify_all() + raise + else: + with c.cond: + c.result = result + c.done = True + c.cond.notify_all() + return result + finally: + with self._mu: + # Only remove if still ours (defensive against re-registration). + if self._inflight.get(key) is c: + del self._inflight[key] + # Follower path. + with c.cond: + while not c.done: + c.cond.wait() + if c.exc is not None: + raise c.exc + return c.result diff --git a/wizard-work/beam/test_singleflight.py b/wizard-work/beam/test_singleflight.py new file mode 100644 index 0000000..cd4734b --- /dev/null +++ b/wizard-work/beam/test_singleflight.py @@ -0,0 +1,163 @@ +import threading +import time +import unittest + +from singleflight import SingleFlight + + +class TestSingleFlight(unittest.TestCase): + def test_single_call_runs_fn(self): + sf = SingleFlight() + calls = {"n": 0} + + def fn(): + calls["n"] += 1 + return 42 + + self.assertEqual(sf.do("k", fn), 42) + self.assertEqual(calls["n"], 1) + + def test_concurrent_coalesces(self): + sf = SingleFlight() + calls = {"n": 0} + lock = threading.Lock() + barrier = threading.Barrier(8) + + def fn(): + with lock: + calls["n"] += 1 + time.sleep(0.2) + return "result" + + results = [None] * 8 + errors = [] + + def worker(i): + barrier.wait() + try: + results[i] = sf.do("k", fn) + except BaseException as e: # noqa + errors.append(e) + + ts = [threading.Thread(target=worker, args=(i,)) for i in range(8)] + for t in ts: + t.start() + for t in ts: + t.join() + self.assertFalse(errors) + self.assertTrue(all(r == "result" for r in results)) + self.assertEqual(calls["n"], 1) + + def test_exception_propagates_to_all(self): + sf = SingleFlight() + calls = {"n": 0} + + def fn(): + calls["n"] += 1 + time.sleep(0.2) + raise ValueError("boom") + + results = [] + errors = [] + barrier = threading.Barrier(5) + + def worker(): + barrier.wait() + try: + sf.do("k", fn) + except Exception as e: # noqa + errors.append(e) + else: + results.append("ok") + + ts = [threading.Thread(target=worker) for _ in range(5)] + for t in ts: + t.start() + for t in ts: + t.join() + self.assertEqual(len(errors), 5) + self.assertTrue(all(isinstance(e, ValueError) and str(e) == "boom" for e in errors)) + self.assertEqual(calls["n"], 1) + # Same exception object shared. + self.assertTrue(len({id(e) for e in errors}) == 1) + + def test_subsequent_call_runs_again(self): + sf = SingleFlight() + calls = {"n": 0} + + def fn(): + calls["n"] += 1 + return calls["n"] + + self.assertEqual(sf.do("k", fn), 1) + self.assertEqual(sf.do("k", fn), 2) + self.assertEqual(calls["n"], 2) + + def test_different_keys_independent(self): + sf = SingleFlight() + calls = {"a": 0, "b": 0} + + def make(key): + def fn(): + calls[key] += 1 + time.sleep(0.1) + return key + return fn + + r = {} + def wa(): + r["a"] = sf.do("a", make("a")) + def wb(): + r["b"] = sf.do("b", make("b")) + ta = threading.Thread(target=wa); tb = threading.Thread(target=wb) + ta.start(); tb.start(); ta.join(); tb.join() + self.assertEqual(r["a"], "a") + self.assertEqual(r["b"], "b") + self.assertEqual(calls["a"], 1) + self.assertEqual(calls["b"], 1) + + def test_inflight_count(self): + sf = SingleFlight() + started = threading.Event() + release = threading.Event() + + def fn(): + started.set() + release.wait() + return 1 + + def leader(): + sf.do("k", fn) + + t = threading.Thread(target=leader) + t.start() + started.wait() + self.assertEqual(sf.inflight_count, 1) + release.set() + t.join() + self.assertEqual(sf.inflight_count, 0) + + def test_leader_crash_does_not_leak_key(self): + sf = SingleFlight() + + def fn(): + raise RuntimeError("crash") + + try: + sf.do("k", fn) + except RuntimeError: + pass + # Key freed: next call runs fn again. + calls = {"n": 0} + + def fn2(): + calls["n"] += 1 + return "ok" + + self.assertEqual(sf.do("k", fn2), "ok") + self.assertEqual(calls["n"], 1) + self.assertEqual(sf.inflight_count, 0) + + +if __name__ == "__main__": + unittest.main()