codekingpro/portable-devtools

1from __future__ import annotations2 3import logging4import os5import signal6import threading7from collections.abc import Callable8from multiprocessing import Pipe9from socket import socket10from typing import Any11 12import click13 14from uvicorn._subprocess import get_subprocess15from uvicorn.config import Config16 17SIGNALS = {18    getattr(signal, f"SIG{x}"): x19    for x in "INT TERM BREAK HUP QUIT TTIN TTOU USR1 USR2 WINCH".split()20    if hasattr(signal, f"SIG{x}")21}22 23logger = logging.getLogger("uvicorn.error")24 25 26class Process:27    def __init__(28        self,29        config: Config,30        target: Callable[[list[socket] | None], None],31        sockets: list[socket],32    ) -> None:33        self.real_target = target34 35        self.parent_conn, self.child_conn = Pipe()36        self.process = get_subprocess(config, self.target, sockets)37 38    def ping(self, timeout: float = 5) -> bool:39        self.parent_conn.send(b"ping")40        if self.parent_conn.poll(timeout):41            self.parent_conn.recv()42            return True43        return False44 45    def pong(self) -> None:46        self.child_conn.recv()47        self.child_conn.send(b"pong")48 49    def always_pong(self) -> None:50        while True:51            self.pong()52 53    def target(self, sockets: list[socket] | None = None) -> Any:  # pragma: no cover54        if os.name == "nt":  # pragma: py-not-win3255            # Windows doesn't support SIGTERM, so we use SIGBREAK instead.56            # And then we raise SIGTERM when SIGBREAK is received.57            # https://learn.microsoft.com/zh-cn/cpp/c-runtime-library/reference/signal?view=msvc-17058            signal.signal(59                signal.SIGBREAK,  # type: ignore[attr-defined]60                lambda sig, frame: signal.raise_signal(signal.SIGTERM),61            )62 63        threading.Thread(target=self.always_pong, daemon=True).start()64        return self.real_target(sockets)65 66    def is_alive(self, timeout: float = 5) -> bool:67        if not self.process.is_alive():68            return False  # pragma: full coverage69 70        return self.ping(timeout)71 72    def start(self) -> None:73        self.process.start()74 75    def terminate(self) -> None:76        if self.process.exitcode is None:  # Process is still running77            assert self.process.pid is not None78            if os.name == "nt":  # pragma: py-not-win3279                # Windows doesn't support SIGTERM.80                # So send SIGBREAK, and then in process raise SIGTERM.81                os.kill(self.process.pid, signal.CTRL_BREAK_EVENT)  # type: ignore[attr-defined]82            else:83                os.kill(self.process.pid, signal.SIGTERM)84            logger.info(f"Terminated child process [{self.process.pid}]")85 86            self.parent_conn.close()87            self.child_conn.close()88 89    def kill(self) -> None:90        # In Windows, the method will call `TerminateProcess` to kill the process.91        # In Unix, the method will send SIGKILL to the process.92        self.process.kill()93 94    def join(self) -> None:95        logger.info(f"Waiting for child process [{self.process.pid}]")96        self.process.join()97 98    @property99    def pid(self) -> int | None:100        return self.process.pid101 102 103class Multiprocess:104    def __init__(105        self,106        config: Config,107        target: Callable[[list[socket] | None], None],108        sockets: list[socket],109    ) -> None:110        self.config = config111        self.target = target112        self.sockets = sockets113 114        self.processes_num = config.workers115        self.processes: list[Process] = []116 117        self.should_exit = threading.Event()118 119        self.signal_queue: list[int] = []120        for sig in SIGNALS:121            signal.signal(sig, lambda sig, frame: self.signal_queue.append(sig))122 123    def init_processes(self) -> None:124        for _ in range(self.processes_num):125            process = Process(self.config, self.target, self.sockets)126            process.start()127            self.processes.append(process)128 129    def terminate_all(self) -> None:130        for process in self.processes:131            process.terminate()132 133    def join_all(self) -> None:134        for process in self.processes:135            process.join()136 137    def restart_all(self) -> None:138        for idx, process in enumerate(self.processes):139            process.terminate()140            process.join()141            new_process = Process(self.config, self.target, self.sockets)142            new_process.start()143            self.processes[idx] = new_process144 145    def run(self) -> None:146        message = f"Started parent process [{os.getpid()}]"147        color_message = "Started parent process [{}]".format(click.style(str(os.getpid()), fg="cyan", bold=True))148        logger.info(message, extra={"color_message": color_message})149 150        self.init_processes()151 152        while not self.should_exit.wait(0.5):153            self.handle_signals()154            self.keep_subprocess_alive()155 156        self.terminate_all()157        self.join_all()158 159        message = f"Stopping parent process [{os.getpid()}]"160        color_message = "Stopping parent process [{}]".format(click.style(str(os.getpid()), fg="cyan", bold=True))161        logger.info(message, extra={"color_message": color_message})162 163    def keep_subprocess_alive(self) -> None:164        if self.should_exit.is_set():165            return  # parent process is exiting, no need to keep subprocess alive166 167        for idx, process in enumerate(self.processes):168            if process.is_alive(timeout=self.config.timeout_worker_healthcheck):169                continue170 171            process.kill()  # process is hung, kill it172            process.join()173 174            if self.should_exit.is_set():175                return  # pragma: full coverage176 177            logger.info(f"Child process [{process.pid}] died")178            process = Process(self.config, self.target, self.sockets)179            process.start()180            self.processes[idx] = process181 182    def handle_signals(self) -> None:183        for sig in tuple(self.signal_queue):184            self.signal_queue.remove(sig)185            sig_name = SIGNALS[sig]186            sig_handler = getattr(self, f"handle_{sig_name.lower()}", None)187            if sig_handler is not None:188                sig_handler()189            else:  # pragma: no cover190                logger.debug(f"Received signal {sig_name}, but no handler is defined for it.")191 192    def handle_int(self) -> None:193        logger.info("Received SIGINT, exiting.")194        self.should_exit.set()195 196    def handle_term(self) -> None:197        logger.info("Received SIGTERM, exiting.")198        self.should_exit.set()199 200    def handle_break(self) -> None:  # pragma: py-not-win32201        logger.info("Received SIGBREAK, exiting.")202        self.should_exit.set()203 204    def handle_hup(self) -> None:  # pragma: py-win32205        logger.info("Received SIGHUP, restarting processes.")206        self.restart_all()207 208    def handle_ttin(self) -> None:  # pragma: py-win32209        logger.info("Received SIGTTIN, increasing the number of processes.")210        self.processes_num += 1211        process = Process(self.config, self.target, self.sockets)212        process.start()213        self.processes.append(process)214 215    def handle_ttou(self) -> None:  # pragma: py-win32216        logger.info("Received SIGTTOU, decreasing number of processes.")217        if self.processes_num <= 1:218            logger.info("Already reached one process, cannot decrease the number of processes anymore.")219            return220        self.processes_num -= 1221        process = self.processes.pop()222        process.terminate()223        process.join()224