fix(events): preserve DB write-lock generation across deletion (#5462)

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
+            # Retire the live-thread pin, but never remove the weak registry
-            # not grow unbounded over the (long-lived, singleton) store's
+            # entry directly. asyncio.Lock.release() clears ``locked()`` before
-            # lifetime. Only pop when no writer is mid-flight; a later write
+            # a queued waiter resumes, so an unlocked check can observe the
-            # recreates the lock lazily and seq restarts correctly from the
+            # handoff window and split one thread onto two lock generations.
-            # now-deleted thread.
+            # Holders/waiters keep the old generation alive until they drain; a
-            lock = self._write_locks.get(thread_id)
+            # later caller therefore resolves that same lock instead of racing
-            if lock is not None and not lock.locked():
+            # it with a fresh one.
-                self._write_locks.pop(thread_id, None)
+            self._write_lock_pins.pop(thread_id, None)
             return count
 
     async def delete_by_run(

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