Throttle background auto-vacuum to at most once per 5 minutes

This commit is contained in:
Yiorgis Gozadinos 2026-06-11 11:01:31 +03:00
parent df8af54298
commit 5d5d87d44c
No known key found for this signature in database
4 changed files with 92 additions and 13 deletions

View file

@ -4,6 +4,7 @@
### Changed ### Changed
- Mutable document attributes (`uri`, `title`, `metadata`, `created_at`, `updated_at`) moved from the `documents` table into a new `document_meta` table (1:1 on `document_id`); metadata/title/`source_revision` updates no longer rewrite the docling blobs. Migration `v0_58_0` relocates existing data and runs a one-time `vacuum` to reclaim prior bloat. - Mutable document attributes (`uri`, `title`, `metadata`, `created_at`, `updated_at`) moved from the `documents` table into a new `document_meta` table (1:1 on `document_id`); metadata/title/`source_revision` updates no longer rewrite the docling blobs. Migration `v0_58_0` relocates existing data and runs a one-time `vacuum` to reclaim prior bloat.
- Background auto-vacuum is throttled to at most once per 5 minutes; a final vacuum on close collapses any throttled writes. Sustained ingestion no longer triggers back-to-back compaction of the `documents` table.
### Fixed ### Fixed

View file

@ -9,6 +9,7 @@ from datetime import datetime
from enum import Enum from enum import Enum
from functools import cached_property from functools import cached_property
from pathlib import Path from pathlib import Path
from time import monotonic
from typing import TYPE_CHECKING, overload from typing import TYPE_CHECKING, overload
from urllib.parse import urlparse from urllib.parse import urlparse
@ -40,6 +41,12 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
# Throttle for the background auto-vacuum: under sustained ingestion, scheduling
# a compaction on every write degenerates into back-to-back optimize() passes
# that churn the blob-bearing documents table. Fire at most one per interval; a
# final vacuum on close collapses anything throttled here.
_VACUUM_MIN_INTERVAL_S = 300.0
class RebuildMode(Enum): class RebuildMode(Enum):
"""Mode for rebuilding the database.""" """Mode for rebuilding the database."""
@ -87,6 +94,8 @@ class HaikuRAG:
self._read_only = read_only self._read_only = read_only
self._before = before self._before = before
self._vacuum_tasks: set[asyncio.Task] = set() self._vacuum_tasks: set[asyncio.Task] = set()
self._last_vacuum_at: float | None = None
self._vacuum_dirty = False
@property @property
def is_read_only(self) -> bool: def is_read_only(self) -> bool:
@ -137,17 +146,20 @@ class HaikuRAG:
return False return False
async def _await_vacuum_tasks(self) -> None: async def _await_vacuum_tasks(self) -> None:
"""Drain background vacuum work before tearing down the connection. """Drain background vacuum work and run a final collapse before teardown.
Each create_document / update_document can schedule its own vacuum task; Writes schedule a throttled background vacuum; many are debounced or skip
all must be awaited, not just the most recently scheduled one. Vacuum because another vacuum holds the lock. The final pass collapses the
skips when another is already running, so the cleanup for the final versions those left behind. It runs whenever writes happened
writes may have been a no-op. Run one more pass once the in-flight tasks (``_vacuum_dirty``) not gated on in-flight tasks remaining, since a
are done to collapse versions created after the last vacuum took the lock. debounced run may have scheduled none but never when nothing was
written (so opening + closing a store still never writes).
""" """
if not self._vacuum_tasks: if self._vacuum_tasks:
await asyncio.gather(*self._vacuum_tasks, return_exceptions=True)
if not self._vacuum_dirty:
return return
await asyncio.gather(*self._vacuum_tasks, return_exceptions=True) self._vacuum_dirty = False
# __aexit__ runs during exception unwinding; a raising vacuum here would # __aexit__ runs during exception unwinding; a raising vacuum here would
# mask the original exception, so the drain stays best-effort. # mask the original exception, so the drain stays best-effort.
try: try:
@ -156,7 +168,19 @@ class HaikuRAG:
logger.debug("Final vacuum on close failed", exc_info=True) logger.debug("Final vacuum on close failed", exc_info=True)
def _schedule_vacuum(self) -> None: def _schedule_vacuum(self) -> None:
"""Schedule a background vacuum and track the task for later awaiting.""" """Schedule a background vacuum, throttled to at most one per
``_VACUUM_MIN_INTERVAL_S``. Sustained writes would otherwise trigger
back-to-back compaction of the blob-bearing documents table. The throttle
only skips the background task ``_vacuum_dirty`` still marks that a
final vacuum on close is owed."""
self._vacuum_dirty = True
now = monotonic()
if (
self._last_vacuum_at is not None
and now - self._last_vacuum_at < _VACUUM_MIN_INTERVAL_S
):
return
self._last_vacuum_at = now
task = asyncio.create_task(self.store.vacuum()) task = asyncio.create_task(self.store.vacuum())
self._vacuum_tasks.add(task) self._vacuum_tasks.add(task)
task.add_done_callback(self._vacuum_tasks.discard) task.add_done_callback(self._vacuum_tasks.discard)

View file

@ -0,0 +1,56 @@
import asyncio
import pytest
import haiku.rag.client as client_mod
from haiku.rag.client import HaikuRAG
@pytest.mark.asyncio
async def test_schedule_vacuum_is_debounced(temp_db_path, monkeypatch):
"""Rapid writes within the throttle window schedule only one background
vacuum; once the interval elapses, a new one is scheduled."""
t = {"now": 1000.0}
monkeypatch.setattr(client_mod, "monotonic", lambda: t["now"])
async with HaikuRAG(temp_db_path, create=True) as client:
calls: list[int] = []
async def fake_vacuum(*_a, **_k):
calls.append(1)
monkeypatch.setattr(client.store, "vacuum", fake_vacuum)
for _ in range(3):
client._schedule_vacuum()
await asyncio.gather(*client._vacuum_tasks)
assert len(calls) == 1 # debounced within the interval
t["now"] += client_mod._VACUUM_MIN_INTERVAL_S + 1
client._schedule_vacuum()
await asyncio.gather(*client._vacuum_tasks)
assert len(calls) == 2 # interval elapsed -> a new vacuum scheduled
@pytest.mark.asyncio
async def test_debounced_writes_still_collapse_on_close(temp_db_path, monkeypatch):
"""Even when scheduled vacuums after the first are debounced, the writes are
marked dirty so the close-time drain runs a final collapse."""
t = {"now": 1000.0}
monkeypatch.setattr(client_mod, "monotonic", lambda: t["now"])
calls: list[int] = []
async with HaikuRAG(temp_db_path, create=True) as client:
async def fake_vacuum(*_a, **_k):
calls.append(1)
monkeypatch.setattr(client.store, "vacuum", fake_vacuum)
client._schedule_vacuum() # schedules the first background pass
client._schedule_vacuum() # debounced (no task)
await client._await_vacuum_tasks()
# one scheduled background pass + one final collapse on drain
assert len(calls) == 2
assert client._vacuum_dirty is False

View file

@ -334,10 +334,8 @@ async def test_close_suppresses_failing_drain_vacuum(temp_db_path, monkeypatch):
calls.append(1) calls.append(1)
raise RuntimeError("vacuum boom") raise RuntimeError("vacuum boom")
# A finished task in the set forces the drain branch to run. # Writes happened, so close owes a final vacuum — force that drain branch.
task = asyncio.create_task(asyncio.sleep(0)) client._vacuum_dirty = True
await task
client._vacuum_tasks.add(task)
monkeypatch.setattr(client.store, "vacuum", boom) monkeypatch.setattr(client.store, "vacuum", boom)
# Must not raise despite the drain vacuum erroring. # Must not raise despite the drain vacuum erroring.