141 lines
5.0 KiB
Python
141 lines
5.0 KiB
Python
import ctypes
|
||
import threading
|
||
import time
|
||
from typing import Dict, Tuple, List
|
||
|
||
from Utils.LogManager import LogManager
|
||
|
||
|
||
def _async_raise(tid: int, exc_type=KeyboardInterrupt) -> bool:
|
||
"""向指定线程抛异常（兜底方案）"""
|
||
try:
|
||
res = ctypes.pythonapi.PyThreadState_SetAsyncExc(
|
||
ctypes.c_long(tid), ctypes.py_object(exc_type)
|
||
)
|
||
if res == 0:
|
||
LogManager.method_info(f"线程 {tid} 不存在", "task")
|
||
return False
|
||
elif res > 1:
|
||
ctypes.pythonapi.PyThreadState_SetAsyncExc(ctypes.c_long(tid), 0)
|
||
LogManager.method_info(f"线程 {tid} 命中多个线程，已回滚", "task")
|
||
return False
|
||
return True
|
||
except Exception as e:
|
||
LogManager.method_error(f"强杀线程失败: {e}", "task")
|
||
return False
|
||
|
||
|
||
class ThreadManager:
|
||
_tasks: Dict[str, Dict] = {}
|
||
_lock = threading.RLock()
|
||
|
||
@classmethod
|
||
def _cleanup_if_dead(cls, udid: str):
|
||
"""如果任务线程已结束，清理占位"""
|
||
obj = cls._tasks.get(udid)
|
||
if obj and not obj["thread"].is_alive():
|
||
cls._tasks.pop(udid, None)
|
||
LogManager.method_info(f"检测到 [{udid}] 线程已结束，自动清理。", "task")
|
||
|
||
@classmethod
|
||
def add(cls, udid: str, thread: threading.Thread, event: threading.Event, force: bool = False) -> Tuple[int, str]:
|
||
with cls._lock:
|
||
cls._cleanup_if_dead(udid)
|
||
|
||
# 已存在任务还在运行
|
||
old = cls._tasks.get(udid)
|
||
if old and old["thread"].is_alive():
|
||
if not force:
|
||
return 1001, "当前设备已存在任务"
|
||
LogManager.method_info(f"[{udid}] 检测到旧任务，尝试强制停止", "task")
|
||
cls._force_stop_locked(udid)
|
||
|
||
# 启动新任务
|
||
try:
|
||
thread.start()
|
||
cls._tasks[udid] = {
|
||
"id": thread.ident,
|
||
"thread": thread,
|
||
"event": event,
|
||
"start_time": time.time(),
|
||
}
|
||
LogManager.method_info(f"创建任务成功 [{udid}]，线程ID={thread.ident}", "task")
|
||
return 200, "创建成功"
|
||
except Exception as e:
|
||
LogManager.method_error(f"线程启动失败: {e}", "task")
|
||
return 1002, f"线程启动失败: {e}"
|
||
|
||
@classmethod
|
||
def stop(cls, udid: str, stop_timeout: float = 5.0, kill_timeout: float = 2.0) -> Tuple[int, str]:
|
||
"""安全停止单个任务"""
|
||
with cls._lock:
|
||
obj = cls._tasks.get(udid)
|
||
if not obj:
|
||
return 200, "任务不存在"
|
||
|
||
thread = obj["thread"]
|
||
event = obj["event"]
|
||
tid = obj["id"]
|
||
|
||
LogManager.method_info(f"请求停止 [{udid}] 线程ID={tid}", "task")
|
||
|
||
# 已经结束
|
||
if not thread.is_alive():
|
||
cls._tasks.pop(udid, None)
|
||
return 200, "已结束"
|
||
|
||
# 1. 协作式停止
|
||
try:
|
||
event.set()
|
||
except Exception as e:
|
||
LogManager.method_error(f"[{udid}] 设置停止事件失败: {e}", "task")
|
||
|
||
thread.join(timeout=stop_timeout)
|
||
if not thread.is_alive():
|
||
cls._tasks.pop(udid, None)
|
||
LogManager.method_info(f"[{udid}] 协作式停止成功", "task")
|
||
return 200, "已停止"
|
||
|
||
# 2. 强杀兜底
|
||
LogManager.method_info(f"[{udid}] 协作式超时，尝试强杀", "task")
|
||
if _async_raise(tid):
|
||
thread.join(timeout=kill_timeout)
|
||
|
||
if not thread.is_alive():
|
||
cls._tasks.pop(udid, None)
|
||
LogManager.method_info(f"[{udid}] 强杀成功", "task")
|
||
return 200, "已停止"
|
||
|
||
# 3. 最终兜底：标记释放占位
|
||
LogManager.method_error(f"[{udid}] 无法停止（线程可能卡死），已释放占位", "task")
|
||
cls._tasks.pop(udid, None)
|
||
return 206, "停止超时，线程可能仍在后台运行"
|
||
|
||
@classmethod
|
||
def _force_stop_locked(cls, udid: str):
|
||
"""内部用，带锁强制停止旧任务"""
|
||
obj = cls._tasks.get(udid)
|
||
if not obj:
|
||
return
|
||
try:
|
||
event = obj["event"]
|
||
event.set()
|
||
obj["thread"].join(timeout=2)
|
||
if obj["thread"].is_alive():
|
||
_async_raise(obj["id"])
|
||
obj["thread"].join(timeout=1)
|
||
except Exception as e:
|
||
LogManager.method_error(f"[{udid}] 强制停止失败: {e}", "task")
|
||
finally:
|
||
cls._tasks.pop(udid, None)
|
||
|
||
@classmethod
|
||
def batch_stop(cls, ids: List[str]) -> Tuple[int, str]:
|
||
failed = []
|
||
for udid in ids:
|
||
code, msg = cls.stop(udid)
|
||
if code != 200:
|
||
failed.append(udid)
|
||
if failed:
|
||
return 207, f"部分任务未成功停止: {failed}"
|
||
return 200, "全部停止成功" |