FinDataHub/Scripts/bin/process_utils.py

# coding=utf-8
"""跨平台进程工具 —— 收敛 POSIX 专属的派生 / 终止 / 存活检测 / 信号处理.

背景
----
采集管线里散落着 POSIX 专属写法：``subprocess.Popen(start_new_session=True)``、
``os.kill(pid, 0)``、``signal.SIGTERM``、``signal.SIGHUP``。这些在 macOS/Linux
上语义正确，在 Windows 上：

* ``start_new_session`` → 直接 ``ValueError``；
* ``os.kill(pid, 0)`` 不是"探测"，Python 在 Windows 上会调 ``TerminateProcess``；
* ``SIGHUP`` / ``SIGUSR1`` 根本不存在，属性访问即 ``AttributeError``；
* ``os.kill(pid, SIGTERM)`` 只能终止目标进程本身，不级联子进程。

本模块提供语义等价的三组能力，macOS 上行为与原实现**逐字一致**（不改变
现网任何可见行为），仅在其他平台切换实现。

用例::

    from process_utils import detach_kwargs, is_running, terminate_pid

    proc = subprocess.Popen(cmd, **detach_kwargs())
    if is_running(pid):
        terminate_pid(pid)
"""

from __future__ import annotations

import os
import signal
import subprocess
import sys

IS_WINDOWS = sys.platform == "win32"

# Windows 进程创建标志（此平台无对应常量时退化为 0）
_CREATE_NEW_PROCESS_GROUP = 0x00000200
_DETACHED_PROCESS = 0x00000008


# ── 派生 ────────────────────────────────────────────────────

def detach_kwargs() -> dict:
    """``subprocess.Popen`` 的"脱离父进程会话"参数。

    POSIX 用 ``start_new_session=True``（等价 setsid）；Windows 用
    ``CREATE_NEW_PROCESS_GROUP | DETACHED_PROCESS``。

    为什么必须脱离：以 Bash/工具方式拉起常驻服务时，父 shell 退出会向
    整个进程组发 SIGHUP/SIGTERM。不脱离的话表现为"刚启动还能查到，
    过一会儿就没了"。
    """
    if IS_WINDOWS:
        return {"creationflags": _CREATE_NEW_PROCESS_GROUP | _DETACHED_PROCESS}
    return {"start_new_session": True}


def hidden_window_kwargs() -> dict:
    """Windows 下后台进程不弹控制台窗口；其他平台返回空 dict."""
    if IS_WINDOWS:
        si = subprocess.STARTUPINFO()
        si.dwFlags |= subprocess.STARTF_USESHOWWINDOW
        return {"startupinfo": si}
    return {}


# ── 存活检测 ────────────────────────────────────────────────

def is_running(pid: int) -> bool:
    """判断进程是否存活（不含僵尸回收语义，仅看是否存在）."""
    if pid <= 0:
        return False
    if IS_WINDOWS:
        return _is_running_windows(pid)
    try:
        os.kill(pid, 0)
    except (OSError, ProcessLookupError):
        return False
    except PermissionError:
        # 进程存在但属于其他用户
        return True
    return True


def _is_running_windows(pid: int) -> bool:
    """用 OpenProcess 探测。

    不能用 ``os.kill(pid, 0)``：Windows 上 Python 会走 TerminateProcess，
    等于"探测即杀人"。
    """
    import ctypes
    from ctypes import wintypes

    PROCESS_QUERY_LIMITED_INFORMATION = 0x1000
    STILL_ACTIVE = 259
    kernel32 = ctypes.windll.kernel32
    handle = kernel32.OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, False, pid)
    if not handle:
        # 拒绝访问说明进程存在（权限不足），其余视为不存在
        return kernel32.GetLastError() == 5
    try:
        code = wintypes.DWORD()
        if kernel32.GetExitCodeProcess(handle, ctypes.byref(code)):
            return code.value == STILL_ACTIVE
        return False
    finally:
        kernel32.CloseHandle(handle)


# ── 终止 ────────────────────────────────────────────────────

def terminate_pid(pid: int, *, force: bool = False, timeout: float = 5.0) -> bool:
    """终止进程，返回是否发出了终止请求。

    POSIX: 先 ``SIGTERM``，``force=True`` 时直接 ``SIGKILL``。
    Windows: ``taskkill /PID x /T``（``/T` 级联子进程），``force=True`` 加 ``/F``。

    注意 Windows 的 taskkill 无"优雅"概念，``/F`` 与否差别只在是否强制；
    因此 Windows 下 ``force=False`` 也可能直接结束进程。
    """
    if pid <= 0:
        return False
    if IS_WINDOWS:
        cmd = ["taskkill", "/PID", str(pid), "/T"]
        if force:
            cmd.append("/F")
        try:
            subprocess.run(cmd, capture_output=True, timeout=timeout, check=False,
                           **hidden_window_kwargs())
            return True
        except (OSError, subprocess.SubprocessError):
            return False
    try:
        os.kill(pid, signal.SIGKILL if force else signal.SIGTERM)
        return True
    except (OSError, ProcessLookupError):
        return False


# ── 信号 ────────────────────────────────────────────────────

def available_signals() -> dict[str, int]:
    """当前平台实际存在的信号子集，名 → 编号。

    采集器注册信号处理器时必须先查这里，否则 Windows 上访问
    ``signal.SIGHUP`` 会直接 AttributeError 崩在启动阶段。
    """
    wanted = ("SIGTERM", "SIGINT", "SIGHUP", "SIGUSR1", "SIGUSR2", "SIGBREAK")
    return {name: getattr(signal, name) for name in wanted if hasattr(signal, name)}


def register_signal_handler(handler, *, names=("SIGTERM", "SIGINT")) -> list[str]:
    """注册信号处理器，跳过当前平台不存在的信号。

    Returns:
        实际注册成功的信号名列表。
    """
    registered = []
    for name in names:
        sig = getattr(signal, name, None)
        if sig is None:
            continue
        try:
            signal.signal(sig, handler)
            registered.append(name)
        except (OSError, ValueError, RuntimeError):
            # 非主线程 / 平台限制：跳过而不是崩
            continue
    return registered


def restore_default_and_die(signum: int) -> None:
    """恢复默认处理并重新触发该信号——保留原始退出码与 core dump 行为.

    用于崩溃留痕：先记录现场，再让进程以"原生方式"死掉，不改变
    外部观察到的终止原因。Windows 上无信号语义，退化为非零退出。
    """
    if IS_WINDOWS:
        os._exit(128 + (signum or 0))
    try:
        signal.signal(signum, signal.SIG_DFL)
        os.kill(os.getpid(), signum)
    except (OSError, ValueError):
        os._exit(1)


# ── 编码 ────────────────────────────────────────────────────

def utf8_subprocess_env(base: dict | None = None) -> dict:
    """返回带 UTF-8 强制的环境变量副本。

    Windows 控制台默认 cp936，``print`` 中文与捕获子进程 stdout 会乱码，
    严重时抛 ``UnicodeEncodeError`` 直接崩掉采集器。采集器 stdout 是
    **JSON 状态协议**，一旦被编码污染，Swift 端解析失败会误判"采集器未运行"。
    """
    env = dict(base if base is not None else os.environ)
    env["PYTHONUTF8"] = "1"
    env["PYTHONIOENCODING"] = "utf-8"
    return env