diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index 0bfbfec7..a39ade8e 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -59,6 +59,11 @@ class StorageConfig(BaseModel): data_dir: Path = Field(default_factory=get_default_data_dir) auto_vacuum: bool = True vacuum_retention_seconds: int = 86400 + # Upper bound (seconds) on a single vacuum/optimize pass. ``None`` leaves it + # unbounded (the historical behavior). A finite value converts a stuck + # compaction into a logged, skipped pass instead of an unbounded hang that + # would also wedge client teardown (which drains background vacuums). + vacuum_timeout_seconds: float | None = None class LanceDBConfig(BaseModel): diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 50d4a73a..8ba21796 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -571,14 +571,29 @@ class Store: # Evaluate config at runtime to allow dynamic changes if retention_seconds is None: retention_seconds = self._config.storage.vacuum_retention_seconds - # Perform maintenance per table using optimize() with configurable retention - retention = timedelta(seconds=retention_seconds) - for table in self._tables().values(): - await table.optimize( - cleanup_older_than=await self._tag_safe_retention( - table, retention + # Bound the pass so a stuck compaction degrades to a logged, + # skipped pass rather than an unbounded await -- which would also + # wedge client teardown (__aexit__ drains background vacuums). + # timeout=None leaves it unbounded (historical behavior). + timeout = self._config.storage.vacuum_timeout_seconds + async with asyncio.timeout(timeout): + # Perform maintenance per table using optimize() with configurable retention + retention = timedelta(seconds=retention_seconds) + for table in self._tables().values(): + await table.optimize( + cleanup_older_than=await self._tag_safe_retention( + table, retention + ) ) - ) + except TimeoutError: + # The pass is best-effort maintenance; a timeout is not fatal. + # NOTE: cancellation only takes effect if the underlying + # operation is cancellable -- a fully-unresponsive native future + # may still linger, but the caller's coroutine is freed. + logger.warning( + "Vacuum timed out after %ss; skipping this pass", + self._config.storage.vacuum_timeout_seconds, + ) except OSError as e: # Resource errors (e.g. disk pressure) skip the pass; lance # errors surface as RuntimeError and must not be swallowed — diff --git a/tests/test_vacuum_debounce.py b/tests/test_vacuum_debounce.py index 83a96dbc..e87633ab 100644 --- a/tests/test_vacuum_debounce.py +++ b/tests/test_vacuum_debounce.py @@ -1,4 +1,5 @@ import asyncio +import logging import pytest @@ -18,6 +19,58 @@ def _docling_doc(name: str, text: str): return doc +@pytest.mark.asyncio +async def test_vacuum_times_out_and_skips(temp_db_path, monkeypatch, caplog): + """A stuck optimize is bounded by storage.vacuum_timeout_seconds: the pass + is skipped (logged), the call returns instead of hanging, and no exception + propagates.""" + + class _SlowTable: + async def optimize(self, **_kwargs): + await asyncio.sleep(30) # far longer than the timeout + + async with HaikuRAG(temp_db_path, create=True) as client: + store = client.store + # setattr (not direct assignment) so the shared global Config is restored. + monkeypatch.setattr(store._config.storage, "vacuum_timeout_seconds", 0.05) + monkeypatch.setattr(store, "_tables", lambda: {"t": _SlowTable()}) + + async def _passthrough_retention(_table, retention): + return retention + + monkeypatch.setattr(store, "_tag_safe_retention", _passthrough_retention) + + with caplog.at_level(logging.WARNING, logger="haiku.rag.store.engine"): + # wait_for guards the test itself from hanging if the bound fails. + await asyncio.wait_for(store.vacuum(), timeout=5) + + assert "Vacuum timed out" in caplog.text + + +@pytest.mark.asyncio +async def test_vacuum_unbounded_when_timeout_none(temp_db_path, monkeypatch): + """timeout=None preserves the historical unbounded behavior (optimize runs + to completion).""" + ran: list[int] = [] + + class _Table: + async def optimize(self, **_kwargs): + ran.append(1) + + async with HaikuRAG(temp_db_path, create=True) as client: + store = client.store + monkeypatch.setattr(store._config.storage, "vacuum_timeout_seconds", None) + monkeypatch.setattr(store, "_tables", lambda: {"t": _Table()}) + + async def _passthrough_retention(_table, retention): + return retention + + monkeypatch.setattr(store, "_tag_safe_retention", _passthrough_retention) + + await store.vacuum() + assert ran == [1] + + @pytest.mark.asyncio async def test_schedule_vacuum_is_debounced(temp_db_path, monkeypatch): """Rapid writes within the throttle window schedule only one background