# -*- coding: utf-8 -*- import threading import ctypes import inspect import time from typing import Dict, Optional, List, Tuple, Any from Utils.LogManager import LogManager def _raise_async_exception(tid: int, exc_type) -> int: if not inspect.isclass(exc_type): raise TypeError("exc_type must be a class") return ctypes.pythonapi.PyThreadState_SetAsyncExc( ctypes.c_long(tid), ctypes.py_object(exc_type) ) def _kill_thread_by_tid(tid: Optional[int]) -> bool: if tid is None: return False res = _raise_async_exception(tid, SystemExit) if res == 0: return False if res > 1: _raise_async_exception(tid, None) raise SystemError("PyThreadState_SetAsyncExc affected multiple threads; reverted.") return True class ThreadManager: """ - add(udid, thread_or_target, *args, **kwargs) -> (code, msg) - stop(udid, join_timeout=2.0, retries=5, wait_step=0.2) -> (code, msg) # 强杀 - batch_stop(udids, join_timeout_each=2.0, retries_each=5, wait_step_each=0.2) -> (code, msg) - get_thread / get_tid / is_running / list_udids """ _threads: Dict[str, threading.Thread] = {} _lock = threading.RLock() # ========== 基础 ========== @classmethod def add(cls, udid: str, thread_or_target: Any, *args, **kwargs) -> Tuple[int, str]: """ 兼容两种用法: 1) add(udid, t) # t 是 threading.Thread 实例 2) add(udid, target, *args, **kwargs) # target 是可调用 返回:(200, "创建任务成功") / (1001, "任务已存在") / (1001, "创建任务失败") """ with cls._lock: exist = cls._threads.get(udid) if exist and exist.is_alive(): return 1001, "任务已存在" if isinstance(thread_or_target, threading.Thread): t = thread_or_target try: t.daemon = True except Exception: pass if not t.name: t.name = f"task-{udid}" # 包装 run,退出时从表移除 orig_run = t.run def run_wrapper(): try: orig_run() finally: with cls._lock: if cls._threads.get(udid) is t: cls._threads.pop(udid, None) t.run = run_wrapper # type: ignore else: target = thread_or_target def _wrapper(): try: target(*args, **kwargs) finally: with cls._lock: cur = cls._threads.get(udid) if cur is threading.current_thread(): cls._threads.pop(udid, None) t = threading.Thread(target=_wrapper, daemon=True, name=f"task-{udid}") try: t.start() except Exception: return 1001, "创建任务失败" cls._threads[udid] = t # 保留你原有的创建成功日志 try: LogManager.method_info(f"创建任务成功 [{udid}],线程ID={t.ident}", "task") except Exception: pass return 200, "创建任务成功" @classmethod def get_thread(cls, udid: str) -> Optional[threading.Thread]: with cls._lock: return cls._threads.get(udid) @classmethod def get_tid(cls, udid: str) -> Optional[int]: t = cls.get_thread(udid) return t.ident if t else None @classmethod def is_running(cls, udid: str) -> bool: t = cls.get_thread(udid) return bool(t and t.is_alive()) @classmethod def list_udids(cls) -> List[str]: with cls._lock: return list(cls._threads.keys()) # ========== 内部:强杀一次 ========== @classmethod def _stop_once(cls, udid: str, join_timeout: float, retries: int, wait_step: float) -> bool: """ 对指定 udid 执行一次强杀流程;返回 True=已停止/不存在,False=仍存活或被拒。 """ with cls._lock: t = cls._threads.get(udid) if not t: return True # 视为已停止 main_tid = threading.main_thread().ident cur_tid = threading.get_ident() if t.ident in (main_tid, cur_tid): return False try: _kill_thread_by_tid(t.ident) except Exception: pass if join_timeout < 0: join_timeout = 0.0 t.join(join_timeout) while t.is_alive() and retries > 0: evt = threading.Event() evt.wait(wait_step) retries -= 1 dead = not t.is_alive() if dead: with cls._lock: if cls._threads.get(udid) is t: cls._threads.pop(udid, None) return dead # ========== 对外:stop / batch_stop(均返回二元组) ========== @classmethod def stop(cls, udid: str, join_timeout: float = 2.0, retries: int = 5, wait_step: float = 0.2) -> Tuple[int, str]: """ 强杀单个:返回 (200, "stopped") 或 (1001, "failed") """ ok = cls._stop_once(udid, join_timeout, retries, wait_step) if ok: return 200, "stopped" else: return 1001, "failed" @classmethod def batch_stop(cls, udids: List[str], join_timeout_each: float = 2.0, retries_each: int = 5, wait_step_each: float = 0.2) -> Tuple[int, str]: """ 先全量执行一遍 -> 记录失败 -> 对失败重试 3 轮(每轮间隔 1 秒) 全部完成后统一返回: 全成功 : (200, "停止任务成功") 仍有失败: (1001, "停止任务失败") """ udids = udids or [] fail: List[str] = [] # 第一轮 for u in udids: ok = cls._stop_once(u, join_timeout_each, retries_each, wait_step_each) if not ok: fail.append(u) # 三轮只对失败重试 for _ in range(3): if not fail: break time.sleep(1.0) remain: List[str] = [] for u in fail: ok = cls._stop_once(u, join_timeout_each, retries_each, wait_step_each) if not ok: remain.append(u) fail = remain if not fail: return 200, "停止任务成功" else: return 1001, "停止任务失败"