259 lines
10 KiB
Python
259 lines
10 KiB
Python
"""持久任务队列（交付报告 P0-3）。
|
||
|
||
用 SQLite 表实现可靠队列语义，替代纯内存线程调度：
|
||
- 租约 lease：worker 领取任务时写入 lease_until，超时未续约视为 worker 死亡，任务可被其他 worker 重新领取；
|
||
- 重试退避：失败后按 attempt 计算退避（复用 hermes_core.jittered_backoff），next_run_at 之前不可领取；
|
||
- 死信队列：超过最大重试次数进入 dead 状态，运营可查看和手动重放；
|
||
- 心跳：worker 定期刷新 lease_until；服务重启时孤儿任务自动回收。
|
||
|
||
与现有 _ACTIVE_RUNS 内存调度共存：队列表是事实源，内存线程只是执行器。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
from datetime import datetime, timedelta, timezone
|
||
|
||
from sqlmodel import Session, select, update
|
||
|
||
from .hermes_core import jittered_backoff
|
||
from .models import now
|
||
|
||
logger = logging.getLogger("yawen.queue")
|
||
|
||
LEASE_SECONDS = 300 # 租约时长：5 分钟无心跳即视为 worker 死亡
|
||
MAX_ATTEMPTS = 3 # 含首次执行在内最多尝试次数
|
||
QUEUE_TABLE = "taskqueue"
|
||
|
||
|
||
def _utcnow() -> datetime:
|
||
"""统一使用无时区本地时间（与 models.now() 一致，队列内部时间戳互相比较）。"""
|
||
return datetime.now()
|
||
|
||
|
||
def _parse(value: str) -> datetime | None:
|
||
try:
|
||
parsed = datetime.fromisoformat(value)
|
||
except (ValueError, TypeError):
|
||
return None
|
||
if parsed.tzinfo is not None:
|
||
parsed = parsed.replace(tzinfo=None)
|
||
return parsed
|
||
|
||
|
||
def enqueue(session: Session, queue_model, *, task_key: str, payload: dict,
|
||
update_existing: bool = False) -> int:
|
||
"""入队（幂等：同 task_key 的活跃任务不重复入队）。返回队列行 id。
|
||
|
||
update_existing 用于 parse/compose/redo 这类确定性 campaign 任务：如果任务
|
||
仍在 pending，用户刚修改了参数，则刷新 payload；running 中的任务不能被改。
|
||
"""
|
||
existing = session.exec(select(queue_model).where(
|
||
queue_model.task_key == task_key,
|
||
queue_model.status.in_(["pending", "leased", "running"]),
|
||
)).first()
|
||
if existing:
|
||
if update_existing and existing.status == "pending":
|
||
existing.payload = payload
|
||
existing.error = ""
|
||
existing.next_run_at = now()
|
||
existing.updated_at = now()
|
||
session.add(existing)
|
||
session.commit()
|
||
return existing.id
|
||
row = queue_model(task_key=task_key, payload=payload, status="pending",
|
||
attempt=0, max_attempts=MAX_ATTEMPTS, next_run_at=now())
|
||
session.add(row)
|
||
session.commit()
|
||
session.refresh(row)
|
||
return row.id
|
||
|
||
|
||
def lease_next(session: Session, queue_model, worker_id: str,
|
||
task_key_prefix: str = "") -> object | None:
|
||
"""领取一个到期任务，并以条件 UPDATE 原子认领。
|
||
|
||
旧实现先读取候选行、再提交修改；两个 worker 在读取与提交之间会
|
||
同时看到同一行，导致重复执行。这里仍保留轻量候选查询以维持优先级，
|
||
但真正认领通过 ``UPDATE ... WHERE`` 校验状态和租约，失败时重试下一个
|
||
候选。SQLite 的单条 UPDATE 在写入时持有数据库写锁，因此不会产生双重
|
||
认领。
|
||
"""
|
||
for _ in range(8):
|
||
current = now()
|
||
# 仅取少量候选，避免队列很大时把全部 payload 读进内存。时间戳使用
|
||
# 固定的 YYYY-MM-DD HH:MM:SS 格式，SQLite 文本比较与 _parse 语义一致。
|
||
query = select(queue_model).where(
|
||
queue_model.status.in_(["pending", "running"]),
|
||
)
|
||
# Worker 可按 task_key 前缀领取自己的任务类型。DSH 内容生产执行器
|
||
# 只处理 campaign-run:*，避免误领取训练导图等其他后台任务。
|
||
if task_key_prefix:
|
||
query = query.where(
|
||
queue_model.task_key.like(task_key_prefix.rstrip("%") + "%")
|
||
)
|
||
rows = session.exec(query.order_by(queue_model.id).limit(16)).all()
|
||
candidate = None
|
||
for row in rows:
|
||
if row.status == "pending":
|
||
next_at = _parse(row.next_run_at)
|
||
if next_at and next_at > _parse(current):
|
||
continue
|
||
else: # running：租约超时才能抢（worker 死亡的孤儿任务）
|
||
lease_until = _parse(row.lease_until or "")
|
||
if lease_until and lease_until > _parse(current):
|
||
continue
|
||
candidate = row
|
||
break
|
||
if candidate is None:
|
||
return None
|
||
|
||
deadline = _lease_deadline()
|
||
# 把候选的当前状态/租约再次放入 WHERE，防止另一个 worker 已经抢走。
|
||
# 对 pending 同时检查 next_run_at，避免失败重试被提前领取。
|
||
# Match the exact timestamp observed in the candidate row instead of
|
||
# comparing strings in SQL. Tests and older rows may include fractional
|
||
# seconds while ``now()`` uses second precision; Python's _parse above
|
||
# handles both formats correctly.
|
||
pending_due = (
|
||
(queue_model.status == "pending")
|
||
& (queue_model.next_run_at == candidate.next_run_at)
|
||
)
|
||
orphan_due = (
|
||
(queue_model.status == "running")
|
||
& (queue_model.lease_until == candidate.lease_until)
|
||
)
|
||
statement = update(queue_model).where(
|
||
queue_model.id == candidate.id,
|
||
pending_due | orphan_due,
|
||
).values(
|
||
status="running",
|
||
attempt=(queue_model.attempt + 1),
|
||
worker_id=worker_id,
|
||
lease_until=deadline,
|
||
updated_at=current,
|
||
)
|
||
result = session.exec(statement)
|
||
# SQLAlchemy's rowcount is available for SQLite UPDATE statements. A
|
||
# zero count means another worker won the race; rollback and retry.
|
||
if result.rowcount != 1:
|
||
session.rollback()
|
||
continue
|
||
session.commit()
|
||
claimed = session.get(queue_model, candidate.id)
|
||
if claimed is not None:
|
||
return claimed
|
||
return None
|
||
|
||
|
||
def _lease_deadline() -> str:
|
||
return (_utcnow() + timedelta(seconds=LEASE_SECONDS)).isoformat()
|
||
|
||
|
||
def heartbeat(session: Session, queue_model, task_id: int, worker_id: str) -> bool:
|
||
"""worker 心跳续租；任务已不在本 worker 名下时返回 False。"""
|
||
row = session.get(queue_model, task_id)
|
||
if not row or row.status != "running" or row.worker_id != worker_id:
|
||
return False
|
||
row.lease_until = _lease_deadline()
|
||
row.updated_at = now()
|
||
session.add(row)
|
||
session.commit()
|
||
return True
|
||
|
||
|
||
def complete(session: Session, queue_model, task_id: int, worker_id: str) -> bool:
|
||
row = session.get(queue_model, task_id)
|
||
if not row or row.status != "running" or row.worker_id != worker_id:
|
||
return False
|
||
row.status = "done"
|
||
row.error = ""
|
||
row.lease_until = ""
|
||
row.updated_at = now()
|
||
session.add(row)
|
||
session.commit()
|
||
return True
|
||
|
||
|
||
def fail(session: Session, queue_model, task_id: int, worker_id: str,
|
||
error: str) -> str:
|
||
"""任务失败：可重试则退避后重新排队；超过次数进死信。返回最终状态。"""
|
||
row = session.get(queue_model, task_id)
|
||
if not row or row.worker_id != worker_id:
|
||
return "unknown"
|
||
row.error = str(error)[:500]
|
||
if (row.attempt or 0) >= (row.max_attempts or MAX_ATTEMPTS):
|
||
row.status = "dead"
|
||
logger.error("任务 %s 进入死信队列：%s", row.task_key, row.error[:120])
|
||
else:
|
||
delay = min(24.0, jittered_backoff(row.attempt))
|
||
retry_at = (_utcnow() + timedelta(seconds=delay)).isoformat()
|
||
row.status = "pending"
|
||
row.next_run_at = retry_at
|
||
row.lease_until = ""
|
||
logger.warning("任务 %s 第 %s 次失败，%.1f 秒后重试：%s",
|
||
row.task_key, row.attempt, delay, row.error[:120])
|
||
row.updated_at = now()
|
||
session.add(row)
|
||
session.commit()
|
||
return row.status
|
||
|
||
|
||
def requeue_dead(session: Session, queue_model, task_id: int) -> bool:
|
||
"""手动重放死信任务。"""
|
||
row = session.get(queue_model, task_id)
|
||
if not row or row.status != "dead":
|
||
return False
|
||
row.status = "pending"
|
||
row.attempt = 0
|
||
row.next_run_at = now()
|
||
row.error = ""
|
||
row.updated_at = now()
|
||
session.add(row)
|
||
session.commit()
|
||
return True
|
||
|
||
|
||
def queue_stats(session: Session, queue_model) -> dict:
|
||
rows = session.exec(select(queue_model)).all()
|
||
by_status: dict[str, int] = {"pending": 0, "running": 0, "done": 0, "dead": 0}
|
||
for row in rows:
|
||
by_status[row.status] = by_status.get(row.status, 0) + 1
|
||
return {"total": len(rows), "by_status": by_status,
|
||
"done": by_status.get("done", 0),
|
||
"dead": by_status.get("dead", 0),
|
||
"pending": by_status.get("pending", 0),
|
||
"running": by_status.get("running", 0)}
|
||
|
||
|
||
__all__ = ["enqueue", "lease_next", "claim_by_key", "heartbeat", "complete", "fail",
|
||
"requeue_dead", "queue_stats", "LEASE_SECONDS", "MAX_ATTEMPTS"]
|
||
|
||
|
||
def claim_by_key(session: Session, queue_model, task_key: str, worker_id: str) -> object | None:
|
||
"""Claim a specific durable job for the inline API worker.
|
||
|
||
A running row is only reused by the same owner or stolen after its lease
|
||
expires; otherwise another live executor owns it.
|
||
"""
|
||
row = session.exec(select(queue_model).where(queue_model.task_key == task_key).order_by(queue_model.id.desc())).first()
|
||
if row is None:
|
||
return None
|
||
current = now()
|
||
expired = False
|
||
if row.status == "running":
|
||
deadline = _parse(row.lease_until or "")
|
||
expired = bool(deadline and deadline <= _parse(current))
|
||
if row.status == "pending" or (row.status == "running" and expired):
|
||
row.status = "running"
|
||
row.attempt = (row.attempt or 0) + 1
|
||
row.worker_id = worker_id
|
||
row.lease_until = _lease_deadline()
|
||
row.updated_at = current
|
||
session.add(row)
|
||
session.commit()
|
||
session.refresh(row)
|
||
return row
|
||
if row.status == "running" and row.worker_id == worker_id:
|
||
return row
|
||
return None
|