iOSAI/Utils/ThreadManager.py

import os
import signal
import sys
import threading
import time
import psutil
import subprocess
from pathlib import Path
from threading import Event, Thread
from typing import Dict, Optional

from Utils.LogManager import LogManager


class ThreadManager:
    """
    对调用方完全透明：
    add(udid, thread_obj, stop_event)  保持原签名
    stop(udid)                         保持原签名
    但内部把 thread_obj 当成“壳”，真正拉起的是子进程。
    """
    _pool: Dict[str, psutil.Process] = {}
    _lock = threading.Lock()

    @classmethod
    def add(cls, udid: str, dummy_thread, dummy_event: Event) -> None:
        LogManager.method_info(f"【1】入口 udid={udid} 长度={len(udid)}", method="task")
        if udid in cls._pool:
            LogManager.method_warning(f"{udid} 仍在运行，先强制清理旧任务", method="task")
            cls.stop(udid)
        LogManager.method_info(f"【2】判断旧任务后 udid={udid} 长度={len(udid)}", method="task")
        port = cls._find_free_port()
        LogManager.method_info(f"【3】找端口后 udid={udid} 长度={len(udid)}", method="task")
        proc = cls._start_worker_process(udid, port)
        LogManager.method_info(f"【4】子进程启动后 udid={udid} 长度={len(udid)}", method="task")
        cls._pool[udid] = proc
        LogManager.method_info(f"【5】已写入字典，udid={udid} 长度={len(udid)}", method="task")

    @classmethod
    def stop(cls, udid: str) -> tuple[int, str]:
        with cls._lock:  # 类级锁
            proc = cls._pool.get(udid)  # 1. 只读，不删
            if proc is None:
                return 1001, f"无此任务 {udid}"

            try:
                proc.terminate()
                gone, alive = psutil.wait_procs([proc], timeout=3)
                if alive:
                    for p in alive:
                        for child in p.children(recursive=True):
                            child.kill()
                        p.kill()
                    psutil.wait_procs(alive, timeout=2)

                # 正常退出
                cls._pool.pop(udid)
                LogManager.method_info("任务停止成功", method="task")
                return 200, f"停止线程成功 {udid}"

            except psutil.NoSuchProcess:  # 精准捕获
                cls._pool.pop(udid)
                LogManager.method_info("进程已自然退出", method="task")
                return 200, f"进程已退出 {udid}"

            except Exception as e:  # 真正的异常
                LogManager.method_error(f"停止异常: {e}", method="task")
                return 1002, f"停止异常 {udid}"

    # ------------------------------------------------------
    # 以下全是内部工具，外部无需调用
    # ------------------------------------------------------
    @staticmethod
    def _find_free_port(start: int = 50000) -> int:
        """找个随机空闲端口，给子进程当通信口（可选）"""
        import socket
        for p in range(start, start + 1000):
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                if s.connect_ex(("127.0.0.1", p)) != 0:
                    return p
        raise RuntimeError("无可用端口")

    @staticmethod
    def _start_worker_process(udid: str, port: int) -> psutil.Process:
        """
        真正拉起子进程：
        打包环境：exe --udid=xxx
        源码环境：python -m Module.Worker --udid=xxx
        """
        exe_path = Path(sys.executable).resolve()
        is_frozen = exe_path.suffix.lower() == ".exe" and exe_path.exists()

        if is_frozen:
            # 打包后
            cmd = [str(exe_path), "--role=worker", f"--udid={udid}", f"--port={port}"]
            cwd = str(exe_path.parent)
        else:
            # 源码运行
            cmd = [sys.executable, "-u", "-m", "Module.Worker", f"--udid={udid}", f"--port={port}"]
            cwd = str(Path(__file__).resolve().parent.parent)

        # 核心：CREATE_NO_WINDOW + 独立会话，父进程死也不影响
        creation_flags = 0x08000000 if os.name == "nt" else 0
        proc = subprocess.Popen(
            cmd,
            stdin=subprocess.DEVNULL,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            encoding="utf-8",
            errors="replace",
            bufsize=1,
            cwd=cwd,
            start_new_session=True,      # 独立进程组
            creationflags=creation_flags
        )
        # 守护线程：把子进程 stdout 实时打到日志
        Thread(target=lambda: ThreadManager._log_stdout(proc, udid), daemon=True).start()
        return psutil.Process(proc.pid)

    @staticmethod
    def _log_stdout(proc: subprocess.Popen, udid: str):
        for line in iter(proc.stdout.readline, ""):
            if line:
                LogManager.info(line.rstrip(), udid)
        proc.stdout.close()