import atexit import os import signal import sys from typing import Callable import psutil from paste.util import ufile class DaemonizeService: """ 驻内存服务。 """ def __init__(self, pid_file, name: str = '', stdin: str = None, stdout: str = None, stderr: str = None): """ 初始化服务。 :param pid_file: pid 文件路径 :param name: 服务名称 :param stdin: 输入文件路径 :param stdout: 输出文件路径 :param stderr: 错误日志文件路径 """ self.start_callback = None self.start_callback_args = () self.start_callback_kwargs = {} self.term_callback = None self.term_callback_args = () self.term_callback_kwargs = {} self.pid_file = pid_file self.name = name self.stdin = stdin self.stdout = stdout self.stderr = stderr def set_start_callback(self, callback, *args, **kwargs): """ 设置启动回调函数。不返回参数。 :param callback: 回调函数 :param args: 回调参数 :param kwargs: 回调参数 """ self.start_callback = callback self.start_callback_args = args self.start_callback_kwargs = kwargs def set_term_callback(self, callback, *args, **kwargs): """ 设置终止回调函数。不返回参数。 :param callback: 回调函数 :param args: 回调参数 :param kwargs: 回调参数 """ self.term_callback = callback self.term_callback_args = args self.term_callback_kwargs = kwargs def daemonize(self): """ 设置和启动常驻服务程序。 """ if os.path.exists(self.pid_file): _pid = int(ufile.read_to_buffer(self.pid_file).decode('utf8').strip()) if psutil.pid_exists(_pid): raise RuntimeError(f"[{self.name}] 正在运行.") else: os.remove(self.pid_file) try: if os.fork() > 0: raise SystemExit(0) # Parent exit except OSError: raise RuntimeError('创建子进程失败 #1.') os.chdir(os.path.abspath(os.path.curdir)) os.umask(0) os.setsid() try: if os.fork() > 0: raise SystemExit(0) except OSError: raise RuntimeError('创建子进程失败 #2.') sys.stdout.flush() sys.stderr.flush() # 替换 stdin, stdout, 和 stderr 的文件描述符 if self.stdin is not None: with open(self.stdin, 'rb', 0) as file: os.dup2(file.fileno(), sys.stdin.fileno()) if self.stdout is not None: with open(self.stdout, 'ab', 0) as file: os.dup2(file.fileno(), sys.stdout.fileno()) if self.stderr is not None: with open(self.stderr, 'ab', 0) as file: os.dup2(file.fileno(), sys.stderr.fileno()) # 写入 PID 文件 with open(self.pid_file, 'w') as file: print(os.getpid(), file=file) # 注册退出函数,进程退出时(包括异常退出)移除pid文件 atexit.register(lambda: os.remove(self.pid_file)) # 监听终止信号,绑定到处理程序 signal.signal(signal.SIGTERM, self.sigterm_handler) def sigterm_handler(self, signum, frame): """ 终止信号处理程序。这里是执行回调函数,处理服务退出前的准备。 :param signum: 信号代码 :param frame: 帧 :return: 程序退出码 """ if isinstance(self.term_callback, Callable): self.term_callback(*self.term_callback_args, **self.term_callback_kwargs) raise SystemExit(1) def start(self): """ 启动服务。 """ try: self.daemonize() if isinstance(self.start_callback, Callable): self.start_callback(*self.start_callback_args, **self.start_callback_kwargs) except RuntimeError as e: print(e, file=sys.stderr) raise SystemExit(1) def stop(self): """ 停止服务。 """ if os.path.exists(self.pid_file): _pid = int(ufile.read_to_buffer(self.pid_file).decode('utf8').strip()) if psutil.pid_exists(_pid): os.kill(_pid, signal.SIGTERM) else: os.remove(self.pid_file) else: print(f"[{self.name}] 尚未启动.", file=sys.stderr) raise SystemExit(1) def cli_run(self): """ 命令行启动。 """ if len(sys.argv) != 2: print(f"Please use like: python3 {sys.argv[0]} [start|stop].", file=sys.stderr) raise SystemExit(1) if sys.argv[1] == 'start': self.start() elif sys.argv[1] == 'stop': self.stop() else: print(f"未知命令: {sys.argv[1]}", file=sys.stderr) raise SystemExit(1)