feat: better shutdowns

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"]

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(

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:

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

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

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)

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

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

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))])

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}"