228 lines
9.4 KiB
Python
228 lines
9.4 KiB
Python
from datetime import datetime
|
||
from typing import List, Dict, Any, Optional
|
||
import logging
|
||
import asyncio
|
||
|
||
|
||
# 配置日志
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||
)
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 心跳暂未启用，直接使用100分钟
|
||
TaskTimeout=6000
|
||
class TaskManager:
|
||
def __init__(self):
|
||
self.tasks: Dict[str, Dict[str, Any]] = {} # task_id -> task_info
|
||
self.heartbeats: Dict[str, datetime] = {} # task_id -> last_heartbeat
|
||
self.task_handles: Dict[str, asyncio.Task] = {} # task_id -> main_task_handle
|
||
self.sub_tasks: Dict[str, List[asyncio.Task]] = {} # task_id -> [sub_task_handles]
|
||
self.stop_events: Dict[str, asyncio.Event] = {} # task_id -> stop_event
|
||
self.lock = asyncio.Lock()
|
||
self.heartbeat_interval = 30
|
||
self.task_timeout = TaskTimeout
|
||
self.cleanup_interval = 10
|
||
|
||
async def start(self):
|
||
"""启动任务管理器"""
|
||
self.cleanup_task = asyncio.create_task(self._cleanup_tasks())
|
||
|
||
async def add_task(self, task_id: str, task_info: Dict[str, Any], main_task_handle: asyncio.Task,
|
||
sub_task_handles: List[asyncio.Task] = None):
|
||
"""添加新任务及其所有子任务"""
|
||
async with self.lock:
|
||
if task_id in self.tasks:
|
||
raise ValueError("Task already exists")
|
||
|
||
# 为每个任务创建stop_event
|
||
stop_event = asyncio.Event()
|
||
|
||
self.tasks[task_id] = {
|
||
"task_info": task_info,
|
||
"status": "running",
|
||
"main_task_handle": main_task_handle,
|
||
"sub_task_handles": sub_task_handles if sub_task_handles else []
|
||
}
|
||
self.heartbeats[task_id] = datetime.now()
|
||
self.task_handles[task_id] = main_task_handle
|
||
self.stop_events[task_id] = stop_event # 存储stop_event
|
||
if sub_task_handles:
|
||
self.sub_tasks[task_id] = sub_task_handles
|
||
logger.info(f"Task {task_id} started")
|
||
#
|
||
async def remove_task(self, task_id: str):
|
||
"""移除任务及其所有子任务，包括MQTT资源"""
|
||
async with self.lock:
|
||
if task_id not in self.tasks:
|
||
return False
|
||
|
||
# 获取stop_event并设置
|
||
if task_id in self.stop_events:
|
||
self.stop_events[task_id].set() # 触发停止信号
|
||
|
||
# 获取所有子任务
|
||
sub_tasks = []
|
||
if task_id in self.sub_tasks:
|
||
sub_tasks = self.sub_tasks[task_id]
|
||
|
||
# 获取主任务
|
||
if task_id in self.task_handles:
|
||
main_task = self.task_handles[task_id]
|
||
if not main_task.done():
|
||
try:
|
||
# 尝试取消主任务
|
||
await asyncio.shield(main_task.cancel())
|
||
except Exception as e:
|
||
logger.warning(f"取消主任务 {task_id} 失败: {e}")
|
||
|
||
# 清理MQTT设备资源
|
||
if task_id in self.tasks:
|
||
task_info = self.tasks[task_id]
|
||
if 'device' in task_info: # 假设任务信息中包含MQTT设备引用
|
||
try:
|
||
for device in task_info['device']:
|
||
await device.stop() # 停止MQTT设备
|
||
except Exception as e:
|
||
logger.warning(f"停止MQTT设备 {task_id} 失败: {e}")
|
||
|
||
# 取消所有子任务
|
||
for sub_task in sub_tasks:
|
||
if not sub_task.done():
|
||
try:
|
||
await asyncio.shield(sub_task.cancel())
|
||
except Exception as e:
|
||
logger.warning(f"取消子任务 {task_id} 失败: {e}")
|
||
|
||
# 清理资源
|
||
del self.tasks[task_id]
|
||
if task_id in self.task_handles:
|
||
del self.task_handles[task_id]
|
||
if task_id in self.sub_tasks:
|
||
del self.sub_tasks[task_id]
|
||
if task_id in self.stop_events:
|
||
del self.stop_events[task_id]
|
||
|
||
logger.info(f"Task {task_id} removed")
|
||
return True
|
||
#
|
||
# async def remove_task(self, task_id: str):
|
||
# """改进的任务移除逻辑"""
|
||
# async with self.lock:
|
||
# if task_id not in self.tasks:
|
||
# return False
|
||
#
|
||
# # 设置取消标志
|
||
# if task_id in self.stop_events:
|
||
# self.stop_events[task_id].set() # 触发停止信号
|
||
#
|
||
# # 获取任务信息
|
||
# task_info = self.tasks[task_id]
|
||
#
|
||
# # 取消主任务（如果是协程）
|
||
# if 'main_task' in task_info and isinstance(task_info['main_task'], asyncio.Task):
|
||
# try:
|
||
# task_info['main_task'].cancel()
|
||
# await asyncio.shield(task_info['main_task']) # 等待任务完成取消
|
||
# except Exception as e:
|
||
# logger.warning(f"取消主任务 {task_id} 失败: {str(e)}")
|
||
#
|
||
# # 取消所有子任务
|
||
# if 'sub_tasks' in task_info and task_id in self.sub_tasks:
|
||
# for sub_task in self.sub_tasks[task_id]:
|
||
# if isinstance(sub_task, asyncio.Task):
|
||
# try:
|
||
# sub_task.cancel()
|
||
# await asyncio.shield(sub_task)
|
||
# except Exception as e:
|
||
# logger.warning(f"取消子任务 {task_id} 失败: {str(e)}")
|
||
#
|
||
# # 清理MQTT资源
|
||
# if 'device' in task_info and isinstance(task_info['device'], MQTTDevice):
|
||
# try:
|
||
# await task_info['device'].stop()
|
||
# except Exception as e:
|
||
# logger.warning(f"停止MQTT设备 {task_id} 失败: {str(e)}")
|
||
#
|
||
# # 清理资源
|
||
# del self.tasks[task_id]
|
||
# if task_id in self.sub_tasks:
|
||
# del self.sub_tasks[task_id]
|
||
# if task_id in self.stop_events:
|
||
# del self.stop_events[task_id]
|
||
#
|
||
# logger.info(f"Task {task_id} removed successfully")
|
||
# return True
|
||
async def mark_all_tasks_as_stopped(self):
|
||
"""标记所有任务为停止状态"""
|
||
async with self.lock:
|
||
for task_id in self.tasks:
|
||
self.tasks[task_id]["status"] = "stopped"
|
||
|
||
async def update_heartbeat(self, task_id: str) -> None:
|
||
"""更新任务心跳时间"""
|
||
if task_id in self.heartbeats:
|
||
self.heartbeats[task_id] = datetime.now()
|
||
logger.debug(f"Heartbeat updated for task {task_id}")
|
||
|
||
async def get_task_info(self, task_id: str) -> Optional[Dict[str, any]]:
|
||
"""获取任务信息"""
|
||
async with self.lock:
|
||
if task_id not in self.tasks:
|
||
return None
|
||
|
||
last_heartbeat = self.heartbeats.get(task_id)
|
||
current_time = datetime.now()
|
||
# elapsed_time = (current_time - last_heartbeat).total_seconds()
|
||
elapsed_time = (current_time - last_heartbeat).total_seconds() if last_heartbeat else None
|
||
|
||
return {
|
||
"task_info": self.tasks[task_id],
|
||
"status": "running" if (elapsed_time is not None and elapsed_time < self.task_timeout) else "timeout",
|
||
"last_heartbeat": last_heartbeat.isoformat(),
|
||
"elapsed_time": elapsed_time
|
||
}
|
||
|
||
async def check_tasks_health(self) -> Dict[str, str]:
|
||
"""检查任务健康状态"""
|
||
unhealthy_tasks = {}
|
||
current_time = datetime.now()
|
||
|
||
async with self.lock:
|
||
for task_id in list(self.tasks.keys()):
|
||
last_heartbeat = self.heartbeats.get(task_id)
|
||
elapsed_seconds = (current_time - last_heartbeat).total_seconds() if last_heartbeat else None
|
||
print(f"当前心跳已经进行{elapsed_seconds}s")
|
||
if elapsed_seconds is not None and elapsed_seconds > self.task_timeout:
|
||
unhealthy_tasks[task_id] = "timeout"
|
||
logger.warning(f"----------------------------Task {task_id} has timed out-----------------------------")
|
||
|
||
return unhealthy_tasks
|
||
|
||
async def _cleanup_tasks(self):
|
||
"""清理超时任务"""
|
||
while True:
|
||
await asyncio.sleep(self.cleanup_interval)
|
||
unhealthy_tasks = await self.check_tasks_health()
|
||
for task_id in unhealthy_tasks:
|
||
await self.remove_task(task_id)
|
||
logger.warning(f"Task {task_id} removed due to timeout")
|
||
|
||
async def get_all_tasks(self) -> Dict[str, Dict[str, any]]:
|
||
"""获取所有任务信息"""
|
||
async with self.lock:
|
||
return {
|
||
task_id: {
|
||
# "task_info": task_info,
|
||
"status": "running" if (datetime.now() - self.heartbeats.get(task_id, datetime.min)).total_seconds() < self.task_timeout
|
||
else "timeout",
|
||
"last_heartbeat": self.heartbeats.get(task_id, datetime.min).isoformat(),
|
||
"elapsed_time": (datetime.now() - self.heartbeats.get(task_id, datetime.min)).total_seconds()
|
||
if task_id in self.heartbeats else 0
|
||
}
|
||
for task_id, task_info in self.tasks.items()
|
||
}
|
||
|
||
|
||
task_manager = TaskManager() |