mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
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:
parent
b731d979f1
commit
002bcc521a
@ -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:
|
||||
|
||||
@ -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())
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user