From bec47a1898883919ad8c5ea41e3a7443a6890e7f Mon Sep 17 00:00:00 2001 From: Manmit Singh Date: Fri, 3 Jul 2026 04:30:55 +0530 Subject: [PATCH] fix(memory): singleflight LocalBackend init to stop cold-start races (#1691) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Running the proxy with `--memory` against a large context throws a bare `AssertionError` (empty message, ~0.1s elapsed, no upstream call) on every request; dropping `--memory` makes it go away. Per-project backends handed out by `BackendRouter._get_or_create_backend` init lazily on first use. `LocalBackend._ensure_initialized` guarded init with a bare `if not self._initialized:` and no `asyncio.Lock`, so concurrent first callers each kicked off a parallel `HierarchicalMemory.create()`. A slow cold-start (>2s on the `pytorch_mps` embedder) cancelled by the outer 2s memory-context `wait_for` left the backend half-built (`_hierarchical_memory` still `None`), so the retry tripped `assert self._hierarchical_memory is not None` (local.py:237/385/...) — the empty-message crash. `MemoryHandler._ensure_initialized` already uses a double-checked `asyncio.Lock`; the per-project `LocalBackend` never got the same treatment. Closes #1678 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - Add a lazily-created `asyncio.Lock` singleflight with a double-checked flag to `LocalBackend._ensure_initialized`, mirroring the existing `MemoryHandler` pattern — concurrent first callers await one init instead of racing N. - On `CancelledError` (e.g. the outer `wait_for` timeout mid cold-start), reset `_hierarchical_memory`/`_graph`/`_initialized` and re-raise, so a cancelled init never leaves a half-built backend for the next request to assert on. - Move the init body verbatim into `_init_locked()` (called with the lock held); the large diff is the dedent, no logic change. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text $ pytest tests/test_local_backend_init_race.py tests/test_memory_handler_concurrent_init.py -q tests/test_local_backend_init_race.py .. [ 20%] tests/test_memory_handler_concurrent_init.py .....s.. [100%] 9 passed, 1 skipped in 0.50s $ ruff check headroom/memory/backends/local.py tests/test_local_backend_init_race.py All checks passed! $ ruff format --check headroom/memory/backends/local.py tests/test_local_backend_init_race.py 2 files already formatted $ mypy headroom/memory/backends/local.py Success: no issues found in 1 source file ``` ## Real Behavior Proof - Environment: macOS (arm64), Python 3.14, local `.venv`. - Exact command / steps: `pytest tests/test_local_backend_init_race.py tests/test_memory_handler_concurrent_init.py -q`. First test spawns 10 concurrent first callers against a `LocalBackend` with a patched slow `HierarchicalMemory.create` and asserts `create` runs exactly once; second cancels a cold-start via an outer `asyncio.wait_for` timeout, asserts state resets to `None`/uninitialized, then a later call re-inits cleanly. - Observed result: both pass; `create` is called once under contention, and a cancelled init leaves no half-built backend. - Not tested: no end-to-end repro of the original `--memory` crash against a real large context / GPU embedder — the race is reproduced deterministically at the unit level instead. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable ## Additional Notes Docs/CHANGELOG unchanged — internal concurrency fix with no user-facing API or behavior change beyond removing the crash. --- headroom/memory/backends/local.py | 136 +++++++++++++++++--------- tests/test_local_backend_init_race.py | 90 +++++++++++++++++ 2 files changed, 180 insertions(+), 46 deletions(-) create mode 100644 tests/test_local_backend_init_race.py diff --git a/headroom/memory/backends/local.py b/headroom/memory/backends/local.py index b86abb36f..a74fa12b8 100644 --- a/headroom/memory/backends/local.py +++ b/headroom/memory/backends/local.py @@ -13,6 +13,7 @@ single-process production deployments. from __future__ import annotations +import asyncio import logging import uuid from dataclasses import dataclass, field @@ -116,6 +117,19 @@ class LocalBackend: self._initialized = False self._hierarchical_memory: HierarchicalMemory | None = None self._graph: InMemoryGraphStore | SQLiteGraphStore | None = None + # Async singleflight guard for lazy init. Per-project backends handed + # out by BackendRouter init lazily on first use; concurrent first + # callers must land on ONE init (double-checked pattern below) instead + # of racing N partial inits that leave ``_hierarchical_memory`` None + # and trip the ``assert`` guards downstream. Created lazily so the + # backend can be constructed before an event loop exists. + self._init_lock: asyncio.Lock | None = None + + def _get_init_lock(self) -> asyncio.Lock: + """Lazily create the init lock bound to the running event loop.""" + if self._init_lock is None: + self._init_lock = asyncio.Lock() + return self._init_lock async def _ensure_initialized(self) -> None: """Ensure the backend is initialized with all components. @@ -123,62 +137,92 @@ class LocalBackend: Creates the HierarchicalMemory system and graph store on first use. Uses SQLiteGraphStore (bounded, persistent) when graph_persist=True, or InMemoryGraphStore (unbounded, volatile) when graph_persist=False. + + Singleflight via ``self._init_lock`` with a double-checked flag: + concurrent first callers await the same cold-start (which can exceed + a second on the ``pytorch_mps`` embedder) rather than each kicking off + a parallel init. If a slow init is cancelled (e.g. an outer + ``asyncio.wait_for`` timeout), state is reset so a later call retries + cleanly and ``CancelledError`` is re-raised rather than leaving the + backend half-built. """ - if not self._initialized: - from headroom.memory import HierarchicalMemory, MemoryConfig - from headroom.memory.config import EmbedderBackend + # Fast path: already initialized, no lock contention. + if self._initialized: + return - # Map string embedder_backend to enum - embedder_backend_map = { - "local": EmbedderBackend.LOCAL, - "onnx": EmbedderBackend.ONNX, - "openai": EmbedderBackend.OPENAI, - "ollama": EmbedderBackend.OLLAMA, - } - embedder_backend = embedder_backend_map.get( - self._config.embedder_backend, EmbedderBackend.LOCAL - ) + lock = self._get_init_lock() + async with lock: + # Double-check after acquiring the lock — another task may have + # finished the init while we were waiting. + if self._initialized: + return + try: + await self._init_locked() + except asyncio.CancelledError: + # Cancellation (e.g. wait_for timeout) can leave a partial + # backend. Reset so the next call re-inits from scratch and + # never sees a half-built ``_hierarchical_memory``. + self._hierarchical_memory = None + self._graph = None + self._initialized = False + raise - mem_config = MemoryConfig( - db_path=Path(self._config.db_path), - embedder_backend=embedder_backend, - embedder_model=self._config.embedder_model, - vector_dimension=self._config.vector_dimension, - openai_api_key=self._config.openai_api_key, - ollama_base_url=self._config.ollama_base_url, - cache_enabled=self._config.cache_enabled, - cache_max_size=self._config.cache_max_size, - ) + async def _init_locked(self) -> None: + """Actual init body. Must be called with ``_init_lock`` held.""" + from headroom.memory import HierarchicalMemory, MemoryConfig + from headroom.memory.config import EmbedderBackend - self._hierarchical_memory = await HierarchicalMemory.create(mem_config) + # Map string embedder_backend to enum + embedder_backend_map = { + "local": EmbedderBackend.LOCAL, + "onnx": EmbedderBackend.ONNX, + "openai": EmbedderBackend.OPENAI, + "ollama": EmbedderBackend.OLLAMA, + } + embedder_backend = embedder_backend_map.get( + self._config.embedder_backend, EmbedderBackend.LOCAL + ) - # Choose graph store based on config - if self._config.graph_persist: - from headroom.memory.adapters.sqlite_graph import SQLiteGraphStore + mem_config = MemoryConfig( + db_path=Path(self._config.db_path), + embedder_backend=embedder_backend, + embedder_model=self._config.embedder_model, + vector_dimension=self._config.vector_dimension, + openai_api_key=self._config.openai_api_key, + ollama_base_url=self._config.ollama_base_url, + cache_enabled=self._config.cache_enabled, + cache_max_size=self._config.cache_max_size, + ) - # Derive graph db path from main db path if not specified - if self._config.graph_db_path: - graph_db_path = self._config.graph_db_path - else: - # "memory.db" -> "memory_graph.db" - db_path = Path(self._config.db_path) - graph_db_path = str(db_path.parent / f"{db_path.stem}_graph{db_path.suffix}") + self._hierarchical_memory = await HierarchicalMemory.create(mem_config) - self._graph = SQLiteGraphStore( - db_path=graph_db_path, - page_cache_size_kb=self._config.graph_cache_size_kb, - ) - logger.info( - f"LocalBackend: Using SQLiteGraphStore at {graph_db_path} " - f"(cache: {self._config.graph_cache_size_kb}KB)" - ) + # Choose graph store based on config + if self._config.graph_persist: + from headroom.memory.adapters.sqlite_graph import SQLiteGraphStore + + # Derive graph db path from main db path if not specified + if self._config.graph_db_path: + graph_db_path = self._config.graph_db_path else: - from headroom.memory.adapters.graph import InMemoryGraphStore + # "memory.db" -> "memory_graph.db" + db_path = Path(self._config.db_path) + graph_db_path = str(db_path.parent / f"{db_path.stem}_graph{db_path.suffix}") - self._graph = InMemoryGraphStore() - logger.info("LocalBackend: Using InMemoryGraphStore (unbounded)") + self._graph = SQLiteGraphStore( + db_path=graph_db_path, + page_cache_size_kb=self._config.graph_cache_size_kb, + ) + logger.info( + f"LocalBackend: Using SQLiteGraphStore at {graph_db_path} " + f"(cache: {self._config.graph_cache_size_kb}KB)" + ) + else: + from headroom.memory.adapters.graph import InMemoryGraphStore - self._initialized = True + self._graph = InMemoryGraphStore() + logger.info("LocalBackend: Using InMemoryGraphStore (unbounded)") + + self._initialized = True # ========================================================================= # Core Memory Operations diff --git a/tests/test_local_backend_init_race.py b/tests/test_local_backend_init_race.py new file mode 100644 index 000000000..ad4b37d6b --- /dev/null +++ b/tests/test_local_backend_init_race.py @@ -0,0 +1,90 @@ +"""Concurrency + cancellation tests for LocalBackend._ensure_initialized. + +Per-project backends handed out by ``BackendRouter`` init lazily on first +use. Without a singleflight guard, concurrent first callers each kick off a +parallel init; a slow cold-start (e.g. the ``pytorch_mps`` embedder, >2s) +that is cancelled by an outer ``asyncio.wait_for`` timeout can leave the +backend half-built (``_hierarchical_memory`` still ``None``), tripping the +bare ``assert self._hierarchical_memory is not None`` guards on the retry. + +Covers: +- N concurrent first callers trigger exactly one HierarchicalMemory.create. +- A cancelled cold-start resets state so a later call re-inits cleanly. +""" + +from __future__ import annotations + +import asyncio +from typing import Any +from unittest.mock import MagicMock + +import pytest + +from headroom.memory.backends.local import LocalBackend, LocalBackendConfig + + +def _backend(tmp_path) -> LocalBackend: + return LocalBackend( + LocalBackendConfig( + db_path=str(tmp_path / "memory.db"), + graph_persist=False, # InMemoryGraphStore — no SQLite/embedder needed + ) + ) + + +@pytest.mark.asyncio +async def test_concurrent_ensure_initialized_runs_init_once(tmp_path, monkeypatch): + hits = {"n": 0} + release = asyncio.Event() + + async def fake_create(config: Any) -> Any: + hits["n"] += 1 + # Simulate a slow cold-start so concurrent callers pile up on the + # lock. Without singleflight, hits would exceed 1. + await release.wait() + return MagicMock(name="HierarchicalMemory") + + monkeypatch.setattr("headroom.memory.HierarchicalMemory.create", fake_create) + + backend = _backend(tmp_path) + tasks = [asyncio.create_task(backend._ensure_initialized()) for _ in range(10)] + await asyncio.sleep(0) # let all tasks reach the lock + release.set() + await asyncio.gather(*tasks) + + assert hits["n"] == 1 + assert backend._initialized is True + assert backend._hierarchical_memory is not None + + +@pytest.mark.asyncio +async def test_cancelled_cold_start_resets_state_and_retries(tmp_path, monkeypatch): + attempts = {"n": 0} + first_started = asyncio.Event() + block_first = asyncio.Event() + + async def fake_create(config: Any) -> Any: + attempts["n"] += 1 + if attempts["n"] == 1: + first_started.set() + await block_first.wait() # never released → this attempt is cancelled + return MagicMock(name="HierarchicalMemory") + + monkeypatch.setattr("headroom.memory.HierarchicalMemory.create", fake_create) + + backend = _backend(tmp_path) + + # First cold-start gets cancelled by an outer timeout mid-init. + with pytest.raises((asyncio.TimeoutError, asyncio.CancelledError)): + await asyncio.wait_for(backend._ensure_initialized(), timeout=0.05) + await first_started.wait() + + # State must be reset — no half-built backend left behind. + assert backend._initialized is False + assert backend._hierarchical_memory is None + + # A subsequent call re-inits cleanly. + await backend._ensure_initialized() + assert backend._initialized is True + assert backend._hierarchical_memory is not None + assert attempts["n"] == 2