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 <noreply@anthropic.com>

---------

Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
黄云龙 2026-07-12 13:03:02 +08:00 committed by GitHub
parent b650456c6d
commit 8be7411da8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 173 additions and 45 deletions

View File

@ -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(

View File

@ -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 --