From 07c5a97929a79243f367ada682661f574cbf7bd2 Mon Sep 17 00:00:00 2001 From: Chris McDonough Date: Mon, 1 Jun 2026 07:00:15 -0400 Subject: [PATCH 1/3] Reuse httpx.AsyncClient across requests in HTTP and WebDAV sources MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit HTTPSource and WebDAVSource previously created a new AsyncClient for every head(), fetch(), and discover() call — no connection reuse, TLS renegotiation on every request, and connection pool churn at scale. Create the client once in __init__ and reuse it for the lifetime of the source. Add aclose() to both sources, called by PollerManager on shutdown to cleanly close the connection pool. --- .../haiku/rag/ingester/pollers/manager.py | 3 + .../haiku/rag/ingester/sources/http.py | 135 +++++++++--------- .../haiku/rag/ingester/sources/webdav.py | 59 ++++---- 3 files changed, 96 insertions(+), 101 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py b/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py index 6339a5c8..bd7638d5 100644 --- a/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py +++ b/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py @@ -96,6 +96,9 @@ class PollerManager: await asyncio.gather(*self._tasks, return_exceptions=True) self._tasks.clear() self._started = False + for source in self.sources: + if hasattr(source, "aclose"): + await source.aclose() @property def pollers(self) -> list[BasePoller]: diff --git a/haiku_rag_slim/haiku/rag/ingester/sources/http.py b/haiku_rag_slim/haiku/rag/ingester/sources/http.py index 83c0bada..b6ff4afa 100644 --- a/haiku_rag_slim/haiku/rag/ingester/sources/http.py +++ b/haiku_rag_slim/haiku/rag/ingester/sources/http.py @@ -39,48 +39,44 @@ class HTTPSource: self.source_id = source_id self.urls = list(urls or []) self.headers = dict(headers or {}) - # transport is for testing — production callers leave it None and httpx - # uses its real transport. - self._transport = transport + self._http = httpx.AsyncClient(headers=self.headers, transport=transport) def supports(self, uri: str) -> bool: return urlparse(uri).scheme in ("http", "https") - def _client(self) -> httpx.AsyncClient: - return httpx.AsyncClient(headers=self.headers, transport=self._transport) + async def aclose(self) -> None: + await self._http.aclose() async def head(self, uri: str) -> str | None: """HEAD probe for the cheap revision short-circuit. Returns the ETag (or Last-Modified) so an unchanged remote URL can skip the full GET. None on HTTP error so the caller falls back to fetch(); network errors propagate and the worker's classifier handles them.""" - async with self._client() as http: - response = await http.head(uri) - if response.is_error: - return None + response = await self._http.head(uri) + if response.is_error: + return None revision, _ = _extract_revision(response.headers) return revision async def fetch(self, uri: str) -> FetchResult: - async with self._client() as http: - response = await http.get(uri) - response.raise_for_status() - body = response.content - content_type = ( - response.headers.get("content-type", "application/octet-stream") - .split(";")[0] - .strip() - .lower() - ) - revision, extra = _extract_revision(response.headers) - return FetchResult( - uri=uri, - body=body, - content_type=content_type, - content_hash=hashlib.md5(body, usedforsecurity=False).hexdigest(), - revision=revision, - extra_metadata=extra, - ) + response = await self._http.get(uri) + response.raise_for_status() + body = response.content + content_type = ( + response.headers.get("content-type", "application/octet-stream") + .split(";")[0] + .strip() + .lower() + ) + revision, extra = _extract_revision(response.headers) + return FetchResult( + uri=uri, + body=body, + content_type=content_type, + content_hash=hashlib.md5(body, usedforsecurity=False).hexdigest(), + revision=revision, + extra_metadata=extra, + ) async def discover( self, @@ -104,53 +100,52 @@ class HTTPSource: now = datetime.now(UTC) configured = set(self.urls) - async with self._client() as http: - for url in self.urls: - try: - head = await http.head(url) - except Exception: - yield SourceEvent( - source_id=self.source_id, - uri=url, - kind=SourceEventKind.UPSERT, - revision=None, - discovered_at=now, - ) - continue - - if head.status_code == 410: - yield SourceEvent( - source_id=self.source_id, - uri=url, - kind=SourceEventKind.DELETE, - revision=None, - discovered_at=now, - ) - continue - - if head.is_error: - yield SourceEvent( - source_id=self.source_id, - uri=url, - kind=SourceEventKind.UPSERT, - revision=None, - discovered_at=now, - ) - continue - - revision, _ = _extract_revision(head.headers) - if revision is not None and snapshot.get(url) == revision: - kind = SourceEventKind.UNCHANGED - else: - kind = SourceEventKind.UPSERT - + for url in self.urls: + try: + head = await self._http.head(url) + except Exception: yield SourceEvent( source_id=self.source_id, uri=url, - kind=kind, - revision=revision, + kind=SourceEventKind.UPSERT, + revision=None, discovered_at=now, ) + continue + + if head.status_code == 410: + yield SourceEvent( + source_id=self.source_id, + uri=url, + kind=SourceEventKind.DELETE, + revision=None, + discovered_at=now, + ) + continue + + if head.is_error: + yield SourceEvent( + source_id=self.source_id, + uri=url, + kind=SourceEventKind.UPSERT, + revision=None, + discovered_at=now, + ) + continue + + revision, _ = _extract_revision(head.headers) + if revision is not None and snapshot.get(url) == revision: + kind = SourceEventKind.UNCHANGED + else: + kind = SourceEventKind.UPSERT + + yield SourceEvent( + source_id=self.source_id, + uri=url, + kind=kind, + revision=revision, + discovered_at=now, + ) # Anything previously known to this source that's no longer in # config emits DELETE so delete_orphans can clean up. Without this, diff --git a/haiku_rag_slim/haiku/rag/ingester/sources/webdav.py b/haiku_rag_slim/haiku/rag/ingester/sources/webdav.py index 7dd6a15e..0bb512f8 100644 --- a/haiku_rag_slim/haiku/rag/ingester/sources/webdav.py +++ b/haiku_rag_slim/haiku/rag/ingester/sources/webdav.py @@ -189,41 +189,39 @@ class WebDAVSource: supported_extensions=self.supported_extensions, ) # transport is for testing — production callers leave it None. - self._transport = transport - - def supports(self, uri: str) -> bool: - return uri.startswith(self.base_url) - - def _client(self) -> httpx.AsyncClient: auth = ( (self.username, self.password) if self.username is not None and self.password is not None else None ) - return httpx.AsyncClient( - auth=auth, headers=self.headers, transport=self._transport + self._http = httpx.AsyncClient( + auth=auth, headers=self.headers, transport=transport ) + def supports(self, uri: str) -> bool: + return uri.startswith(self.base_url) + + async def aclose(self) -> None: + await self._http.aclose() + async def head(self, uri: str) -> str | None: - async with self._client() as http: - response = await http.request( - "PROPFIND", - uri, - headers={"Depth": "0", "Content-Type": "application/xml"}, - content=_PROPFIND_BODY, - ) - if response.is_error: - return None - entries = _parse_multistatus(response.content) + response = await self._http.request( + "PROPFIND", + uri, + headers={"Depth": "0", "Content-Type": "application/xml"}, + content=_PROPFIND_BODY, + ) + if response.is_error: + return None + entries = _parse_multistatus(response.content) if not entries: return None return entries[0].revision async def fetch(self, uri: str) -> FetchResult: - async with self._client() as http: - response = await http.get(uri) - response.raise_for_status() - body = response.content + response = await self._http.get(uri) + response.raise_for_status() + body = response.content content_type = ( response.headers.get("content-type", "application/octet-stream") .split(";")[0] @@ -261,15 +259,14 @@ class WebDAVSource: now = datetime.now(UTC) seen: set[str] = set() - async with self._client() as http: - response = await http.request( - "PROPFIND", - self.base_url, - headers={"Depth": "infinity", "Content-Type": "application/xml"}, - content=_PROPFIND_BODY, - ) - response.raise_for_status() - entries = _parse_multistatus(response.content) + response = await self._http.request( + "PROPFIND", + self.base_url, + headers={"Depth": "infinity", "Content-Type": "application/xml"}, + content=_PROPFIND_BODY, + ) + response.raise_for_status() + entries = _parse_multistatus(response.content) for entry in entries: if entry.is_collection: From 22ab79c49223ad665a1320eff7af75151c098653 Mon Sep 17 00:00:00 2001 From: Chris McDonough Date: Mon, 1 Jun 2026 09:55:39 -0400 Subject: [PATCH 2/3] 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 From c5814b31bf176ad1553e0a658a40f0f3f0e3ff39 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 1 Jun 2026 17:29:53 +0300 Subject: [PATCH 3/3] Test source clients close after the worker pool stops Workers share the pollers' Source instances for fetch(), so the httpx clients must be closed only after the pool has stopped. Guards both run_batch() and serve() against reintroducing the shutdown-order bug. Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/ingester/test_run_batch.py | 58 ++++++++++++++++++++++++++++++++ 1 file changed, 58 insertions(+) diff --git a/tests/ingester/test_run_batch.py b/tests/ingester/test_run_batch.py index c646d5d6..8edd461d 100644 --- a/tests/ingester/test_run_batch.py +++ b/tests/ingester/test_run_batch.py @@ -23,6 +23,7 @@ from haiku.rag.config import ( WorkerConfig, ) from haiku.rag.ingester.app import IngesterApp +from haiku.rag.ingester.pollers.manager import PollerManager from haiku.rag.ingester.workers.pool import WorkerPool from haiku.rag.store.models.document import Document @@ -305,3 +306,60 @@ async def test_stop_pool_warns_when_shutdown_grace_elapses(tmp_path, caplog): assert pool.released == 1 assert "Shutdown grace" in caplog.text + + +def _record_close_order(monkeypatch) -> list[str]: + """Record the order of _stop_pool and PollerManager.close_sources. + + Workers share the pollers' Source instances for fetch(), so the source + clients must be closed only after the pool has stopped — otherwise an + in-flight fetch during the shutdown grace hits a closed client. + """ + order: list[str] = [] + orig_stop_pool = IngesterApp._stop_pool + orig_close = PollerManager.close_sources + + async def rec_stop(self): + order.append("stop_pool") + await orig_stop_pool(self) + + async def rec_close(self): + order.append("close_sources") + await orig_close(self) + + monkeypatch.setattr(IngesterApp, "_stop_pool", rec_stop) + monkeypatch.setattr(PollerManager, "close_sources", rec_close) + return order + + +@pytest.mark.asyncio +async def test_run_batch_closes_sources_after_pool_stops( + tmp_path, use_client, monkeypatch +): + (tmp_path / "a.md").write_text("hello") + use_client(_mock_client()) + app = IngesterApp(config=_config(tmp_path), db_path=tmp_path / "db.lancedb") + order = _record_close_order(monkeypatch) + + await app.run_batch() + + assert order == ["stop_pool", "close_sources"] + + +@pytest.mark.asyncio +async def test_serve_closes_sources_after_pool_stops(tmp_path, use_client, monkeypatch): + use_client(_mock_client()) + config = _config(tmp_path) + config.ingester.api = APIConfig(enabled=False) + app = IngesterApp(config=config, db_path=tmp_path / "db.lancedb") + order = _record_close_order(monkeypatch) + + task = asyncio.create_task(app.serve(api=False)) + try: + await _wait_until(lambda: app._pool is not None and app._pool.live_workers > 0) + finally: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=5.0) + + assert order == ["stop_pool", "close_sources"]