from __future__ import annotations import os from collections.abc import MutableMapping from datetime import datetime, timezone from inspect import signature from multiprocessing.context import BaseContext from signal import SIGINT from threading import Thread from time import sleep from typing import Any from sanic.log import Colors, logger from sanic.worker.constants import ProcessState, RestartOrder def get_now(): now = datetime.now(tz=timezone.utc) return now class WorkerProcess: """A worker process.""" THRESHOLD = 300 # == 30 seconds SERVER_LABEL = "Server" SERVER_IDENTIFIER = "Srv" def __init__( self, factory, name, ident, target, kwargs, worker_state, restartable: bool = False, ): self.state = ProcessState.IDLE self.factory = factory self.name = name self.ident = ident self.target = target self.kwargs = kwargs self.worker_state = worker_state self.restartable = restartable if self.name not in self.worker_state: self.worker_state[self.name] = { "server": self.SERVER_LABEL in self.name } self.spawn() def set_state(self, state: ProcessState, force=False): if not force and state < self.state: raise Exception("...") self.state = state self.worker_state[self.name] = { **self.worker_state[self.name], "state": self.state.name, } def start(self): os.environ["SANIC_WORKER_NAME"] = self.name os.environ["SANIC_WORKER_IDENTIFIER"] = self.ident logger.debug( f"{Colors.BLUE}Starting a process: {Colors.BOLD}" f"{Colors.SANIC}%s{Colors.END}", self.name, ) self.set_state(ProcessState.STARTING) self._current_process.start() self.set_state(ProcessState.STARTED) if not self.worker_state[self.name].get("starts"): self.worker_state[self.name] = { **self.worker_state[self.name], "pid": self.pid, "start_at": get_now(), "starts": 1, } del os.environ["SANIC_WORKER_NAME"] del os.environ["SANIC_WORKER_IDENTIFIER"] def join(self): self.set_state(ProcessState.JOINED) self._current_process.join() def exit(self): limit = 100 while self.is_alive() and limit > 0: sleep(0.1) limit -= 1 if not self.is_alive(): try: del self.worker_state[self.name] except ( BrokenPipeError, ConnectionRefusedError, ConnectionResetError, EOFError, ): logger.debug("Monitor process has already exited.") except KeyError: logger.debug("Could not find worker state to delete.") def terminate(self): if self.state is not ProcessState.TERMINATED: logger.debug( f"{Colors.BLUE}Terminating a process: " f"{Colors.BOLD}{Colors.SANIC}" f"%s {Colors.BLUE}[%s]{Colors.END}", self.name, self.pid, ) try: self.set_state(ProcessState.TERMINATED, force=True) except (BrokenPipeError, ConnectionResetError, EOFError): pass try: os.kill(self.pid, SIGINT) except (KeyError, AttributeError, ProcessLookupError): ... def restart(self, restart_order=RestartOrder.SHUTDOWN_FIRST, **kwargs): logger.debug( f"{Colors.BLUE}Restarting a process: {Colors.BOLD}{Colors.SANIC}" f"%s {Colors.BLUE}[%s]{Colors.END}", self.name, self.pid, ) self.set_state(ProcessState.RESTARTING, force=True) if restart_order is RestartOrder.SHUTDOWN_FIRST: self._terminate_now() else: self._old_process = self._current_process if self._add_config(): self.kwargs.update( {"config": {k.upper(): v for k, v in kwargs.items()}} ) try: self.spawn() self.start() except AttributeError: raise RuntimeError("Restart failed") if restart_order is RestartOrder.STARTUP_FIRST: self._terminate_soon() self.worker_state[self.name] = { **self.worker_state[self.name], "pid": self.pid, "starts": self.worker_state[self.name]["starts"] + 1, "restart_at": get_now(), } def is_alive(self): try: return self._current_process.is_alive() except AssertionError: return False def spawn(self): if self.state not in (ProcessState.IDLE, ProcessState.RESTARTING): raise Exception("Cannot spawn a worker process until it is idle.") self._current_process = self.factory( name=self.name, target=self.target, kwargs=self.kwargs, daemon=True, ) @property def pid(self): return self._current_process.pid @property def exitcode(self): return self._current_process.exitcode def _terminate_now(self): if not self._current_process.is_alive(): return logger.debug( f"{Colors.BLUE}Begin restart termination: " f"{Colors.BOLD}{Colors.SANIC}" f"%s {Colors.BLUE}[%s]{Colors.END}", self.name, self._current_process.pid, ) self._current_process.terminate() def _terminate_soon(self): logger.debug( f"{Colors.BLUE}Begin restart termination: " f"{Colors.BOLD}{Colors.SANIC}" f"%s {Colors.BLUE}[%s]{Colors.END}", self.name, self._current_process.pid, ) termination_thread = Thread(target=self._wait_to_terminate) termination_thread.start() def _wait_to_terminate(self): logger.debug( f"{Colors.BLUE}Waiting for process to be acked: " f"{Colors.BOLD}{Colors.SANIC}" f"%s {Colors.BLUE}[%s]{Colors.END}", self.name, self._old_process.pid, ) misses = 0 while self.state is not ProcessState.ACKED: sleep(0.1) misses += 1 if misses > self.THRESHOLD: raise TimeoutError( f"Worker {self.name} failed to come ack within " f"{self.THRESHOLD / 10} seconds" ) else: logger.debug( f"{Colors.BLUE}Process acked. Terminating: " f"{Colors.BOLD}{Colors.SANIC}" f"%s {Colors.BLUE}[%s]{Colors.END}", self.name, self._old_process.pid, ) self._old_process.terminate() delattr(self, "_old_process") def _add_config(self) -> bool: sig = signature(self.target) if "config" in sig.parameters or any( param.kind == param.VAR_KEYWORD for param in sig.parameters.values() ): return True return False class Worker: WORKER_PREFIX = "Sanic" def __init__( self, ident: str, name: str, serve, server_settings, context: BaseContext, worker_state: MutableMapping[str, Any], num: int = 1, restartable: bool = False, tracked: bool = True, auto_start: bool = True, ): self.ident = ident self.name = name self.num = num self.context = context self.serve = serve self.server_settings = server_settings self.worker_state = worker_state self.processes: set[WorkerProcess] = set() self.restartable = restartable self.tracked = tracked self.auto_start = auto_start for _ in range(num): self.create_process() def create_process(self) -> WorkerProcess: process = WorkerProcess( # Need to ignore this typing error - The problem is the # BaseContext itself has no Process. But, all of its # implementations do. We can safely ignore as it is a typing # issue in the standard lib. factory=self.context.Process, # type: ignore name="-".join( [self.WORKER_PREFIX, self.name, str(len(self.processes))] ), ident=self.ident, target=self.serve, kwargs={**self.server_settings}, worker_state=self.worker_state, restartable=self.restartable, ) self.processes.add(process) return process def has_alive_processes(self) -> bool: return any(process.is_alive() for process in self.processes)