85 lines
2.8 KiB
Python
85 lines
2.8 KiB
Python
"""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
|