From dd2817ff6d4a47d98b645f823445c1455d245a02 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 16 Jul 2026 10:56:43 +0300 Subject: [PATCH] Protect tag-creation cleanup from cancellation --- haiku_rag_slim/haiku/rag/store/engine.py | 91 +++++++++++++++--------- tests/store/test_tags.py | 87 ++++++++++++++++++++++ 2 files changed, 144 insertions(+), 34 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 0b53fc1f..a5a26c3f 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -1,6 +1,7 @@ import asyncio import json import logging +from collections.abc import Coroutine from dataclasses import dataclass from datetime import UTC, datetime, timedelta from enum import Enum @@ -200,6 +201,27 @@ RESTORE_TABLE_ORDER: tuple[str, ...] = tuple( ) + ("documents",) +async def _wait_protected[T](coro: Coroutine[Any, Any, T]) -> tuple[T, bool]: + """Await a recovery coroutine that a cancellation cannot interrupt. + + Runs the coroutine as a task and keeps waiting for it even if this + coroutine is cancelled, so a Ctrl-C cannot leave recovery half applied. + Returns the result and whether a cancellation was absorbed; the caller + must re-deliver an absorbed cancellation. + """ + task = asyncio.ensure_future(coro) + cancelled = False + while True: + try: + return await asyncio.shield(task), cancelled + except asyncio.CancelledError: + if task.done(): + # The recovery coroutine itself ended cancelled; there is + # nothing left to wait for. + raise + cancelled = True + + def _safety_tag_name(existing: set[str]) -> str: """Collision-resistant name for the pre-restore safety tag.""" base = f"before-restore-{datetime.now(UTC).strftime('%Y%m%dT%H%M%SZ')}" @@ -926,26 +948,47 @@ class Store: ) versions = await self.current_table_versions() - created: list[str] = [] try: for table_name, table in tables.items(): await table.tags.create(name, versions[table_name]) - created.append(table_name) - except Exception as exc: - failed_cleanup: list[str] = [] - for table_name in created: - try: - await tables[table_name].tags.delete(name) - except Exception: - failed_cleanup.append(table_name) + except BaseException as exc: + # BaseException: cancellation must also trigger cleanup, and the + # cleanup itself is protected from further cancellation. The + # sweep covers all tables, not only the recorded ones: a + # cancellation can land after lance committed a table's tag but + # before this attempt recorded it, and preflight guarantees the + # name was unused, so any occurrence belongs to this attempt. + (_, failed_cleanup), cancelled = await _wait_protected( + self._delete_tag_locked(name) + ) if failed_cleanup: raise RuntimeError( - f"Tag '{name}' creation failed ({exc}) and cleanup " + f"Tag '{name}' creation failed ({exc!r}) and cleanup " f"failed on: {', '.join(failed_cleanup)}. A partial " "tag may remain; delete it with delete_tag." ) from exc + if cancelled and not isinstance(exc, asyncio.CancelledError): + raise asyncio.CancelledError() raise + async def _delete_tag_locked(self, name: str) -> tuple[bool, list[str]]: + """Delete the tag from every table that has it; the caller must + hold the write lock. + + Returns whether the tag was found anywhere and the tables where + listing or deletion failed. + """ + found = False + failed: list[str] = [] + for table_name, table in self._tables().items(): + try: + if name in await table.tags.list(): + found = True + await table.tags.delete(name) + except Exception: + failed.append(table_name) + return found, failed + async def list_tags(self) -> dict[str, TagInfo]: """Aggregate per-table tags into database-level tags. @@ -977,15 +1020,7 @@ class Store: self._assert_writable() self._assert_not_rebuilding() async with self._rebuild_lock, self._write_lock: - found = False - failed: list[str] = [] - for table_name, table in self._tables().items(): - try: - if name in await table.tags.list(): - found = True - await table.tags.delete(name) - except Exception: - failed.append(table_name) + found, failed = await self._delete_tag_locked(name) if failed: # A listing failure obscures whether the tag exists on that # table, so failures take precedence over not-found. @@ -1020,22 +1055,10 @@ class Store: ) -> tuple[list[tuple[str, Exception]], bool]: """Best-effort rollback that a cancellation cannot interrupt. - Runs the rollback as a task and keeps waiting for it even if this - coroutine is cancelled, so a Ctrl-C cannot leave the rollback half - applied. Returns the rollback failures and whether a cancellation - was absorbed; the caller must re-deliver an absorbed cancellation. + Returns the rollback failures and whether a cancellation was + absorbed; the caller must re-deliver an absorbed cancellation. """ - task = asyncio.ensure_future(self._restore_tables(snapshot, best_effort=True)) - cancelled = False - while True: - try: - return await asyncio.shield(task), cancelled - except asyncio.CancelledError: - if task.done(): - # The rollback coroutine itself ended cancelled; there is - # nothing left to wait for. - raise - cancelled = True + return await _wait_protected(self._restore_tables(snapshot, best_effort=True)) async def restore_tag(self, name: str) -> str: """Restore every table to the versions of a complete tag. diff --git a/tests/store/test_tags.py b/tests/store/test_tags.py index 74288ada..9b20f10a 100644 --- a/tests/store/test_tags.py +++ b/tests/store/test_tags.py @@ -445,3 +445,90 @@ async def test_delete_tag_reports_listing_failures(temp_db_path, monkeypatch): await store.delete_tag("release-1") assert await store.list_tags() == {} + + +@pytest.mark.asyncio +async def test_create_tag_cancellation_cleans_up(temp_db_path, monkeypatch): + """Cancellation during per-table tag creation must not leave a partial + tag behind: cleanup runs before the cancellation propagates.""" + async with Store(temp_db_path, create=True) as store: + real_create = AsyncTags.create + calls = {"n": 0} + + async def cancelled_create(self, name: str, version: int) -> None: + calls["n"] += 1 + if calls["n"] == 4: + raise asyncio.CancelledError() + await real_create(self, name, version) + + monkeypatch.setattr(AsyncTags, "create", cancelled_create) + + with pytest.raises(asyncio.CancelledError): + await store.create_tag("broken") + + monkeypatch.undo() + assert await store.list_tags() == {} + + +@pytest.mark.asyncio +async def test_create_tag_cleanup_survives_cancellation(temp_db_path, monkeypatch): + """Cancelling create_tag while it cleans up a failed creation does not + interrupt the cleanup: no partial tag remains and the cancellation is + delivered afterwards.""" + async with Store(temp_db_path, create=True) as store: + real_create = AsyncTags.create + real_delete = AsyncTags.delete + create_calls = {"n": 0} + cleanup_started = asyncio.Event() + release = asyncio.Event() + + async def flaky_create(self, name: str, version: int) -> None: + create_calls["n"] += 1 + if create_calls["n"] == 4: + raise RuntimeError("create boom") + await real_create(self, name, version) + + async def slow_delete(self, name: str) -> None: + cleanup_started.set() + await release.wait() + await real_delete(self, name) + + monkeypatch.setattr(AsyncTags, "create", flaky_create) + monkeypatch.setattr(AsyncTags, "delete", slow_delete) + + task = asyncio.create_task(store.create_tag("broken")) + await cleanup_started.wait() + task.cancel() + release.set() + + with pytest.raises(asyncio.CancelledError): + await task + + monkeypatch.undo() + assert await store.list_tags() == {} + + +@pytest.mark.asyncio +async def test_create_tag_cancellation_after_commit_cleans_committed_tag( + temp_db_path, monkeypatch +): + """Cancellation arriving after lance committed a table's tag but before + the attempt recorded it must still clean that table: cleanup sweeps all + tables, relying on the preflight guarantee that the name was unused.""" + async with Store(temp_db_path, create=True) as store: + real_create = AsyncTags.create + calls = {"n": 0} + + async def committing_cancelled_create(self, name: str, version: int) -> None: + calls["n"] += 1 + await real_create(self, name, version) + if calls["n"] == 4: + raise asyncio.CancelledError() + + monkeypatch.setattr(AsyncTags, "create", committing_cancelled_create) + + with pytest.raises(asyncio.CancelledError): + await store.create_tag("broken") + + monkeypatch.undo() + assert await store.list_tags() == {}