From 8be7411da8c30295a6594bfc0ab592b4bdc96a4a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=BB=84=E4=BA=91=E9=BE=99?= <76432572+nankingjing@users.noreply.github.com> Date: Sun, 12 Jul 2026 13:03:02 +0800 Subject: [PATCH] fix(runtime): serialize SQLite event-store writes to prevent per-thread seq collisions (#4077) * fix(run-events): serialize seq assignment with a per-thread asyncio lock put() and put_batch() read max(seq) and then INSERT seq+1 in separate awaits. Two coroutines writing the same thread in one process could interleave between the read and the insert and assign the same seq, colliding on SQLite where the DB-level FOR UPDATE lock is weaker than Postgres. Add a per-thread asyncio.Lock (_write_locks / _get_write_lock) around the read-assign-insert critical section in both methods. * address review: evict orphaned per-thread write-lock in delete_by_thread _write_locks accumulated one asyncio.Lock per thread ever seen and never released, leaking in the long-lived DbRunEventStore singleton. Evict the entry after delete_by_thread when no writer holds it (lock recreated lazily on the next write). Per @willem-bd review on #4077. Co-Authored-By: Claude --------- Co-authored-by: Claude --- .../deerflow/runtime/events/store/db.py | 115 +++++++++++------- backend/tests/test_run_event_store.py | 103 ++++++++++++++++ 2 files changed, 173 insertions(+), 45 deletions(-) diff --git a/backend/packages/harness/deerflow/runtime/events/store/db.py b/backend/packages/harness/deerflow/runtime/events/store/db.py index 4190bf440..719c45ddb 100644 --- a/backend/packages/harness/deerflow/runtime/events/store/db.py +++ b/backend/packages/harness/deerflow/runtime/events/store/db.py @@ -6,6 +6,7 @@ at ``max_trace_content`` bytes to avoid bloating the database. from __future__ import annotations +import asyncio import json import logging from datetime import UTC, datetime @@ -26,6 +27,20 @@ class DbRunEventStore(RunEventStore): def __init__(self, session_factory: async_sessionmaker[AsyncSession], *, max_trace_content: int = 10240): self._sf = session_factory self._max_trace_content = max_trace_content + # Per-thread asyncio locks serialize seq assignment for concurrent + # in-process writers on the same thread. The DB-level FOR UPDATE / + # 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] = {} + + def _get_write_lock(self, thread_id: str) -> asyncio.Lock: + """Return (creating if needed) the per-thread seq-assignment lock.""" + lock = self._write_locks.get(thread_id) + if lock is None: + lock = asyncio.Lock() + self._write_locks[thread_id] = lock + return lock @staticmethod def _row_to_dict(row: RunEventRow) -> dict: @@ -123,23 +138,24 @@ class DbRunEventStore(RunEventStore): content, metadata = self._truncate_trace(category, content, metadata) db_content, metadata = self._content_to_db(content, metadata) user_id = self._user_id_from_context() - async with self._sf() as session: - async with session.begin(): - max_seq = await self._max_seq_for_thread(session, thread_id) - seq = (max_seq or 0) + 1 - row = RunEventRow( - thread_id=thread_id, - run_id=run_id, - user_id=user_id, - event_type=event_type, - category=category, - content=db_content, - event_metadata=metadata, - seq=seq, - created_at=datetime.fromisoformat(created_at) if created_at else datetime.now(UTC), - ) - session.add(row) - return self._row_to_dict(row) + async with self._get_write_lock(thread_id): + async with self._sf() as session: + async with session.begin(): + max_seq = await self._max_seq_for_thread(session, thread_id) + seq = (max_seq or 0) + 1 + row = RunEventRow( + thread_id=thread_id, + run_id=run_id, + user_id=user_id, + event_type=event_type, + category=category, + content=db_content, + event_metadata=metadata, + seq=seq, + created_at=datetime.fromisoformat(created_at) if created_at else datetime.now(UTC), + ) + session.add(row) + return self._row_to_dict(row) async def put_batch(self, events): if not events: @@ -148,34 +164,35 @@ class DbRunEventStore(RunEventStore): if len(thread_ids) > 1: raise ValueError(f"put_batch requires all events to belong to the same thread; got {thread_ids!r}") user_id = self._user_id_from_context() - async with self._sf() as session: - async with session.begin(): - # All events belong to the same thread (validated above). - thread_id = events[0]["thread_id"] - max_seq = await self._max_seq_for_thread(session, thread_id) - seq = max_seq or 0 - rows = [] - for e in events: - seq += 1 - content = e.get("content", "") - category = e.get("category", "trace") - metadata = e.get("metadata") - content, metadata = self._truncate_trace(category, content, metadata) - db_content, metadata = self._content_to_db(content, metadata) - row = RunEventRow( - thread_id=e["thread_id"], - run_id=e["run_id"], - user_id=e.get("user_id", user_id), - event_type=e["event_type"], - category=category, - content=db_content, - event_metadata=metadata, - seq=seq, - created_at=datetime.fromisoformat(e["created_at"]) if e.get("created_at") else datetime.now(UTC), - ) - session.add(row) - rows.append(row) - return [self._row_to_dict(r) for r in rows] + # All events belong to the same thread (validated above). + thread_id = events[0]["thread_id"] + async with self._get_write_lock(thread_id): + async with self._sf() as session: + async with session.begin(): + max_seq = await self._max_seq_for_thread(session, thread_id) + seq = max_seq or 0 + rows = [] + for e in events: + seq += 1 + content = e.get("content", "") + category = e.get("category", "trace") + metadata = e.get("metadata") + content, metadata = self._truncate_trace(category, content, metadata) + db_content, metadata = self._content_to_db(content, metadata) + row = RunEventRow( + thread_id=e["thread_id"], + run_id=e["run_id"], + user_id=e.get("user_id", user_id), + event_type=e["event_type"], + category=category, + content=db_content, + event_metadata=metadata, + seq=seq, + created_at=datetime.fromisoformat(e["created_at"]) if e.get("created_at") else datetime.now(UTC), + ) + session.add(row) + rows.append(row) + return [self._row_to_dict(r) for r in rows] async def list_messages( self, @@ -304,6 +321,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) return count async def delete_by_run( diff --git a/backend/tests/test_run_event_store.py b/backend/tests/test_run_event_store.py index 7c409c711..6480ef2ee 100644 --- a/backend/tests/test_run_event_store.py +++ b/backend/tests/test_run_event_store.py @@ -498,6 +498,109 @@ class TestDbRunEventStore: await close_engine() +class TestDbRunEventStoreWriteLock: + """Per-thread seq-assignment lock (fixes SQLite UNIQUE(thread_id, seq) races). + + Two in-process coroutines writing to the same thread can interleave between + the ``max(seq)`` read and the INSERT, both computing the same next seq and + colliding. A per-thread ``asyncio.Lock`` serializes seq assignment. + """ + + def test_get_write_lock_same_thread_returns_same_lock(self): + import asyncio + from unittest.mock import MagicMock + + from deerflow.runtime.events.store.db import DbRunEventStore + + # The lock accessor does not touch the session factory, so a stub is fine. + store = DbRunEventStore(MagicMock()) + + lock = store._get_write_lock("thread-1") + assert isinstance(lock, asyncio.Lock) + assert store._get_write_lock("thread-1") is lock + + def test_get_write_lock_distinct_threads_get_distinct_locks(self): + from unittest.mock import MagicMock + + from deerflow.runtime.events.store.db import DbRunEventStore + + store = DbRunEventStore(MagicMock()) + + assert store._get_write_lock("thread-1") is not store._get_write_lock("thread-2") + + @pytest.mark.anyio + async def test_concurrent_put_batch_same_thread_has_no_seq_collision(self, tmp_path): + import asyncio + + from deerflow.persistence.engine import close_engine, get_session_factory, init_engine + from deerflow.runtime.events.store.db import DbRunEventStore + + url = f"sqlite+aiosqlite:///{tmp_path / 'test.db'}" + await init_engine("sqlite", url=url, sqlite_dir=str(tmp_path)) + s = DbRunEventStore(get_session_factory()) + + def _batch(run_id: str): + return [{"thread_id": "t1", "run_id": run_id, "event_type": "trace", "category": "trace"} for _ in range(20)] + + # Fire two concurrent batches at the same thread; without the per-thread + # lock this races on seq and raises IntegrityError / duplicates seq. + results = await asyncio.gather(s.put_batch(_batch("r1")), s.put_batch(_batch("r2"))) + + all_seqs = [r["seq"] for batch in results for r in batch] + assert len(all_seqs) == 40 + # Seq values are unique and contiguous 1..40 across both writers. + assert sorted(all_seqs) == list(range(1, 41)) + + await close_engine() + + @pytest.mark.anyio + async def test_delete_by_thread_evicts_orphaned_write_lock(self, tmp_path): + from deerflow.persistence.engine import close_engine, get_session_factory, init_engine + from deerflow.runtime.events.store.db import DbRunEventStore + + url = f"sqlite+aiosqlite:///{tmp_path / 'test.db'}" + await init_engine("sqlite", url=url, sqlite_dir=str(tmp_path)) + s = DbRunEventStore(get_session_factory()) + + # A write materializes the per-thread lock in the registry. + await s.put_batch([{"thread_id": "t1", "run_id": "r1", "event_type": "trace", "category": "trace"}]) + assert "t1" in s._write_locks + + # Deleting the thread must evict the now-orphaned lock so the registry + # does not grow unbounded across the singleton store's lifetime. + await s.delete_by_thread("t1") + assert "t1" not in s._write_locks + + # A subsequent write recreates a fresh lock and seq restarts from 1. + result = await s.put_batch([{"thread_id": "t1", "run_id": "r2", "event_type": "trace", "category": "trace"}]) + assert "t1" in s._write_locks + assert result[0]["seq"] == 1 + + await close_engine() + + @pytest.mark.anyio + async def test_delete_by_thread_keeps_lock_held_by_inflight_writer(self, tmp_path): + from deerflow.persistence.engine import close_engine, get_session_factory, init_engine + from deerflow.runtime.events.store.db import DbRunEventStore + + url = f"sqlite+aiosqlite:///{tmp_path / 'test.db'}" + await init_engine("sqlite", url=url, sqlite_dir=str(tmp_path)) + s = DbRunEventStore(get_session_factory()) + + # Simulate a writer mid-flight by holding the lock; the eviction must + # not drop a lock another coroutine is actively using. + lock = s._get_write_lock("t1") + await lock.acquire() + try: + await s.delete_by_thread("t1") + assert "t1" in s._write_locks + assert s._write_locks["t1"] is lock + finally: + lock.release() + + await close_engine() + + # -- Factory tests --