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
|