danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #395 from mcdonc/perf/reuse-httpx-clients

Browse source 
perf: reuse httpx.AsyncClient in HTTP and WebDAV sources
This commit is contained in:
Yiorgis Gozadinos 2026-06-01 17:38:05 +03:00 committed by GitHub
commit e9875fc842
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 174 additions and 101 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
haiku_rag_slim/haiku/rag/ingester/app.py
View file

@ -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

7
haiku_rag_slim/haiku/rag/ingester/pollers/manager.py
View file

@ -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)

5
haiku_rag_slim/haiku/rag/ingester/sources/base.py
View file

@ -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(

3
haiku_rag_slim/haiku/rag/ingester/sources/fs.py
View file

@ -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():

135
haiku_rag_slim/haiku/rag/ingester/sources/http.py
View file

@ -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,

3
haiku_rag_slim/haiku/rag/ingester/sources/s3.py
View file

@ -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]

59
haiku_rag_slim/haiku/rag/ingester/sources/webdav.py
View file

@ -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:

58
tests/ingester/test_run_batch.py
View file

@ -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"]

3
tests/ingester/test_sources_base.py
View file

@ -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