fix(skills): let skill_manage lock waiters wake across event loops (#7020)

* fix(skills): let skill_manage lock waiters wake across event loops

The sync tool wrapper runs every call under its own asyncio.run, and
ToolNode runs same-turn calls on separate threads. The per-(user, skill)
asyncio.Lock bound itself to the first contending loop, so a release from
another thread never woke that waiter, and a third caller failed with
"bound to a different event loop".

Reuse the read-before-write gate lock: each waiter parks on its own
loop without holding an executor thread, and release wakes every waiter
to re-contend, so a cancelled waiter cannot strand the skill. Keep the
gate's waiters in arrival order, so same-turn calls on one skill still
take the lock in the order they reached it, as with asyncio.Lock. Guard
the registry lookup, which the sync path runs on several threads.

Fixes #5846

* fix(skills): hand the skill lock to the oldest waiter

Reusing the read-before-write gate let a later same-skill call overtake
one already waiting. Release unlocked the gate and woke every waiter, so
a call whose storage lookup finished in that same loop step took the
free lock first. asyncio.Lock on main never lets a new caller pass a
queued one.

Give skill_manage its own small lock. Release hands the lock straight to
the oldest waiter and keeps it locked, so a later caller queues behind
it. Each waiter still parks on its own future and wakes on its own loop,
so waiting holds no executor thread. Claiming the future settles a race
with its cancellation, and a waiter cancelled after the handoff passes
the lock on. Releasing an unlocked skill lock raises, as with
asyncio.Lock. read_before_write_middleware.py is back to main.

---------

Co-authored-by: Totoro-qaq <279883115+Totoro-qaq@users.noreply.github.com>
This commit is contained in:
Totoro 2026-10-18 08:03:13 +08:00 committed by GitHub
parent b731d979f1
commit 002bcc521a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 326 additions and 7 deletions

View File

@ -7,6 +7,9 @@ import json
import logging
import shutil
import tempfile
import threading
from collections import deque
from concurrent.futures import Future
from pathlib import Path
from typing import Any, NoReturn
from weakref import WeakValueDictionary
@ -30,17 +33,76 @@ from deerflow.tools.types import Runtime
logger = logging.getLogger(__name__)
class _SkillLock:
"""Per-(user, skill) lock shared by every event loop that runs this tool.
The sync wrapper gives each call its own ``asyncio.run`` loop; an
``asyncio.Lock`` binds to one of them and strands waiters on the others
(#5846). Each waiter here parks on its own future, woken on its own loop,
so waiting holds no executor thread. Like ``asyncio.Lock``, release hands
the lock straight to the oldest waiter, so a later call cannot overtake it.
"""
def __init__(self) -> None:
self._state = threading.Lock()
self._locked = False
self._waiters: deque[tuple[Future[None], asyncio.AbstractEventLoop]] = deque()
def locked(self) -> bool:
with self._state:
return self._locked
async def __aenter__(self) -> None:
with self._state:
if not self._locked:
self._locked = True
return
waiter: Future[None] = Future()
self._waiters.append((waiter, asyncio.get_running_loop()))
try:
await asyncio.wrap_future(waiter)
except BaseException:
# cancel() fails only once release has handed us the lock: pass it
# on, or the callers queued behind us would wait forever.
if not waiter.cancel():
self._release()
raise
async def __aexit__(self, *_exc: object) -> None:
self._release()
def _release(self) -> None:
with self._state:
if not self._locked:
raise RuntimeError("release unlocked skill lock")
while self._waiters:
waiter, loop = self._waiters.popleft()
# Claiming the future makes a racing cancellation a no-op, so
# the lock is either handed over here or skipped to the next.
if not loop.is_closed() and waiter.set_running_or_notify_cancel():
break
else:
self._locked = False
return
# Outside the state lock: the result callback schedules the waiter's
# wakeup on its own loop.
waiter.set_result(None)
# Lock granularity: (user_id, skill_name) to avoid cross-user blocking.
_skill_locks: WeakValueDictionary[tuple[str, str], asyncio.Lock] = WeakValueDictionary()
_skill_locks: WeakValueDictionary[tuple[str, str], _SkillLock] = WeakValueDictionary()
_skill_locks_guard = threading.Lock()
def _get_lock(user_id: str, name: str) -> asyncio.Lock:
def _get_lock(user_id: str, name: str) -> _SkillLock:
key = (user_id, name)
lock = _skill_locks.get(key)
if lock is None:
lock = asyncio.Lock()
_skill_locks[key] = lock
return lock
with _skill_locks_guard:
lock = _skill_locks.get(key)
if lock is None:
lock = _SkillLock()
_skill_locks[key] = lock
return lock
def _get_thread_id(runtime: Runtime | None) -> str | None:

View File

@ -1,4 +1,8 @@
import asyncio
import gc
import importlib
import threading
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from types import SimpleNamespace
@ -495,3 +499,256 @@ def test_scan_or_raise_does_not_attach_model_tracing(monkeypatch, tmp_path):
assert result["decision"] == "allow"
assert create_kwargs["attach_tracing"] is False
# --- per-(user, skill) lock across event loops (#5846) ---
class _ParkingStorage:
"""Skill storage whose first delete parks inside the skill lock until released."""
def __init__(self) -> None:
self.entered = threading.Event()
self.release = threading.Event()
self.deleted: list[str] = []
self.callers: list[str] = []
self.max_active = 0
self._active = 0
self._guard = threading.Lock()
def delete_custom_skill(self, name: str, *, history_meta: dict) -> None:
with self._guard:
self._active += 1
self.max_active = max(self.max_active, self._active)
try:
if not self.entered.is_set():
self.entered.set()
# Bounded so a broken lock cannot hang pytest teardown.
assert self.release.wait(5), "test never released the parked holder"
self.deleted.append(name)
self.callers.append(history_meta["thread_id"])
finally:
with self._guard:
self._active -= 1
class _ObservedSkillLock:
"""Report waits on a held skill lock without replacing its locking behavior."""
def __init__(self, lock, arrived) -> None:
self.lock = lock
self.arrived = arrived
async def __aenter__(self):
if self.lock.locked():
self.arrived()
return await self.lock.__aenter__()
async def __aexit__(self, *exc):
return await self.lock.__aexit__(*exc)
def _park_skill_mutations(monkeypatch, arrived) -> _ParkingStorage:
storage = _ParkingStorage()
monkeypatch.setattr(skill_manage_module, "get_or_new_user_skill_storage", lambda _user_id: storage)
async def _refresh(_user_id: str) -> None:
return None
monkeypatch.setattr(skill_manage_module, "refresh_user_skills_system_prompt_cache_async", _refresh)
get_lock = skill_manage_module._get_lock
monkeypatch.setattr(skill_manage_module, "_get_lock", lambda user_id, name: _ObservedSkillLock(get_lock(user_id, name), arrived))
return storage
def _delete(user_id: str, name: str, *, thread_id: str = "thread-1"):
return skill_manage_module.skill_manage_tool.coroutine(_make_runtime(thread_id=thread_id, user_id=user_id), "delete", name)
def test_skill_manage_same_skill_sync_calls_do_not_hang_across_event_loops(monkeypatch):
# ToolNode runs same-turn sync calls on separate threads; the sync wrapper
# gives each one its own asyncio.run loop.
arrivals = threading.Semaphore(0)
storage = _park_skill_mutations(monkeypatch, arrivals.release)
results = {}
def call(tag):
try:
results[tag] = skill_manage_module.skill_manage_tool.func(runtime=_make_runtime(user_id="alice"), action="delete", name="cross-loop-skill")
except Exception as exc:
results[tag] = exc
threads = [threading.Thread(target=call, args=("holder",), daemon=True)]
threads[0].start()
try:
assert storage.entered.wait(5)
for tag in ("waiter-1", "waiter-2"):
threads.append(threading.Thread(target=call, args=(tag,), daemon=True))
threads[-1].start()
assert arrivals.acquire(timeout=5), f"{tag} did not wait for the held skill lock"
finally:
storage.release.set()
for thread in threads:
thread.join(5)
assert not any(thread.is_alive() for thread in threads), "a same-skill sync call never returned"
assert results == dict.fromkeys(("holder", "waiter-1", "waiter-2"), "Deleted custom skill 'cross-loop-skill'.")
assert storage.max_active == 1
gc.collect()
assert ("alice", "cross-loop-skill") not in skill_manage_module._skill_locks
def test_skill_manage_lock_waiters_leave_workers_for_other_skills(monkeypatch):
async def scenario():
asyncio.get_running_loop().set_default_executor(ThreadPoolExecutor(max_workers=2))
waiters_parked = asyncio.Event()
arrived = 0
def on_arrival():
nonlocal arrived
arrived += 1
if arrived == 2:
waiters_parked.set()
storage = _park_skill_mutations(monkeypatch, on_arrival)
holder = asyncio.create_task(_delete("alice", "busy-skill"))
tasks = [holder]
try:
assert await asyncio.to_thread(storage.entered.wait, 5)
tasks += [asyncio.create_task(_delete("alice", "busy-skill")) for _ in range(2)]
await asyncio.wait_for(waiters_parked.wait(), 5)
# The holder parks one worker; the waiters must leave the other free
# for another skill and for the same skill of another user.
assert await asyncio.wait_for(_delete("alice", "other-skill"), 5) == "Deleted custom skill 'other-skill'."
assert await asyncio.wait_for(_delete("bob", "busy-skill"), 5) == "Deleted custom skill 'busy-skill'."
assert not holder.done()
finally:
storage.release.set()
results = await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), 5)
assert results == ["Deleted custom skill 'busy-skill'."] * 3
asyncio.run(scenario())
def test_skill_manage_cancelled_lock_waiter_neither_waits_nor_strands_the_skill(monkeypatch):
async def scenario():
lock = skill_manage_module._get_lock("alice", "cancel-skill")
await lock.__aenter__()
held = True
all_parked = asyncio.Event()
arrived = 0
def on_arrival():
nonlocal arrived
arrived += 1
if arrived == 3:
all_parked.set()
storage = _park_skill_mutations(monkeypatch, on_arrival)
storage.entered.set()
waiters = [asyncio.create_task(_delete("alice", "cancel-skill")) for _ in range(3)]
try:
await asyncio.wait_for(all_parked.wait(), 5)
waiters[0].cancel()
# asyncio.wait, not wait_for: a waiter that shields its wait from
# cancellation must fail this test, not deadlock against the lock we hold.
await asyncio.wait({waiters[0]}, timeout=5)
assert waiters[0].cancelled(), "cancellation waited for the holder"
# Cancel the next waiter in the same step as the release, before it can re-acquire.
held = False
await lock.__aexit__(None, None, None)
waiters[1].cancel()
await asyncio.wait({waiters[1]}, timeout=5)
assert waiters[1].cancelled()
assert await asyncio.wait_for(waiters[2], 5) == "Deleted custom skill 'cancel-skill'."
finally:
if held:
await lock.__aexit__(None, None, None)
for task in waiters:
task.cancel()
await asyncio.wait_for(asyncio.gather(*waiters, return_exceptions=True), 5)
assert storage.deleted == ["cancel-skill"]
asyncio.run(scenario())
gc.collect()
assert ("alice", "cancel-skill") not in skill_manage_module._skill_locks
def test_skill_manage_lock_rejects_release_when_unlocked():
# Like asyncio.Lock: releasing a lock nobody holds is an error, not a no-op.
lock = skill_manage_module._get_lock("alice", "unlocked-skill")
with pytest.raises(RuntimeError):
asyncio.run(lock.__aexit__(None, None, None))
def test_skill_manage_same_skill_waiters_keep_arrival_order(monkeypatch):
# Same-turn calls on one skill can depend on each other (patch a->b, then
# b->c). The gateway runs them on one loop, where they must take the lock in
# the order they reached it, as asyncio.Lock did.
async def scenario():
lock = skill_manage_module._get_lock("alice", "ordered-skill")
await lock.__aenter__()
arrived = asyncio.Event()
storage = _park_skill_mutations(monkeypatch, arrived.set)
storage.entered.set()
# The thread id only tags each call in the storage record.
calls = [f"call-{i}" for i in range(8)]
waiters = []
try:
for call in calls:
arrived.clear()
waiters.append(asyncio.create_task(_delete("alice", "ordered-skill", thread_id=call)))
await asyncio.wait_for(arrived.wait(), 5)
finally:
await lock.__aexit__(None, None, None)
await asyncio.wait_for(asyncio.gather(*waiters, return_exceptions=True), 5)
assert storage.callers == calls
asyncio.run(scenario())
def test_skill_manage_later_call_does_not_overtake_a_parked_waiter(monkeypatch):
# The storage lookup runs before the lock, so a later call can reach the
# lock in the same loop step as a release. As with asyncio.Lock, it must
# queue behind the call already waiting instead of taking the lock first.
async def scenario():
lock = skill_manage_module._get_lock("alice", "fair-skill")
await lock.__aenter__()
arrived = asyncio.Event()
storage = _park_skill_mutations(monkeypatch, arrived.set)
storage.entered.set()
to_thread = skill_manage_module._to_thread
lookup_parked = asyncio.Event()
finish_lookup = asyncio.Event()
late = None
async def gated_to_thread(func, /, *args, **kwargs):
if func is skill_manage_module.get_or_new_user_skill_storage and asyncio.current_task() is late:
lookup_parked.set()
await finish_lookup.wait()
# Resume straight into the lock, with no thread hop in between.
return func(*args, **kwargs)
return await to_thread(func, *args, **kwargs)
monkeypatch.setattr(skill_manage_module, "_to_thread", gated_to_thread)
waiter = asyncio.create_task(_delete("alice", "fair-skill", thread_id="queued"))
tasks = [waiter]
held = True
try:
await asyncio.wait_for(arrived.wait(), 5)
late = asyncio.create_task(_delete("alice", "fair-skill", thread_id="late"))
tasks.append(late)
await asyncio.wait_for(lookup_parked.wait(), 5)
# Release while the later call's lookup completes in the same step.
finish_lookup.set()
held = False
await lock.__aexit__(None, None, None)
await asyncio.wait_for(asyncio.gather(*tasks), 5)
finally:
if held:
await lock.__aexit__(None, None, None)
finish_lookup.set()
await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), 5)
assert storage.callers == ["queued", "late"]
asyncio.run(scenario())