33 lines
1.3 KiB
Python
33 lines
1.3 KiB
Python
 
|
"""Per-event-loop semaphores for bounded concurrency (refactor §4, §7).
|
||
|
|
|
||
|
|
Connection and parse caps are process-local but must be bound to the running
|
||
|
|
event loop: an :class:`asyncio.Semaphore` created on one loop cannot be awaited
|
||
|
|
on another (a real hazard under pytest, which uses a fresh loop per test). These
|
||
|
|
helpers hand back a semaphore keyed by ``(group, key)`` for the current loop,
|
||
|
|
creating it on first use and dropping the whole set when the loop is collected.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import weakref
|
||
|
|
|
||
|
|
# group -> {loop: {key: Semaphore}}. WeakKeyDictionary drops a loop's semaphores
|
||
|
|
# once the loop is gone, so no state leaks across tests or restarts.
|
||
|
|
_registry: dict[
|
||
|
|
str, weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, dict[str, asyncio.Semaphore]]
|
||
|
|
] = {}
|
||
|
|
|
||
|
|
|
||
|
|
def loop_semaphore(group: str, key: str, cap: int) -> asyncio.Semaphore:
|
||
|
|
"""Return the current loop's semaphore for ``(group, key)`` (cap on first use)."""
|
||
|
|
loop = asyncio.get_running_loop()
|
||
|
|
per_group = _registry.setdefault(group, weakref.WeakKeyDictionary())
|
||
|
|
per_loop = per_group.get(loop)
|
||
|
|
if per_loop is None:
|
||
|
|
per_loop = per_group[loop] = {}
|
||
|
|
sem = per_loop.get(key)
|
||
|
|
if sem is None:
|
||
|
|
sem = per_loop[key] = asyncio.Semaphore(cap)
|
||
|
|
return sem
|