418 lines
17 KiB
Python
418 lines
17 KiB
Python
 
|
"""跨插件的「任务槽」：长任务的进度上报 + 同槽去重 + SSE 推送。
|
|||
|
|
|
|||
|
|
起因（2026-10-10）：video_analysis 的 temp 清理是个几十秒级的阻塞任务，直接放在
|
|||
|
|
HTTP 请求里会把页面交互一起卡住，用户还可能连点几次。这类需求（清缓存、批量下载、
|
|||
|
|
批量生成…）会越来越多，所以在 core 出一份通用机制，插件按需注册。
|
|||
|
|
|
|||
|
|
三个概念：
|
|||
|
|
|
|||
|
|
- **槽（slot）**：字符串标识，约定 `插件id:动作`（如 `video_analysis:temp_cleanup`）。
|
|||
|
|
**同一槽同时只跑一个任务** —— 重复 `submit` 不新建，直接把你交给正在跑的那个并返回
|
|||
|
|
`reused=True`，前端据此禁用按钮 / 提示"已在运行"（服务端是最终防线，不能只靠前端）。
|
|||
|
|
- **执行**：任务体是**同步可调用对象**，由本模块丢线程池跑（`asyncio.to_thread`），
|
|||
|
|
所以清理/下载那类阻塞 IO 不会卡事件循环；函数签名固定为 `func(report) -> 结果`。
|
|||
|
|
- **进度**：`report(done=…, total=…, phase=…, message=…, warning=…)`，在**任务体所在的
|
|||
|
|
线程**里调用，内部用 `loop.call_soon_threadsafe` 切回事件循环改状态。
|
|||
|
|
推送按"百分比桶"节流（同一百分点内不重复发），上万个文件也不会灌爆 SSE。
|
|||
|
|
|
|||
|
|
`report` 的**返回值**是协作式取消信号：`False` 表示有人请求取消，任务体应尽快收尾
|
|||
|
|
（清理这种天然可中断的任务在循环里判断即可，不要硬杀线程）。
|
|||
|
|
|
|||
|
|
推送方式：`subscribe()` 拿一个 `asyncio.Queue`，此后**任何**任务变化都会往队列里塞一份
|
|||
|
|
全量快照；hub 的 `/api/tasks/stream` 直接搭在它上面。任务少、变化不频繁，全量快照比
|
|||
|
|
增量协议省心得多（代价是 `result` 之类的字段要克制，别往里塞大对象）。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import asyncio
|
|||
|
|
import time
|
|||
|
|
import uuid
|
|||
|
|
from dataclasses import dataclass, field
|
|||
|
|
from typing import Any
|
|||
|
|
from collections.abc import Callable
|
|||
|
|
|
|||
|
|
from nonebot.log import logger
|
|||
|
|
|
|||
|
|
#: 任务体：接收 report 回调，返回任意结果（会放进 task.result）
|
|||
|
|
TaskFunc = Callable[[Callable[..., bool]], Any]
|
|||
|
|
|
|||
|
|
#: 终态任务保留时长（秒）与总量上限（超出先丢最旧的终态任务）
|
|||
|
|
TERMINAL_TTL = 30 * 60
|
|||
|
|
MAX_TASKS = 50
|
|||
|
|
|
|||
|
|
#: 状态取值：running → done / failed / canceled
|
|||
|
|
STATUS_RUNNING = "running"
|
|||
|
|
STATUS_DONE = "done"
|
|||
|
|
STATUS_FAILED = "failed"
|
|||
|
|
STATUS_CANCELED = "canceled"
|
|||
|
|
|
|||
|
|
|
|||
|
|
@dataclass
|
|||
|
|
class Task:
|
|||
|
|
"""一条任务记录；`to_dict()` 就是 SSE / REST 的负载。
|
|||
|
|
|
|||
|
|
字段分四组：**身份**（id / slot / title / plugin_id）、**来源**（trigger / actor /
|
|||
|
|
group_id，出问题时能查到是谁点的）、**进度**（phase / done / total / unit /
|
|||
|
|
message，to_dict 里归成 `progress` 对象）、**结果**（result / warnings / error）
|
|||
|
|
与**时间**（created / started / updated / finished）。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
# ── 身份 ──────────────────────────────────────────
|
|||
|
|
id: str
|
|||
|
|
#: 槽名：同槽互斥，约定 `插件id:动作`
|
|||
|
|
slot: str
|
|||
|
|
#: 展示名，如"清理 temp（保留 7 天）"
|
|||
|
|
title: str
|
|||
|
|
plugin_id: str = ""
|
|||
|
|
|
|||
|
|
# ── 来源（审计） ───────────────────────────────────
|
|||
|
|
#: 触发方式：web / command / schedule
|
|||
|
|
trigger: str = ""
|
|||
|
|
#: web 触发是登录用户名，command 触发是 QQ 号
|
|||
|
|
actor: str = ""
|
|||
|
|
#: command 触发时的群号
|
|||
|
|
group_id: str = ""
|
|||
|
|
|
|||
|
|
# ── 状态 ──────────────────────────────────────────
|
|||
|
|
status: str = STATUS_RUNNING
|
|||
|
|
#: 有人请求取消（任务体通过 `report()` 的返回值感知）
|
|||
|
|
cancel_requested: bool = False
|
|||
|
|
|
|||
|
|
# ── 进度 ──────────────────────────────────────────
|
|||
|
|
#: 阶段文案（"扫描中" / "清理中" / "回收空目录"）
|
|||
|
|
phase: str = ""
|
|||
|
|
done: int = 0
|
|||
|
|
#: 0 = 总数未知（扫描阶段没有进度条）
|
|||
|
|
total: int = 0
|
|||
|
|
#: 进度单位（"个文件"），前端拼 "3/7 个文件"
|
|||
|
|
unit: str = ""
|
|||
|
|
#: 当前细节（"已删除 2 个"）
|
|||
|
|
message: str = ""
|
|||
|
|
|
|||
|
|
# ── 结果 ──────────────────────────────────────────
|
|||
|
|
result: Any = None
|
|||
|
|
#: 部分失败（跳过的被占用文件等）——不算失败，但要让人看见
|
|||
|
|
warnings: list[str] = field(default_factory=list)
|
|||
|
|
error: str = ""
|
|||
|
|
|
|||
|
|
# ── 时间 ──────────────────────────────────────────
|
|||
|
|
created_at: float = field(default_factory=time.time)
|
|||
|
|
started_at: float = 0.0
|
|||
|
|
updated_at: float = field(default_factory=time.time)
|
|||
|
|
finished_at: float = 0.0
|
|||
|
|
#: 结束信号（不出 API）：命令层要"等任务跑完再回复"时用 `TASKS.wait(task)`
|
|||
|
|
_done: asyncio.Event = field(
|
|||
|
|
default_factory=asyncio.Event, repr=False, compare=False
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def active(self) -> bool:
|
|||
|
|
return self.status == STATUS_RUNNING
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def percent(self) -> int | None:
|
|||
|
|
"""0-100；总数未知时 None（前端显示不确定进度）。"""
|
|||
|
|
if not self.active:
|
|||
|
|
return 100
|
|||
|
|
if self.total <= 0:
|
|||
|
|
return None
|
|||
|
|
return max(0, min(100, int(self.done * 100 / self.total)))
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def duration(self) -> float:
|
|||
|
|
"""耗时（秒）：跑着就是"已运行多久"，结束就是总用时，没开始是 0。"""
|
|||
|
|
if not self.started_at:
|
|||
|
|
return 0.0
|
|||
|
|
end = self.finished_at or (time.time() if self.active else 0.0)
|
|||
|
|
return max(0.0, end - self.started_at)
|
|||
|
|
|
|||
|
|
def to_dict(self) -> dict[str, Any]:
|
|||
|
|
return {
|
|||
|
|
"id": self.id,
|
|||
|
|
"slot": self.slot,
|
|||
|
|
"title": self.title,
|
|||
|
|
"plugin_id": self.plugin_id,
|
|||
|
|
"trigger": self.trigger,
|
|||
|
|
"actor": self.actor,
|
|||
|
|
"group_id": self.group_id,
|
|||
|
|
"status": self.status,
|
|||
|
|
"active": self.active,
|
|||
|
|
"cancel_requested": self.cancel_requested,
|
|||
|
|
"progress": {
|
|||
|
|
"percent": self.percent,
|
|||
|
|
"done": self.done,
|
|||
|
|
"total": self.total,
|
|||
|
|
"unit": self.unit,
|
|||
|
|
"phase": self.phase,
|
|||
|
|
"message": self.message,
|
|||
|
|
},
|
|||
|
|
"result": self.result,
|
|||
|
|
"warnings": list(self.warnings),
|
|||
|
|
"error": self.error,
|
|||
|
|
"created_at": self.created_at,
|
|||
|
|
"started_at": self.started_at,
|
|||
|
|
"updated_at": self.updated_at,
|
|||
|
|
"finished_at": self.finished_at,
|
|||
|
|
"duration": self.duration,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
class TaskSlots:
|
|||
|
|
"""任务注册表（进程内单例 `TASKS`）。
|
|||
|
|
|
|||
|
|
只在事件循环里用：`submit` 需要 running loop 来建后台任务，`subscribe` 的队列
|
|||
|
|
也绑定在事件循环上。任务体本身跑在线程池里，但改状态一律切回事件循环。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self, *, max_tasks: int = MAX_TASKS, ttl: float = TERMINAL_TTL) -> None:
|
|||
|
|
self._tasks: dict[str, Task] = {}
|
|||
|
|
self._subscribers: set[asyncio.Queue[list[dict[str, Any]]]] = set()
|
|||
|
|
self._loop: asyncio.AbstractEventLoop | None = None
|
|||
|
|
self._bg: set[asyncio.Task[None]] = set()
|
|||
|
|
self.max_tasks = max_tasks
|
|||
|
|
self.ttl = ttl
|
|||
|
|
#: 节流用：task.id -> 上次推送的百分比桶（-1 表示"未定进度"）
|
|||
|
|
self._last_bucket: dict[str, int] = {}
|
|||
|
|
|
|||
|
|
# ── 查询 ──────────────────────────────────────────────
|
|||
|
|
|
|||
|
|
def list(self) -> list[Task]:
|
|||
|
|
"""全部任务，新的在前（顺带回收超龄/超量的终态任务）。"""
|
|||
|
|
self._prune()
|
|||
|
|
return sorted(self._tasks.values(), key=lambda t: t.created_at, reverse=True)
|
|||
|
|
|
|||
|
|
def get(self, task_id: str) -> Task | None:
|
|||
|
|
return self._tasks.get(task_id)
|
|||
|
|
|
|||
|
|
def active(self, slot: str) -> Task | None:
|
|||
|
|
"""该槽正在跑的任务（没有则 None）—— 去重判定就是它。"""
|
|||
|
|
for task in self._tasks.values():
|
|||
|
|
if task.slot == slot and task.active:
|
|||
|
|
return task
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
# ── 提交 / 取消 ───────────────────────────────────────
|
|||
|
|
|
|||
|
|
def submit(
|
|||
|
|
self,
|
|||
|
|
*,
|
|||
|
|
slot: str,
|
|||
|
|
title: str,
|
|||
|
|
func: TaskFunc,
|
|||
|
|
plugin_id: str = "",
|
|||
|
|
trigger: str = "",
|
|||
|
|
actor: str = "",
|
|||
|
|
group_id: str = "",
|
|||
|
|
) -> tuple[Task, bool]:
|
|||
|
|
"""提交一个任务；同槽已有在跑的 → 不新建，返回 (已有任务, True)。
|
|||
|
|
|
|||
|
|
立即返回（不等任务跑完），任务体在线程池里执行。
|
|||
|
|
"""
|
|||
|
|
running = self.active(slot)
|
|||
|
|
if running is not None:
|
|||
|
|
logger.info(f"任务槽 {slot} 已有任务在跑（{running.title}），本次提交复用")
|
|||
|
|
return running, True
|
|||
|
|
|
|||
|
|
self._loop = asyncio.get_running_loop()
|
|||
|
|
task = Task(
|
|||
|
|
id=uuid.uuid4().hex[:12],
|
|||
|
|
slot=slot,
|
|||
|
|
title=title,
|
|||
|
|
plugin_id=plugin_id,
|
|||
|
|
trigger=trigger,
|
|||
|
|
actor=str(actor or ""),
|
|||
|
|
group_id=str(group_id or ""),
|
|||
|
|
)
|
|||
|
|
self._tasks[task.id] = task
|
|||
|
|
self._prune()
|
|||
|
|
# 审计信息进日志：任务表本身是内存态（重启即清），但"谁在什么时候点了什么"
|
|||
|
|
# 要能在 /hub 日志页翻到
|
|||
|
|
who = trigger or "unknown"
|
|||
|
|
if actor:
|
|||
|
|
who += f" · {actor}"
|
|||
|
|
if group_id:
|
|||
|
|
who += f"@群{group_id}"
|
|||
|
|
logger.info(f"任务槽 {slot} 启动任务：{title}（{who}）")
|
|||
|
|
|
|||
|
|
bg = asyncio.create_task(self._run(task, func))
|
|||
|
|
self._bg.add(bg)
|
|||
|
|
bg.add_done_callback(self._bg.discard)
|
|||
|
|
|
|||
|
|
self._publish()
|
|||
|
|
return task, False
|
|||
|
|
|
|||
|
|
def cancel(self, task_id: str) -> bool:
|
|||
|
|
"""请求取消：只置标志，**任务体自己决定何时收尾**（协作式，不硬杀线程）。
|
|||
|
|
|
|||
|
|
任务体在 `report(...)` 的返回值里看到 False 就该尽快结束并返回已有结果，
|
|||
|
|
任务随即落到 `canceled` 状态（结果与 warnings 保留，方便看"取消前清了多少"）。
|
|||
|
|
"""
|
|||
|
|
task = self._tasks.get(task_id)
|
|||
|
|
if task is None or not task.active:
|
|||
|
|
return False
|
|||
|
|
task.cancel_requested = True
|
|||
|
|
task.updated_at = time.time()
|
|||
|
|
logger.info(f"任务槽 {task.slot} 收到取消请求：{task.title}")
|
|||
|
|
self._publish()
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
async def wait(self, task: Task) -> Task:
|
|||
|
|
"""等任务跑完并返回它（已结束则立即返回）。
|
|||
|
|
|
|||
|
|
命令层用："跑完再回复结果"；Web 层不用（提交完就返回，进度走 SSE）。
|
|||
|
|
"""
|
|||
|
|
await task._done.wait()
|
|||
|
|
return task
|
|||
|
|
|
|||
|
|
def dismiss(self, task_id: str) -> bool:
|
|||
|
|
"""从列表里删掉一条**已结束**的任务，返回是否删掉了。
|
|||
|
|
|
|||
|
|
在跑的任务删不得（删了就再也没人回传它的状态了），只能先取消、等它收尾；
|
|||
|
|
所以这里对 active 的任务直接返回 False，让调用方去提示。
|
|||
|
|
"""
|
|||
|
|
task = self._tasks.get(task_id)
|
|||
|
|
if task is None or task.active:
|
|||
|
|
return False
|
|||
|
|
self._tasks.pop(task_id, None)
|
|||
|
|
self._last_bucket.pop(task_id, None)
|
|||
|
|
self._publish()
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
def clear_finished(self) -> int:
|
|||
|
|
"""清掉所有终态任务（在跑的不动），返回清掉的条数。"""
|
|||
|
|
finished = [tid for tid, t in self._tasks.items() if not t.active]
|
|||
|
|
for tid in finished:
|
|||
|
|
self._tasks.pop(tid, None)
|
|||
|
|
self._last_bucket.pop(tid, None)
|
|||
|
|
if finished:
|
|||
|
|
self._publish()
|
|||
|
|
return len(finished)
|
|||
|
|
|
|||
|
|
# ── 订阅（SSE） ───────────────────────────────────────
|
|||
|
|
|
|||
|
|
def subscribe(self) -> asyncio.Queue[list[dict[str, Any]]]:
|
|||
|
|
queue: asyncio.Queue[list[dict[str, Any]]] = asyncio.Queue()
|
|||
|
|
self._subscribers.add(queue)
|
|||
|
|
return queue
|
|||
|
|
|
|||
|
|
def unsubscribe(self, queue: asyncio.Queue[list[dict[str, Any]]]) -> None:
|
|||
|
|
self._subscribers.discard(queue)
|
|||
|
|
|
|||
|
|
def snapshot(self) -> list[dict[str, Any]]:
|
|||
|
|
return [t.to_dict() for t in self.list()]
|
|||
|
|
|
|||
|
|
# ── 内部：执行与推送 ──────────────────────────────────
|
|||
|
|
|
|||
|
|
async def _run(self, task: Task, func: TaskFunc) -> None:
|
|||
|
|
"""线程池里跑任务体；进度回调切回事件循环；异常落进 task.error。"""
|
|||
|
|
task.started_at = time.time()
|
|||
|
|
task.updated_at = task.started_at
|
|||
|
|
|
|||
|
|
def report(
|
|||
|
|
done: int | None = None,
|
|||
|
|
total: int | None = None,
|
|||
|
|
phase: str | None = None,
|
|||
|
|
unit: str | None = None,
|
|||
|
|
message: str | None = None,
|
|||
|
|
warning: str | None = None,
|
|||
|
|
) -> bool:
|
|||
|
|
"""任务体所在线程调用；返回 False 表示"请尽快收尾"（协作式取消）。"""
|
|||
|
|
loop = self._loop
|
|||
|
|
if loop is not None:
|
|||
|
|
loop.call_soon_threadsafe(
|
|||
|
|
self._apply_progress, task, done, total, phase, unit, message, warning
|
|||
|
|
)
|
|||
|
|
# cancel_requested 由事件循环线程写、工作线程读：CPython 下读 bool 是安全的
|
|||
|
|
return not task.cancel_requested
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
result = await asyncio.to_thread(func, report)
|
|||
|
|
except Exception as e: # noqa: BLE001 - 任务边界：失败即任务状态，不外抛
|
|||
|
|
logger.exception(f"任务失败（{task.slot}）：{task.title}")
|
|||
|
|
task.status = STATUS_FAILED
|
|||
|
|
task.error = str(e) or type(e).__name__
|
|||
|
|
else:
|
|||
|
|
task.status = STATUS_CANCELED if task.cancel_requested else STATUS_DONE
|
|||
|
|
task.result = result
|
|||
|
|
warn = f"，{len(task.warnings)} 条警告" if task.warnings else ""
|
|||
|
|
logger.info(
|
|||
|
|
f"任务{'取消' if task.status == STATUS_CANCELED else '完成'}"
|
|||
|
|
f"（{task.slot}）：{task.title}，用时 {task.duration:.1f}s{warn}"
|
|||
|
|
)
|
|||
|
|
task.finished_at = time.time()
|
|||
|
|
task.updated_at = task.finished_at
|
|||
|
|
task._done.set()
|
|||
|
|
self._publish()
|
|||
|
|
|
|||
|
|
def _apply_progress(
|
|||
|
|
self,
|
|||
|
|
task: Task,
|
|||
|
|
done: int | None,
|
|||
|
|
total: int | None,
|
|||
|
|
phase: str | None,
|
|||
|
|
unit: str | None,
|
|||
|
|
message: str | None,
|
|||
|
|
warning: str | None = None,
|
|||
|
|
) -> None:
|
|||
|
|
"""在事件循环里更新进度；按百分比桶节流推送。"""
|
|||
|
|
if not task.active: # 任务已结束（被 clear / 已失败）→ 丢弃迟到的进度
|
|||
|
|
return
|
|||
|
|
if done is not None:
|
|||
|
|
task.done = max(0, int(done))
|
|||
|
|
if total is not None:
|
|||
|
|
task.total = max(0, int(total))
|
|||
|
|
if phase is not None:
|
|||
|
|
task.phase = phase
|
|||
|
|
if unit is not None:
|
|||
|
|
task.unit = unit
|
|||
|
|
if message is not None:
|
|||
|
|
task.message = message
|
|||
|
|
if warning:
|
|||
|
|
task.warnings.append(str(warning))
|
|||
|
|
task.updated_at = time.time()
|
|||
|
|
|
|||
|
|
bucket = task.percent if task.percent is not None else -1
|
|||
|
|
noteworthy = phase is not None or message is not None or bool(warning)
|
|||
|
|
if bucket != self._last_bucket.get(task.id) or noteworthy:
|
|||
|
|
self._last_bucket[task.id] = bucket
|
|||
|
|
self._publish()
|
|||
|
|
|
|||
|
|
def _publish(self) -> None:
|
|||
|
|
"""给所有订阅者塞一份全量快照（队列满则丢掉最旧的，宁可丢进度也不阻塞）。"""
|
|||
|
|
if not self._subscribers:
|
|||
|
|
return
|
|||
|
|
snapshot = self.snapshot()
|
|||
|
|
for queue in list(self._subscribers):
|
|||
|
|
if queue.full():
|
|||
|
|
try:
|
|||
|
|
queue.get_nowait()
|
|||
|
|
except asyncio.QueueEmpty: # pragma: no cover - full 与 empty 之间被取走
|
|||
|
|
pass
|
|||
|
|
try:
|
|||
|
|
queue.put_nowait(snapshot)
|
|||
|
|
except asyncio.QueueFull: # pragma: no cover - 上面已腾位
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
def _prune(self) -> None:
|
|||
|
|
"""先清超龄终态任务，再按上限砍最旧的终态任务（在跑的不动）。"""
|
|||
|
|
now = time.time()
|
|||
|
|
for tid, task in list(self._tasks.items()):
|
|||
|
|
if not task.active and now - task.updated_at > self.ttl:
|
|||
|
|
self._tasks.pop(tid, None)
|
|||
|
|
self._last_bucket.pop(tid, None)
|
|||
|
|
|
|||
|
|
overflow = len(self._tasks) - self.max_tasks
|
|||
|
|
if overflow <= 0:
|
|||
|
|
return
|
|||
|
|
finished = sorted(
|
|||
|
|
(t for t in self._tasks.values() if not t.active),
|
|||
|
|
key=lambda t: t.created_at,
|
|||
|
|
)
|
|||
|
|
for task in finished[:overflow]:
|
|||
|
|
self._tasks.pop(task.id, None)
|
|||
|
|
self._last_bucket.pop(task.id, None)
|
|||
|
|
|
|||
|
|
|
|||
|
|
#: 全局单例（hub 路由与各插件共用）
|
|||
|
|
TASKS = TaskSlots()
|