From 22ab79c49223ad665a1320eff7af75151c098653 Mon Sep 17 00:00:00 2001 From: Chris McDonough Date: Mon, 1 Jun 2026 09:55:39 -0400 Subject: [PATCH] Fix shutdown-order bug: close source clients after workers stop Workers share the same Source instances as pollers and use them for fetch(). PollerManager.stop() was closing httpx clients before the worker pool drained, so in-flight fetches during the shutdown grace hit a closed client. - Move source closing out of stop() into a separate close_sources() - Call close_sources() after _stop_pool() in both serve() and run_batch() - Promote aclose() to the Source protocol with no-op defaults for FS and S3, removing the hasattr duck-typing --- haiku_rag_slim/haiku/rag/ingester/app.py | 2 ++ haiku_rag_slim/haiku/rag/ingester/pollers/manager.py | 8 ++++++-- haiku_rag_slim/haiku/rag/ingester/sources/base.py | 5 +++++ haiku_rag_slim/haiku/rag/ingester/sources/fs.py | 3 +++ haiku_rag_slim/haiku/rag/ingester/sources/s3.py | 3 +++ tests/ingester/test_sources_base.py | 3 +++ 6 files changed, 22 insertions(+), 2 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/ingester/app.py b/haiku_rag_slim/haiku/rag/ingester/app.py index 7b64556f..70c5a673 100644 --- a/haiku_rag_slim/haiku/rag/ingester/app.py +++ b/haiku_rag_slim/haiku/rag/ingester/app.py @@ -181,6 +181,7 @@ class IngesterApp: await asyncio.gather(api_task, return_exceptions=True) await self._pollers.stop() await self._stop_pool() + await self._pollers.close_sources() async def run_batch(self) -> BatchReport: """Run one discover() sweep across every configured source, drain the @@ -213,6 +214,7 @@ class IngesterApp: ) finally: await self._stop_pool() + await self._pollers.close_sources() async def _maybe_start_api(self, api: bool): """Spin up the FastAPI control plane on an asyncio task. Returns diff --git a/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py b/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py index bd7638d5..0e390693 100644 --- a/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py +++ b/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py @@ -96,9 +96,13 @@ class PollerManager: await asyncio.gather(*self._tasks, return_exceptions=True) self._tasks.clear() self._started = False + + async def close_sources(self) -> None: + """Close all source adapters (e.g. HTTP connection pools). Must be + called after the worker pool has fully stopped so in-flight fetches + don't hit a closed client.""" for source in self.sources: - if hasattr(source, "aclose"): - await source.aclose() + await source.aclose() @property def pollers(self) -> list[BasePoller]: diff --git a/haiku_rag_slim/haiku/rag/ingester/sources/base.py b/haiku_rag_slim/haiku/rag/ingester/sources/base.py index 32099b35..c282ece9 100644 --- a/haiku_rag_slim/haiku/rag/ingester/sources/base.py +++ b/haiku_rag_slim/haiku/rag/ingester/sources/base.py @@ -61,6 +61,11 @@ class Source(Protocol): """ ... + async def aclose(self) -> None: + """Release any resources held by the source (e.g. HTTP connection + pools). Called once during shutdown, after all workers have stopped.""" + ... + async def fetch(self, uri: str) -> FetchResult: ... def discover( diff --git a/haiku_rag_slim/haiku/rag/ingester/sources/fs.py b/haiku_rag_slim/haiku/rag/ingester/sources/fs.py index 8b944e04..ba97f120 100644 --- a/haiku_rag_slim/haiku/rag/ingester/sources/fs.py +++ b/haiku_rag_slim/haiku/rag/ingester/sources/fs.py @@ -75,6 +75,9 @@ class FSSource: return False return self._resolve_within_root(uri) is not None + async def aclose(self) -> None: # pragma: no cover - no resources to release + pass + async def head(self, uri: str) -> str | None: path = self._resolve_within_root(uri) if path is None or not path.exists(): diff --git a/haiku_rag_slim/haiku/rag/ingester/sources/s3.py b/haiku_rag_slim/haiku/rag/ingester/sources/s3.py index 0b0b229f..f3c8a446 100644 --- a/haiku_rag_slim/haiku/rag/ingester/sources/s3.py +++ b/haiku_rag_slim/haiku/rag/ingester/sources/s3.py @@ -63,6 +63,9 @@ class S3Source: def supports(self, uri: str) -> bool: return uri.startswith(self.uri_prefix) + async def aclose(self) -> None: # pragma: no cover - no resources to release + pass + async def head(self, uri: str) -> str | None: import obstore # type: ignore[import-not-found] diff --git a/tests/ingester/test_sources_base.py b/tests/ingester/test_sources_base.py index b2ffe995..ebaaebae 100644 --- a/tests/ingester/test_sources_base.py +++ b/tests/ingester/test_sources_base.py @@ -59,6 +59,9 @@ def test_source_protocol_runtime_checkable(): def supports(self, uri: str) -> bool: return True + async def aclose(self) -> None: + pass + async def head(self, uri: str): return None