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

Reuse httpx.AsyncClient across requests in HTTP and WebDAV sources

Browse source 
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.
This commit is contained in:
Chris McDonough 2026-06-01 07:00:15 -04:00
parent d5e5733f67
commit 07c5a97929
3 changed files with 96 additions and 101 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -96,6 +96,9 @@ class PollerManager:
await asyncio.gather(*self._tasks, return_exceptions=True) await asyncio.gather(*self._tasks, return_exceptions=True)
self._tasks.clear() self._tasks.clear()
self._started = False self._started = False
for source in self.sources:
if hasattr(source, "aclose"):
await source.aclose()
@property @property
def pollers(self) -> list[BasePoller]: def pollers(self) -> list[BasePoller]:

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

@ -39,48 +39,44 @@ class HTTPSource:
self.source_id = source_id self.source_id = source_id
self.urls = list(urls or []) self.urls = list(urls or [])
self.headers = dict(headers or {}) self.headers = dict(headers or {})
# transport is for testing — production callers leave it None and httpx self._http = httpx.AsyncClient(headers=self.headers, transport=transport)
# uses its real transport.
self._transport = transport
def supports(self, uri: str) -> bool: def supports(self, uri: str) -> bool:
return urlparse(uri).scheme in ("http", "https") return urlparse(uri).scheme in ("http", "https")
def _client(self) -> httpx.AsyncClient: async def aclose(self) -> None:
return httpx.AsyncClient(headers=self.headers, transport=self._transport) await self._http.aclose()
async def head(self, uri: str) -> str | None: async def head(self, uri: str) -> str | None:
"""HEAD probe for the cheap revision short-circuit. Returns the """HEAD probe for the cheap revision short-circuit. Returns the
ETag (or Last-Modified) so an unchanged remote URL can skip 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(); full GET. None on HTTP error so the caller falls back to fetch();
network errors propagate and the worker's classifier handles them.""" network errors propagate and the worker's classifier handles them."""
async with self._client() as http: response = await self._http.head(uri)
response = await http.head(uri) if response.is_error:
if response.is_error: return None
return None
revision, _ = _extract_revision(response.headers) revision, _ = _extract_revision(response.headers)
return revision return revision
async def fetch(self, uri: str) -> FetchResult: async def fetch(self, uri: str) -> FetchResult:
async with self._client() as http: response = await self._http.get(uri)
response = await http.get(uri) response.raise_for_status()
response.raise_for_status() body = response.content
body = response.content content_type = (
content_type = ( response.headers.get("content-type", "application/octet-stream")
response.headers.get("content-type", "application/octet-stream") .split(";")[0]
.split(";")[0] .strip()
.strip() .lower()
.lower() )
) revision, extra = _extract_revision(response.headers)
revision, extra = _extract_revision(response.headers) return FetchResult(
return FetchResult( uri=uri,
uri=uri, body=body,
body=body, content_type=content_type,
content_type=content_type, content_hash=hashlib.md5(body, usedforsecurity=False).hexdigest(),
content_hash=hashlib.md5(body, usedforsecurity=False).hexdigest(), revision=revision,
revision=revision, extra_metadata=extra,
extra_metadata=extra, )
)
async def discover( async def discover(
self, self,
@ -104,53 +100,52 @@ class HTTPSource:
now = datetime.now(UTC) now = datetime.now(UTC)
configured = set(self.urls) configured = set(self.urls)
async with self._client() as http: for url in self.urls:
for url in self.urls: try:
try: head = await self._http.head(url)
head = await http.head(url) except Exception:
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
yield SourceEvent( yield SourceEvent(
source_id=self.source_id, source_id=self.source_id,
uri=url, uri=url,
kind=kind, kind=SourceEventKind.UPSERT,
revision=revision, revision=None,
discovered_at=now, 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 # Anything previously known to this source that's no longer in
# config emits DELETE so delete_orphans can clean up. Without this, # config emits DELETE so delete_orphans can clean up. Without this,

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

@ -189,41 +189,39 @@ class WebDAVSource:
supported_extensions=self.supported_extensions, supported_extensions=self.supported_extensions,
) )
# transport is for testing — production callers leave it None. # 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 = ( auth = (
(self.username, self.password) (self.username, self.password)
if self.username is not None and self.password is not None if self.username is not None and self.password is not None
else None else None
) )
return httpx.AsyncClient( self._http = httpx.AsyncClient(
auth=auth, headers=self.headers, transport=self._transport 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 def head(self, uri: str) -> str | None:
async with self._client() as http: response = await self._http.request(
response = await http.request( "PROPFIND",
"PROPFIND", uri,
uri, headers={"Depth": "0", "Content-Type": "application/xml"},
headers={"Depth": "0", "Content-Type": "application/xml"}, content=_PROPFIND_BODY,
content=_PROPFIND_BODY, )
) if response.is_error:
if response.is_error: return None
return None entries = _parse_multistatus(response.content)
entries = _parse_multistatus(response.content)
if not entries: if not entries:
return None return None
return entries[0].revision return entries[0].revision
async def fetch(self, uri: str) -> FetchResult: async def fetch(self, uri: str) -> FetchResult:
async with self._client() as http: response = await self._http.get(uri)
response = await http.get(uri) response.raise_for_status()
response.raise_for_status() body = response.content
body = response.content
content_type = ( content_type = (
response.headers.get("content-type", "application/octet-stream") response.headers.get("content-type", "application/octet-stream")
.split(";")[0] .split(";")[0]
@ -261,15 +259,14 @@ class WebDAVSource:
now = datetime.now(UTC) now = datetime.now(UTC)
seen: set[str] = set() seen: set[str] = set()
async with self._client() as http: response = await self._http.request(
response = await http.request( "PROPFIND",
"PROPFIND", self.base_url,
self.base_url, headers={"Depth": "infinity", "Content-Type": "application/xml"},
headers={"Depth": "infinity", "Content-Type": "application/xml"}, content=_PROPFIND_BODY,
content=_PROPFIND_BODY, )
) response.raise_for_status()
response.raise_for_status() entries = _parse_multistatus(response.content)
entries = _parse_multistatus(response.content)
for entry in entries: for entry in entries:
if entry.is_collection: if entry.is_collection: