From e21245fd5b6a3e1e5131d50030ace0d6e8edb5c1 Mon Sep 17 00:00:00 2001 From: Jun <84921700+Amazingjun-j@users.noreply.github.com> Date: Fri, 4 Sep 2026 20:20:28 +0800 Subject: [PATCH] fix(runtime): add waiter-safe keyed lock reclamation (#5176) * 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 --- .../packages/harness/deerflow/runtime/goal.py | 26 +-- .../harness/deerflow/runtime/keyed_lock.py | 87 ++++++++ .../harness/deerflow/runtime/runs/worker.py | 25 +-- backend/tests/test_runtime_keyed_lock.py | 200 ++++++++++++++++++ 4 files changed, 298 insertions(+), 40 deletions(-) create mode 100644 backend/packages/harness/deerflow/runtime/keyed_lock.py create mode 100644 backend/tests/test_runtime_keyed_lock.py diff --git a/backend/packages/harness/deerflow/runtime/goal.py b/backend/packages/harness/deerflow/runtime/goal.py index 3f603a42b..abb0cea0e 100644 --- a/backend/packages/harness/deerflow/runtime/goal.py +++ b/backend/packages/harness/deerflow/runtime/goal.py @@ -14,10 +14,7 @@ import inspect import json import logging import os -import threading -import weakref -from collections.abc import AsyncIterator -from contextlib import asynccontextmanager +from contextlib import AbstractAsyncContextManager from typing import Any, Literal, NamedTuple from langchain_core.messages import HumanMessage, SystemMessage @@ -26,6 +23,7 @@ from langgraph.checkpoint.base import empty_checkpoint, uuid6 import deerflow.utils.llm_text as llm_text from deerflow.agents.goal_state import GoalBlocker, GoalEvaluation, GoalState from deerflow.models import create_chat_model +from deerflow.runtime.keyed_lock import AsyncKeyedLockTable from deerflow.tracing import inject_langfuse_metadata from deerflow.utils.messages import message_to_text from deerflow.utils.time import now_iso @@ -56,30 +54,16 @@ _extract_response_text = llm_text.extract_response_text _strip_markdown_code_fence = llm_text.strip_markdown_code_fence _strip_think_blocks = llm_text.strip_think_blocks -_goal_locks_guard = threading.Lock() -_goal_locks_by_loop: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, dict[str, asyncio.Lock]] = weakref.WeakKeyDictionary() +_goal_locks = AsyncKeyedLockTable[str]() class GoalWriteConflict(RuntimeError): """Raised when a goal write is based on a stale checkpoint.""" -@asynccontextmanager -async def goal_thread_lock(thread_id: str) -> AsyncIterator[None]: +def goal_thread_lock(thread_id: str) -> AbstractAsyncContextManager[None]: """Serialize goal read-modify-write sequences within the current event loop.""" - loop = asyncio.get_running_loop() - with _goal_locks_guard: - locks = _goal_locks_by_loop.get(loop) - if locks is None: - locks = {} - _goal_locks_by_loop[loop] = locks - lock = locks.get(thread_id) - if lock is None: - lock = asyncio.Lock() - locks[thread_id] = lock - - async with lock: - yield + return _goal_locks.hold(thread_id) class GoalCommand(NamedTuple): diff --git a/backend/packages/harness/deerflow/runtime/keyed_lock.py b/backend/packages/harness/deerflow/runtime/keyed_lock.py new file mode 100644 index 000000000..c6a3d56dc --- /dev/null +++ b/backend/packages/harness/deerflow/runtime/keyed_lock.py @@ -0,0 +1,87 @@ +"""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) diff --git a/backend/packages/harness/deerflow/runtime/runs/worker.py b/backend/packages/harness/deerflow/runtime/runs/worker.py index 8334d5b23..1685fd8cc 100644 --- a/backend/packages/harness/deerflow/runtime/runs/worker.py +++ b/backend/packages/harness/deerflow/runtime/runs/worker.py @@ -25,8 +25,8 @@ import sys import threading import time import weakref -from collections.abc import AsyncIterator, Callable, Coroutine, Mapping -from contextlib import asynccontextmanager +from collections.abc import Callable, Coroutine, Mapping +from contextlib import AbstractAsyncContextManager from contextvars import Context from dataclasses import dataclass, field from datetime import datetime @@ -73,6 +73,7 @@ from deerflow.runtime.goal import ( visible_conversation_signature, write_thread_goal, ) +from deerflow.runtime.keyed_lock import AsyncKeyedLockTable from deerflow.runtime.serialization import serialize from deerflow.runtime.stream_bridge import StreamBridge from deerflow.runtime.stream_modes import normalize_stream_modes, to_langgraph_stream_modes @@ -90,8 +91,7 @@ from .schemas import RunStatus logger = logging.getLogger(__name__) -_checkpoint_locks_guard = threading.Lock() -_checkpoint_locks_by_loop: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, dict[str, asyncio.Lock]] = weakref.WeakKeyDictionary() +_checkpoint_locks = AsyncKeyedLockTable[str]() # Completed LangGraph runs can leave callback Contexts and AsyncPregelLoop # instances in unreachable reference cycles. They are collectable, but a busy @@ -229,22 +229,9 @@ def _release_run_scoped_references( runtime_context.pop(key, None) -@asynccontextmanager -async def _checkpoint_thread_lock(thread_id: str) -> AsyncIterator[None]: +def _checkpoint_thread_lock(thread_id: str) -> AbstractAsyncContextManager[None]: """Serialize checkpoint mutations for one thread without blocking goal commands.""" - loop = asyncio.get_running_loop() - with _checkpoint_locks_guard: - locks = _checkpoint_locks_by_loop.get(loop) - if locks is None: - locks = {} - _checkpoint_locks_by_loop[loop] = locks - lock = locks.get(thread_id) - if lock is None: - lock = asyncio.Lock() - locks[thread_id] = lock - - async with lock: - yield + return _checkpoint_locks.hold(thread_id) _DELIVERY_RECEIPT_RETRY_DELAYS_SECONDS = (0.1, 0.5) diff --git a/backend/tests/test_runtime_keyed_lock.py b/backend/tests/test_runtime_keyed_lock.py new file mode 100644 index 000000000..34d9e1855 --- /dev/null +++ b/backend/tests/test_runtime_keyed_lock.py @@ -0,0 +1,200 @@ +from __future__ import annotations + +import asyncio +import gc +import threading +import weakref +from concurrent.futures import ThreadPoolExecutor +from uuid import uuid4 + +import pytest + +from deerflow.runtime.goal import goal_thread_lock +from deerflow.runtime.runs.worker import _checkpoint_thread_lock + + +class _WeakThreadId(str): + pass + + +class _WeakKey: + pass + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "lock_factory", + [goal_thread_lock, _checkpoint_thread_lock], + ids=["goal", "checkpoint"], +) +async def test_runtime_thread_lock_releases_idle_thread_id(lock_factory) -> None: + thread_id = _WeakThreadId(f"retention-{uuid4().hex}") + thread_id_ref = weakref.ref(thread_id) + + async with lock_factory(thread_id): + pass + + del thread_id + gc.collect() + + assert thread_id_ref() is None + + +@pytest.mark.asyncio +async def test_goal_and_checkpoint_lock_domains_remain_independent() -> None: + release_checkpoint = asyncio.Event() + checkpoint_entered = asyncio.Event() + goal_entered = asyncio.Event() + thread_id = f"independent-{uuid4().hex}" + + async def hold_checkpoint() -> None: + async with _checkpoint_thread_lock(thread_id): + checkpoint_entered.set() + await release_checkpoint.wait() + + async def hold_goal() -> None: + async with goal_thread_lock(thread_id): + goal_entered.set() + + checkpoint_task = asyncio.create_task(hold_checkpoint()) + await checkpoint_entered.wait() + goal_task = asyncio.create_task(hold_goal()) + try: + await asyncio.wait_for(goal_entered.wait(), timeout=1) + finally: + release_checkpoint.set() + await checkpoint_task + await goal_task + + +@pytest.mark.asyncio +async def test_late_arrival_cannot_bypass_queued_waiter() -> None: + from deerflow.runtime.keyed_lock import AsyncKeyedLockTable + + table = AsyncKeyedLockTable[str]() + release_first = asyncio.Event() + release_second = asyncio.Event() + first_entered = asyncio.Event() + second_started = asyncio.Event() + second_entered = asyncio.Event() + third_started = asyncio.Event() + third_entered = asyncio.Event() + active = 0 + max_active = 0 + order: list[str] = [] + + async def participant( + name: str, + started: asyncio.Event | None, + entered: asyncio.Event, + release: asyncio.Event | None, + ) -> None: + nonlocal active, max_active + if started is not None: + started.set() + async with table.hold("thread"): + active += 1 + max_active = max(max_active, active) + order.append(name) + entered.set() + try: + if release is not None: + await release.wait() + finally: + active -= 1 + + first = asyncio.create_task(participant("first", None, first_entered, release_first)) + await first_entered.wait() + + second = asyncio.create_task(participant("second", second_started, second_entered, release_second)) + await second_started.wait() + + release_first.set() + await second_entered.wait() + + third = asyncio.create_task(participant("third", third_started, third_entered, None)) + await third_started.wait() + + assert not third_entered.is_set() + assert max_active == 1 + + release_second.set() + await asyncio.gather(first, second, third) + + assert order == ["first", "second", "third"] + assert max_active == 1 + + +@pytest.mark.asyncio +async def test_cancelled_waiter_releases_its_participation() -> None: + from deerflow.runtime.keyed_lock import AsyncKeyedLockTable + + table = AsyncKeyedLockTable[_WeakKey]() + key = _WeakKey() + key_ref = weakref.ref(key) + release_holder = asyncio.Event() + holder_entered = asyncio.Event() + waiter_started = asyncio.Event() + + async def holder(lock_key: _WeakKey) -> None: + async with table.hold(lock_key): + holder_entered.set() + await release_holder.wait() + + async def waiter(lock_key: _WeakKey) -> None: + waiter_started.set() + async with table.hold(lock_key): + raise AssertionError("cancelled waiter entered the critical section") + + holder_task = asyncio.create_task(holder(key)) + await holder_entered.wait() + waiter_task = asyncio.create_task(waiter(key)) + await waiter_started.wait() + + waiter_task.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter_task + + release_holder.set() + await holder_task + + del holder_task, waiter_task, key + gc.collect() + assert key_ref() is None + + +@pytest.mark.asyncio +async def test_many_unique_keys_are_reclaimed() -> None: + from deerflow.runtime.keyed_lock import AsyncKeyedLockTable + + table = AsyncKeyedLockTable[_WeakKey]() + keys = [_WeakKey() for _ in range(1000)] + key_refs = [weakref.ref(key) for key in keys] + + for key in keys: + async with table.hold(key): + pass + + del key, keys + gc.collect() + + assert all(key_ref() is None for key_ref in key_refs) + + +def test_same_key_is_independent_across_event_loops() -> None: + from deerflow.runtime.keyed_lock import AsyncKeyedLockTable + + table = AsyncKeyedLockTable[str]() + barrier = threading.Barrier(2, timeout=2) + + def run_loop() -> None: + async def run() -> None: + async with table.hold("thread"): + await asyncio.to_thread(barrier.wait) + + asyncio.run(run()) + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [executor.submit(run_loop) for _ in range(2)] + for future in futures: + future.result(timeout=3)