From bfd09f6880c0a4972bc1919cb28b05469ec80793 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 28 Aug 2026 12:02:07 +0300 Subject: [PATCH] Leave no database reading when a federated fan-out fails MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `asyncio.gather` propagates the first failure while its siblings run on, and the caller unwinding from that closes the set through `async with` — so a sibling still reading reads through a closed session. Seven fan-outs were affected: lookup, search, image enrichment, multimodal picture loading, context expansion, document listing and counting, and the sandbox's document load. `gather_all` cancels and drains the rest, then re-raises the original exception. A `TaskGroup` would drain them too but raise an `ExceptionGroup`, which every caller and both CLIs' exception handlers would have to unwrap. The two `return_exceptions=True` gathers in session opening and teardown already drain their children and are left alone. --- haiku_rag_slim/haiku/rag/client/__init__.py | 8 +-- haiku_rag_slim/haiku/rag/client/search.py | 10 ++-- haiku_rag_slim/haiku/rag/sandbox/sandbox.py | 3 +- haiku_rag_slim/haiku/rag/utils.py | 18 +++++++ tests/multi_db/test_lifecycle.py | 35 ++++++++++++- tests/test_utils.py | 58 ++++++++++++++++++++- 6 files changed, 120 insertions(+), 12 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 68d305e0..1289d8c1 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -43,7 +43,7 @@ from haiku.rag.store.repositories.chunk import ChunkRepository from haiku.rag.store.repositories.document import DocumentRepository from haiku.rag.store.repositories.document_item import DocumentItemRepository from haiku.rag.store.repositories.settings import SettingsRepository -from haiku.rag.utils import escape_sql_string +from haiku.rag.utils import escape_sql_string, gather_all if TYPE_CHECKING: from docling_core.types.doc.document import DoclingDocument @@ -69,7 +69,7 @@ async def all_found( asked at once. Asking in turn would cost a round trip per database for an identifier that is missing or held by the last of them. """ - found_by_client = await asyncio.gather(*(lookup(client) for client in clients)) + found_by_client = await gather_all(*(lookup(client) for client in clients)) return [ (client, found) for client, found in zip(clients, found_by_client, strict=True) @@ -784,7 +784,7 @@ class HaikuRAG: # the window is applied to the merged listing: a limit means that # many documents in total, not that many per database. wanted = None if limit is None else limit + (offset or 0) - groups = await asyncio.gather( + groups = await gather_all( *( owner.list_documents( limit=wanted, filter=filter, include_content=include_content @@ -813,7 +813,7 @@ class HaikuRAG: Number of documents matching the criteria. """ if self.covers_multiple: - counts = await asyncio.gather( + counts = await gather_all( *( owner.count_documents(filter=filter) for owner in await self.clients_covering() diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index 133bac85..2d963038 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -1,4 +1,3 @@ -import asyncio import base64 from collections.abc import Sequence from typing import TYPE_CHECKING @@ -10,6 +9,7 @@ from haiku.rag.store.models.chunk import ( qualified_id, ) from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX +from haiku.rag.utils import gather_all if TYPE_CHECKING: from PIL import Image as PILImage @@ -109,7 +109,7 @@ async def search_sources( fetch_limit = _fetch_limit(client, query, limit) query_vector = await _embed_query(selected[0], query, resolved) text = query if isinstance(query, str) else "" - per_source = await asyncio.gather( + per_source = await gather_all( *( c.chunk_repository.search( query=text, @@ -137,7 +137,7 @@ async def search_sources( if result.source: by_owner.setdefault(result.source, []).append(result) owners = await client.clients_for(list(by_owner)) - await asyncio.gather( + await gather_all( *( _populate_image_data(owner, by_owner[name]) for name, owner in zip(by_owner, owners, strict=True) @@ -176,7 +176,7 @@ async def _fuse( if reranker is not None: chunks = [chunk for _, chunk, _ in owned] if federator._config.reranking.multimodal: - await asyncio.gather( + await gather_all( *( _attach_picture_data( c, [chunk for owner, chunk, _ in owned if owner is c] @@ -438,7 +438,7 @@ async def expand_sources( unsourced.append(result) names = list(by_source) sessions = await federated.sessions_for(names) - expanded_groups = await asyncio.gather( + expanded_groups = await gather_all( *( expand_context(session, by_source[name]) for name, session in zip(names, sessions, strict=True) diff --git a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py index 24f260f2..d88a60f0 100644 --- a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py +++ b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py @@ -20,6 +20,7 @@ from haiku.rag.config.models import AppConfig from haiku.rag.sandbox.dependencies import AnalysisContext from haiku.rag.store.models.chunk import SearchResult from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX, DocumentItem +from haiku.rag.utils import gather_all if TYPE_CHECKING: from pathlib import Path, PurePosixPath @@ -280,7 +281,7 @@ class Sandbox: owners = await rag.clients_covering(self._context.sources) # Resolve owners under the shared-client lock; owner sessions perform # reads independently. - groups = await asyncio.gather( + groups = await gather_all( *(owner.list_documents(filter=self._context.filter) for owner in owners) ) # Interleaved, not concatenated: code that prints the listing is read diff --git a/haiku_rag_slim/haiku/rag/utils.py b/haiku_rag_slim/haiku/rag/utils.py index 50f859bb..30b07084 100644 --- a/haiku_rag_slim/haiku/rag/utils.py +++ b/haiku_rag_slim/haiku/rag/utils.py @@ -1,5 +1,7 @@ +import asyncio import math import sys +from collections.abc import Awaitable from importlib import metadata from pathlib import Path from typing import TYPE_CHECKING, Any, NoReturn, cast @@ -16,6 +18,22 @@ if TYPE_CHECKING: from haiku.rag.store.models.citation import Citation +async def gather_all[T](*awaitables: Awaitable[T]) -> list[T]: + """Run `awaitables` concurrently, leaving none of them running when one fails. + + `asyncio.gather` leaves its siblings running; `TaskGroup` wraps the failure in + an `ExceptionGroup`. + """ + tasks = [asyncio.ensure_future(awaitable) for awaitable in awaitables] + try: + return await asyncio.gather(*tasks) + except BaseException: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + raise + + def parse_model_option(value: str) -> "ModelConfig": """Parse a 'provider:name' string into a ModelConfig.""" from haiku.rag.config.models import ModelConfig diff --git a/tests/multi_db/test_lifecycle.py b/tests/multi_db/test_lifecycle.py index 1c12906f..2f7e65bc 100644 --- a/tests/multi_db/test_lifecycle.py +++ b/tests/multi_db/test_lifecycle.py @@ -7,10 +7,11 @@ from docling_core.types.doc.document import DoclingDocument from docling_core.types.doc.labels import DocItemLabel from haiku.rag.client import HaikuRAG -from haiku.rag.client.session import FederatedSession +from haiku.rag.client.session import FederatedSession, SingleDatabaseSession from haiku.rag.config import get_config from haiku.rag.store.exceptions import ( AmbiguousDatabaseError, + ReadOnlyError, SourceUnavailableError, ) from haiku.rag.store.models import Chunk @@ -93,6 +94,38 @@ class TestOpeningDatabases: assert not alpha.store.db.is_open() + @pytest.mark.asyncio + async def test_a_failing_read_leaves_no_sibling_reading( + self, tmp_path, monkeypatch + ): + """Unwinding through `async with` closes every session, so a sibling still + reading reads through a closed one.""" + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha document about cats"]) + await _seed(config, "beta", ["beta document about cats"]) + + reading = asyncio.Event() + unwound = asyncio.Event() + + async def listing(self, *args, **kwargs): + if self.source == "beta": + await reading.wait() + raise ReadOnlyError("beta is read-only") + reading.set() + try: + await asyncio.sleep(60) + finally: + unwound.set() + return [] + + monkeypatch.setattr(SingleDatabaseSession, "list_documents", listing) + + async with HaikuRAG(config=config) as rag: + with pytest.raises(ReadOnlyError, match="beta"): + await rag.list_documents() + + assert unwound.is_set() + @pytest.mark.asyncio async def test_a_database_named_twice_is_opened_once(self, tmp_path): """Fusion would count a repeated database as two rank lists.""" diff --git a/tests/test_utils.py b/tests/test_utils.py index cd08144b..28fcf629 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,3 +1,4 @@ +import asyncio import importlib.util from unittest.mock import AsyncMock @@ -7,7 +8,8 @@ from pydantic_ai.models.openai import OpenAIChatModel from haiku.rag.config import get_config from haiku.rag.config.models import ModelConfig from haiku.rag.converters import get_converter -from haiku.rag.utils import get_model +from haiku.rag.store.exceptions import ReadOnlyError +from haiku.rag.utils import gather_all, get_model # Check for optional dependencies HAS_ANTHROPIC = importlib.util.find_spec("anthropic") is not None @@ -1104,3 +1106,57 @@ def test_get_model_api_key_rejected_on_unplumbed_provider(): get_model( ModelConfig(provider="anthropic", name="claude-sonnet-4-5", api_key="sk-x") ) + + +class TestGatherAll: + """A fan-out leaves nothing running: a caller unwinding from a failure closes + the sessions its siblings are still reading through.""" + + @pytest.mark.asyncio + async def test_results_arrive_in_the_order_asked_for(self): + async def slow(value): + await asyncio.sleep(0.01) + return value + + async def fast(value): + return value + + assert await gather_all(slow("a"), fast("b"), slow("c")) == ["a", "b", "c"] + + @pytest.mark.asyncio + async def test_a_failure_leaves_no_sibling_running(self): + started = asyncio.Event() + unwound = asyncio.Event() + + async def sibling(): + started.set() + try: + await asyncio.sleep(60) + finally: + unwound.set() + + async def failing(): + await started.wait() + raise ReadOnlyError("beta is read-only") + + before = asyncio.all_tasks() + with pytest.raises(ReadOnlyError, match="beta"): + await gather_all(sibling(), failing()) + + assert unwound.is_set() + assert asyncio.all_tasks() - before == set() + + @pytest.mark.asyncio + async def test_the_failure_arrives_as_itself(self): + """A `TaskGroup` drains the siblings too, but raises an `ExceptionGroup`.""" + + async def failing(): + raise ReadOnlyError("beta is read-only") + + async def sibling(): + await asyncio.sleep(60) + + with pytest.raises(ReadOnlyError) as raised: + await gather_all(sibling(), failing()) + + assert not isinstance(raised.value, BaseExceptionGroup)