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

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):

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)

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)

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)