257 lines
10 KiB
Python
257 lines
10 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。"""
|
||
|
||
|
||
class JobQueueFullError(RuntimeError):
|
||
"""在途任务达到上限，拒绝继续创建线程。"""
|
||
|
||
|
||
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
|
||
tenant_uuid: str = "platform"
|
||
project_id: str = "default"
|
||
|
||
def raise_if_cancelled(self) -> None:
|
||
if self.cancel_requested:
|
||
raise JobCancelled()
|
||
|
||
|
||
class JobQueue:
|
||
"""轻量异步任务队列：提交 / 轮询 / 取消 / 结果 / 统计。"""
|
||
|
||
def __init__(self, *, max_finished: int = 200, max_concurrent: int = 2,
|
||
max_inflight: int = 8) -> None:
|
||
self._lock = threading.RLock()
|
||
self._jobs: dict[str, JobRecord] = {}
|
||
self._seq = itertools.count(1)
|
||
self.max_finished = max_finished
|
||
self.max_concurrent = max(1, int(max_concurrent))
|
||
self.max_inflight = max(self.max_concurrent, int(max_inflight))
|
||
self._slots = threading.BoundedSemaphore(self.max_concurrent)
|
||
self._inflight_threads = 0
|
||
|
||
# ---------------- 提交 ----------------
|
||
def submit(
|
||
self,
|
||
kind: str,
|
||
fn: Callable[[JobRecord], Any],
|
||
*,
|
||
params: dict[str, Any] | None = None,
|
||
actor: str = "planner",
|
||
tenant_uuid: str = "platform",
|
||
project_id: str = "default",
|
||
) -> str:
|
||
"""提交任务：立即起 daemon 线程执行，返回 job_id。"""
|
||
job_id = f"job-{next(self._seq):04d}-{uuid.uuid4().hex[:6]}"
|
||
with self._lock:
|
||
if self._inflight_threads >= self.max_inflight:
|
||
raise JobQueueFullError(f"异步任务在途已达上限 {self.max_inflight}")
|
||
record = JobRecord(
|
||
job_id=job_id, kind=kind, params=params or {}, actor=actor,
|
||
tenant_uuid=tenant_uuid, project_id=project_id, created_at=_now(),
|
||
)
|
||
self._jobs[job_id] = record
|
||
self._inflight_threads += 1
|
||
thread = threading.Thread(
|
||
target=self._worker, args=(job_id, fn),
|
||
name=f"aps-job-{job_id}", daemon=True,
|
||
)
|
||
try:
|
||
thread.start()
|
||
except BaseException:
|
||
with self._lock:
|
||
self._jobs.pop(job_id, None)
|
||
self._inflight_threads = max(0, self._inflight_threads - 1)
|
||
raise
|
||
return job_id
|
||
|
||
def _worker(self, job_id: str, fn: Callable[[JobRecord], Any]) -> None:
|
||
try:
|
||
with self._slots:
|
||
with self._lock:
|
||
record = self._jobs[job_id]
|
||
if record.status == "cancelled" or record.cancel_requested:
|
||
self._finish(job_id, "cancelled")
|
||
return
|
||
record.status = "running"; record.started_at = _now()
|
||
try:
|
||
result = fn(record)
|
||
with self._lock:
|
||
current = self._jobs[job_id]
|
||
if current.cancel_requested:
|
||
self._finish(job_id, "cancelled")
|
||
else:
|
||
current.result = result; current.progress = 100
|
||
self._finish(job_id, "done")
|
||
except JobCancelled:
|
||
self._finish(job_id, "cancelled")
|
||
except Exception as exc: # noqa: BLE001
|
||
with self._lock:
|
||
current = self._jobs[job_id]
|
||
if current.cancel_requested:
|
||
self._finish(job_id, "cancelled")
|
||
else:
|
||
current.error = str(exc)
|
||
self._finish(job_id, "failed")
|
||
finally:
|
||
with self._lock:
|
||
self._inflight_threads = max(0, self._inflight_threads - 1)
|
||
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, *, tenant_uuid: str | None = None,
|
||
project_id: str | None = None,
|
||
actor: str | None = None) -> dict[str, Any] | None:
|
||
"""查询任务快照（含 result/error）；不存在返回 None。"""
|
||
with self._lock:
|
||
record = self._jobs.get(job_id)
|
||
if record and tenant_uuid is not None and record.tenant_uuid != tenant_uuid: return None
|
||
if record and project_id is not None and record.project_id != project_id: return None
|
||
if record and actor is not None and record.actor != actor: return None
|
||
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, *, tenant_uuid: str | None = None,
|
||
project_id: str | None = None,
|
||
actor: str | None = None) -> 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
|
||
if tenant_uuid is not None and record.tenant_uuid != tenant_uuid: return False
|
||
if project_id is not None and record.project_id != project_id: return False
|
||
if actor is not None and record.actor != actor: 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, *, tenant_uuid: str | None = None,
|
||
project_id: str | None = None,
|
||
actor: 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]
|
||
if tenant_uuid is not None: items = [i for i in items if i["tenant_uuid"] == tenant_uuid]
|
||
if project_id is not None: items = [i for i in items if i["project_id"] == project_id]
|
||
if actor is not None: items = [i for i in items if i["actor"] == actor]
|
||
items.sort(key=lambda i: i["created_at"] or "", reverse=True)
|
||
return items
|
||
|
||
def stats(self, *, tenant_uuid: str | None = None,
|
||
project_id: str | None = None,
|
||
actor: str | None = None) -> 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():
|
||
if tenant_uuid is not None and r.tenant_uuid != tenant_uuid: continue
|
||
if project_id is not None and r.project_id != project_id: continue
|
||
if actor is not None and r.actor != actor: continue
|
||
counts[r.status] = counts.get(r.status, 0) + 1
|
||
counts["total"] = sum(counts[s] for s in ("pending", "running", "done", "failed", "cancelled"))
|
||
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
|