205 lines
7.1 KiB
Python
205 lines
7.1 KiB
Python
|
# ============================================================
|
|||
|
|
# 轻量异步任务队列（moduleId: core-async-jobs, 可重生 ✅）
|
|||
|
|
# 矩阵 88 行剩余：敏感性 / 蒙特卡洛重算放后台（提交 / 轮询 / 取消 / 结果）
|
|||
|
|
# 线程模型：daemon worker 线程执行；fn 签名 fn(record)，可检查
|
|||
|
|
# record.cancel_requested 提前退出（或抛 JobCancelled）实现可取消长任务。
|
|||
|
|
# 与网关解耦：本模块只做队列与状态机，业务闭包由调用方注入。
|
|||
|
|
# ============================================================
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import itertools
|
|||
|
|
import threading
|
|||
|
|
import time
|
|||
|
|
import uuid
|
|||
|
|
from collections.abc import Callable
|
|||
|
|
from datetime import UTC, datetime
|
|||
|
|
from typing import Any, Literal
|
|||
|
|
|
|||
|
|
from pydantic import BaseModel, Field
|
|||
|
|
|
|||
|
|
from server.timeutil import fmt_dt
|
|||
|
|
|
|||
|
|
JobStatus = Literal["pending", "running", "done", "failed", "cancelled"]
|
|||
|
|
_TERMINAL = ("done", "failed", "cancelled")
|
|||
|
|
|
|||
|
|
|
|||
|
|
class JobCancelled(Exception):
|
|||
|
|
"""任务主动取消：worker 捕获后状态归为 cancelled。"""
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _now() -> str:
|
|||
|
|
# DTZ005：统一取本地时区的 aware 时间，墙钟与仓库其余模块一致
|
|||
|
|
return fmt_dt(datetime.now(UTC).astimezone())
|
|||
|
|
|
|||
|
|
|
|||
|
|
class JobRecord(BaseModel):
|
|||
|
|
"""任务记录：状态机 pending → running → done/failed/cancelled。"""
|
|||
|
|
|
|||
|
|
job_id: str
|
|||
|
|
kind: str
|
|||
|
|
status: JobStatus = "pending"
|
|||
|
|
progress: int = 0
|
|||
|
|
params: dict[str, Any] = Field(default_factory=dict)
|
|||
|
|
actor: str = "planner"
|
|||
|
|
result: Any = None
|
|||
|
|
error: str | None = None
|
|||
|
|
created_at: str = ""
|
|||
|
|
started_at: str | None = None
|
|||
|
|
finished_at: str | None = None
|
|||
|
|
cancel_requested: bool = False
|
|||
|
|
|
|||
|
|
|
|||
|
|
class JobQueue:
|
|||
|
|
"""轻量异步任务队列：提交 / 轮询 / 取消 / 结果 / 统计。"""
|
|||
|
|
|
|||
|
|
def __init__(self, *, max_finished: int = 200) -> None:
|
|||
|
|
self._lock = threading.RLock()
|
|||
|
|
self._jobs: dict[str, JobRecord] = {}
|
|||
|
|
self._seq = itertools.count(1)
|
|||
|
|
self.max_finished = max_finished
|
|||
|
|
|
|||
|
|
# ---------------- 提交 ----------------
|
|||
|
|
def submit(
|
|||
|
|
self,
|
|||
|
|
kind: str,
|
|||
|
|
fn: Callable[[JobRecord], Any],
|
|||
|
|
*,
|
|||
|
|
params: dict[str, Any] | None = None,
|
|||
|
|
actor: str = "planner",
|
|||
|
|
) -> str:
|
|||
|
|
"""提交任务：立即起 daemon 线程执行，返回 job_id。"""
|
|||
|
|
job_id = f"job-{next(self._seq):04d}-{uuid.uuid4().hex[:6]}"
|
|||
|
|
record = JobRecord(
|
|||
|
|
job_id=job_id, kind=kind, params=params or {}, actor=actor, created_at=_now(),
|
|||
|
|
)
|
|||
|
|
with self._lock:
|
|||
|
|
self._jobs[job_id] = record
|
|||
|
|
thread = threading.Thread(
|
|||
|
|
target=self._worker, args=(job_id, fn),
|
|||
|
|
name=f"aps-job-{job_id}", daemon=True,
|
|||
|
|
)
|
|||
|
|
thread.start()
|
|||
|
|
return job_id
|
|||
|
|
|
|||
|
|
def _worker(self, job_id: str, fn: Callable[[JobRecord], Any]) -> None:
|
|||
|
|
with self._lock:
|
|||
|
|
record = self._jobs[job_id]
|
|||
|
|
record.status = "running"
|
|||
|
|
record.started_at = _now()
|
|||
|
|
cancel = record.cancel_requested
|
|||
|
|
if cancel:
|
|||
|
|
self._finish(job_id, "cancelled")
|
|||
|
|
return
|
|||
|
|
try:
|
|||
|
|
result = fn(record)
|
|||
|
|
with self._lock:
|
|||
|
|
cancelled = self._jobs[job_id].cancel_requested
|
|||
|
|
if cancelled:
|
|||
|
|
self._finish(job_id, "cancelled")
|
|||
|
|
else:
|
|||
|
|
with self._lock:
|
|||
|
|
self._jobs[job_id].result = result
|
|||
|
|
self._jobs[job_id].progress = 100
|
|||
|
|
self._finish(job_id, "done")
|
|||
|
|
except JobCancelled:
|
|||
|
|
self._finish(job_id, "cancelled")
|
|||
|
|
except Exception as exc: # noqa: BLE001 - worker 兜底任意业务异常，转 failed 状态
|
|||
|
|
with self._lock:
|
|||
|
|
self._jobs[job_id].error = str(exc)
|
|||
|
|
self._finish(job_id, "failed")
|
|||
|
|
finally:
|
|||
|
|
self._prune()
|
|||
|
|
|
|||
|
|
def _finish(self, job_id: str, status: JobStatus) -> None:
|
|||
|
|
with self._lock:
|
|||
|
|
record = self._jobs[job_id]
|
|||
|
|
record.status = status
|
|||
|
|
record.finished_at = _now()
|
|||
|
|
if status == "cancelled":
|
|||
|
|
record.result = None
|
|||
|
|
record.error = "cancelled"
|
|||
|
|
|
|||
|
|
# ---------------- 轮询 / 结果 ----------------
|
|||
|
|
def poll(self, job_id: str) -> dict[str, Any] | None:
|
|||
|
|
"""查询任务快照（含 result/error）；不存在返回 None。"""
|
|||
|
|
with self._lock:
|
|||
|
|
record = self._jobs.get(job_id)
|
|||
|
|
return record.model_dump(mode="json") if record else None
|
|||
|
|
|
|||
|
|
def wait(self, job_id: str, timeout: float = 15.0) -> dict[str, Any]:
|
|||
|
|
"""阻塞轮询至终态（done/failed/cancelled）。"""
|
|||
|
|
deadline = time.time() + timeout
|
|||
|
|
while True:
|
|||
|
|
record = self.poll(job_id)
|
|||
|
|
if record is None:
|
|||
|
|
raise KeyError(job_id)
|
|||
|
|
if record["status"] in _TERMINAL:
|
|||
|
|
return record
|
|||
|
|
if time.time() >= deadline:
|
|||
|
|
raise TimeoutError(f"任务 {job_id} 等待超时（{timeout}s）")
|
|||
|
|
time.sleep(0.05)
|
|||
|
|
|
|||
|
|
# ---------------- 取消 ----------------
|
|||
|
|
def cancel(self, job_id: str) -> bool:
|
|||
|
|
"""取消任务：pending 立即取消；running 置取消标记（worker 收尾时落 cancelled）；
|
|||
|
|
已终态或不存在返回 False。"""
|
|||
|
|
with self._lock:
|
|||
|
|
record = self._jobs.get(job_id)
|
|||
|
|
if record is None or record.status in _TERMINAL:
|
|||
|
|
return False
|
|||
|
|
record.cancel_requested = True
|
|||
|
|
if record.status == "pending":
|
|||
|
|
record.status = "cancelled"
|
|||
|
|
record.finished_at = _now()
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
# ---------------- 列表 / 统计 ----------------
|
|||
|
|
def list(self, status: str | None = None) -> list[dict[str, Any]]:
|
|||
|
|
with self._lock:
|
|||
|
|
items = [r.model_dump(mode="json") for r in self._jobs.values()]
|
|||
|
|
if status:
|
|||
|
|
items = [i for i in items if i["status"] == status]
|
|||
|
|
items.sort(key=lambda i: i["created_at"] or "", reverse=True)
|
|||
|
|
return items
|
|||
|
|
|
|||
|
|
def stats(self) -> dict[str, int]:
|
|||
|
|
with self._lock:
|
|||
|
|
counts: dict[str, int] = {
|
|||
|
|
"total": len(self._jobs), "pending": 0, "running": 0,
|
|||
|
|
"done": 0, "failed": 0, "cancelled": 0,
|
|||
|
|
}
|
|||
|
|
for r in self._jobs.values():
|
|||
|
|
counts[r.status] = counts.get(r.status, 0) + 1
|
|||
|
|
return counts
|
|||
|
|
|
|||
|
|
def _prune(self) -> None:
|
|||
|
|
"""只保留最多 max_finished 条终态任务（按完成时间淘汰最旧）。"""
|
|||
|
|
with self._lock:
|
|||
|
|
finished = sorted(
|
|||
|
|
(r for r in self._jobs.values() if r.status in _TERMINAL),
|
|||
|
|
key=lambda r: r.finished_at or "",
|
|||
|
|
)
|
|||
|
|
overflow = len(finished) - self.max_finished
|
|||
|
|
for r in finished[:overflow]:
|
|||
|
|
self._jobs.pop(r.job_id, None)
|
|||
|
|
|
|||
|
|
def clear(self) -> None:
|
|||
|
|
with self._lock:
|
|||
|
|
self._jobs.clear()
|
|||
|
|
|
|||
|
|
|
|||
|
|
_queue: JobQueue | None = None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_job_queue() -> JobQueue:
|
|||
|
|
global _queue
|
|||
|
|
if _queue is None:
|
|||
|
|
_queue = JobQueue()
|
|||
|
|
return _queue
|
|||
|
|
|
|||
|
|
|
|||
|
|
def reset_job_queue() -> None:
|
|||
|
|
"""测试用：丢弃单例，下次重建。"""
|
|||
|
|
global _queue
|
|||
|
|
_queue = None
|