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 asyncio.gather(api_task, return_exceptions=True)
await self._pollers.stop() await self._pollers.stop()
await self._stop_pool() await self._stop_pool()
await self._pollers.close_sources()
async def run_batch(self) -> BatchReport: async def run_batch(self) -> BatchReport:
"""Run one discover() sweep across every configured source, drain the """Run one discover() sweep across every configured source, drain the
@ -213,6 +214,7 @@ class IngesterApp:
) )
finally: finally:
await self._stop_pool() await self._stop_pool()
await self._pollers.close_sources()
async def _maybe_start_api(self, api: bool): async def _maybe_start_api(self, api: bool):
"""Spin up the FastAPI control plane on an asyncio task. Returns """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) await asyncio.gather(*self._tasks, return_exceptions=True)
self._tasks.clear() self._tasks.clear()
self._started = False 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: for source in self.sources:
if hasattr(source, "aclose"): await source.aclose()
await source.aclose()
@property @property
def pollers(self) -> list[BasePoller]: 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: ... async def fetch(self, uri: str) -> FetchResult: ...
def discover( def discover(

View file

@ -75,6 +75,9 @@ class FSSource:
return False return False
return self._resolve_within_root(uri) is not None 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: async def head(self, uri: str) -> str | None:
path = self._resolve_within_root(uri) path = self._resolve_within_root(uri)
if path is None or not path.exists(): if path is None or not path.exists():

View file

@ -63,6 +63,9 @@ class S3Source:
def supports(self, uri: str) -> bool: def supports(self, uri: str) -> bool:
return uri.startswith(self.uri_prefix) 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: async def head(self, uri: str) -> str | None:
import obstore # type: ignore[import-not-found] 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: def supports(self, uri: str) -> bool:
return True return True
async def aclose(self) -> None:
pass
async def head(self, uri: str): async def head(self, uri: str):
return None return None