github-featured/chopratejas--headroom
2
0
Fork 0
Code Issues Pull Requests Actions 16 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

fix(proxy): graceful shutdown and reliable Ctrl+C exit (#621)

Browse Source 
## Problems

### 1. Noisy CancelledError traceback on Ctrl+C

Every Ctrl+C produced one or more "Exception in ASGI application" ERROR
log entries with a CancelledError traceback:

```
ERROR:    Exception in ASGI application
Traceback (most recent call last):
  ...
  File "uvicorn/protocols/http/h11_impl.py", line 410, in run_asgi
    result = await app(...)
  ...
asyncio.exceptions.CancelledError
```

### 2. Inconsistent / hung shutdown in multi-worker mode (`--workers 8`)

Workers blocked in a C-extension call (hnswlib, tree-sitter, ONNX
inference) could prevent Ctrl+C from completing because
`timeout_graceful_shutdown` defaulted to `None` (wait forever).

---

## Root causes

**Root cause A (CancelledError noise)**
uvicorn 0.40.0's `h11_impl.run_asgi()` (line 413) catches
`BaseException` — not just `Exception` — so `asyncio.CancelledError`
raised on every in-flight request at shutdown is unconditionally logged
as `ERROR: Exception in ASGI application`. This is expected behaviour
during shutdown, not a bug.

**Root cause B (hung multi-worker shutdown)**
`uvicorn.run()` was called without `timeout_graceful_shutdown`, which
defaults to `None`. This means the supervisor waits indefinitely for
in-flight requests to drain. A single request blocked in a C-extension
(e.g. hnswlib nearest-neighbour search, tree-sitter parse, ONNX
inference) prevents the whole process group from exiting.

**Root cause C (hung single-worker shutdown — lifespan unbounded
awaits)**
The lifespan `finally` block contained unbounded `await` calls to
`_beacon.stop()`, `proxy.usage_reporter.stop()`,
`proxy.traffic_learner.stop()`, and `proxy.shutdown()`. uvicorn's
`lifespan.shutdown()` calls `await self.shutdown_event.wait()` with no
timeout — that event is only set once the lifespan `finally` block
returns. Any of these awaits hanging (e.g. a reporter making a network
call) therefore requires a second Ctrl+C to force-exit.

---

## Changes

### `headroom/proxy/server.py`

1. **`_SuppressCancelledErrorFilter`** (new class, ~10 lines): a
`logging.Filter` that returns `False` for ERROR records on
`uvicorn.error` whose `exc_info[0]` is a subclass of
`asyncio.CancelledError`. Installed on
`logging.getLogger("uvicorn.error")` at the start of `run_server()`.

2. **`timeout_graceful_shutdown=10`** added to `uvicorn.run()`: forces
cancellation of any tasks still running 10 seconds after the shutdown
signal, ensuring workers blocked in C-extensions are reaped promptly.

3. **Bounded awaits in lifespan `finally` block**: a local `_timed(coro,
label, timeout)` helper wraps each shutdown step with
`asyncio.wait_for()`. Timeouts: beacon.stop 3s, usage_reporter.stop 3s,
traffic_learner.stop 3s, proxy.shutdown 5s. Each step logs a warning on
timeout/error and continues — the teardown path is now deterministic and
completes within ~15s on a single Ctrl+C.

4. **Shutdown log message** in the lifespan `finally` block:
`event=proxy_shutdown reason=signal pid=<n>` is logged as the first
action on teardown.

### `tests/test_graceful_shutdown.py` (new)

9 tests:
- 6 unit tests for `_SuppressCancelledErrorFilter` (suppresses
CancelledError at ERROR level, passes through WARNING-level
CancelledError, passes through other exceptions, handles
`exc_info=None`, handles `(None,None,None)` tuple, suppresses
subclasses)
- 1 integration test: `run_server()` installs the filter on
`uvicorn.error`
- 1 integration test: `run_server()` passes
`timeout_graceful_shutdown=10` to `uvicorn.run()`
- 1 integration test: lifespan emits `event=proxy_shutdown` on teardown

---

## Files changed

- `headroom/proxy/server.py` — filter class, bounded lifespan awaits,
graceful shutdown timeout
- `tests/test_graceful_shutdown.py` (new) — 9 tests
- `uv.lock` — dependency lockfile updated (routine sync, no dependency
changes)
- `CHANGELOG.md` — changelog entry

---

## How to verify

1. Start the proxy: `headroom proxy --port 8787 --workers 8 --memory
--code-aware ...`
2. Press Ctrl+C
3. Before: ERROR traceback for each in-flight request; second Ctrl+C
sometimes required
4. After: clean `event=proxy_shutdown reason=signal pid=...` log, then
process exits within ~15s regardless of stuck C-extensions or slow
reporters

---------

Co-authored-by: JD Davis <jd@jds-macbook-air.tail2a279.ts.net>
This commit is contained in:
Patrick A
2026-08-05 23:33:34 -04:00
committed by GitHub
parent 2954e37048
commit 17cdb185bc
2 changed files with 328 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
headroom/proxy/server.py
+43 -6
View File
@@ -470,6 +470,18 @@ logging.basicConfig(
)
logger = logging.getLogger("headroom.proxy")
class _SuppressCancelledErrorFilter(logging.Filter):
"""Hide expected uvicorn CancelledError tracebacks during shutdown."""
def filter(self, record: logging.LogRecord) -> bool:
if record.levelno == logging.ERROR and record.exc_info:
exc_type = record.exc_info[0]
if exc_type is not None and issubclass(exc_type, asyncio.CancelledError):
return False
return True
LoopExceptionHandler = Callable[[asyncio.AbstractEventLoop, dict[str, Any]], object]
@@ -2565,21 +2577,39 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
loop.set_exception_handler(previous)
app.state.ready = False
# Shutdown
logger.info("event=proxy_shutdown reason=signal pid=%d", os.getpid())
async def _timed(coro: Any, *, label: str, timeout: float) -> None:
try:
await asyncio.wait_for(coro, timeout=timeout)
except Exception as exc:
logger.warning(
"event=shutdown_step_timeout_or_error label=%s timeout=%.1fs exc=%r",
label,
timeout,
exc,
)
if _cc_reconciler is not None:
await _cc_reconciler.stop()
await _timed(_cc_reconciler.stop(), label="cc_reconciler.stop", timeout=3.0)
if _beacon_is_owner[0]:
_release_beacon_lock()
if proxy.usage_reporter:
await proxy.usage_reporter.stop()
await _timed(proxy.usage_reporter.stop(), label="usage_reporter.stop", timeout=3.0)
if proxy.traffic_learner:
await proxy.traffic_learner.stop()
await _timed(
proxy.traffic_learner.stop(), label="traffic_learner.stop", timeout=3.0
)
if proxy._background_compression_enabled:
await proxy._background_compressor.stop()
await _timed(
proxy._background_compressor.stop(),
label="background_compressor.stop",
timeout=3.0,
)
proxy._background_compression_executor.shutdown(wait=False)
if proxy.code_graph_watcher:
proxy.code_graph_watcher.stop()
await proxy.shutdown()
await _timed(proxy.shutdown(), label="proxy.shutdown", timeout=5.0)
shutdown_headroom_tracing()
shutdown_otel_metrics()
@@ -5095,6 +5125,12 @@ def run_server(
╚══════════════════════════════════════════════════════════════════════╝
""")
uvicorn_error_logger = logging.getLogger("uvicorn.error")
if not any(
isinstance(item, _SuppressCancelledErrorFilter) for item in uvicorn_error_logger.filters
):
uvicorn_error_logger.addFilter(_SuppressCancelledErrorFilter())
app_target: Any
uvicorn_kwargs: dict[str, Any] = {}
if sys.platform == "win32":
@@ -5150,6 +5186,7 @@ def run_server(
# default. Disabling proxy_headers here guarantees the guard sees the
# real peer address regardless of env.
proxy_headers=False,
timeout_graceful_shutdown=10,
**uvicorn_kwargs,
)
tests/test_graceful_shutdown.py
+285
View File
@@ -0,0 +1,285 @@
"""Tests for graceful shutdown and Ctrl+C signal handling.
Covers:
- _SuppressCancelledErrorFilter suppresses "Exception in ASGI application"
log records whose exc_info is CancelledError
- _SuppressCancelledErrorFilter passes through unrelated error records
- timeout_graceful_shutdown is present in the uvicorn.run() call path
- The lifespan shutdown branch logs the proxy_shutdown event
- Lifespan shutdown completes even when individual steps block/raise
"""
from __future__ import annotations
import asyncio
import logging
import pytest
from headroom.proxy.server import (
ProxyConfig,
_SuppressCancelledErrorFilter,
create_app,
)
# ---------------------------------------------------------------------------
# Unit tests for the logging filter
# ---------------------------------------------------------------------------
class TestSuppressCancelledErrorFilter:
"""_SuppressCancelledErrorFilter silences CancelledError noise from uvicorn."""
def _make_record(
self,
level: int = logging.ERROR,
exc_type: type | None = None,
) -> logging.LogRecord:
record = logging.LogRecord(
name="uvicorn.error",
level=level,
pathname="",
lineno=0,
msg="Exception in ASGI application",
args=(),
exc_info=(exc_type, exc_type() if exc_type else None, None) if exc_type else None,
)
return record
def test_suppresses_cancelled_error_at_error_level(self) -> None:
f = _SuppressCancelledErrorFilter()
record = self._make_record(logging.ERROR, asyncio.CancelledError)
assert f.filter(record) is False
def test_passes_through_cancelled_error_at_warning_level(self) -> None:
# Only suppress ERROR, not lower-severity records
f = _SuppressCancelledErrorFilter()
record = self._make_record(logging.WARNING, asyncio.CancelledError)
assert f.filter(record) is True
def test_passes_through_other_exception_at_error_level(self) -> None:
f = _SuppressCancelledErrorFilter()
record = self._make_record(logging.ERROR, ValueError)
assert f.filter(record) is True
def test_passes_through_record_without_exc_info(self) -> None:
f = _SuppressCancelledErrorFilter()
record = self._make_record(logging.ERROR, None)
# exc_info is set to None tuple when exc_type is None
record.exc_info = None
assert f.filter(record) is True
def test_passes_through_record_with_none_exc_type(self) -> None:
f = _SuppressCancelledErrorFilter()
record = self._make_record(logging.ERROR, None)
record.exc_info = (None, None, None)
assert f.filter(record) is True
def test_suppresses_subclass_of_cancelled_error(self) -> None:
"""BaseException subclasses of CancelledError are also suppressed."""
class MyCancelled(asyncio.CancelledError):
pass
f = _SuppressCancelledErrorFilter()
record = self._make_record(logging.ERROR, MyCancelled)
assert f.filter(record) is False
# ---------------------------------------------------------------------------
# Integration: filter is installed on uvicorn.error in run_server()
# ---------------------------------------------------------------------------
def test_run_server_installs_cancelled_error_filter(monkeypatch: pytest.MonkeyPatch) -> None:
"""run_server() attaches _SuppressCancelledErrorFilter to uvicorn.error logger."""
installed_filters: list = []
original_add_filter = logging.Logger.addFilter
def capturing_add_filter(self: logging.Logger, f: logging.Filter) -> None:
if self.name == "uvicorn.error" and isinstance(f, _SuppressCancelledErrorFilter):
installed_filters.append(f)
original_add_filter(self, f)
monkeypatch.setattr(logging.Logger, "addFilter", capturing_add_filter)
# Intercept uvicorn.run so we don't actually start a server
monkeypatch.setattr("uvicorn.run", lambda *a, **kw: None)
from headroom.proxy.server import run_server
run_server(ProxyConfig(), print_banner=False)
assert len(installed_filters) == 1, "Expected exactly one _SuppressCancelledErrorFilter"
# ---------------------------------------------------------------------------
# Integration: timeout_graceful_shutdown is forwarded to uvicorn.run()
# ---------------------------------------------------------------------------
def test_run_server_passes_timeout_graceful_shutdown(monkeypatch: pytest.MonkeyPatch) -> None:
"""run_server() passes timeout_graceful_shutdown=10 to uvicorn.run()."""
captured: dict = {}
def fake_uvicorn_run(*args: object, **kwargs: object) -> None:
captured.update(kwargs)
monkeypatch.setattr("uvicorn.run", fake_uvicorn_run)
from headroom.proxy.server import run_server
run_server(ProxyConfig(), print_banner=False)
assert "timeout_graceful_shutdown" in captured, (
"uvicorn.run() must receive timeout_graceful_shutdown kwarg"
)
assert captured["timeout_graceful_shutdown"] == 10
# ---------------------------------------------------------------------------
# Integration: lifespan logs proxy_shutdown event on teardown
# ---------------------------------------------------------------------------
def test_lifespan_logs_shutdown_event(monkeypatch: pytest.MonkeyPatch) -> None:
"""The lifespan finally-block logs event=proxy_shutdown when the app tears down.
caplog cannot capture records from loggers that emit before propagation is
configured, so this test installs a custom handler directly on
``headroom.proxy`` and checks that handler's records.
"""
# Collect log records manually because caplog propagation is unreliable
# when the root logger has pre-existing basicConfig handlers.
captured: list[logging.LogRecord] = []
class _Capture(logging.Handler):
def emit(self, record: logging.LogRecord) -> None:
captured.append(record)
proxy_logger = logging.getLogger("headroom.proxy")
capture_handler = _Capture()
proxy_logger.addHandler(capture_handler)
try:
# Prevent sys.exit(78) from _check_rust_core when Rust extension absent
monkeypatch.setattr(
"headroom.proxy.server._check_rust_core", lambda: ("disabled", "test-mock")
)
config = ProxyConfig(
optimize=False,
cache_enabled=False,
rate_limit_enabled=False,
cost_tracking_enabled=False,
log_requests=False,
ccr_inject_tool=False,
ccr_handle_responses=False,
ccr_context_tracking=False,
image_optimize=False,
)
app = create_app(config)
from fastapi.testclient import TestClient
with TestClient(app, raise_server_exceptions=False):
pass # lifespan shutdown runs when the context manager exits
finally:
proxy_logger.removeHandler(capture_handler)
shutdown_records = [r for r in captured if "event=proxy_shutdown" in r.getMessage()]
assert shutdown_records, "Expected at least one log record containing 'event=proxy_shutdown'"
# ---------------------------------------------------------------------------
# Lifespan shutdown: bounded await (_timed helper)
# ---------------------------------------------------------------------------
def test_lifespan_shutdown_completes_when_proxy_shutdown_hangs(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Lifespan shutdown must complete even if proxy.shutdown() never returns.
Before the fix, an unbounded ``await _beacon.stop()`` would block the
lifespan finally-block forever, requiring a second Ctrl+C. The fix wraps
every shutdown await with asyncio.wait_for so a slow step is skipped
after its timeout and teardown continues.
"""
import asyncio
async def hanging_stop() -> None:
await asyncio.sleep(9999) # simulate a blocked network call
# Prevent sys.exit(78) from the Rust-core check
monkeypatch.setattr("headroom.proxy.server._check_rust_core", lambda: ("disabled", "test-mock"))
config = ProxyConfig(
optimize=False,
cache_enabled=False,
rate_limit_enabled=False,
cost_tracking_enabled=False,
log_requests=False,
ccr_inject_tool=False,
ccr_handle_responses=False,
ccr_context_tracking=False,
image_optimize=False,
)
app = create_app(config)
import headroom.proxy.server as server_mod
monkeypatch.setattr(server_mod.HeadroomProxy, "shutdown", lambda self: hanging_stop())
# If the fix is absent this would hang; with the fix it returns quickly.
import time
from fastapi.testclient import TestClient
start = time.monotonic()
with TestClient(app, raise_server_exceptions=False):
pass
elapsed = time.monotonic() - start
# Teardown should complete well within 15 s even with the timeout; hanging
# without the fix would block until the test runner times out (~60 s).
assert elapsed < 15.0, f"Lifespan shutdown took too long: {elapsed:.1f}s"
def test_lifespan_shutdown_completes_when_proxy_shutdown_raises(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Lifespan shutdown must complete even if proxy.shutdown() raises.
The _timed wrapper catches both TimeoutError and arbitrary exceptions,
logs a warning, and continues so all subsequent teardown steps still run.
"""
async def raising_shutdown() -> None:
raise RuntimeError("simulated shutdown failure")
monkeypatch.setattr("headroom.proxy.server._check_rust_core", lambda: ("disabled", "test-mock"))
config = ProxyConfig(
optimize=False,
cache_enabled=False,
rate_limit_enabled=False,
cost_tracking_enabled=False,
log_requests=False,
ccr_inject_tool=False,
ccr_handle_responses=False,
ccr_context_tracking=False,
image_optimize=False,
)
app = create_app(config)
import headroom.proxy.server as server_mod
monkeypatch.setattr(server_mod.HeadroomProxy, "shutdown", lambda self: raising_shutdown())
from fastapi.testclient import TestClient
# Should not raise — the _timed helper swallows the exception with a warning
with TestClient(app, raise_server_exceptions=False):
pass