mirror of
https://github.com/rachelos/we-mp-rss.git
synced 2026-11-05 02:14:07 +08:00
512 lines
19 KiB
Python
512 lines
19 KiB
Python
import queue
|
|
import threading
|
|
import time
|
|
import gc
|
|
import json
|
|
from typing import Callable, Any, Optional
|
|
from datetime import datetime
|
|
from dataclasses import dataclass, field, asdict
|
|
from core.print import print_error, print_info, print_warning, print_success
|
|
|
|
# Redis 键前缀
|
|
REDIS_KEY_PREFIX = "werss:queue"
|
|
REDIS_KEY_PENDING = f"{REDIS_KEY_PREFIX}:pending"
|
|
REDIS_KEY_CURRENT = f"{REDIS_KEY_PREFIX}:current"
|
|
REDIS_KEY_HISTORY = f"{REDIS_KEY_PREFIX}:history"
|
|
REDIS_KEY_STATUS = f"{REDIS_KEY_PREFIX}:status"
|
|
|
|
# 全局 Redis 客户端（延迟初始化）
|
|
_redis_client = None
|
|
|
|
|
|
def _get_redis():
|
|
"""获取 Redis 客户端（延迟初始化）"""
|
|
global _redis_client
|
|
if _redis_client is not None:
|
|
return _redis_client
|
|
|
|
try:
|
|
import redis
|
|
from core.config import cfg
|
|
|
|
redis_url = cfg.get("redis.url", "")
|
|
if not redis_url:
|
|
print_info("Redis URL 未配置，队列状态将仅在进程内共享")
|
|
return None
|
|
|
|
_redis_client = redis.from_url(
|
|
redis_url,
|
|
decode_responses=True,
|
|
socket_connect_timeout=5,
|
|
socket_timeout=5,
|
|
retry_on_timeout=True
|
|
)
|
|
# 测试连接
|
|
_redis_client.ping()
|
|
print_info("Queue Redis 连接成功")
|
|
return _redis_client
|
|
except Exception as e:
|
|
print_error(f"Queue Redis 连接失败: {e}")
|
|
_redis_client = None
|
|
return None
|
|
|
|
|
|
@dataclass
|
|
class TaskItem:
|
|
"""待执行任务项"""
|
|
task_name: str
|
|
args: tuple = field(default_factory=tuple)
|
|
kwargs: dict = field(default_factory=dict)
|
|
max_retries: int = 3
|
|
|
|
def _serialize_value(self, value: Any) -> Any:
|
|
"""安全序列化值，处理不可 JSON 序列化的对象"""
|
|
if value is None:
|
|
return None
|
|
elif isinstance(value, (str, int, float, bool)):
|
|
return value
|
|
elif isinstance(value, (list, tuple)):
|
|
return [self._serialize_value(v) for v in value]
|
|
elif isinstance(value, dict):
|
|
return {k: self._serialize_value(v) for k, v in value.items()}
|
|
else:
|
|
# 对于复杂对象，返回其字符串表示
|
|
return repr(value)
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
'task_name': self.task_name,
|
|
'args': self._serialize_value(list(self.args) if self.args else []),
|
|
'kwargs': self._serialize_value(self.kwargs),
|
|
'max_retries': self.max_retries
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class TaskRecord:
|
|
"""任务执行记录"""
|
|
task_name: str
|
|
start_time: str
|
|
end_time: Optional[str] = None
|
|
duration: Optional[float] = None
|
|
status: str = "running" # running, completed, failed
|
|
error: Optional[str] = None
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
'task_name': self.task_name,
|
|
'start_time': self.start_time,
|
|
'end_time': self.end_time or '',
|
|
'duration': self.duration or 0,
|
|
'status': self.status,
|
|
'error': self.error or ''
|
|
}
|
|
|
|
|
|
class TaskQueueManager:
|
|
"""任务队列管理器，用于管理和执行排队任务（支持多进程）"""
|
|
|
|
def __init__(self, maxsize=0, tag=""):
|
|
"""初始化任务队列"""
|
|
self._queue = queue.Queue(maxsize=maxsize)
|
|
self._thread_lock = threading.Lock()
|
|
self._is_running = False
|
|
self.tag = tag
|
|
# 任务历史记录（最近100条）
|
|
self._history: list[TaskRecord] = []
|
|
self._history_max_size = 100
|
|
# 当前执行的任务
|
|
self._current_task: Optional[TaskRecord] = None
|
|
# 待执行任务列表（用于展示）
|
|
self._pending_items: list[TaskItem] = []
|
|
self._instance_id = id(self)
|
|
print_info(f"TaskQueueManager initialized, instance id: {self._instance_id}")
|
|
|
|
def _save_pending_to_redis(self):
|
|
"""保存待执行任务到 Redis"""
|
|
redis_client = _get_redis()
|
|
if not redis_client:
|
|
return
|
|
|
|
try:
|
|
pending_list = [item.to_dict() for item in self._pending_items]
|
|
redis_client.delete(REDIS_KEY_PENDING)
|
|
if pending_list:
|
|
for item in pending_list:
|
|
redis_client.rpush(REDIS_KEY_PENDING, json.dumps(item, ensure_ascii=False))
|
|
except Exception as e:
|
|
print_error(f"保存待执行任务到 Redis 失败: {e}")
|
|
|
|
def _save_current_task_to_redis(self, task_record: Optional[TaskRecord]):
|
|
"""保存当前任务到 Redis"""
|
|
redis_client = _get_redis()
|
|
if not redis_client:
|
|
return
|
|
|
|
try:
|
|
if task_record:
|
|
redis_client.hset(REDIS_KEY_CURRENT, mapping=task_record.to_dict())
|
|
else:
|
|
redis_client.delete(REDIS_KEY_CURRENT)
|
|
except Exception as e:
|
|
print_error(f"保存当前任务到 Redis 失败: {e}")
|
|
|
|
def _save_history_to_redis(self, task_record: TaskRecord):
|
|
"""保存历史记录到 Redis"""
|
|
redis_client = _get_redis()
|
|
if not redis_client:
|
|
return
|
|
|
|
try:
|
|
redis_client.lpush(REDIS_KEY_HISTORY, json.dumps(task_record.to_dict(), ensure_ascii=False))
|
|
redis_client.ltrim(REDIS_KEY_HISTORY, 0, self._history_max_size - 1)
|
|
except Exception as e:
|
|
print_error(f"保存历史记录到 Redis 失败: {e}")
|
|
|
|
def _clear_history_from_redis(self):
|
|
"""清空 Redis 历史记录"""
|
|
redis_client = _get_redis()
|
|
if not redis_client:
|
|
return
|
|
|
|
try:
|
|
redis_client.delete(REDIS_KEY_HISTORY)
|
|
except Exception as e:
|
|
print_error(f"清空 Redis 历史记录失败: {e}")
|
|
|
|
def _save_status_to_redis(self):
|
|
"""保存队列状态到 Redis"""
|
|
redis_client = _get_redis()
|
|
if not redis_client:
|
|
return
|
|
|
|
try:
|
|
redis_client.hset(REDIS_KEY_STATUS, mapping={
|
|
'is_running': str(self._is_running).lower(),
|
|
'tag': self.tag or '',
|
|
'pending_count': self._queue.qsize(),
|
|
'history_count': len(self._history)
|
|
})
|
|
except Exception as e:
|
|
print_error(f"保存队列状态到 Redis 失败: {e}")
|
|
|
|
def _get_pending_from_redis(self) -> list:
|
|
"""从 Redis 获取待执行任务"""
|
|
redis_client = _get_redis()
|
|
if not redis_client:
|
|
return []
|
|
|
|
try:
|
|
items = redis_client.lrange(REDIS_KEY_PENDING, 0, -1)
|
|
return [json.loads(item) for item in items]
|
|
except Exception as e:
|
|
print_error(f"从 Redis 获取待执行任务失败: {e}")
|
|
return []
|
|
|
|
def _get_current_task_from_redis(self) -> Optional[dict]:
|
|
"""从 Redis 获取当前任务"""
|
|
redis_client = _get_redis()
|
|
if not redis_client:
|
|
return None
|
|
|
|
try:
|
|
data = redis_client.hgetall(REDIS_KEY_CURRENT)
|
|
if data and 'task_name' in data:
|
|
return data
|
|
return None
|
|
except Exception as e:
|
|
print_error(f"从 Redis 获取当前任务失败: {e}")
|
|
return None
|
|
|
|
def _get_history_from_redis(self, limit: int = 20) -> list:
|
|
"""从 Redis 获取历史记录"""
|
|
redis_client = _get_redis()
|
|
if not redis_client:
|
|
return []
|
|
|
|
try:
|
|
items = redis_client.lrange(REDIS_KEY_HISTORY, 0, limit - 1)
|
|
return [json.loads(item) for item in items]
|
|
except Exception as e:
|
|
print_error(f"从 Redis 获取历史记录失败: {e}")
|
|
return []
|
|
|
|
def _get_history_count_from_redis(self) -> int:
|
|
"""从 Redis 获取历史记录数量"""
|
|
redis_client = _get_redis()
|
|
if not redis_client:
|
|
return 0
|
|
|
|
try:
|
|
return redis_client.llen(REDIS_KEY_HISTORY)
|
|
except Exception as e:
|
|
print_error(f"从 Redis 获取历史记录数量失败: {e}")
|
|
return 0
|
|
|
|
def add_task(self, task: Callable[..., Any], *args: Any, max_retries: int = 3, **kwargs: Any) -> bool:
|
|
"""添加任务到队列
|
|
|
|
Args:
|
|
task: 要执行的任务函数
|
|
*args: 任务函数的参数
|
|
max_retries: 最大重试次数，默认3次
|
|
**kwargs: 任务函数的关键字参数
|
|
|
|
Returns:
|
|
bool: 是否成功添加到队列
|
|
"""
|
|
with self._thread_lock:
|
|
# 检查队列是否已满（如果设置了maxsize）
|
|
try:
|
|
self._queue.put_nowait((task, args, kwargs, max_retries))
|
|
except queue.Full:
|
|
print_error(f"{self.tag}队列已满，任务添加失败")
|
|
return False
|
|
|
|
# 记录待执行任务
|
|
task_name = getattr(task, '__name__', str(task))
|
|
self._pending_items.append(TaskItem(
|
|
task_name=task_name,
|
|
args=args,
|
|
kwargs=kwargs,
|
|
max_retries=max_retries
|
|
))
|
|
|
|
# 保存到 Redis
|
|
self._save_pending_to_redis()
|
|
self._save_status_to_redis()
|
|
|
|
print_success(f"{self.tag}队列任务添加成功\n")
|
|
return True
|
|
|
|
def run_task_background(self)->None:
|
|
threading.Thread(target=self.run_tasks, daemon=True).start()
|
|
print_warning(f"队列任务后台运行, instance id: {self._instance_id}")
|
|
|
|
def run_tasks(self, timeout: float = 1.0) -> None:
|
|
"""执行队列中的所有任务，并持续运行以接收新任务
|
|
|
|
Args:
|
|
timeout: 等待新任务的超时时间(秒)
|
|
"""
|
|
with self._thread_lock:
|
|
if self._is_running:
|
|
return
|
|
self._is_running = True
|
|
self._save_status_to_redis()
|
|
|
|
try:
|
|
while self._is_running:
|
|
time.sleep(0.1) # 避免过于频繁的任务获取
|
|
try:
|
|
# 阻塞获取任务，避免CPU空转
|
|
task_item = self._queue.get(timeout=timeout)
|
|
|
|
# 兼容新旧格式：新格式4个元素，旧格式3个元素
|
|
if len(task_item) == 4:
|
|
task, args, kwargs, max_retries = task_item
|
|
else:
|
|
task, args, kwargs = task_item
|
|
max_retries = 3
|
|
|
|
# 从待执行列表中移除
|
|
with self._thread_lock:
|
|
if self._pending_items:
|
|
self._pending_items.pop(0)
|
|
self._save_pending_to_redis()
|
|
|
|
# 记录任务开始
|
|
task_name = getattr(task, '__name__', str(task))
|
|
with self._thread_lock:
|
|
self._current_task = TaskRecord(
|
|
task_name=task_name,
|
|
start_time=datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
|
)
|
|
# 保存当前任务到 Redis
|
|
self._save_current_task_to_redis(self._current_task)
|
|
|
|
print_info(f"Task started: {task_name}")
|
|
|
|
retry_count = 0
|
|
success = False
|
|
last_error = None
|
|
start_time = time.time() # 提前定义，确保异常时可用
|
|
|
|
while retry_count <= max_retries and not success:
|
|
if retry_count > 0:
|
|
print_warning(f"任务 [{task_name}] 第 {retry_count} 次重试...")
|
|
time.sleep(2 ** retry_count) # 指数退避
|
|
|
|
try:
|
|
# 记录任务开始时间
|
|
start_time = time.time()
|
|
task(*args, **kwargs)
|
|
# 记录任务执行时间
|
|
duration = time.time() - start_time
|
|
|
|
# 更新当前任务记录
|
|
with self._thread_lock:
|
|
if self._current_task:
|
|
self._current_task.end_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
|
self._current_task.duration = duration
|
|
self._current_task.status = "completed"
|
|
|
|
print_info(f"\n任务执行完成，耗时: {duration:.2f}秒")
|
|
success = True
|
|
|
|
except Exception as e:
|
|
last_error = e
|
|
retry_count += 1
|
|
|
|
if retry_count <= max_retries:
|
|
print_warning(f"任务 [{task_name}] 执行失败: {e}，准备重试 ({retry_count}/{max_retries})")
|
|
else:
|
|
# 达到最大重试次数，记录失败
|
|
with self._thread_lock:
|
|
if self._current_task:
|
|
self._current_task.end_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
|
self._current_task.duration = time.time() - start_time
|
|
self._current_task.status = "failed"
|
|
self._current_task.error = str(e)
|
|
|
|
print_error(f"任务 [{task_name}] 执行失败，已达到最大重试次数: {e}")
|
|
|
|
# 保存到历史记录
|
|
with self._thread_lock:
|
|
if self._current_task:
|
|
self._history.append(self._current_task)
|
|
# 保存到 Redis
|
|
self._save_history_to_redis(self._current_task)
|
|
print_info(f"Task completed and saved to history. Local count: {len(self._history)}")
|
|
# 限制历史记录大小
|
|
if len(self._history) > self._history_max_size:
|
|
self._history = self._history[-self._history_max_size:]
|
|
self._current_task = None
|
|
# 清除 Redis 当前任务
|
|
self._save_current_task_to_redis(None)
|
|
self._save_status_to_redis()
|
|
|
|
# 确保任务完成标记和资源释放
|
|
self._queue.task_done()
|
|
# 强制垃圾回收
|
|
gc.collect()
|
|
|
|
except queue.Empty:
|
|
# 超时无任务，继续检查运行状态
|
|
continue
|
|
|
|
finally:
|
|
# 确保停止状态设置和资源清理
|
|
with self._thread_lock:
|
|
self._is_running = False
|
|
self._save_status_to_redis()
|
|
# 清理可能残留的资源
|
|
gc.collect()
|
|
|
|
def stop(self) -> None:
|
|
"""停止任务执行"""
|
|
with self._thread_lock:
|
|
self._is_running = False
|
|
self._save_status_to_redis()
|
|
|
|
def get_queue_info(self) -> dict:
|
|
"""
|
|
获取队列的当前状态信息
|
|
|
|
返回:
|
|
dict: 包含队列信息的字典，包括:
|
|
- is_running: 队列是否正在运行
|
|
- pending_tasks: 等待执行的任务数量
|
|
"""
|
|
with self._thread_lock:
|
|
return {
|
|
'is_running': self._is_running,
|
|
'pending_tasks': self._queue.qsize()
|
|
}
|
|
|
|
def get_detailed_status(self) -> dict:
|
|
"""
|
|
获取队列的详细状态信息（从 Redis 读取，支持多进程）
|
|
|
|
返回:
|
|
dict: 包含详细队列信息的字典
|
|
"""
|
|
# 从 Redis 获取数据
|
|
pending_list = self._get_pending_from_redis()
|
|
history_list = self._get_history_from_redis(20)
|
|
history_count = self._get_history_count_from_redis()
|
|
current_task = self._get_current_task_from_redis()
|
|
|
|
# 获取运行状态
|
|
redis_client = _get_redis()
|
|
is_running = self._is_running
|
|
if redis_client:
|
|
try:
|
|
status = redis_client.hgetall(REDIS_KEY_STATUS)
|
|
if status and 'is_running' in status:
|
|
is_running = status['is_running'] == 'true'
|
|
except Exception:
|
|
pass
|
|
|
|
return {
|
|
'tag': self.tag,
|
|
'is_running': is_running,
|
|
'pending_count': len(pending_list),
|
|
'pending_tasks': pending_list,
|
|
'current_task': current_task,
|
|
'history_count': history_count,
|
|
'recent_history': history_list
|
|
}
|
|
|
|
def clear_history(self) -> None:
|
|
"""清空任务历史记录"""
|
|
with self._thread_lock:
|
|
self._history.clear()
|
|
self._clear_history_from_redis()
|
|
self._save_status_to_redis()
|
|
print_success("任务历史记录已清空")
|
|
|
|
def clear_queue(self) -> None:
|
|
"""清空队列中的所有任务"""
|
|
with self._thread_lock:
|
|
while not self._queue.empty():
|
|
try:
|
|
self._queue.get_nowait()
|
|
self._queue.task_done()
|
|
except queue.Empty:
|
|
break
|
|
self._pending_items.clear()
|
|
self._save_pending_to_redis()
|
|
self._save_status_to_redis()
|
|
print_success("队列已清空")
|
|
|
|
def delete_queue(self) -> None:
|
|
"""删除队列(停止并清空所有任务)"""
|
|
with self._thread_lock:
|
|
self._is_running = False
|
|
while not self._queue.empty():
|
|
try:
|
|
self._queue.get_nowait()
|
|
self._queue.task_done()
|
|
except queue.Empty:
|
|
break
|
|
self._save_status_to_redis()
|
|
print_success("队列已删除")
|
|
|
|
# 创建全局单例（模块级变量，确保唯一实例）
|
|
TaskQueue = TaskQueueManager(tag="默认队列")
|
|
print_info(f"TaskQueue singleton created, instance id: {TaskQueue._instance_id}")
|
|
TaskQueue.run_task_background()
|
|
|
|
if __name__ == "__main__":
|
|
def task1():
|
|
print("执行任务1")
|
|
|
|
def task2(name):
|
|
print(f"执行任务2，参数: {name}")
|
|
|
|
manager = TaskQueueManager()
|
|
manager.add_task(task1)
|
|
manager.add_task(task2, "测试任务")
|
|
manager.run_tasks() # 按顺序执行任务1和任务2
|