137 lines
4.4 KiB
Python
137 lines
4.4 KiB
Python
"""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)
|