diff --git a/application/Dockerfile b/application/Dockerfile index 0eb037d3..18eb943c 100644 --- a/application/Dockerfile +++ b/application/Dockerfile @@ -93,15 +93,18 @@ EXPOSE 7091 # Switch to non-root user USER appuser +# BoundedDrainUvicornWorker makes max_requests recycles safe with held-open SSE +# connections (see application/gunicorn_worker.py); with recycles now safe, +# --max-requests is raised (kept for memory hygiene) to cut churn. CMD ["gunicorn", \ "-w", "1", \ - "-k", "uvicorn_worker.UvicornWorker", \ + "-k", "application.gunicorn_worker.BoundedDrainUvicornWorker", \ "--bind", "0.0.0.0:7091", \ "--timeout", "180", \ "--graceful-timeout", "120", \ "--keep-alive", "5", \ "--worker-tmp-dir", "/dev/shm", \ - "--max-requests", "1000", \ - "--max-requests-jitter", "100", \ + "--max-requests", "5000", \ + "--max-requests-jitter", "500", \ "--config", "application/gunicorn_conf.py", \ "application.asgi:asgi_app"] diff --git a/application/api/devices/session.py b/application/api/devices/session.py index 8f15aabc..04ed30c5 100644 --- a/application/api/devices/session.py +++ b/application/api/devices/session.py @@ -13,6 +13,7 @@ from flask import Response, jsonify, make_response, request, stream_with_context from application.api.devices.auth import verify_device_session from application.core.settings import settings +from application.core.shutdown import is_shutting_down from application.devices.broker import get_broker from application.storage.db.repositories.device_audit_log import ( DeviceAuditLogRepository, @@ -107,6 +108,9 @@ def session_events(session_id: str) -> Response: try: last_keepalive = time.time() while not sess.closed.is_set(): + # Break promptly on shutdown (see application/core/shutdown.py). + if is_shutting_down(): + break now = time.time() if now - sess.last_activity_at > idle_seconds: yield _sse_event( diff --git a/application/api/events/routes.py b/application/api/events/routes.py index 039259f1..9080a58a 100644 --- a/application/api/events/routes.py +++ b/application/api/events/routes.py @@ -18,6 +18,7 @@ from flask import Blueprint, Response, jsonify, make_response, request, stream_w from application.cache import get_redis_instance from application.core.settings import settings +from application.core.shutdown import is_shutting_down from application.events.keys import ( connection_counter_key, replay_budget_key, @@ -411,6 +412,11 @@ def stream_events() -> Response: on_subscribe=_on_subscribe_callback, poll_timeout=SUBSCRIBE_POLL_INTERVAL_SECONDS, ): + # Break promptly on shutdown — this a2wsgi thread can't be + # cancelled by asyncio (see application/core/shutdown.py). + if is_shutting_down(): + break + # Flush snapshot on the first iteration after the SUBSCRIBE # callback ran. This runs at most once per connection. if replay_done and replay_lines: diff --git a/application/core/settings.py b/application/core/settings.py index a5b88cf9..681ead64 100644 --- a/application/core/settings.py +++ b/application/core/settings.py @@ -215,6 +215,10 @@ class Settings(BaseSettings): # Per-user durable backlog cap (~entries). At typical event rates this # gives ~24h of replay; tune up for verbose feeds, down for memory. EVENTS_STREAM_MAXLEN: int = 1000 + # Bounds uvicorn's graceful-shutdown drain (uvicorn_worker doesn't forward + # --graceful-timeout). Keep below the gunicorn --timeout (180) watchdog. + # Used by gunicorn_worker.BoundedDrainUvicornWorker. + GRACEFUL_SHUTDOWN_TIMEOUT_SECONDS: int = 30 # SSE keepalive comment cadence. Must sit under Cloudflare's 100s idle # close and iOS Safari's ~60s — 15s gives generous headroom. SSE_KEEPALIVE_SECONDS: int = 15 diff --git a/application/core/shutdown.py b/application/core/shutdown.py new file mode 100644 index 00000000..be1e4345 --- /dev/null +++ b/application/core/shutdown.py @@ -0,0 +1,28 @@ +"""Process-wide graceful-shutdown flag. + +Raised by ``BoundedDrainUvicornWorker`` when a drain starts and polled by the +SSE/long-poll generators (which run in a2wsgi threads asyncio can't cancel) so +they return promptly instead of hanging the worker until the ``--timeout`` +watchdog kills it. Standalone module so the worker can import it without the app. +""" + +from __future__ import annotations + +import threading + +_shutting_down = threading.Event() + + +def begin_shutdown() -> None: + """Mark the process as shutting down so streaming generators stop looping.""" + _shutting_down.set() + + +def is_shutting_down() -> bool: + """Return ``True`` once the server has begun a graceful shutdown.""" + return _shutting_down.is_set() + + +def reset_shutdown() -> None: + """Clear the shutdown flag. Intended for tests only.""" + _shutting_down.clear() diff --git a/application/gunicorn_worker.py b/application/gunicorn_worker.py new file mode 100644 index 00000000..6a6dc42d --- /dev/null +++ b/application/gunicorn_worker.py @@ -0,0 +1,57 @@ +"""Gunicorn worker that bounds uvicorn's graceful-shutdown drain. + +``uvicorn_worker`` doesn't forward gunicorn's ``--graceful-timeout`` to uvicorn's +``timeout_graceful_shutdown``, so after a ``max_requests`` recycle the drain is +unbounded: a held-open SSE/long-poll connection (a Flask/WSGI generator in +a2wsgi's thread pool, which asyncio can't cancel) hangs the worker until the +``--timeout`` watchdog SIGKILLs it (mislabeled "Perhaps out of memory?"). + +This worker bounds the drain (``timeout_graceful_shutdown`` from settings) and +raises the shutdown flag at drain start — including the signal-less +``max_requests`` path — so generators stop within one poll tick. Wire in via +``-k application.gunicorn_worker.BoundedDrainUvicornWorker``. +""" + +from __future__ import annotations + +import socket +import sys +from typing import Any + +from gunicorn.arbiter import Arbiter +from uvicorn.server import Server +from uvicorn_worker import UvicornWorker + +from application.core.settings import settings +from application.core.shutdown import begin_shutdown + + +class _ShutdownAwareServer(Server): + """uvicorn ``Server`` that raises the shutdown flag when a drain starts. + + ``serve()`` always calls ``shutdown()`` when its loop exits — including the + signal-less ``max_requests`` recycle — so this is the reliable hook point. + """ + + async def shutdown(self, sockets: list[socket.socket] | None = None) -> None: + begin_shutdown() + await super().shutdown(sockets=sockets) + + +class BoundedDrainUvicornWorker(UvicornWorker): + """UvicornWorker with a bounded drain and shutdown-aware streaming.""" + + CONFIG_KWARGS: dict[str, Any] = { + "loop": "auto", + "http": "auto", + "timeout_graceful_shutdown": settings.GRACEFUL_SHUTDOWN_TIMEOUT_SECONDS, + } + + async def _serve(self) -> None: + # Mirrors UvicornWorker._serve but swaps in the shutdown-aware Server. + self.config.app = self.wsgi + server = _ShutdownAwareServer(config=self.config) + self._install_sigquit_handler() + await server.serve(sockets=self.sockets) + if not server.started: + sys.exit(Arbiter.WORKER_BOOT_ERROR) diff --git a/tests/core/test_gunicorn_worker.py b/tests/core/test_gunicorn_worker.py new file mode 100644 index 00000000..0cf16979 --- /dev/null +++ b/tests/core/test_gunicorn_worker.py @@ -0,0 +1,48 @@ +"""Unit tests for the bounded-drain gunicorn worker.""" + +from unittest import mock + +import pytest + +from application.core import shutdown +from application.core.settings import settings +from application.gunicorn_worker import ( + BoundedDrainUvicornWorker, + _ShutdownAwareServer, +) + + +@pytest.fixture(autouse=True) +def _reset_flag(): + shutdown.reset_shutdown() + yield + shutdown.reset_shutdown() + + +@pytest.mark.unit +def test_worker_bounds_graceful_shutdown(): + cfg = BoundedDrainUvicornWorker.CONFIG_KWARGS + assert cfg["timeout_graceful_shutdown"] == settings.GRACEFUL_SHUTDOWN_TIMEOUT_SECONDS + assert isinstance(cfg["timeout_graceful_shutdown"], int) + # Must sit under the gunicorn --timeout (180) or the watchdog fires first. + assert 0 < cfg["timeout_graceful_shutdown"] < 180 + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_server_shutdown_raises_flag(): + from uvicorn.config import Config + + async def _dummy_asgi(scope, receive, send): # pragma: no cover - never called + pass + + server = _ShutdownAwareServer(Config(app=_dummy_asgi)) + assert shutdown.is_shutting_down() is False + + # Stub the heavy uvicorn drain so the test only exercises our override. + with mock.patch("uvicorn.server.Server.shutdown", new=mock.AsyncMock()) as sup: + await server.shutdown() + + # The flag is raised before delegating to uvicorn's real shutdown. + assert shutdown.is_shutting_down() is True + sup.assert_awaited_once() diff --git a/tests/core/test_shutdown_flag.py b/tests/core/test_shutdown_flag.py new file mode 100644 index 00000000..f5ff4a17 --- /dev/null +++ b/tests/core/test_shutdown_flag.py @@ -0,0 +1,37 @@ +"""Unit tests for the process-wide graceful-shutdown flag.""" + +import pytest + +from application.core import shutdown + + +@pytest.fixture(autouse=True) +def _reset_flag(): + shutdown.reset_shutdown() + yield + shutdown.reset_shutdown() + + +@pytest.mark.unit +def test_flag_starts_clear(): + assert shutdown.is_shutting_down() is False + + +@pytest.mark.unit +def test_begin_shutdown_sets_flag(): + shutdown.begin_shutdown() + assert shutdown.is_shutting_down() is True + + +@pytest.mark.unit +def test_begin_shutdown_is_idempotent(): + shutdown.begin_shutdown() + shutdown.begin_shutdown() + assert shutdown.is_shutting_down() is True + + +@pytest.mark.unit +def test_reset_clears_flag(): + shutdown.begin_shutdown() + shutdown.reset_shutdown() + assert shutdown.is_shutting_down() is False diff --git a/tests/integration/_drain_harness_app.py b/tests/integration/_drain_harness_app.py new file mode 100644 index 00000000..743b75c1 --- /dev/null +++ b/tests/integration/_drain_harness_app.py @@ -0,0 +1,51 @@ +"""Minimal ASGI app for the worker graceful-shutdown drain e2e. + +Mirrors production (a Flask SSE generator behind a2wsgi's thread pool) without +Postgres/Redis. ``/sse`` holds the connection open like an idle subscriber; +``DRAIN_HARNESS_COOPERATIVE=1`` makes it poll the real +``application.core.shutdown.is_shutting_down`` flag (the fix), else it +reproduces the pre-fix hang. +""" + +from __future__ import annotations + +import os +import time + +from a2wsgi import WSGIMiddleware +from flask import Flask, Response +from starlette.applications import Starlette +from starlette.routing import Mount + +from application.core.shutdown import is_shutting_down + +_COOPERATIVE = os.environ.get("DRAIN_HARNESS_COOPERATIVE") == "1" +_POLL_SECONDS = 1.0 +_MAX_HOLD_SECONDS = 120.0 + +flask_app = Flask(__name__) + + +@flask_app.get("/health") +def health() -> Response: + return Response('{"ok": true}', mimetype="application/json") + + +@flask_app.get("/sse") +def sse() -> Response: + def generate(): + # Emit headers immediately (like the real ": connected" frame) so the + # client establishes the stream, then hold without writing further. + yield ": connected\n\n" + deadline = time.monotonic() + _MAX_HOLD_SECONDS + while time.monotonic() < deadline: + # Cooperative variant bails on the flag; non-cooperative pins the + # a2wsgi thread to the deadline (the pre-fix hang). + if _COOPERATIVE and is_shutting_down(): + break + time.sleep(_POLL_SECONDS) + + return Response(generate(), mimetype="text/event-stream") + + +asgi_app = Starlette(routes=[Mount("/", app=WSGIMiddleware(flask_app, workers=8))]) diff --git a/tests/integration/test_worker_drain_e2e.py b/tests/integration/test_worker_drain_e2e.py new file mode 100644 index 00000000..e22d298a --- /dev/null +++ b/tests/integration/test_worker_drain_e2e.py @@ -0,0 +1,206 @@ +"""End-to-end regression tests for the gunicorn graceful-shutdown drain fix. + +Boots real gunicorn against a minimal a2wsgi+Flask SSE app, holds an SSE +connection open, trips a ``max_requests`` recycle, and checks the worker log: +RED proves the stock worker is ``WORKER TIMEOUT``'d; GREEN proves the +bounded-drain worker recycles cleanly. Opt-in (slow). Run with: + python -m pytest tests/integration/test_worker_drain_e2e.py -o addopts="" +""" + +from __future__ import annotations + +import os +import signal +import socket +import subprocess +import sys +import threading +import time +import urllib.request +from pathlib import Path + +import pytest + +pytestmark = [pytest.mark.integration, pytest.mark.slow] + +_ROOT = Path(__file__).resolve().parents[2] +_APP = "tests.integration._drain_harness_app:asgi_app" +_GUNICORN = Path(sys.executable).with_name("gunicorn") + +# Shrunk timing so a full recycle takes seconds, not the production minutes. +_MAX_REQUESTS = "5" +_TIMEOUT = "6" # gunicorn worker-timeout watchdog +_GRACEFUL = "30" # gunicorn --graceful-timeout (not the lever; here for parity) +_GUNICORN_CONF = _ROOT / "application" / "gunicorn_conf.py" + +pytestmark.append( + pytest.mark.skipif(not _GUNICORN.exists(), reason="gunicorn binary not found") +) + + +def _free_port() -> int: + s = socket.socket() + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + s.close() + return port + + +def _boot(worker_class: str, port: int, logpath: Path, extra_env: dict) -> tuple: + env = dict(os.environ) + env.update( + { + "OTEL_SDK_DISABLED": "true", + "AUTO_MIGRATE": "false", + "AUTO_CREATE_DB": "false", + "GRACEFUL_SHUTDOWN_TIMEOUT_SECONDS": "3", + **extra_env, + } + ) + cmd = [ + str(_GUNICORN), + "-w", "1", + "-k", worker_class, + "--bind", f"127.0.0.1:{port}", + "--timeout", _TIMEOUT, + "--graceful-timeout", _GRACEFUL, + "--keep-alive", "2", + "--max-requests", _MAX_REQUESTS, + "--max-requests-jitter", "0", + "--pythonpath", str(_ROOT), + "--config", str(_GUNICORN_CONF), + _APP, + ] + fh = open(logpath, "w") + proc = subprocess.Popen( + cmd, cwd=str(_ROOT), env=env, stdout=fh, stderr=subprocess.STDOUT, + start_new_session=True, + ) + return proc, fh + + +def _reap(proc, fh) -> None: + try: + os.killpg(proc.pid, signal.SIGTERM) + except Exception: + proc.terminate() + try: + proc.wait(timeout=5) + except Exception: + try: + os.killpg(proc.pid, signal.SIGKILL) + except Exception: + proc.kill() + finally: + try: + fh.close() + except Exception: + pass + + +def _get(url: str, timeout: float = 3.0) -> int: + try: + with urllib.request.urlopen(url, timeout=timeout) as r: + r.read() + return r.status + except Exception: + return 0 + + +def _wait_http(base: str, timeout: float = 30.0) -> bool: + end = time.monotonic() + timeout + while time.monotonic() < end: + if _get(base + "/health", timeout=2) == 200: + return True + time.sleep(0.2) + return False + + +def _hold_sse(url: str, hold: float = 25.0) -> threading.Thread: + def _run(): + try: + urllib.request.urlopen(url, timeout=hold).read() + except Exception: + pass + + t = threading.Thread(target=_run, daemon=True) + t.start() + return t + + +def _trip_recycle(base: str, n: int = 8) -> None: + # Fire more than --max-requests completed responses so the recycle trips + # deterministically (the held SSE response never completes, so it doesn't + # count, and some requests race the drain once it starts). + for _ in range(n): + _get(base + "/health", timeout=3) + + +def _wait_for(logpath: Path, needle: str, timeout: float) -> str: + end = time.monotonic() + timeout + text = "" + while time.monotonic() < end: + text = logpath.read_text() + if needle in text: + return text + time.sleep(0.3) + return text + + +@pytest.mark.integration +@pytest.mark.slow +def test_stock_worker_is_force_killed_with_held_sse(tmp_path): + """RED: stock worker + non-cooperative SSE -> drain hangs -> WORKER TIMEOUT.""" + port = _free_port() + log = tmp_path / "stock.log" + base = f"http://127.0.0.1:{port}" + proc, fh = _boot("uvicorn_worker.UvicornWorker", port, log, {}) + try: + assert _wait_http(base), f"app never came up:\n{log.read_text()}" + _hold_sse(base + "/sse", hold=25) + time.sleep(1.0) # let the SSE request register before tripping + _trip_recycle(base) + text = _wait_for(log, "WORKER TIMEOUT", timeout=15) + finally: + _reap(proc, fh) + + assert "Maximum request limit" in text, text + # The held SSE hangs the unbounded drain until the watchdog force-kills it. + assert "WORKER TIMEOUT" in text, text + + +@pytest.mark.integration +@pytest.mark.slow +def test_bounded_drain_worker_exits_cleanly_with_held_sse(tmp_path): + """GREEN: bounded-drain worker + cooperative SSE -> clean recycle, no kill.""" + port = _free_port() + log = tmp_path / "fixed.log" + base = f"http://127.0.0.1:{port}" + proc, fh = _boot( + "application.gunicorn_worker.BoundedDrainUvicornWorker", + port, + log, + {"DRAIN_HARNESS_COOPERATIVE": "1"}, + ) + try: + assert _wait_http(base), f"app never came up:\n{log.read_text()}" + _hold_sse(base + "/sse", hold=25) + time.sleep(1.0) + _trip_recycle(base) + # Clean recycle => replacement worker reaches "Application startup complete" twice. + end = time.monotonic() + 14 + text = "" + while time.monotonic() < end: + text = log.read_text() + if text.count("Application startup complete") >= 2: + break + time.sleep(0.3) + finally: + _reap(proc, fh) + + assert "Maximum request limit" in text, text + assert text.count("Application startup complete") >= 2, ( + f"replacement worker did not boot cleanly:\n{text}" + ) + for bad in ("WORKER TIMEOUT", "was sent SIGABRT", "was sent SIGKILL"): + assert bad not in text, f"unexpected watchdog kill ({bad}):\n{text}"