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
This commit is contained in:
Chris McDonough 2026-06-01 09:55:39 -04:00
parent 07c5a97929
commit 22ab79c492
6 changed files with 22 additions and 2 deletions

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

View file

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

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(

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

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]

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