Add RSS-based worker recycling

CHANGELOG.md

@@ -1,5 +1,10 @@
 # Version history
 
+## 0.20.2
+
+- Add a `max_rss` config flag and `--max-rss` CLI option to gracefully recycle workers
+  whose current RSS exceeds a configured MiB limit.
+
 ## 0.20.1
 
 - Drain in-flight requests until `graceful_timeout` on shutdown.

README.md

@@ -35,3 +35,18 @@ anyio.run(serve, app, Config())
 See Hypercorn's
 [documentation](https://hypercorn.readthedocs.io/en/latest/how_to_guides/api_usage.html) for more
 details.
+
+## Worker recycling
+
+Workers can be recycled after a fixed number of requests with `--max-requests`, or after their
+current resident set size exceeds a limit with `--max-rss`:
+
+```bash
+anycorn --workers 4 --max-rss 512 module:app
+```
+
+`max_rss` is specified in MiB and is disabled by default. Anycorn samples the worker RSS
+periodically and uses the same graceful shutdown path as request-count recycling, so in-flight
+requests can drain until `graceful_timeout` before the master process respawns the worker. When
+`--workers 0` is used, no master process exists to respawn the worker; exceeding `max_rss` stops
+that single worker, matching `max_requests` behavior.

pyproject.toml

@@ -43,6 +43,7 @@ dependencies = [
     "typing_extensions; python_version<'3.11'",
     "wsproto >=0.14.0",
     "rich-click >=1.8.3,<2.0.0",
+    "psutil>=7.2.2",
 ]
 
 [project.optional-dependencies]

src/anycorn/__main__.py

@@ -102,6 +102,11 @@ def _load_config(config_path: str | None) -> Config:
     help="Maximum number of requests a worker will process before restarting",
     type=int,
 )
+@click.option(
+    "--max-rss",
+    help="Maximum RSS in MiB a worker may use before restarting",
+    type=int,
+)
 @click.option(
     "--max-requests-jitter",
     help="This jitter causes the max-requests per worker to be "
@@ -247,6 +252,7 @@ def main(  # noqa: C901 PLR0913 PLR0912 PLR0915
     graceful_timeout: int | None,
     read_timeout: int | None,
     max_requests: int | None,
+    max_rss: int | None,
     max_requests_jitter: int | None,
     group: int | None,
     worker_class: str | None,
@@ -312,6 +318,8 @@ def main(  # noqa: C901 PLR0913 PLR0912 PLR0915
         cfg.logconfig = log_config
     if max_requests is not None:
         cfg.max_requests = max_requests
+    if max_rss is not None:
+        cfg.max_rss = max_rss
     if max_requests_jitter is not None:
         cfg.max_requests_jitter = max_requests_jitter
     if pid is not None:

src/anycorn/config.py

@@ -139,6 +139,7 @@ def __init__(self) -> None:
     logger_class = Logger
     loglevel: str = "INFO"
     max_app_queue_size: int = 10
+    max_rss: int | None = None
     max_requests: int | None = None
     max_requests_jitter: int = 0
     pid_path: str | None = None

src/anycorn/run.py

@@ -18,6 +18,7 @@
 import anyio
 import anyio.abc
 import anyio.streams.tls
+import psutil
 import sniffio
 
 from .datagram import wrap_datagram_socket
@@ -39,6 +40,9 @@
 )
 from .worker_context import WorkerContext
 
+MAX_RSS_CHECK_INTERVAL = 30
+_BYTES_PER_MIB = 1024 * 1024
+
 if TYPE_CHECKING:
     from collections.abc import Awaitable, Callable
     from multiprocessing.context import BaseContext
@@ -51,6 +55,34 @@
     from exceptiongroup import BaseExceptionGroup
 
 
+def _read_current_rss() -> int:
+    """Return this process's current RSS in bytes."""
+    return psutil.Process().memory_info().rss
+
+
+async def _watch_max_rss(
+    context: WorkerContext,
+    config: Config,
+    read_current_rss: Callable[[], int] | None = None,
+) -> None:
+    if config.max_rss is None:
+        return
+
+    read_current_rss = _read_current_rss if read_current_rss is None else read_current_rss
+    max_rss = config.max_rss * _BYTES_PER_MIB
+    while True:
+        await context.sleep(MAX_RSS_CHECK_INTERVAL)
+        current_rss = read_current_rss()
+        if current_rss > max_rss:
+            await config.log.info(
+                "Worker RSS %.1f MiB exceeded max_rss %d MiB; recycling worker",
+                current_rss / _BYTES_PER_MIB,
+                config.max_rss,
+            )
+            await context.terminate.set()
+            return
+
+
 def run(config: Config) -> int:  # noqa: C901, PLR0912, PLR0915
     """Start the server, blocking until it exits, and return an exit code."""
     if config.pid_path is not None:
@@ -315,6 +347,7 @@ async def worker_serve(  # noqa: C901, PLR0912, PLR0915
                     if shutdown_trigger is not None:
                         tg.start_soon(raise_shutdown, shutdown_trigger, context.terminated.set)
                     tg.start_soon(raise_shutdown, context.terminate.wait, context.terminated.set)
+                    tg.start_soon(_watch_max_rss, context, config)
 
                     for udp_server in udp_servers:
                         await tg.start(udp_server.run)

tests/test___main__.py

@@ -52,6 +52,7 @@ def test_load_config(monkeypatch: MonkeyPatch) -> None:
         ("--ciphers", "DHE-RSA-AES128-SHA", "ciphers"),
         ("--keep-alive", 20, "keep_alive_timeout"),
         ("--keyfile", "/path", "keyfile"),
+        ("--max-rss", 128, "max_rss"),
         ("--pid", "/path", "pid_path"),
         ("--root-path", "/path", "root_path"),
         ("--workers", 2, "workers"),

tests/test_config.py

@@ -20,6 +20,7 @@
 
 access_log_format = "bob"
 h11_max_incomplete_size = 4
+max_rss = 128
 
 # The reuse option _create_sockets sets is platform-specific: SO_EXCLUSIVEADDRUSE on
 # Windows (so a second server on a busy port fails), SO_REUSEADDR elsewhere. getattr
@@ -64,6 +65,18 @@ def test_config_from_toml() -> None:
     _check_standard_config(config)
 
 
+def test_config_from_toml_sets_max_rss(tmp_path: Path) -> None:
+    path = tmp_path / "config.toml"
+    path.write_text(f"max_rss = {max_rss}\n")
+    config = Config.from_toml(path)
+    assert config.max_rss == max_rss
+
+
+def test_config_from_mapping_sets_max_rss() -> None:
+    config = Config.from_mapping(max_rss=max_rss)
+    assert config.max_rss == max_rss
+
+
 def test_create_ssl_context() -> None:
     path = str(Path(__file__).parent / "assets/config_ssl.py")
     config = Config.from_pyfile(path)

tests/test_run.py

@@ -140,6 +140,44 @@ async def run(
     assert observed == [True]
 
 
+@pytest.mark.anyio
+async def test_watch_max_rss_terminates_worker_when_over_limit(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    config = Config()
+    config.max_rss = 1
+    context = WorkerContext(None)
+    logged: list[tuple[str, tuple[object, ...]]] = []
+
+    async def info(message: str, *args: object) -> None:
+        logged.append((message, args))
+
+    monkeypatch.setattr(config.log, "info", info)
+    monkeypatch.setattr(anycorn.run, "MAX_RSS_CHECK_INTERVAL", 0)
+
+    await anycorn.run._watch_max_rss(context, config, lambda: 2 * 1024 * 1024)
+
+    assert context.terminate.is_set()
+    assert logged == [
+        (
+            "Worker RSS %.1f MiB exceeded max_rss %d MiB; recycling worker",
+            (2.0, 1),
+        )
+    ]
+
+
+@pytest.mark.anyio
+async def test_worker_serve_terminates_when_over_max_rss(monkeypatch: pytest.MonkeyPatch) -> None:
+    config = Config()
+    config.bind = ["127.0.0.1:0"]
+    config.max_rss = 1
+    monkeypatch.setattr(anycorn.run, "MAX_RSS_CHECK_INTERVAL", 0)
+    monkeypatch.setattr(anycorn.run, "_read_current_rss", lambda: 2 * 1024 * 1024)
+
+    with anyio.fail_after(5):
+        await worker_serve(wrap_app(app, config.wsgi_max_body_size, None), config)
+
+
 @pytest.mark.anyio
 async def test_udp_server_serialises_concurrent_sends() -> None:
     """QUIC sends from several tasks at once must not collide on the socket.