beam: add single-flight (request coalescing) primitive with tests

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

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()