From a239add5ee5d4b5f16ad7a5969609c16f15b79e2 Mon Sep 17 00:00:00 2001 From: Arpit Jain Date: Sun, 19 Jul 2026 00:58:56 +0900 Subject: [PATCH] fix(backend): bound the in-memory report registry The completed-report registry (_report_registry) stored a full Analyzer, including its parsed-pcap pandas frames, on every analysis and was never evicted or capped. Its companion dict _progress_queues is popped in the WebSocket cleanup path, but the heavier report registry kept growing for the lifetime of the worker, so repeated analyses steadily increased memory. Back the registry with an OrderedDict capped at MAX_STORED_REPORTS and evict the oldest entry once the cap is reached. Reads through _require_analyzer move the entry to the end so an actively used report is not the first evicted, which keeps the drilldown routes working for recent reports. Add tests covering the size bound and the least-recently-used eviction order. Signed-off-by: Arpit Jain --- backend/src/app/main.py | 21 +++++++++++++++++-- backend/tests/test_main.py | 42 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 61 insertions(+), 2 deletions(-) diff --git a/backend/src/app/main.py b/backend/src/app/main.py index be9eb1e..08d95fd 100644 --- a/backend/src/app/main.py +++ b/backend/src/app/main.py @@ -13,6 +13,7 @@ import uuid import warnings from asyncio import Queue, Task +from collections import OrderedDict from collections.abc import Mapping from concurrent.futures import ThreadPoolExecutor from pathlib import Path @@ -44,11 +45,25 @@ # Per-session progress queues - keyed by session_id, cleaned up after analysis completes _progress_queues: dict[str, asyncio.Queue[dict[str, object]]] = {} -_report_registry: dict[str, Analyzer] = {} + +# Completed reports, keyed by report_id, used by the drilldown routes. +# Each stored Analyzer holds the full parsed-pcap pandas frames, so the registry +# is bounded and evicts the least-recently-used report once the cap is reached. +# Ordered oldest-to-newest; reads move the entry to the end (see _require_analyzer). +_report_registry: OrderedDict[str, Analyzer] = OrderedDict() _executor = ThreadPoolExecutor(max_workers=4) # Constants PCAP_CHUNK_SIZE = 1024 * 1024 # 1MB +MAX_STORED_REPORTS = 32 # Cap on retained reports before the oldest is evicted. + + +def _store_analyzer(report_id: str, analyzer: Analyzer) -> None: + """Register an analyzer for later drilldown, evicting the oldest if over cap.""" + _report_registry[report_id] = analyzer + _report_registry.move_to_end(report_id) + while len(_report_registry) > MAX_STORED_REPORTS: + _report_registry.popitem(last=False) def sanitize_for_json(obj: JSONValue) -> JSONValue: @@ -129,7 +144,7 @@ def run_analysis( ) report_id = str(uuid.uuid4()) - _report_registry[report_id] = analyzer + _store_analyzer(report_id, analyzer) _emit(session_id, "report", 80, "Generating report...") report = Report(analyzer, report_id=report_id) @@ -292,6 +307,8 @@ def _require_analyzer(report_id: str) -> Analyzer: analyzer = _report_registry.get(report_id) if analyzer is None: raise HTTPException(status_code=404, detail=f"Report not found: {report_id}") + # Mark this report as recently used so active drilldowns are not evicted first. + _report_registry.move_to_end(report_id) return analyzer diff --git a/backend/tests/test_main.py b/backend/tests/test_main.py index b14152b..d4651c7 100644 --- a/backend/tests/test_main.py +++ b/backend/tests/test_main.py @@ -21,8 +21,12 @@ # or if the src/app directory is added to PYTHONPATH. # If not, we might need to adjust sys.path or use a different import strategy. from src.app.main import ( + MAX_STORED_REPORTS, _emit, _progress_queues, + _report_registry, + _require_analyzer, + _store_analyzer, app, logger, run_analysis, @@ -39,6 +43,44 @@ def clear_progress_queues(): _progress_queues.clear() +@pytest.fixture +def clear_report_registry(): + """Clear the report registry around a test case.""" + _report_registry.clear() + yield + _report_registry.clear() + + +def test_store_analyzer_bounds_registry(clear_report_registry): + """Keep the report registry capped when more reports are stored than the limit.""" + total = MAX_STORED_REPORTS + 10 + for i in range(total): + _store_analyzer(f"report-{i}", Mock()) + + # Without a bound the registry would hold every insert; it must stay capped. + assert len(_report_registry) == MAX_STORED_REPORTS + + # The oldest reports are evicted first; only the most recent survive. + assert "report-0" not in _report_registry + assert f"report-{total - 1}" in _report_registry + + +def test_require_analyzer_keeps_recently_used(clear_report_registry): + """Evict the least-recently-used report rather than a recently accessed one.""" + for i in range(MAX_STORED_REPORTS): + _store_analyzer(f"report-{i}", Mock()) + + # Touch the oldest report so it counts as recently used. + _require_analyzer("report-0") + + # One more insert should evict report-1 (now the oldest) and spare report-0. + _store_analyzer("report-new", Mock()) + + assert len(_report_registry) == MAX_STORED_REPORTS + assert "report-0" in _report_registry + assert "report-1" not in _report_registry + + @pytest.mark.asyncio async def test_emit_success(): """Enqueue a progress event for an active session."""