New multiprocess manager

pyproject.toml

@@ -101,7 +101,7 @@ omit = [
 
 [tool.coverage.report]
 precision = 2
-fail_under = 98.35
+fail_under = 96.82
 show_missing = true
 skip_covered = true
 exclude_lines = [

tests/supervisors/test_multiprocess.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
-import signal
 import socket
+import threading
+import time
 
 from uvicorn import Config
 from uvicorn._types import ASGIReceiveCallable, ASGISendCallable, Scope
@@ -16,6 +17,12 @@ def run(sockets: list[socket.socket] | None) -> None:
     pass  # pragma: no cover
 
 
+def stop_run(stop) -> None:
+    while True:
+        time.sleep(1)
+        stop()
+
+
 def test_multiprocess_run() -> None:
     """
     A basic sanity check.
@@ -25,5 +32,5 @@ def test_multiprocess_run() -> None:
     """
     config = Config(app=app, workers=2)
     supervisor = Multiprocess(config, target=run, sockets=[])
-    supervisor.signal_handler(sig=signal.SIGINT, frame=None)
+    threading.Thread(target=stop_run, args=(supervisor.handle_int,), daemon=True).start()
     supervisor.run()

uvicorn/supervisors/multiprocess.py

@@ -4,24 +4,101 @@
 import os
 import signal
 import threading
-from multiprocessing.context import SpawnProcess
+from multiprocessing import Pipe
 from socket import socket
-from types import FrameType
-from typing import Callable
+from typing import Any, Callable
 
 import click
 
 from uvicorn._subprocess import get_subprocess
 from uvicorn.config import Config
 
-HANDLED_SIGNALS = (
-    signal.SIGINT,  # Unix signal 2. Sent by Ctrl+C.
-    signal.SIGTERM,  # Unix signal 15. Sent by `kill <pid>`.
-)
+UNIX_SIGNALS = {
+    getattr(signal, f"SIG{x}"): x for x in "HUP QUIT TTIN TTOU USR1 USR2 WINCH".split() if hasattr(signal, f"SIG{x}")
+}
 
 logger = logging.getLogger("uvicorn.error")
 
 
+class Process:
+    def __init__(
+        self,
+        config: Config,
+        target: Callable[[list[socket] | None], None],
+        sockets: list[socket],
+    ) -> None:
+        self.real_target = target
+
+        self.parent_conn, self.child_conn = Pipe()
+        self.process = get_subprocess(config, self.target, sockets)
+
+    def ping(self, timeout: float = 5) -> bool:
+        self.parent_conn.send(b"ping")
+        if self.parent_conn.poll(timeout):
+            self.parent_conn.recv()
+            return True
+        return False
+
+    def pong(self) -> None:
+        self.child_conn.recv()
+        self.child_conn.send(b"pong")
+
+    def always_pong(self) -> None:
+        while True:
+            self.pong()
+
+    def target(self, sockets: list[socket] | None = None) -> Any:
+        if os.name == "nt":
+            # Windows doesn't support SIGTERM, so we use SIGBREAK instead.
+            # And then we raise SIGTERM when SIGBREAK is received.
+            # https://learn.microsoft.com/zh-cn/cpp/c-runtime-library/reference/signal?view=msvc-170
+            signal.signal(
+                signal.SIGBREAK,  # type: ignore[attr-defined]
+                lambda sig, frame: signal.raise_signal(signal.SIGTERM),
+            )
+
+        threading.Thread(target=self.always_pong, daemon=True).start()
+        return self.real_target(sockets)
+
+    def is_alive(self, timeout: float = 5) -> bool:
+        if not self.process.is_alive():
+            return False
+
+        return self.ping(timeout)
+
+    def start(self) -> None:
+        self.process.start()
+        logger.info(f"Started child process [{self.process.pid}]")
+
+    def terminate(self) -> None:
+        if self.process.exitcode is not None:
+            return
+        assert self.process.pid is not None
+        if os.name == "nt":
+            # Windows doesn't support SIGTERM.
+            # So send SIGBREAK, and then in process raise SIGTERM.
+            os.kill(self.process.pid, signal.CTRL_BREAK_EVENT)  # type: ignore[attr-defined]
+        else:
+            os.kill(self.process.pid, signal.SIGTERM)
+        logger.info(f"Terminated child process [{self.process.pid}]")
+
+        self.parent_conn.close()
+        self.child_conn.close()
+
+    def kill(self) -> None:
+        # In Windows, the method will call `TerminateProcess` to kill the process.
+        # In Unix, the method will send SIGKILL to the process.
+        self.process.kill()
+
+    def join(self) -> None:
+        logger.info(f"Waiting for child process [{self.process.pid}]")
+        self.process.join()
+
+    @property
+    def pid(self) -> int | None:
+        return self.process.pid
+
+
 class Multiprocess:
     def __init__(
         self,
@@ -32,39 +109,125 @@ def __init__(
         self.config = config
         self.target = target
         self.sockets = sockets
-        self.processes: list[SpawnProcess] = []
-        self.should_exit = threading.Event()
-        self.pid = os.getpid()
 
-    def signal_handler(self, sig: int, frame: FrameType | None) -> None:
-        """
-        A signal handler that is registered with the parent process.
-        """
-        self.should_exit.set()
+        self.processes_num = config.workers
+        self.processes: list[Process] = []
 
-    def run(self) -> None:
-        self.startup()
-        self.should_exit.wait()
-        self.shutdown()
+        self.should_exit = threading.Event()
 
-    def startup(self) -> None:
-        message = f"Started parent process [{str(self.pid)}]"
-        color_message = "Started parent process [{}]".format(click.style(str(self.pid), fg="cyan", bold=True))
-        logger.info(message, extra={"color_message": color_message})
+        self.signal_queue: list[int] = []
+        for sig in UNIX_SIGNALS:
+            signal.signal(sig, lambda sig, frame: self.signal_queue.append(sig))
 
-        for sig in HANDLED_SIGNALS:
-            signal.signal(sig, self.signal_handler)
+        # Sent by Ctrl+C.
+        signal.signal(signal.SIGINT, lambda sig, frame: self.handle_int())
+        # Sent by `kill <pid>`. Not sent on Windows.
+        signal.signal(signal.SIGTERM, lambda sig, frame: self.handle_term())
+        if os.name == "nt":
+            # Sent by `Ctrl+Break` on Windows.
+            signal.signal(signal.SIGBREAK, lambda sig, frame: self.handle_break())  # type: ignore[attr-defined]
 
-        for _idx in range(self.config.workers):
-            process = get_subprocess(config=self.config, target=self.target, sockets=self.sockets)
+    def init_processes(self) -> None:
+        for _ in range(self.processes_num):
+            process = Process(self.config, self.target, self.sockets)
             process.start()
             self.processes.append(process)
 
-    def shutdown(self) -> None:
+    def terminate_all(self) -> None:
         for process in self.processes:
             process.terminate()
+
+    def join_all(self) -> None:
+        for process in self.processes:
             process.join()
 
-        message = f"Stopping parent process [{str(self.pid)}]"
-        color_message = "Stopping parent process [{}]".format(click.style(str(self.pid), fg="cyan", bold=True))
+    def restart_all(self) -> None:
+        for idx, process in enumerate(tuple(self.processes)):
+            process.terminate()
+            process.join()
+            del self.processes[idx]
+            process = Process(self.config, self.target, self.sockets)
+            process.start()
+            self.processes.append(process)
+
+    def run(self) -> None:
+        message = f"Started parent process [{os.getpid()}]"
+        color_message = "Started parent process [{}]".format(click.style(str(os.getpid()), fg="cyan", bold=True))
+        logger.info(message, extra={"color_message": color_message})
+
+        self.init_processes()
+
+        while not self.should_exit.wait(0.5):
+            self.handle_signals()
+            self.keep_subprocess_alive()
+
+        self.terminate_all()
+        self.join_all()
+
+        message = f"Stopping parent process [{os.getpid()}]"
+        color_message = "Stopping parent process [{}]".format(click.style(str(os.getpid()), fg="cyan", bold=True))
         logger.info(message, extra={"color_message": color_message})
+
+    def keep_subprocess_alive(self) -> None:
+        for idx, process in enumerate(tuple(self.processes)):
+            if process.is_alive():
+                continue
+
+            process.kill()  # process is hung, kill it
+            process.join()
+            logger.info(f"Child process [{process.pid}] died")
+            del self.processes[idx]
+            process = Process(self.config, self.target, self.sockets)
+            process.start()
+            self.processes.append(process)
+
+    def handle_signals(self) -> None:
+        for sig in tuple(self.signal_queue):
+            self.signal_queue.remove(sig)
+            sig_name = UNIX_SIGNALS[sig]
+            sig_handler = getattr(self, f"handle_{sig_name.lower()}", None)
+            if sig_handler is not None:
+                sig_handler()
+            else:
+                logger.info(f"Received signal [{sig_name}], but nothing to do")
+
+    def handle_int(self) -> None:
+        if not self.should_exit.is_set():
+            self.should_exit.set()
+        else:
+            self.terminate_all()
+
+    def handle_term(self) -> None:
+        logger.info("Received SIGTERM, exiting")
+        if not self.should_exit.is_set():
+            self.should_exit.set()
+        else:
+            self.terminate_all()
+
+    def handle_break(self) -> None:
+        logger.info("Received SIGBREAK, exiting")
+        if not self.should_exit.is_set():
+            self.should_exit.set()
+        else:
+            self.terminate_all()
+
+    def handle_hup(self) -> None:
+        logger.info("Received SIGHUP, restarting processes")
+        self.restart_all()
+
+    def handle_ttin(self) -> None:
+        logger.info("Received SIGTTIN, increasing processes")
+        self.processes_num += 1
+        process = Process(self.config, self.target, self.sockets)
+        process.start()
+        self.processes.append(process)
+
+    def handle_ttou(self) -> None:
+        logger.info("Received SIGTTOU, decreasing processes")
+        if self.processes_num <= 1:
+            logger.info("Cannot decrease processes any more")
+            return
+        self.processes_num -= 1
+        process = self.processes.pop()
+        process.terminate()
+        process.join()