mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-19 19:16:17 +00:00
79 lines
2.3 KiB
Python
79 lines
2.3 KiB
Python
import asyncio
|
|
import weakref
|
|
|
|
import pytest
|
|
|
|
from deerflow.runtime.events.store.db import DbRunEventStore
|
|
|
|
|
|
class _PausedDeleteSession:
|
|
def __init__(self, scalar_started: asyncio.Event, allow_scalar: asyncio.Event) -> None:
|
|
self._scalar_started = scalar_started
|
|
self._allow_scalar = allow_scalar
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc, tb):
|
|
return False
|
|
|
|
async def scalar(self, _stmt):
|
|
self._scalar_started.set()
|
|
await self._allow_scalar.wait()
|
|
return 1
|
|
|
|
async def execute(self, _stmt):
|
|
return None
|
|
|
|
async def commit(self) -> None:
|
|
return None
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_delete_waiter_handoff_keeps_one_write_lock_generation():
|
|
scalar_started = asyncio.Event()
|
|
allow_scalar = asyncio.Event()
|
|
session = _PausedDeleteSession(scalar_started, allow_scalar)
|
|
store = DbRunEventStore(lambda: session)
|
|
|
|
old_lock = store._get_write_lock("t1")
|
|
old_lock_ref = weakref.ref(old_lock)
|
|
await old_lock.acquire()
|
|
|
|
waiter_resolved = asyncio.Event()
|
|
waiter_entered = asyncio.Event()
|
|
release_waiter = asyncio.Event()
|
|
|
|
async def queued_writer() -> None:
|
|
lock = store._get_write_lock("t1")
|
|
waiter_resolved.set()
|
|
async with lock:
|
|
waiter_entered.set()
|
|
await release_waiter.wait()
|
|
|
|
waiter_task = asyncio.create_task(queued_writer())
|
|
await waiter_resolved.wait()
|
|
await asyncio.sleep(0)
|
|
assert not waiter_entered.is_set()
|
|
|
|
delete_task = asyncio.create_task(store.delete_by_thread("t1", user_id=None))
|
|
await scalar_started.wait()
|
|
|
|
# Resume deletion first, then release the holder. asyncio.Lock.release()
|
|
# marks the lock unlocked before the queued waiter resumes, so the old
|
|
# implementation can evict the registry entry in that handoff window.
|
|
allow_scalar.set()
|
|
old_lock.release()
|
|
del old_lock
|
|
|
|
await delete_task
|
|
await waiter_entered.wait()
|
|
|
|
try:
|
|
# The waiter still owns the old generation. Any later writer must
|
|
# resolve that exact lock rather than creating a concurrent generation.
|
|
assert store._get_write_lock("t1") is old_lock_ref()
|
|
finally:
|
|
release_waiter.set()
|
|
await waiter_task
|