mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-11 06:28:58 +00:00
e21245fd5b
* fix(runtime): reclaim idle keyed locks safely Replace the per-loop thread lock registries with a waiter-aware keyed lock table. Count holders and queued waiters before acquisition so idle entries can be reclaimed without allowing a late caller to bypass an existing waiter. Add regression coverage for runtime call-site reclamation, goal/checkpoint domain independence, queued-waiter ordering, cancellation cleanup, high-cardinality key reclamation, and cross-event-loop isolation. Fixes #5171 * style(runtime): format keyed lock helper
88 lines
2.9 KiB
Python
88 lines
2.9 KiB
Python
"""Async per-key serialization with waiter-aware entry reclamation."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import threading
|
|
import weakref
|
|
from collections.abc import AsyncIterator, Hashable
|
|
from contextlib import asynccontextmanager
|
|
from dataclasses import dataclass
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class _Entry:
|
|
lock: asyncio.Lock
|
|
participants: int = 0 # current holder plus queued waiters
|
|
|
|
|
|
class AsyncKeyedLockTable[KeyT: Hashable]:
|
|
"""Serialize same-key work without retaining idle keys.
|
|
|
|
A table may be shared by event loops running in different threads. Each
|
|
loop receives its own entries because ``asyncio.Lock`` instances become
|
|
loop-affine once contended. The thread lock protects only the registry;
|
|
async critical sections never hold it.
|
|
|
|
Participants are counted before awaiting the lock. This keeps an entry
|
|
discoverable until its final holder or waiter leaves, so a new caller
|
|
cannot create a second lock and bypass an already queued waiter. Cancelled
|
|
waiters check their participation back in through the same ``finally``
|
|
path.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self._guard = threading.Lock()
|
|
self._entries_by_loop: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, dict[KeyT, _Entry]] = weakref.WeakKeyDictionary()
|
|
|
|
@asynccontextmanager
|
|
async def hold(self, key: KeyT) -> AsyncIterator[None]:
|
|
"""Hold the current event loop's lock for ``key``."""
|
|
loop = asyncio.get_running_loop()
|
|
entries, entry = self._checkout(loop, key)
|
|
acquired = False
|
|
try:
|
|
await entry.lock.acquire()
|
|
acquired = True
|
|
yield
|
|
finally:
|
|
try:
|
|
if acquired:
|
|
entry.lock.release()
|
|
finally:
|
|
self._checkin(loop, entries, key, entry)
|
|
|
|
def _checkout(
|
|
self,
|
|
loop: asyncio.AbstractEventLoop,
|
|
key: KeyT,
|
|
) -> tuple[dict[KeyT, _Entry], _Entry]:
|
|
with self._guard:
|
|
entries = self._entries_by_loop.get(loop)
|
|
if entries is None:
|
|
entries = {}
|
|
self._entries_by_loop[loop] = entries
|
|
entry = entries.get(key)
|
|
if entry is None:
|
|
entry = _Entry(lock=asyncio.Lock())
|
|
entries[key] = entry
|
|
entry.participants += 1
|
|
return entries, entry
|
|
|
|
def _checkin(
|
|
self,
|
|
loop: asyncio.AbstractEventLoop,
|
|
entries: dict[KeyT, _Entry],
|
|
key: KeyT,
|
|
entry: _Entry,
|
|
) -> None:
|
|
with self._guard:
|
|
entry.participants -= 1
|
|
if entry.participants != 0 or entry.lock.locked():
|
|
return
|
|
if entries.get(key) is not entry:
|
|
return
|
|
entries.pop(key)
|
|
if not entries and self._entries_by_loop.get(loop) is entries:
|
|
self._entries_by_loop.pop(loop, None)
|