From e6333719ee37d79a62ba8ad52db89ea7ecf7e6af Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Mon, 23 Mar 2026 16:06:04 +0100 Subject: [PATCH] fix(entity_resolver): prevent _pending_stats/_pending_cooccurrences memory leak (#662) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(entity_resolver): prevent _pending_stats/_pending_cooccurrences memory leak Add discard_pending_stats() to EntityResolver to clean up both pending dicts for the current task key. Call it at the start of each _run_db_work attempt so that exceptions between accumulation and flush_pending_stats() — including deadlock retries — never leave stale entries keyed by recycled task IDs. Fixes #660 * test(entity_resolver): add unit tests for discard_pending_stats() Covers: clears both dicts for current task, is idempotent when empty, and does not touch entries belonging to other task keys. No database required — purely in-memory logic. --- .../hindsight_api/engine/entity_resolver.py | 13 ++++++ .../engine/retain/orchestrator.py | 4 ++ .../tests/test_entity_resolver.py | 44 +++++++++++++++++++ 3 files changed, 61 insertions(+) diff --git a/hindsight-api-slim/hindsight_api/engine/entity_resolver.py b/hindsight-api-slim/hindsight_api/engine/entity_resolver.py index 734adb2c..47ee5be3 100644 --- a/hindsight-api-slim/hindsight_api/engine/entity_resolver.py +++ b/hindsight-api-slim/hindsight_api/engine/entity_resolver.py @@ -86,6 +86,19 @@ class EntityResolver: task = asyncio.current_task() return id(task) if task is not None else 0 + def discard_pending_stats(self) -> None: + """ + Discard accumulated entity stats and co-occurrence counts for the current task. + + Call this on any exception path between resolve_entities_batch / + link_units_to_entities_batch and flush_pending_stats() to prevent the + per-task dicts from growing unbounded when tasks fail before flushing. + Safe to call even if no entries exist for the current task. + """ + key = self._task_key() + self._pending_stats.pop(key, None) + self._pending_cooccurrences.pop(key, None) + async def flush_pending_stats(self) -> None: """ Flush accumulated entity stats and co-occurrence counts for the current task. diff --git a/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py b/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py index 3bc54c77..f5c0cd3a 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py @@ -292,6 +292,10 @@ async def retain_batch( pf.document_id = None pf.chunk_id = None + # Discard any leftover pending stats from a previous failed attempt so + # retries don't double-count or accumulate unbounded state. + entity_resolver.discard_pending_stats() + async with acquire_with_retry(pool) as conn: async with conn.transaction(): # Handle document tracking for all documents diff --git a/hindsight-api-slim/tests/test_entity_resolver.py b/hindsight-api-slim/tests/test_entity_resolver.py index 988abaa5..4b31da01 100644 --- a/hindsight-api-slim/tests/test_entity_resolver.py +++ b/hindsight-api-slim/tests/test_entity_resolver.py @@ -11,6 +11,50 @@ import pytest from hindsight_api.engine.entity_resolver import EntityResolver from hindsight_api.pg0 import resolve_database_url +# --------------------------------------------------------------------------- +# Unit tests for discard_pending_stats() — no database required +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_discard_pending_stats_clears_both_dicts(): + """discard_pending_stats() must remove entries for the current task from + both _pending_stats and _pending_cooccurrences.""" + resolver = EntityResolver(pool=None) # type: ignore[arg-type] + key = resolver._task_key() + + resolver._pending_stats[key] = [object()] # type: ignore[list-item] + resolver._pending_cooccurrences[key] = [object()] # type: ignore[list-item] + + resolver.discard_pending_stats() + + assert key not in resolver._pending_stats + assert key not in resolver._pending_cooccurrences + + +@pytest.mark.asyncio +async def test_discard_pending_stats_is_idempotent(): + """Calling discard_pending_stats() when nothing is pending must not raise.""" + resolver = EntityResolver(pool=None) # type: ignore[arg-type] + resolver.discard_pending_stats() + resolver.discard_pending_stats() # second call — still safe + + +@pytest.mark.asyncio +async def test_discard_pending_stats_does_not_affect_other_task_keys(): + """discard_pending_stats() must only remove the current task's entries, + leaving entries keyed under other task IDs untouched.""" + resolver = EntityResolver(pool=None) # type: ignore[arg-type] + other_key = -1 # A fake key that can never be a real task id + + resolver._pending_stats[other_key] = [object()] # type: ignore[list-item] + resolver._pending_cooccurrences[other_key] = [object()] # type: ignore[list-item] + + resolver.discard_pending_stats() # discards current task's key only + + assert other_key in resolver._pending_stats, "other task's stats must be preserved" + assert other_key in resolver._pending_cooccurrences, "other task's cooccurrences must be preserved" + @pytest.mark.asyncio async def test_resolve_entities_batch_handles_unicode_lower_conflicts(pg0_db_url):