126 lines
4.9 KiB
Python
126 lines
4.9 KiB
Python
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() |