d3i-szct/paste/service/daemonize.py

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)