"""Keep the listening socket alive while replacement ASGI workers become ready.""" from __future__ import annotations import asyncio import logging import multiprocessing import signal import threading import time from dataclasses import dataclass from functools import partial from multiprocessing.connection import Connection from multiprocessing.process import BaseProcess from pathlib import Path import uvicorn from uvicorn._subprocess import get_subprocess from uvicorn.importer import import_from_string logger = logging.getLogger("uvicorn.error") class _ReadinessGate: def __init__(self, app, check): self.app, self.check = app, check async def __call__(self, scope, receive, send): if scope["type"] != "lifespan": return await self.app(scope, receive, send) async def checked_send(message): if message["type"] == "lifespan.startup.complete" and not await self.check(): await send({"type": "lifespan.startup.failed", "message": "API readiness check failed"}) else: await send(message) await self.app(scope, receive, checked_send) class _ReadyServer(uvicorn.Server): def __init__(self, config: uvicorn.Config, control: Connection): super().__init__(config) self.control = control async def startup(self, sockets=None) -> None: await super().startup(sockets=sockets) if self.started and not self.should_exit: self.control.send("ready") async def shutdown(self, sockets=None) -> None: # Acknowledge that the old worker has actually stopped accepting. Its # established streams may still run for minutes after this handshake. for listener in self.servers: listener.close() for sock in sockets or []: sock.close() self.control.send("draining") await super().shutdown(sockets=sockets) def _serve_worker(*, config: uvicorn.Config, control: Connection, readiness_check: str | None, sockets) -> None: config.reload = False config.workers = 1 if readiness_check: config.load() config.loaded_app = _ReadinessGate(config.loaded_app, import_from_string(readiness_check)) server = _ReadyServer(config, control) async def receive_commands() -> None: while True: if control.poll(): try: command = control.recv() except EOFError: command = "drain" if command == "drain": server.should_exit = True return await asyncio.sleep(0.05) async def serve() -> None: receiver = asyncio.create_task(receive_commands()) try: await server.serve(sockets=sockets) finally: receiver.cancel() await asyncio.gather(receiver, return_exceptions=True) control.close() config.setup_event_loop() asyncio.run(serve()) @dataclass class Worker: process: BaseProcess control: Connection def drain(self, *, wait_timeout: float = 0) -> bool: if self.process.is_alive(): try: self.control.send("drain") if wait_timeout: deadline = time.monotonic() + wait_timeout while time.monotonic() < deadline: if self.control.poll(0.05) and self.control.recv() == "draining": return True if not self.process.is_alive(): return True return False except (BrokenPipeError, EOFError, OSError): return not self.process.is_alive() return True class HandoffSupervisor: """Start the replacement before draining the worker with existing connections.""" def __init__( self, config: uvicorn.Config, *, startup_timeout: float = 60, reload_dirs: list[Path] | None = None, poll_interval: float = 0.25, readiness_check: str | None = None, ): self.config = config self.startup_timeout = startup_timeout self.reload_dirs = reload_dirs or [Path("app"), Path("samples/peripheral")] self.poll_interval = poll_interval self.readiness_check = readiness_check self.current: Worker | None = None self.retiring: list[Worker] = [] self.socket = config.bind_socket() self.socket.listen(config.backlog) self.should_exit = threading.Event() self._files: dict[Path, tuple[int, int]] = {} def _candidate(self) -> Worker | None: parent = child = process = None ready = False try: parent, child = multiprocessing.get_context("spawn").Pipe() process = get_subprocess( config=self.config, target=partial(_serve_worker, config=self.config, control=child, readiness_check=self.readiness_check), sockets=[self.socket], ) worker = Worker(process, parent) process.start() child.close() deadline = time.monotonic() + self.startup_timeout while not self.should_exit.is_set() and time.monotonic() < deadline: if parent.poll(0.05): try: ready = parent.recv() == "ready" and process.is_alive() except EOFError: break if ready: return worker if not process.is_alive(): break logger.error("Replacement worker did not become ready; keeping the current worker") return None except Exception: logger.exception("Replacement startup or readiness handshake failed; keeping the current worker") return None finally: if child is not None: child.close() if not ready: # Only an unready candidate is terminated. The serving worker # and workers draining established connections are untouched. if process is not None and process.pid is not None: if process.is_alive(): process.terminate() process.join(timeout=2) if process.is_alive(): process.kill() process.join() if parent is not None: parent.close() def start(self) -> bool: self._files = self._snapshot() self.current = self._candidate() return self.current is not None def reload(self) -> bool: self.reap() candidate = self._candidate() if candidate is None: return False old, self.current = self.current, candidate logger.info("Worker %s is ready to accept new connections", candidate.process.pid) if old is not None: acknowledged = old.drain(wait_timeout=2) self.retiring.append(old) if not acknowledged: logger.error("Worker %s has not yet acknowledged draining", old.process.pid) return False logger.info("Worker %s is draining existing connections", old.process.pid) return True def reap(self) -> None: remaining = [] for worker in self.retiring: if worker.process.is_alive(): remaining.append(worker) else: worker.process.join() worker.control.close() self.retiring = remaining def _snapshot(self) -> dict[Path, tuple[int, int]]: files = {} for directory in self.reload_dirs: for path in directory.rglob("*.py"): try: stat = path.stat() files[path.resolve()] = (stat.st_mtime_ns, stat.st_size) except OSError: continue return files def run(self) -> None: signals = (signal.SIGINT, signal.SIGTERM) original = {sig: signal.signal(sig, lambda *_: self.should_exit.set()) for sig in signals} try: if not self.start(): raise RuntimeError("The initial API worker failed to start") pending_reload = False while not self.should_exit.wait(self.poll_interval): self.reap() files = self._snapshot() if files != self._files: self._files = files pending_reload = True active_failed = self.current is None or not self.current.process.is_alive() if active_failed: logger.error("The active worker exited; starting a replacement") pending_reload = True # Coalesce rapid edits while old SSE connections drain, rather # than keeping an unbounded number of application generations. if pending_reload and (active_failed or len(self.retiring) < 2): pending_reload = False self.reload() finally: try: self.close() finally: for sig, handler in original.items(): signal.signal(sig, handler) def close(self) -> None: self.should_exit.set() workers = [*self.retiring, *([self.current] if self.current else [])] for worker in workers: worker.drain() self.socket.close() for worker in workers: worker.process.join() worker.control.close() self.current = None self.retiring = []