beam: thread-safe bulkhead (bounded concurrency) with tests

wizard-work/beam/CHECK

@@ -1,6 +1,5 @@
 #!/bin/bash
-# CHECK: verify sha256_manifest.py works correctly
+# CHECK: verify all beam reliability primitives
 cd /root/troupe-standup/wizard-work/beam
-python3 -m unittest test_sha256_manifest -v
+python3 -m unittest discover -p 'test_*.py' -v
 exit $?
-" path="/root/troupe-standup/wizard-work/beam/CHECK

wizard-work/beam/bulkhead.py

@@ -0,0 +1,136 @@
+"""Thread-safe bulkhead: bounded concurrency with an optional wait queue.
+
+A bulkhead isolates a resource by capping how many calls may run at once.
+Callers beyond the cap either queue (up to max_queue) and wait for a slot,
+or are rejected immediately. This prevents a slow downstream from exhausting
+all threads/connections of the caller.
+
+Design notes:
+- State guarded by a single lock + condition variable.
+- active: number of in-flight calls.
+- queued: number of callers currently waiting for a slot.
+- max_concurrent >= 1, max_queue >= 0.
+- acquire(timeout=None) blocks up to timeout seconds; returns a release handle
+  (call release() or use as context manager). Raises BulkheadFull if the queue
+  is full or the wait times out; raises TimeoutError on timeout.
+- try_acquire() is non-blocking: returns handle or None.
+- Releasing a slot wakes one waiter (FIFO-ish, OS scheduling).
+"""
+from __future__ import annotations
+
+import threading
+import time
+from contextlib import contextmanager
+from typing import Optional
+
+
+class BulkheadFull(Exception):
+    """Raised when the bulkhead cannot accept another caller."""
+
+
+class _Handle:
+    __slots__ = ("_bh", "_released")
+
+    def __init__(self, bh: "Bulkhead"):
+        self._bh = bh
+        self._released = False
+
+    def release(self) -> None:
+        if self._released:
+            return
+        self._released = True
+        self._bh._release()
+
+    def __enter__(self) -> "_Handle":
+        return self
+
+    def __exit__(self, *exc) -> bool:
+        self.release()
+        return False
+
+
+class Bulkhead:
+    def __init__(self, max_concurrent: int, max_queue: int = 0):
+        if not isinstance(max_concurrent, int) or max_concurrent < 1:
+            raise ValueError("max_concurrent must be an int >= 1")
+        if not isinstance(max_queue, int) or max_queue < 0:
+            raise ValueError("max_queue must be an int >= 0")
+        self._max_concurrent = max_concurrent
+        self._max_queue = max_queue
+        self._active = 0
+        self._queued = 0
+        self._lock = threading.Lock()
+        self._cond = threading.Condition(self._lock)
+
+    @property
+    def max_concurrent(self) -> int:
+        return self._max_concurrent
+
+    @property
+    def max_queue(self) -> int:
+        return self._max_queue
+
+    def active(self) -> int:
+        with self._lock:
+            return self._active
+
+    def queued(self) -> int:
+        with self._lock:
+            return self._queued
+
+    def try_acquire(self) -> Optional[_Handle]:
+        with self._lock:
+            if self._active < self._max_concurrent:
+                self._active += 1
+                return _Handle(self)
+            return None
+
+    def acquire(self, timeout: Optional[float] = None) -> _Handle:
+        if timeout is not None and timeout < 0:
+            raise ValueError("timeout must be >= 0")
+        deadline = None if timeout is None else (time.monotonic() + timeout)
+        with self._lock:
+            # fast path
+            if self._active < self._max_concurrent:
+                self._active += 1
+                return _Handle(self)
+            # need to queue
+            if self._queued >= self._max_queue:
+                raise BulkheadFull("bulkhead queue full")
+            self._queued += 1
+            try:
+                while self._active >= self._max_concurrent:
+                    if deadline is None:
+                        self._cond.wait()
+                    else:
+                        remaining = deadline - time.monotonic()
+                        if remaining <= 0:
+                            raise TimeoutError("bulkhead acquire timed out")
+                        self._cond.wait(timeout=remaining)
+                # slot available
+                self._queued -= 1
+                self._active += 1
+                return _Handle(self)
+            except BaseException:
+                # ensure queued counter consistency on any exit
+                if self._queued > 0:
+                    self._queued -= 1
+                raise
+
+    def _release(self) -> None:
+        with self._lock:
+            if self._active > 0:
+                self._active -= 1
+            self._cond.notify(1)
+
+    @contextmanager
+    def slot(self, timeout: Optional[float] = None):
+        h = self.acquire(timeout=timeout)
+        try:
+            yield h
+        finally:
+            h.release()
+
+    def run(self, fn, *args, timeout: Optional[float] = None, **kwargs):
+        with self.slot(timeout=timeout):
+            return fn(*args, **kwargs)

wizard-work/beam/test_bulkhead.py

@@ -0,0 +1,126 @@
+import threading
+import time
+import unittest
+
+from bulkhead import Bulkhead, BulkheadFull
+
+
+class TestBulkheadBasic(unittest.TestCase):
+    def test_invalid_args(self):
+        with self.assertRaises(ValueError):
+            Bulkhead(0)
+        with self.assertRaises(ValueError):
+            Bulkhead(-1)
+        with self.assertRaises(ValueError):
+            Bulkhead(1, -1)
+        with self.assertRaises(ValueError):
+            Bulkhead(1, max_queue=-1)
+
+    def test_try_acquire_fast_path(self):
+        bh = Bulkhead(2)
+        h1 = bh.try_acquire()
+        h2 = bh.try_acquire()
+        self.assertIsNotNone(h1)
+        self.assertIsNotNone(h2)
+        self.assertIsNone(bh.try_acquire())
+        self.assertEqual(bh.active(), 2)
+        h1.release()
+        h2.release()
+        self.assertEqual(bh.active(), 0)
+
+    def test_acquire_release_context(self):
+        bh = Bulkhead(1)
+        with bh.slot():
+            self.assertEqual(bh.active(), 1)
+            self.assertIsNone(bh.try_acquire())
+        self.assertEqual(bh.active(), 0)
+
+    def test_run(self):
+        bh = Bulkhead(1)
+        self.assertEqual(bh.run(lambda x: x + 1, 4), 5)
+
+    def test_queue_full_raises(self):
+        bh = Bulkhead(1, max_queue=0)
+        h = bh.acquire()
+        with self.assertRaises(BulkheadFull):
+            bh.acquire()
+        h.release()
+
+    def test_queued_counter(self):
+        bh = Bulkhead(1, max_queue=2)
+        h = bh.acquire()
+        # two waiters queue
+        results = []
+        def worker(i):
+            try:
+                hh = bh.acquire(timeout=5)
+                results.append(("ok", i))
+                time.sleep(0.2)
+                hh.release()
+            except Exception as e:
+                results.append(("err", i, type(e).__name__))
+        ts = [threading.Thread(target=worker, args=(i,)) for i in range(2)]
+        for t in ts: t.start()
+        time.sleep(0.1)
+        self.assertEqual(bh.queued(), 2)
+        h.release()
+        for t in ts: t.join(timeout=5)
+        self.assertEqual(bh.active(), 0)
+        self.assertEqual(len([r for r in results if r[0] == "ok"]), 2)
+
+    def test_timeout(self):
+        bh = Bulkhead(1, max_queue=1)
+        h = bh.acquire()
+        t0 = time.monotonic()
+        with self.assertRaises(TimeoutError):
+            bh.acquire(timeout=0.1)
+        elapsed = time.monotonic() - t0
+        self.assertLess(elapsed, 1.0)
+        # queued counter restored
+        self.assertEqual(bh.queued(), 0)
+        h.release()
+
+    def test_timeout_negative(self):
+        bh = Bulkhead(1)
+        with self.assertRaises(ValueError):
+            bh.acquire(timeout=-1)
+
+    def test_double_release_idempotent(self):
+        bh = Bulkhead(1)
+        h = bh.acquire()
+        h.release()
+        h.release()
+        self.assertEqual(bh.active(), 0)
+        # slot reusable
+        h2 = bh.try_acquire()
+        self.assertIsNotNone(h2)
+        h2.release()
+
+    def test_concurrent_stress(self):
+        bh = Bulkhead(4, max_queue=100)
+        active_peak = [0]
+        lock = threading.Lock()
+        def work():
+            with bh.slot():
+                with lock:
+                    # approximate peak via active()
+                    pass
+                time.sleep(0.01)
+        ts = [threading.Thread(target=work) for _ in range(50)]
+        for t in ts: t.start()
+        for t in ts: t.join(timeout=10)
+        self.assertEqual(bh.active(), 0)
+        self.assertEqual(bh.queued(), 0)
+
+    def test_handle_context_exit_releases_on_exception(self):
+        bh = Bulkhead(1)
+        try:
+            with bh.slot():
+                raise ValueError("boom")
+        except ValueError:
+            pass
+        self.assertEqual(bh.active(), 0)
+
+
+if __name__ == "__main__":
+    unittest.main()