From 78117354b2e1e6864cd230311ccc9e360a20945d Mon Sep 17 00:00:00 2001 From: NanPan <111261006+poijygfdyy@users.noreply.github.com> Date: Thu, 17 Sep 2026 22:06:03 +0800 Subject: [PATCH] fix(events): preserve DB write-lock generation across deletion (#5462) --- .../deerflow/runtime/events/store/db.py | 29 ++++--- .../test_db_event_store_lock_lifecycle.py | 78 +++++++++++++++++++ 2 files changed, 98 insertions(+), 9 deletions(-) create mode 100644 backend/tests/test_db_event_store_lock_lifecycle.py diff --git a/backend/packages/harness/deerflow/runtime/events/store/db.py b/backend/packages/harness/deerflow/runtime/events/store/db.py index 6736d8bc3..9b3d1266a 100644 --- a/backend/packages/harness/deerflow/runtime/events/store/db.py +++ b/backend/packages/harness/deerflow/runtime/events/store/db.py @@ -10,6 +10,7 @@ import asyncio import json import logging import re +import weakref from datetime import UTC, datetime from typing import Any @@ -34,7 +35,14 @@ class DbRunEventStore(RunEventStore): # advisory lock guards cross-process races; this guards the common # single-process case where two coroutines interleave between the # max(seq) read and the INSERT and would otherwise collide on seq. - self._write_locks: dict[str, asyncio.Lock] = {} + # + # The weak registry preserves one lock generation while an admitted + # holder/waiter still references it. A separate pin keeps the historical + # one-lock-per-live-thread behavior until delete_by_thread() explicitly + # retires that thread; after retirement, outstanding users alone keep + # the generation alive until they drain. + self._write_locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary() + self._write_lock_pins: dict[str, asyncio.Lock] = {} def _get_write_lock(self, thread_id: str) -> asyncio.Lock: """Return (creating if needed) the per-thread seq-assignment lock.""" @@ -42,6 +50,9 @@ class DbRunEventStore(RunEventStore): if lock is None: lock = asyncio.Lock() self._write_locks[thread_id] = lock + # A fresh caller after deletion makes the thread live again. Repin the + # current generation so normal live-thread registry lifetime is stable. + self._write_lock_pins[thread_id] = lock return lock @staticmethod @@ -486,14 +497,14 @@ class DbRunEventStore(RunEventStore): if count > 0: await session.execute(delete(RunEventRow).where(*count_conditions)) await session.commit() - # Evict the per-thread seq-assignment lock so ``_write_locks`` does - # not grow unbounded over the (long-lived, singleton) store's - # lifetime. Only pop when no writer is mid-flight; a later write - # recreates the lock lazily and seq restarts correctly from the - # now-deleted thread. - lock = self._write_locks.get(thread_id) - if lock is not None and not lock.locked(): - self._write_locks.pop(thread_id, None) + # Retire the live-thread pin, but never remove the weak registry + # entry directly. asyncio.Lock.release() clears ``locked()`` before + # a queued waiter resumes, so an unlocked check can observe the + # handoff window and split one thread onto two lock generations. + # Holders/waiters keep the old generation alive until they drain; a + # later caller therefore resolves that same lock instead of racing + # it with a fresh one. + self._write_lock_pins.pop(thread_id, None) return count async def delete_by_run( diff --git a/backend/tests/test_db_event_store_lock_lifecycle.py b/backend/tests/test_db_event_store_lock_lifecycle.py new file mode 100644 index 000000000..ba3bee647 --- /dev/null +++ b/backend/tests/test_db_event_store_lock_lifecycle.py @@ -0,0 +1,78 @@ +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