421 lines
16 KiB
Python
421 lines
16 KiB
Python
import sys
|
|
import time
|
|
import subprocess
|
|
import uuid
|
|
import select
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from threading import Lock
|
|
import threading
|
|
from dataclasses import dataclass, field
|
|
from typing import Dict, Optional, Callable, Any
|
|
|
|
# ========== 1. 定义常量和类型别名 ==========
|
|
DEFAULT_MAX_WORKERS = 5
|
|
TaskCallback = Optional[Callable[[str, Dict[str, Any]], None]] # (TaskID, 结果字典) -> None
|
|
|
|
# 适配不同系统的select支持
|
|
if sys.platform == 'win32':
|
|
import msvcrt
|
|
def _win32_readline(process):
|
|
"""Windows下非阻塞读取子进程输出"""
|
|
lines = []
|
|
while True:
|
|
if msvcrt.kbhit():
|
|
break
|
|
try:
|
|
line = process.stdout.readline()
|
|
if not line:
|
|
break
|
|
lines.append(line)
|
|
except:
|
|
break
|
|
return lines
|
|
else:
|
|
def _unix_readline(process):
|
|
"""Unix/Linux下非阻塞读取子进程输出"""
|
|
lines = []
|
|
ready, _, _ = select.select([process.stdout], [], [], 0.1)
|
|
if ready:
|
|
for line in iter(process.stdout.readline, ''):
|
|
if line:
|
|
lines.append(line)
|
|
else:
|
|
break
|
|
return lines
|
|
|
|
# ========== 2. 任务类（记录任务状态） ==========
|
|
@dataclass
|
|
class Task:
|
|
TaskID: str # 任务唯一ID（帕斯卡命名）
|
|
Command: str # 要执行的命令
|
|
Process: Optional[subprocess.Popen] = None # 命令进程
|
|
IsRunning: bool = False # 任务运行状态
|
|
IsStarted: bool = False # 任务是否已启动过（用于重复启动）
|
|
StartTime: Optional[float] = None # 启动时间
|
|
Callback: TaskCallback = None # 任务结束回调函数
|
|
ExitCode: Optional[int] = None # 任务退出码
|
|
ErrorMsg: Optional[str] = None # 错误信息
|
|
Lock: threading.Lock = field(default_factory=Lock) # 单个任务的锁
|
|
StopFlag: bool = False # 任务停止标记
|
|
|
|
# ========== 3. 任务池核心类 ==========
|
|
class CommandTaskPool:
|
|
def __init__(self, maxWorkers: int = DEFAULT_MAX_WORKERS):
|
|
"""
|
|
初始化命令任务池
|
|
:param maxWorkers: 最大并发工作线程数
|
|
"""
|
|
# 线程池（maxWorkers限制并发任务数）
|
|
self.Executor = ThreadPoolExecutor(max_workers=maxWorkers)
|
|
# 任务字典：key=TaskID, value=Task对象
|
|
self.Tasks: Dict[str, Task] = {}
|
|
# 全局锁（防止多线程修改任务字典冲突）
|
|
self.GlobalLock = Lock()
|
|
# 日志前缀（区分不同任务的输出）
|
|
self.LogPrefix = "[TaskPool] "
|
|
|
|
def _Log(self, msg: str, taskID: str = "") -> None:
|
|
"""
|
|
日志输出（带时间戳和任务ID）
|
|
:param msg: 日志信息
|
|
:param taskID: 任务ID
|
|
"""
|
|
timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
|
|
taskTag = f"[{taskID}] " if taskID else ""
|
|
logMsg = f"{self.LogPrefix}{timestamp} {taskTag}{msg}\n"
|
|
sys.stdout.write(logMsg)
|
|
sys.stdout.flush() # 强制刷新，避免日志缓存
|
|
|
|
def _ExecuteCallback(self, task: Task) -> None:
|
|
"""
|
|
执行任务回调函数
|
|
:param task: 任务对象
|
|
"""
|
|
if task.Callback:
|
|
try:
|
|
result = {
|
|
"Command": task.Command,
|
|
"IsRunning": task.IsRunning,
|
|
"ExitCode": task.ExitCode,
|
|
"ErrorMsg": task.ErrorMsg,
|
|
"RuntimeSeconds": round(time.time() - task.StartTime, 2) if task.StartTime else 0
|
|
}
|
|
task.Callback(task.TaskID, result)
|
|
except Exception as e:
|
|
self._Log(f"回调函数执行失败: {str(e)}", task.TaskID)
|
|
|
|
def _RunCommand(self, task: Task) -> None:
|
|
"""
|
|
执行命令核心逻辑（实时输出日志，修复卡死问题）
|
|
:param task: 任务对象
|
|
"""
|
|
with task.Lock:
|
|
if task.IsRunning:
|
|
self._Log("任务已在运行中，跳过重复执行", task.TaskID)
|
|
return
|
|
task.IsRunning = True
|
|
task.IsStarted = True
|
|
task.StartTime = time.time()
|
|
task.ExitCode = None
|
|
task.ErrorMsg = None
|
|
task.StopFlag = False
|
|
|
|
try:
|
|
self._Log(f"启动命令: {task.Command}", task.TaskID)
|
|
# 启动子进程执行命令（修改为非阻塞模式）
|
|
task.Process = subprocess.Popen(
|
|
task.Command,
|
|
shell=True,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT, # 将stderr合并到stdout
|
|
text=True, # 文本模式读取
|
|
bufsize=0, # 无缓冲，避免数据堆积
|
|
universal_newlines=True,
|
|
creationflags=subprocess.CREATE_NEW_PROCESS_GROUP if sys.platform == 'win32' else 0
|
|
)
|
|
|
|
# 非阻塞读取输出（核心修复：避免死锁）
|
|
while True:
|
|
# 检查是否需要停止
|
|
with task.Lock:
|
|
if task.StopFlag:
|
|
break
|
|
# 检查进程是否已退出
|
|
if task.Process.poll() is not None:
|
|
break
|
|
|
|
# 根据系统选择读取方式
|
|
try:
|
|
if sys.platform == 'win32':
|
|
lines = _win32_readline(task.Process)
|
|
else:
|
|
lines = _unix_readline(task.Process)
|
|
|
|
for line in lines:
|
|
if line:
|
|
sys.stdout.write(f"[{task.TaskID}] {line}")
|
|
sys.stdout.flush()
|
|
except Exception as e:
|
|
self._Log(f"读取输出异常: {str(e)}", task.TaskID)
|
|
break
|
|
|
|
time.sleep(0.01) # 降低CPU占用
|
|
|
|
# 等待命令结束并获取退出码
|
|
task.ExitCode = task.Process.wait(timeout=2) if task.Process.poll() is None else task.Process.poll()
|
|
with task.Lock:
|
|
task.IsRunning = False
|
|
if task.ExitCode == 0:
|
|
self._Log(f"命令正常结束，退出码: {task.ExitCode}", task.TaskID)
|
|
else:
|
|
self._Log(f"命令异常结束，退出码: {task.ExitCode}", task.TaskID)
|
|
|
|
except subprocess.TimeoutExpired:
|
|
with task.Lock:
|
|
task.IsRunning = False
|
|
task.ExitCode = -2
|
|
self._Log(f"任务超时终止", task.TaskID)
|
|
except Exception as e:
|
|
errorMsg = str(e)
|
|
with task.Lock:
|
|
task.IsRunning = False
|
|
task.ErrorMsg = errorMsg
|
|
task.ExitCode = -1
|
|
self._Log(f"命令执行失败: {errorMsg}", task.TaskID)
|
|
|
|
finally:
|
|
# 确保进程被终止
|
|
if task.Process and task.Process.poll() is None:
|
|
try:
|
|
task.Process.kill()
|
|
except:
|
|
pass
|
|
# 执行回调函数
|
|
self._ExecuteCallback(task)
|
|
|
|
def GenerateTaskID(self, prefix: str = "TASK") -> str:
|
|
"""
|
|
生成唯一任务ID（前缀+UUID）
|
|
:param prefix: 任务ID前缀
|
|
:return: 唯一任务ID
|
|
"""
|
|
return f"{prefix}_{uuid.uuid4().hex[:8].upper()}"
|
|
|
|
def AddTask(self, command: str, taskID: Optional[str] = None,
|
|
callback: TaskCallback = None, autoStart: bool = True) -> str:
|
|
"""
|
|
添加任务到任务池
|
|
:param command: 要执行的命令
|
|
:param taskID: 自定义任务ID（为空则自动生成）
|
|
:param callback: 任务结束回调函数
|
|
:param autoStart: 是否添加后自动启动
|
|
:return: 任务ID（添加失败返回空字符串）
|
|
"""
|
|
with self.GlobalLock:
|
|
# 自动生成任务ID
|
|
taskID = taskID or self.GenerateTaskID()
|
|
|
|
# 检查任务ID是否已存在
|
|
if taskID in self.Tasks:
|
|
self._Log(f"任务ID已存在: {taskID}")
|
|
return ""
|
|
|
|
# 创建任务对象
|
|
task = Task(
|
|
TaskID=taskID,
|
|
Command=command,
|
|
Callback=callback
|
|
)
|
|
self.Tasks[taskID] = task
|
|
|
|
# 自动启动任务
|
|
if autoStart:
|
|
self.StartTask(taskID)
|
|
|
|
self._Log(f"任务已添加 {'(自动启动)' if autoStart else '(未启动)'}", taskID)
|
|
return taskID
|
|
|
|
def StartTask(self, taskID: str) -> bool:
|
|
"""
|
|
手动启动任务（支持重复启动已停止的任务）
|
|
:param taskID: 任务ID
|
|
:return: 启动成功/失败
|
|
"""
|
|
with self.GlobalLock:
|
|
if taskID not in self.Tasks:
|
|
self._Log(f"任务不存在: {taskID}")
|
|
return False
|
|
|
|
task = self.Tasks[taskID]
|
|
with task.Lock:
|
|
if task.IsRunning:
|
|
self._Log(f"任务已在运行: {taskID}")
|
|
return False
|
|
task.StopFlag = False # 重置停止标记
|
|
|
|
# 提交任务到线程池
|
|
self.Executor.submit(self._RunCommand, task)
|
|
self._Log(f"任务启动成功", taskID)
|
|
return True
|
|
|
|
def StopTask(self, taskID: str) -> bool:
|
|
"""
|
|
停止指定任务（修复卡死问题）
|
|
:param taskID: 任务ID
|
|
:return: 停止成功/失败
|
|
"""
|
|
with self.GlobalLock:
|
|
if taskID not in self.Tasks:
|
|
self._Log(f"任务不存在: {taskID}")
|
|
return False
|
|
|
|
task = self.Tasks[taskID]
|
|
with task.Lock:
|
|
if not task.IsRunning or task.Process is None:
|
|
self._Log(f"任务未运行: {taskID}")
|
|
return False
|
|
|
|
# 设置停止标记
|
|
task.StopFlag = True
|
|
# 终止子进程
|
|
try:
|
|
if sys.platform == 'win32':
|
|
# Windows下终止进程组
|
|
subprocess.call(['taskkill', '/F', '/T', '/PID', str(task.Process.pid)])
|
|
else:
|
|
task.Process.terminate()
|
|
|
|
task.IsRunning = False
|
|
task.ExitCode = -1 # 标记为手动终止
|
|
self._Log(f"任务已终止", taskID)
|
|
return True
|
|
except Exception as e:
|
|
self._Log(f"终止任务失败: {str(e)}", taskID)
|
|
return False
|
|
|
|
def ListTasks(self) -> Dict[str, dict]:
|
|
"""
|
|
查看所有任务状态
|
|
:return: 任务状态字典
|
|
"""
|
|
with self.GlobalLock:
|
|
taskStatus = {}
|
|
currentTime = time.time()
|
|
for taskID, task in self.Tasks.items():
|
|
with task.Lock:
|
|
runtime = round(currentTime - task.StartTime, 2) if task.StartTime else 0
|
|
taskStatus[taskID] = {
|
|
"Command": task.Command,
|
|
"IsRunning": task.IsRunning,
|
|
"IsStarted": task.IsStarted,
|
|
"RuntimeSeconds": runtime,
|
|
"ExitCode": task.ExitCode,
|
|
"ErrorMsg": task.ErrorMsg
|
|
}
|
|
return taskStatus
|
|
|
|
def RemoveTask(self, taskID: str) -> bool:
|
|
"""
|
|
移除任务（需先停止）
|
|
:param taskID: 任务ID
|
|
:return: 移除成功/失败
|
|
"""
|
|
with self.GlobalLock:
|
|
if taskID not in self.Tasks:
|
|
self._Log(f"任务不存在: {taskID}")
|
|
return False
|
|
|
|
task = self.Tasks[taskID]
|
|
with task.Lock:
|
|
if task.IsRunning:
|
|
self._Log(f"任务正在运行，无法移除: {taskID}")
|
|
return False
|
|
|
|
del self.Tasks[taskID]
|
|
self._Log(f"任务已移除", taskID)
|
|
return True
|
|
|
|
def Shutdown(self, wait: bool = False) -> None:
|
|
"""
|
|
关闭任务池，停止所有任务（修改wait默认值为False避免卡死）
|
|
:param wait: 是否等待所有任务结束
|
|
"""
|
|
self._Log("开始关闭任务池...")
|
|
|
|
# 停止所有运行中的任务
|
|
with self.GlobalLock:
|
|
for taskID in list(self.Tasks.keys()):
|
|
self.StopTask(taskID)
|
|
|
|
# 关闭线程池（非阻塞模式）
|
|
self.Executor.shutdown(wait=wait, cancel_futures=True)
|
|
self._Log("任务池已关闭")
|
|
|
|
# ========== 4. 测试示例 ==========
|
|
def TaskCallbackExample(taskID: str, result: dict) -> None:
|
|
"""示例回调函数"""
|
|
print(f"\n=== 任务 {taskID} 回调结果 ===")
|
|
print(f"命令: {result['Command']}")
|
|
print(f"运行状态: {result['IsRunning']}")
|
|
print(f"退出码: {result['ExitCode']}")
|
|
print(f"运行时长: {result['RuntimeSeconds']}秒")
|
|
if result['ErrorMsg']:
|
|
print(f"错误信息: {result['ErrorMsg']}")
|
|
|
|
if __name__ == "__main__":
|
|
# 创建任务池（最多并发5个任务）
|
|
taskPool = CommandTaskPool(maxWorkers=5)
|
|
|
|
# 示例1：添加任务（自动启动，带回调）
|
|
pingTaskID = taskPool.AddTask(
|
|
command="ping -t 127.0.0.1", # Windows
|
|
# command="ping 127.0.0.1 -c 100", # Linux/Mac
|
|
callback=TaskCallbackExample
|
|
)
|
|
print(f"添加Ping任务，ID: {pingTaskID}")
|
|
|
|
# 示例2：添加任务（不自动启动）
|
|
frpTaskID = taskPool.AddTask(
|
|
command="frps -c /etc/frp/frps.ini",
|
|
taskID="FRPS_SERVICE", # 自定义任务ID
|
|
autoStart=False # 不自动启动
|
|
)
|
|
print(f"添加FRP任务，ID: {frpTaskID}")
|
|
|
|
# 查看初始任务状态
|
|
time.sleep(1)
|
|
print("\n=== 初始任务状态 ===")
|
|
for taskID, status in taskPool.ListTasks().items():
|
|
print(f"任务ID: {taskID}")
|
|
print(f" 命令: {status['Command']}")
|
|
print(f" 运行中: {status['IsRunning']}")
|
|
print(f" 已启动: {status['IsStarted']}")
|
|
|
|
# 手动启动FRP任务
|
|
print("\n=== 手动启动FRP任务 ===")
|
|
taskPool.StartTask(frpTaskID)
|
|
time.sleep(1)
|
|
frpStatus = taskPool.ListTasks().get(frpTaskID, {})
|
|
print(f"FRP任务状态: {frpStatus.get('IsRunning', False)}")
|
|
|
|
# 模拟运行3秒后停止Ping任务
|
|
time.sleep(3)
|
|
print("\n=== 停止Ping任务 ===")
|
|
taskPool.StopTask(pingTaskID)
|
|
|
|
# 重新启动Ping任务（支持重复启动）
|
|
print("\n=== 重新启动Ping任务 ===")
|
|
taskPool.StartTask(pingTaskID)
|
|
time.sleep(2)
|
|
|
|
# 查看最终任务状态
|
|
print("\n=== 最终任务状态 ===")
|
|
for taskID, status in taskPool.ListTasks().items():
|
|
print(f"任务ID: {taskID}")
|
|
print(f" 运行中: {status['IsRunning']}")
|
|
print(f" 运行时长: {status['RuntimeSeconds']}秒")
|
|
print(f" 退出码: {status['ExitCode']}")
|
|
|
|
# 关闭任务池（非阻塞模式）
|
|
taskPool.Shutdown(wait=False)
|
|
print("\n任务池已关闭，程序正常退出") |