Shut down the multiprocess supervisor when a worker dies before startup

Author: Kludex

tests/supervisors/test_multiprocess.py

@@ -3,7 +3,6 @@
 import functools
 import os
 import signal
-import socket
 import threading
 import time
 from collections.abc import Callable
@@ -13,6 +12,7 @@
 
 from uvicorn import Config
 from uvicorn._types import ASGIReceiveCallable, ASGISendCallable, Scope
+from uvicorn.server import STARTUP_FAILURE
 from uvicorn.supervisors import Multiprocess
 from uvicorn.supervisors.multiprocess import Process
 
@@ -41,22 +41,32 @@ async def app(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallable
     pass  # pragma: no cover
 
 
-def run(sockets: list[socket.socket] | None) -> None:
-    while True:  # pragma: no cover
-        time.sleep(1)
+async def lifespan_failure_app(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallable) -> None:
+    assert scope["type"] == "lifespan"
+    message = await receive()
+    if message["type"] == "lifespan.startup":
+        raise RuntimeError("Startup failed")
 
 
 def test_process_ping_pong() -> None:
-    process = Process(Config(app=app), target=lambda x: None, sockets=[])
+    process = Process(Config(app=app), sockets=[])
     threading.Thread(target=process.always_pong, daemon=True).start()
     assert process.ping()
 
 
 def test_process_ping_pong_timeout() -> None:
-    process = Process(Config(app=app), target=lambda x: None, sockets=[])
+    process = Process(Config(app=app), sockets=[])
     assert not process.ping(0.1)
 
 
+def test_process_ready() -> None:
+    process = Process(Config(app=app), sockets=[])
+    assert not process.ready
+    process.notify_started()
+    assert process.ready
+    assert process.ready
+
+
 @new_console_in_windows
 def test_multiprocess_run() -> None:
     """
@@ -66,7 +76,7 @@ def test_multiprocess_run() -> None:
     quit immediately.
     """
     config = Config(app=app, workers=2)
-    supervisor = Multiprocess(config, target=run, sockets=[])
+    supervisor = Multiprocess(config, sockets=[])
     threading.Thread(target=supervisor.run, daemon=True).start()
     supervisor.signal_queue.append(signal.SIGINT)
     supervisor.join_all()
@@ -78,9 +88,12 @@ def test_multiprocess_health_check() -> None:
     Ensure that the health check works as expected.
     """
     config = Config(app=app, workers=2)
-    supervisor = Multiprocess(config, target=run, sockets=[])
+    supervisor = Multiprocess(config, sockets=[])
     threading.Thread(target=supervisor.run, daemon=True).start()
-    time.sleep(1)
+    deadline = time.monotonic() + 10
+    while len(supervisor.processes) < 2 or not all(p.ready for p in supervisor.processes):  # pragma: no cover
+        assert time.monotonic() < deadline, "Timed out waiting for processes to become ready"
+        time.sleep(0.1)
     process = supervisor.processes[0]
     process.kill()
     assert not process.is_alive()
@@ -92,13 +105,32 @@ def test_multiprocess_health_check() -> None:
     supervisor.join_all()
 
 
+def test_multiprocess_stops_on_startup_failure() -> None:
+    """A worker that dies before startup completes shuts the supervisor down instead of
+    being restarted forever (https://github.com/encode/uvicorn/discussions/2440)."""
+    config = Config(app="tests.supervisors.test_multiprocess:app_does_not_exist", workers=2)
+    supervisor = Multiprocess(config, sockets=[])
+    with pytest.raises(SystemExit) as exc_info:
+        supervisor.run()
+    assert exc_info.value.code == STARTUP_FAILURE
+
+
+def test_multiprocess_stops_on_lifespan_startup_failure() -> None:
+    """A lifespan startup failure is a startup failure too, not a reason to restart the worker."""
+    config = Config(app=lifespan_failure_app, lifespan="on", workers=2)
+    supervisor = Multiprocess(config, sockets=[])
+    with pytest.raises(SystemExit) as exc_info:
+        supervisor.run()
+    assert exc_info.value.code == STARTUP_FAILURE
+
+
 @new_console_in_windows
 def test_multiprocess_sigterm() -> None:
     """
     Ensure that the SIGTERM signal is handled as expected.
     """
     config = Config(app=app, workers=2)
-    supervisor = Multiprocess(config, target=run, sockets=[])
+    supervisor = Multiprocess(config, sockets=[])
     threading.Thread(target=supervisor.run, daemon=True).start()
     time.sleep(1)
     supervisor.signal_queue.append(signal.SIGTERM)
@@ -112,7 +144,7 @@ def test_multiprocess_sigbreak() -> None:  # pragma: py-not-win32
     Ensure that the SIGBREAK signal is handled as expected.
     """
     config = Config(app=app, workers=2)
-    supervisor = Multiprocess(config, target=run, sockets=[])
+    supervisor = Multiprocess(config, sockets=[])
     threading.Thread(target=supervisor.run, daemon=True).start()
     time.sleep(1)
     supervisor.signal_queue.append(getattr(signal, "SIGBREAK"))
@@ -125,7 +157,7 @@ def test_multiprocess_sighup() -> None:
     Ensure that the SIGHUP signal is handled as expected.
     """
     config = Config(app=app, workers=2)
-    supervisor = Multiprocess(config, target=run, sockets=[])
+    supervisor = Multiprocess(config, sockets=[])
     threading.Thread(target=supervisor.run, daemon=True).start()
     time.sleep(1)
     pids = [p.pid for p in supervisor.processes]
@@ -148,7 +180,7 @@ def test_multiprocess_sigttin() -> None:
     Ensure that the SIGTTIN signal is handled as expected.
     """
     config = Config(app=app, workers=2)
-    supervisor = Multiprocess(config, target=run, sockets=[])
+    supervisor = Multiprocess(config, sockets=[])
     threading.Thread(target=supervisor.run, daemon=True).start()
     supervisor.signal_queue.append(signal.SIGTTIN)
     time.sleep(1)
@@ -163,7 +195,7 @@ def test_multiprocess_sigttou() -> None:
     Ensure that the SIGTTOU signal is handled as expected.
     """
     config = Config(app=app, workers=2)
-    supervisor = Multiprocess(config, target=run, sockets=[])
+    supervisor = Multiprocess(config, sockets=[])
     threading.Thread(target=supervisor.run, daemon=True).start()
     supervisor.signal_queue.append(signal.SIGTTOU)
     time.sleep(1)

tests/test_main.py

@@ -1,3 +1,4 @@
+import asyncio
 import importlib
 import inspect
 import socket
@@ -167,6 +168,20 @@ def test_run_match_config_params() -> None:
     assert config_params == run_params
 
 
+async def test_server_on_started_callback(unused_tcp_port: int) -> None:
+    config = Config(app=app, loop="asyncio", port=unused_tcp_port)
+    server = Server(config=config)
+    started_events: list[bool] = []
+    server.on_started = lambda: started_events.append(server.started)
+
+    task = asyncio.create_task(server.serve())
+    while not server.started:
+        await asyncio.sleep(0.01)
+    assert started_events == [True]
+    await server.shutdown()
+    task.cancel()
+
+
 async def test_exit_on_create_server_with_invalid_host() -> None:
     with pytest.raises(SystemExit) as exc_info:
         config = Config(app=app, host="illegal_host")

uvicorn/main.py

@@ -28,7 +28,7 @@
     LoopFactoryType,
     WSProtocolType,
 )
-from uvicorn.server import Server
+from uvicorn.server import STARTUP_FAILURE, Server
 from uvicorn.supervisors import ChangeReload, Multiprocess
 
 LEVEL_CHOICES = click.Choice(list(LOG_LEVELS.keys()))
@@ -40,8 +40,6 @@ def _metavar_from_type(_type: Any) -> str:
     return f"[{'|'.join(key for key in get_args(_type) if key != 'none')}]"
 
 
-STARTUP_FAILURE = 3
-
 logger = logging.getLogger("uvicorn.error")
 
 
@@ -615,7 +613,7 @@ def run(
             ChangeReload(config, target=server.run, sockets=[sock]).run()
         elif config.workers > 1:
             sock = config.bind_socket()
-            Multiprocess(config, target=server.run, sockets=[sock]).run()
+            Multiprocess(config, sockets=[sock]).run()
         else:
             server.run()
     except KeyboardInterrupt:  # pragma: full coverage

uvicorn/server.py

@@ -12,7 +12,7 @@
 import sys
 import threading
 import time
-from collections.abc import Generator, Sequence
+from collections.abc import Callable, Generator, Sequence
 from email.utils import formatdate
 from types import FrameType
 from typing import TYPE_CHECKING, TypeAlias
@@ -40,6 +40,8 @@
 
 logger = logging.getLogger("uvicorn.error")
 
+STARTUP_FAILURE = 3
+
 
 class ServerState:
     """
@@ -62,6 +64,7 @@ def __init__(self, config: Config) -> None:
         self.should_exit = False
         self.force_exit = False
         self.last_notified = 0.0
+        self.on_started: Callable[[], None] | None = None
 
         self._captured_signals: list[int] = []
 
@@ -193,6 +196,8 @@ def _share_socket(
             pass  # pragma: full coverage
 
         self.started = True
+        if self.on_started is not None:
+            self.on_started()
 
     def _log_started_message(self, listeners: Sequence[socket.SocketType]) -> None:
         config = self.config

uvicorn/supervisors/multiprocess.py

@@ -3,16 +3,16 @@
 import logging
 import os
 import signal
+import sys
 import threading
-from collections.abc import Callable
 from multiprocessing import Pipe
 from socket import socket
-from typing import Any
 
 import click
 
 from uvicorn._subprocess import get_subprocess
 from uvicorn.config import Config
+from uvicorn.server import STARTUP_FAILURE, Server
 
 SIGNALS = {
     getattr(signal, f"SIG{x}"): x
@@ -27,12 +27,13 @@ class Process:
     def __init__(
         self,
         config: Config,
-        target: Callable[[list[socket] | None], None],
         sockets: list[socket],
     ) -> None:
-        self.real_target = target
+        self.config = config
+        self._ready = False
 
         self.parent_conn, self.child_conn = Pipe()
+        self.ready_parent_conn, self.ready_child_conn = Pipe(duplex=False)
         self.process = get_subprocess(config, self.target, sockets)
 
     def ping(self, timeout: float = 5) -> bool:
@@ -50,7 +51,7 @@ def always_pong(self) -> None:
         while True:
             self.pong()
 
-    def target(self, sockets: list[socket] | None = None) -> Any:  # pragma: no cover
+    def target(self, sockets: list[socket] | None = None) -> None:  # pragma: no cover
         if os.name == "nt":  # pragma: py-not-win32
             # Windows doesn't support SIGTERM, so we use SIGBREAK instead.
             # And then we raise SIGTERM when SIGBREAK is received.
@@ -61,7 +62,20 @@ def target(self, sockets: list[socket] | None = None) -> Any:  # pragma: no cove
             )
 
         threading.Thread(target=self.always_pong, daemon=True).start()
-        return self.real_target(sockets)
+
+        server = Server(config=self.config)
+        server.on_started = self.notify_started
+        server.run(sockets)
+
+    def notify_started(self) -> None:
+        self.ready_child_conn.send(b"ready")
+
+    @property
+    def ready(self) -> bool:
+        if not self._ready and self.ready_parent_conn.poll(0):
+            self.ready_parent_conn.recv()
+            self._ready = True
+        return self._ready
 
     def is_alive(self, timeout: float = 5) -> bool:
         if not self.process.is_alive():
@@ -85,6 +99,8 @@ def terminate(self) -> None:
 
             self.parent_conn.close()
             self.child_conn.close()
+            self.ready_parent_conn.close()
+            self.ready_child_conn.close()
 
     def kill(self) -> None:
         # In Windows, the method will call `TerminateProcess` to kill the process.
@@ -104,25 +120,24 @@ class Multiprocess:
     def __init__(
         self,
         config: Config,
-        target: Callable[[list[socket] | None], None],
         sockets: list[socket],
     ) -> None:
         self.config = config
-        self.target = target
         self.sockets = sockets
 
         self.processes_num = config.workers
         self.processes: list[Process] = []
 
         self.should_exit = threading.Event()
+        self.startup_failed = False
 
         self.signal_queue: list[int] = []
         for sig in SIGNALS:
             signal.signal(sig, lambda sig, frame: self.signal_queue.append(sig))
 
     def init_processes(self) -> None:
         for _ in range(self.processes_num):
-            process = Process(self.config, self.target, self.sockets)
+            process = Process(self.config, self.sockets)
             process.start()
             self.processes.append(process)
 
@@ -138,7 +153,7 @@ def restart_all(self) -> None:
         for idx, process in enumerate(self.processes):
             process.terminate()
             process.join()
-            new_process = Process(self.config, self.target, self.sockets)
+            new_process = Process(self.config, self.sockets)
             new_process.start()
             self.processes[idx] = new_process
 
@@ -160,6 +175,9 @@ def run(self) -> None:
         color_message = "Stopping parent process [{}]".format(click.style(str(os.getpid()), fg="cyan", bold=True))
         logger.info(message, extra={"color_message": color_message})
 
+        if self.startup_failed:
+            sys.exit(STARTUP_FAILURE)
+
     def keep_subprocess_alive(self) -> None:
         if self.should_exit.is_set():
             return  # parent process is exiting, no need to keep subprocess alive
@@ -174,8 +192,14 @@ def keep_subprocess_alive(self) -> None:
             if self.should_exit.is_set():
                 return  # pragma: full coverage
 
+            if not process.ready:
+                logger.error(f"Child process [{process.pid}] died before startup completed, shutting down.")
+                self.startup_failed = True
+                self.should_exit.set()
+                return
+
             logger.info(f"Child process [{process.pid}] died")
-            process = Process(self.config, self.target, self.sockets)
+            process = Process(self.config, self.sockets)
             process.start()
             self.processes[idx] = process
 
@@ -208,7 +232,7 @@ def handle_hup(self) -> None:  # pragma: py-win32
     def handle_ttin(self) -> None:  # pragma: py-win32
         logger.info("Received SIGTTIN, increasing the number of processes.")
         self.processes_num += 1
-        process = Process(self.config, self.target, self.sockets)
+        process = Process(self.config, self.sockets)
         process.start()
         self.processes.append(process)