mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-27 00:17:53 +00:00
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:
parent
b650456c6d
commit
8be7411da8
@ -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(
|
||||
|
||||
@ -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 --
|
||||
|
||||
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user