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 6339a5c8..0e390693 100644 --- a/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py +++ b/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py @@ -97,6 +97,13 @@ class PollerManager: 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: + await source.aclose() + @property def pollers(self) -> list[BasePoller]: return list(self._pollers) 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 8de87469..a76f17bc 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/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/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/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: 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"] 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