perf(runtime): index MemoryRunEventStore events by run_id to avoid O(n) scans (#3686)

* perf(runtime): index MemoryRunEventStore events by run_id to avoid O(n) scans

list_messages was already served from a thread-wide messages projection
(#3531), but list_events and list_messages_by_run still scanned the whole
thread's event log (every run, every category) to return one run's events --
O(N_thread) on every run-scoped /messages page-load and /events request.

Add _events_by_run / _messages_by_run projections (same dict objects, kept
in lockstep in _put_one / delete_by_run / delete_by_thread), so both reads
are O(M_run), with bisect cursor pagination for messages. Semantics are
unchanged, pinned by a brute-force parity test over interleaved traces and
both cursors.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* style: ruff format test_run_event_store_by_run_index.py

Clears the lint-backend (ruff format --check) failure on the PR; the original
commit ran `ruff check` but not `ruff format`. No behavior change (test fixture
formatting only).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: ly-wang19 <ly-wang19@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
ly-wang19 2026-06-24 17:47:04 +08:00 committed by GitHub
parent 4a8f94eb28
commit 820560e576
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 162 additions and 13 deletions

View File

@ -19,6 +19,14 @@ class MemoryRunEventStore(RunEventStore):
# kept in seq order so message pagination is O(log m + page) via bisect
# instead of re-scanning every event on each request.
self._messages: dict[str, list[dict]] = {} # thread_id -> seq-sorted message list
# Run-keyed projections of the two lists above (same dict objects, no
# copies), kept in seq order. Per-run reads then cost O(events-in-run)
# instead of O(events-in-thread): without these, ``list_events`` and
# ``list_messages_by_run`` re-scan the whole thread's event log on every
# request even though one run holds only a handful of events. This is
# the per-run analogue of the thread-wide ``_messages`` projection.
self._events_by_run: dict[str, dict[str, list[dict]]] = {} # thread_id -> run_id -> seq-sorted events
self._messages_by_run: dict[str, dict[str, list[dict]]] = {} # thread_id -> run_id -> seq-sorted messages
self._seq_counters: dict[str, int] = {} # thread_id -> last assigned seq
def _next_seq(self, thread_id: str) -> int:
@ -50,8 +58,10 @@ class MemoryRunEventStore(RunEventStore):
"created_at": created_at or datetime.now(UTC).isoformat(),
}
self._events.setdefault(thread_id, []).append(record)
self._events_by_run.setdefault(thread_id, {}).setdefault(run_id, []).append(record)
if category == "message":
self._messages.setdefault(thread_id, []).append(record)
self._messages_by_run.setdefault(thread_id, {}).setdefault(run_id, []).append(record)
return record
async def put(
@ -100,23 +110,27 @@ class MemoryRunEventStore(RunEventStore):
return messages[-limit:]
async def list_events(self, thread_id, run_id, *, event_types=None, limit=500):
all_events = self._events.get(thread_id, [])
filtered = [e for e in all_events if e["run_id"] == run_id]
# ``_events_by_run`` is already scoped to this run and seq-ordered, so we
# touch only this run's events instead of scanning the whole thread.
run_events = self._events_by_run.get(thread_id, {}).get(run_id, [])
if event_types is not None:
filtered = [e for e in filtered if e["event_type"] in event_types]
return filtered[:limit]
run_events = [e for e in run_events if e["event_type"] in event_types]
return run_events[:limit]
async def list_messages_by_run(self, thread_id, run_id, *, limit=50, before_seq=None, after_seq=None):
all_events = self._events.get(thread_id, [])
filtered = [e for e in all_events if e["run_id"] == run_id and e["category"] == "message"]
if before_seq is not None:
filtered = [e for e in filtered if e["seq"] < before_seq]
# Per-run, messages-only, seq-sorted: the seq window is a contiguous
# slice located with bisect (O(log m_run)) over only this run's
# messages, instead of re-scanning the whole thread's event log.
messages = self._messages_by_run.get(thread_id, {}).get(run_id, [])
lo = 0 if after_seq is None else bisect.bisect_right(messages, after_seq, key=lambda e: e["seq"])
hi = len(messages) if before_seq is None else bisect.bisect_left(messages, before_seq, key=lambda e: e["seq"])
window = messages[lo:hi]
# An ``after_seq`` cursor pages forward (first ``limit``); otherwise
# return the last ``limit`` (the latest page, or the page ending just
# before ``before_seq``). Matches the prior filter-based semantics.
if after_seq is not None:
filtered = [e for e in filtered if e["seq"] > after_seq]
if after_seq is not None:
return filtered[:limit]
else:
return filtered[-limit:] if len(filtered) > limit else filtered
return window[:limit]
return window[-limit:]
async def count_messages(self, thread_id):
return len(self._messages.get(thread_id, []))
@ -124,6 +138,8 @@ class MemoryRunEventStore(RunEventStore):
async def delete_by_thread(self, thread_id):
events = self._events.pop(thread_id, [])
self._messages.pop(thread_id, None)
self._events_by_run.pop(thread_id, None)
self._messages_by_run.pop(thread_id, None)
self._seq_counters.pop(thread_id, None)
return len(events)
@ -136,4 +152,7 @@ class MemoryRunEventStore(RunEventStore):
self._events[thread_id] = remaining
# Keep the message projection in lockstep (same surviving dict objects).
self._messages[thread_id] = [e for e in remaining if e["category"] == "message"]
# Drop the deleted run from the run-keyed projections.
self._events_by_run.get(thread_id, {}).pop(run_id, None)
self._messages_by_run.get(thread_id, {}).pop(run_id, None)
return removed

View File

@ -0,0 +1,130 @@
"""Regression tests for MemoryRunEventStore's run-keyed event/message index.
``list_events`` and ``list_messages_by_run`` are served from per-run
projections so a single run's reads cost O(events-in-run) instead of
re-scanning O(events-in-thread). These tests pin the indexed implementation to
the exact semantics of a brute-force full-thread scan -- including interleaved
trace events (non-contiguous message seqs), both cursors supplied at once, and
index upkeep after ``delete_by_run`` -- so the optimization can never silently
drift from the reference behavior.
"""
import pytest
from deerflow.runtime.events.store.memory import MemoryRunEventStore
def _ref_messages_by_run(records, thread_id, run_id, *, limit=50, before_seq=None, after_seq=None):
"""Brute-force reference: the pre-index full-thread scan it replaced."""
filtered = [e for e in records if e["thread_id"] == thread_id and e["run_id"] == run_id and e["category"] == "message"]
if before_seq is not None:
filtered = [e for e in filtered if e["seq"] < before_seq]
if after_seq is not None:
filtered = [e for e in filtered if e["seq"] > after_seq]
if after_seq is not None:
return filtered[:limit]
return filtered[-limit:] if len(filtered) > limit else filtered
def _ref_events(records, thread_id, run_id, *, event_types=None, limit=500):
filtered = [e for e in records if e["thread_id"] == thread_id and e["run_id"] == run_id]
if event_types is not None:
filtered = [e for e in filtered if e["event_type"] in event_types]
return filtered[:limit]
async def _seed(store):
"""Two runs interleaved within one thread; messages and traces mixed so
each run's message seqs are non-contiguous (the bisect must handle gaps)."""
plan = [
("run-a", "message"),
("run-a", "trace"),
("run-b", "message"),
("run-a", "message"),
("run-b", "trace"),
("run-b", "message"),
("run-a", "trace"),
("run-a", "message"),
("run-b", "message"),
("run-a", "message"),
("run-b", "message"),
("run-a", "message"),
]
records = []
for i, (run_id, category) in enumerate(plan):
rec = await store.put(thread_id="t1", run_id=run_id, event_type=f"e{i}", category=category, content=str(i))
records.append(rec)
return records
@pytest.mark.anyio
async def test_list_messages_by_run_matches_reference_across_cursors():
store = MemoryRunEventStore()
records = await _seed(store)
seqs = [r["seq"] for r in records]
cursors = [None, 0, *seqs, max(seqs) + 1]
for run_id in ("run-a", "run-b", "run-missing"):
for limit in (1, 2, 3, 50):
for before_seq in cursors:
for after_seq in cursors:
got = await store.list_messages_by_run("t1", run_id, limit=limit, before_seq=before_seq, after_seq=after_seq)
want = _ref_messages_by_run(records, "t1", run_id, limit=limit, before_seq=before_seq, after_seq=after_seq)
assert got == want, (run_id, limit, before_seq, after_seq)
@pytest.mark.anyio
async def test_list_events_matches_reference_with_filters():
store = MemoryRunEventStore()
records = await _seed(store)
all_types = sorted({r["event_type"] for r in records})
for run_id in ("run-a", "run-b", "run-missing"):
assert await store.list_events("t1", run_id) == _ref_events(records, "t1", run_id)
assert await store.list_events("t1", run_id, limit=2) == _ref_events(records, "t1", run_id, limit=2)
for et in all_types:
assert await store.list_events("t1", run_id, event_types=[et]) == _ref_events(records, "t1", run_id, event_types=[et])
@pytest.mark.anyio
async def test_run_keyed_index_partitions_every_event():
"""Every stored event is filed under exactly its (thread, run), each run's
list is seq-ordered, and the union reconstructs the flat event log."""
store = MemoryRunEventStore()
records = await _seed(store)
indexed = [e for run_events in store._events_by_run["t1"].values() for e in run_events]
assert sorted(e["seq"] for e in indexed) == sorted(r["seq"] for r in records)
for run_id, run_events in store._events_by_run["t1"].items():
assert all(e["run_id"] == run_id for e in run_events)
assert [e["seq"] for e in run_events] == sorted(e["seq"] for e in run_events)
for run_id, run_msgs in store._messages_by_run["t1"].items():
assert all(e["run_id"] == run_id and e["category"] == "message" for e in run_msgs)
@pytest.mark.anyio
async def test_run_index_stays_in_lockstep_after_delete_by_run():
store = MemoryRunEventStore()
await _seed(store)
removed = await store.delete_by_run("t1", "run-a")
assert removed == 7 # run-a: 5 messages + 2 traces
# The deleted run vanishes from both per-run reads.
assert await store.list_events("t1", "run-a") == []
assert await store.list_messages_by_run("t1", "run-a") == []
assert "run-a" not in store._events_by_run.get("t1", {})
assert "run-a" not in store._messages_by_run.get("t1", {})
# The surviving run is untouched, and the thread-wide projection agrees.
msgs_b = await store.list_messages_by_run("t1", "run-b")
assert len(msgs_b) == 4
assert all(m["run_id"] == "run-b" for m in msgs_b)
assert all(m["run_id"] == "run-b" for m in await store.list_messages("t1"))
@pytest.mark.anyio
async def test_delete_by_thread_clears_run_indexes():
store = MemoryRunEventStore()
await _seed(store)
await store.delete_by_thread("t1")
assert "t1" not in store._events_by_run
assert "t1" not in store._messages_by_run
assert await store.list_events("t1", "run-a") == []
assert await store.list_messages_by_run("t1", "run-b") == []