mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-19 02:56:17 +00:00
fix(subagents): preserve capacity release across repeated cancellation (#5477)
* fix(subagents): preserve capacity release across cancellation * test(subagents): cover repeated cancellation during slot release * chore(subagents): remove unreachable release branch
This commit is contained in:
parent
0f2195e994
commit
8e94cc3432
@ -105,13 +105,32 @@ class SubagentExecutionCapacity:
|
||||
async with self._lock:
|
||||
self._release_locked()
|
||||
|
||||
async def _release_cancellation_safe(self) -> None:
|
||||
"""Release an acquired slot before propagating repeated cancellation."""
|
||||
release_task = asyncio.create_task(self._release())
|
||||
cancellation: asyncio.CancelledError | None = None
|
||||
while not release_task.done():
|
||||
try:
|
||||
await asyncio.shield(release_task)
|
||||
except asyncio.CancelledError as exc:
|
||||
if cancellation is None:
|
||||
cancellation = exc
|
||||
|
||||
if cancellation is not None:
|
||||
try:
|
||||
release_task.result()
|
||||
except Exception as exc:
|
||||
raise cancellation from exc
|
||||
raise cancellation
|
||||
release_task.result()
|
||||
|
||||
@asynccontextmanager
|
||||
async def slot(self) -> AsyncIterator[None]:
|
||||
await self._acquire()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
await self._release()
|
||||
await self._release_cancellation_safe()
|
||||
|
||||
|
||||
_config = SubagentRuntimeConfig()
|
||||
|
||||
@ -57,3 +57,43 @@ async def test_snapshot_reports_running_and_queued_waiters():
|
||||
queued.cancel()
|
||||
holder.cancel()
|
||||
await asyncio.gather(queued, holder, return_exceptions=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_repeated_cancellation_cannot_leak_running_slot(monkeypatch):
|
||||
"""A second cancellation during cleanup must not strand process capacity."""
|
||||
capacity = SubagentExecutionCapacity(SubagentRuntimeConfig(max_running=1, max_queued=1, queue_timeout_seconds=1))
|
||||
entered = asyncio.Event()
|
||||
release_started = asyncio.Event()
|
||||
never = asyncio.Event()
|
||||
original_release = capacity._release
|
||||
|
||||
async def observed_release():
|
||||
release_started.set()
|
||||
await original_release()
|
||||
|
||||
monkeypatch.setattr(capacity, "_release", observed_release)
|
||||
|
||||
async def hold_slot():
|
||||
async with capacity.slot():
|
||||
entered.set()
|
||||
await never.wait()
|
||||
|
||||
holder = asyncio.create_task(hold_slot())
|
||||
await asyncio.wait_for(entered.wait(), timeout=1)
|
||||
await capacity._lock.acquire()
|
||||
try:
|
||||
holder.cancel()
|
||||
await asyncio.wait_for(release_started.wait(), timeout=1)
|
||||
holder.cancel()
|
||||
await asyncio.sleep(0)
|
||||
finally:
|
||||
capacity._lock.release()
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await holder
|
||||
|
||||
assert capacity.snapshot().running == 0
|
||||
async with asyncio.timeout(0.5):
|
||||
async with capacity.slot():
|
||||
pass
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user