From 346164d740ef08c131b0d65722a2448d3b402476 Mon Sep 17 00:00:00 2001 From: ly-wang19 <94427531+ly-wang19@users.noreply.github.com> Date: Wed, 24 Jun 2026 10:57:45 +0800 Subject: [PATCH] perf(runtime): index MemoryRunStore by thread_id to avoid O(n) scans (#3562) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * perf(runtime): index MemoryRunStore by thread_id to avoid O(n) scans MemoryRunStore is the default run backend (database.backend=memory) and backs RunManager.list_by_thread, which calls it on every thread-runs query to hydrate persisted runs. list_by_thread scanned every run in the store (O(total runs)) to filter by thread_id, so listing one thread's runs got linearly slower as unrelated runs accumulated across all threads. Add a thread_id -> insertion-ordered run_id set secondary index, maintained in lockstep with _runs in put()/delete(), and use it in list_by_thread for an O(runs-in-thread) lookup. This mirrors the index RunManager already keeps over its own in-memory records (#3499); the store extraction — whose docstring notes it is "Equivalent to the original RunManager._runs dict behavior" — did not carry the index across. Behavior is unchanged: same user_id filtering, newest-first ordering, and limit. The store runs each method without awaits on the event loop, so the index and _runs stay consistent without a lock. Extends tests/test_persistence_scaffold.py::TestMemoryRunStore with coverage for unknown-thread, newest-first ordering, limit, and index cleanup on delete (including empty-bucket removal). * perf(runtime): route aggregate_tokens_by_thread through the thread index too list_by_thread already uses the _runs_by_thread index this PR adds, but aggregate_tokens_by_thread (the /token-usage endpoint) still scanned every run in the process to pick out one thread's runs. Route it through the same index for an O(runs-in-thread) lookup, completing the thread-scoped read coverage. Co-Authored-By: Claude Opus 4.8 (1M context) --------- Co-authored-by: ly-wang19 Co-authored-by: Claude Opus 4.8 (1M context) --- .../deerflow/runtime/runs/store/memory.py | 35 +++++++- backend/tests/test_persistence_scaffold.py | 90 +++++++++++++++++++ 2 files changed, 122 insertions(+), 3 deletions(-) diff --git a/backend/packages/harness/deerflow/runtime/runs/store/memory.py b/backend/packages/harness/deerflow/runtime/runs/store/memory.py index 0cb4e82cc..902f847ee 100644 --- a/backend/packages/harness/deerflow/runtime/runs/store/memory.py +++ b/backend/packages/harness/deerflow/runtime/runs/store/memory.py @@ -14,6 +14,23 @@ from deerflow.runtime.runs.store.base import RunStore class MemoryRunStore(RunStore): def __init__(self) -> None: self._runs: dict[str, dict[str, Any]] = {} + # Secondary index: thread_id -> insertion-ordered run_id set (a dict is + # used as an ordered set), maintained in lockstep with ``_runs`` so + # per-thread queries avoid O(total in-memory runs) full scans. Mirrors + # the index ``RunManager`` keeps over its own in-memory records. + self._runs_by_thread: dict[str, dict[str, None]] = {} + + def _index_run(self, run_id: str, thread_id: str) -> None: + """Register *run_id* under *thread_id* in the secondary index.""" + self._runs_by_thread.setdefault(thread_id, {})[run_id] = None + + def _unindex_run(self, run_id: str, thread_id: str) -> None: + """Drop *run_id* from the *thread_id* bucket, removing the bucket when empty.""" + bucket = self._runs_by_thread.get(thread_id) + if bucket is not None: + bucket.pop(run_id, None) + if not bucket: + self._runs_by_thread.pop(thread_id, None) async def put( self, @@ -45,6 +62,7 @@ class MemoryRunStore(RunStore): "created_at": created_at or now, "updated_at": now, } + self._index_run(run_id, thread_id) async def get(self, run_id, *, user_id=None): run = self._runs.get(run_id) @@ -55,7 +73,13 @@ class MemoryRunStore(RunStore): return run async def list_by_thread(self, thread_id, *, user_id=None, limit=100): - results = [r for r in self._runs.values() if r["thread_id"] == thread_id and (user_id is None or r.get("user_id") == user_id)] + # Use the thread index for an O(runs-in-thread) lookup instead of + # scanning every run. ``self._runs.get`` is defense-in-depth: it drops a + # stale id still in the index but already gone from ``_runs``. + run_ids = self._runs_by_thread.get(thread_id) + if not run_ids: + return [] + results = [run for run_id in run_ids if (run := self._runs.get(run_id)) is not None and (user_id is None or run.get("user_id") == user_id)] results.sort(key=lambda r: r["created_at"], reverse=True) return results[:limit] @@ -74,7 +98,9 @@ class MemoryRunStore(RunStore): self._runs[run_id]["updated_at"] = datetime.now(UTC).isoformat() async def delete(self, run_id): - self._runs.pop(run_id, None) + run = self._runs.pop(run_id, None) + if run is not None: + self._unindex_run(run_id, run["thread_id"]) async def update_run_completion(self, run_id, *, status, **kwargs): if run_id in self._runs: @@ -107,7 +133,10 @@ class MemoryRunStore(RunStore): async def aggregate_tokens_by_thread(self, thread_id: str, *, include_active: bool = False) -> dict[str, Any]: statuses = ("success", "error", "running") if include_active else ("success", "error") - completed = [r for r in self._runs.values() if r["thread_id"] == thread_id and r.get("status") in statuses] + # Use the thread index for an O(runs-in-thread) lookup instead of + # scanning every run in the process (mirrors ``list_by_thread``). + run_ids = self._runs_by_thread.get(thread_id) or () + completed = [run for run_id in run_ids if (run := self._runs.get(run_id)) is not None and run.get("status") in statuses] by_model: dict[str, dict] = {} for r in completed: usage_by_model = r.get("token_usage_by_model") or {} diff --git a/backend/tests/test_persistence_scaffold.py b/backend/tests/test_persistence_scaffold.py index e0bca3eef..baba88c4b 100644 --- a/backend/tests/test_persistence_scaffold.py +++ b/backend/tests/test_persistence_scaffold.py @@ -131,6 +131,96 @@ class TestMemoryRunStore: async def test_delete_nonexistent_is_noop(self, store): await store.delete("nope") # should not raise + @pytest.mark.anyio + async def test_list_by_thread_unknown_thread_is_empty(self, store): + await store.put("r1", thread_id="t1") + assert await store.list_by_thread("missing") == [] + + @pytest.mark.anyio + async def test_list_by_thread_newest_first(self, store): + await store.put("r1", thread_id="t1", created_at="2024-01-01T00:00:00+00:00") + await store.put("r2", thread_id="t1", created_at="2024-01-03T00:00:00+00:00") + await store.put("r3", thread_id="t1", created_at="2024-01-02T00:00:00+00:00") + rows = await store.list_by_thread("t1") + assert [r["run_id"] for r in rows] == ["r2", "r3", "r1"] + + @pytest.mark.anyio + async def test_list_by_thread_respects_limit(self, store): + for i in range(5): + await store.put(f"r{i}", thread_id="t1", created_at=f"2024-01-0{i + 1}T00:00:00+00:00") + rows = await store.list_by_thread("t1", limit=2) + assert [r["run_id"] for r in rows] == ["r4", "r3"] + + @pytest.mark.anyio + async def test_delete_keeps_thread_index_consistent(self, store): + await store.put("r1", thread_id="t1") + await store.put("r2", thread_id="t1") + await store.delete("r1") + rows = await store.list_by_thread("t1") + assert [r["run_id"] for r in rows] == ["r2"] + # deleting the last run in a thread drops the now-empty index bucket + await store.delete("r2") + assert await store.list_by_thread("t1") == [] + assert "t1" not in store._runs_by_thread + + @pytest.mark.anyio + async def test_aggregate_tokens_by_thread_scopes_to_thread(self, store): + await store.put("r1", thread_id="t1") + await store.update_run_completion("r1", status="success", model_name="m-a", total_tokens=100) + await store.put("r2", thread_id="t1") + await store.update_run_completion("r2", status="error", model_name="m-a", total_tokens=20) + await store.put("r3", thread_id="t2") + await store.update_run_completion("r3", status="success", model_name="m-b", total_tokens=999) + + agg = await store.aggregate_tokens_by_thread("t1") + assert agg["total_tokens"] == 120 # the other thread's run is excluded + assert agg["total_runs"] == 2 + assert agg["by_model"]["m-a"] == {"tokens": 120, "runs": 2} + assert "m-b" not in agg["by_model"] + + @pytest.mark.anyio + async def test_aggregate_tokens_by_thread_excludes_active_unless_requested(self, store): + await store.put("r1", thread_id="t1") + await store.update_run_completion("r1", status="success", total_tokens=10) + await store.put("r2", thread_id="t1") + await store.update_run_completion("r2", status="running", total_tokens=5) + + assert (await store.aggregate_tokens_by_thread("t1"))["total_tokens"] == 10 + assert (await store.aggregate_tokens_by_thread("t1", include_active=True))["total_tokens"] == 15 + + @pytest.mark.anyio + async def test_aggregate_tokens_by_thread_unknown_thread_is_zero(self, store): + await store.put("r1", thread_id="t1") + await store.update_run_completion("r1", status="success", total_tokens=10) + agg = await store.aggregate_tokens_by_thread("missing") + assert agg["total_tokens"] == 0 + assert agg["total_runs"] == 0 + assert agg["by_model"] == {} + + @pytest.mark.anyio + async def test_aggregate_tokens_by_thread_matches_full_scan_reference(self, store): + plan = [ + ("r0", "t1", "success", "m-a", 10), + ("r1", "t1", "error", "m-b", 20), + ("r2", "t1", "running", "m-a", 7), + ("r3", "t2", "success", "m-a", 999), + ("r4", "t1", "pending", "m-a", 3), + ] + for run_id, thread_id, status, model, tokens in plan: + await store.put(run_id, thread_id=thread_id) + await store.update_run_completion(run_id, status=status, model_name=model, total_tokens=tokens) + + def _reference(thread_id, include_active): + statuses = ("success", "error", "running") if include_active else ("success", "error") + completed = [r for r in store._runs.values() if r["thread_id"] == thread_id and r.get("status") in statuses] + return len(completed), sum(r.get("total_tokens", 0) for r in completed) + + for thread_id in ("t1", "t2", "missing"): + for include_active in (False, True): + agg = await store.aggregate_tokens_by_thread(thread_id, include_active=include_active) + ref_runs, ref_tokens = _reference(thread_id, include_active) + assert (agg["total_runs"], agg["total_tokens"]) == (ref_runs, ref_tokens), (thread_id, include_active) + @pytest.mark.anyio async def test_list_pending(self, store): await store.put("r1", thread_id="t1", status="pending")