mirror of
https://github.com/MemeMeow-Studio/MemeMeow.git
synced 2026-10-29 06:06:47 +08:00
76 lines
3.6 KiB
Python
76 lines
3.6 KiB
Python
# 持久 Task 观测：从已有 Session 的状态变化生成提交后的终态指标。
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
|
|
from opentelemetry import metrics
|
|
from sqlalchemy import event, inspect
|
|
from sqlalchemy.orm import Session, sessionmaker
|
|
from runtime_telemetry import telemetry_snapshot
|
|
|
|
from backend.persistence.models import Task
|
|
|
|
_PENDING = "task_terminal_telemetry"
|
|
_ENABLED = "task_terminal_telemetry_enabled"
|
|
_TERMINALS = {"succeeded", "failed", "cancelled"}
|
|
|
|
|
|
def terminal_snapshot(task_type: str, status: str, error: object) -> dict[str, str] | None:
|
|
"""将已有业务字段转成有界指标标签，非终态返回空值。"""
|
|
if status not in _TERMINALS:
|
|
return None
|
|
reason = error.get("error") or error.get("code") if isinstance(error, dict) else None
|
|
if not isinstance(reason, str) or re.fullmatch(r"[a-z][a-z0-9_]{0,127}", reason) is None:
|
|
reason = "none" if status == "succeeded" else "unspecified"
|
|
result = "cancelled" if status == "failed" and reason == "task_cancelled" else status
|
|
return {"task_type": task_type, "result": result, "reason": reason}
|
|
|
|
|
|
def stage_terminal(session: Session, identifier: object, snapshot: dict[str, str] | None) -> None:
|
|
"""在当前业务事务内暂存内存快照，SQL 更新路径复用同一个入口。"""
|
|
if not session.info.get(_ENABLED):
|
|
return
|
|
transaction = session.get_nested_transaction() or session.get_transaction()
|
|
if transaction is None:
|
|
raise RuntimeError("task_telemetry_transaction_missing")
|
|
session.info.setdefault(_PENDING, {}).setdefault(transaction, {})[identifier] = snapshot
|
|
|
|
|
|
def _before_flush(session: Session, flush_context: object, instances: object) -> None:
|
|
"""在既有事务保存新 Task 关联及终态快照，不执行查询或创建 Session。"""
|
|
for record in session.new | session.dirty:
|
|
if isinstance(record, Task) and (record in session.new or "_telemetry" not in (record.payload or {})):
|
|
record.payload = {**(record.payload or {}), "_telemetry": telemetry_snapshot()}
|
|
if isinstance(record, Task) and (record in session.new or inspect(record).attrs.status.history.has_changes()):
|
|
stage_terminal(session, record.id, terminal_snapshot(record.task_type, record.status, record.error))
|
|
|
|
|
|
def _after_commit(session: Session) -> None:
|
|
"""保存点合入父事务，最外层提交成功后增加真实终态计数。"""
|
|
transaction = session.get_nested_transaction() or session.get_transaction()
|
|
pending = session.info.get(_PENDING, {})
|
|
snapshots = pending.pop(transaction, {})
|
|
if transaction is not None and transaction.parent is not None:
|
|
pending.setdefault(transaction.parent, {}).update(snapshots)
|
|
return
|
|
counter = metrics.get_meter("mememeow.tasks").create_counter("task_terminals", unit="{task}")
|
|
for snapshot in snapshots.values():
|
|
if snapshot is not None:
|
|
counter.add(1, snapshot)
|
|
|
|
|
|
def _after_transaction_end(session: Session, transaction: object) -> None:
|
|
"""事务结束后移除未提交快照，回滚不会产生成功指标。"""
|
|
pending = session.info.get(_PENDING, {})
|
|
pending.pop(transaction, None)
|
|
if not pending:
|
|
session.info.pop(_PENDING, None)
|
|
|
|
|
|
def instrument_task_sessions(factory: sessionmaker[Session]) -> None:
|
|
"""仅为所属资源工厂注册标准 Session 事件，覆盖取消、恢复及 Worker 终态。"""
|
|
factory.configure(info={**factory.kw.get("info", {}), _ENABLED: True})
|
|
event.listen(factory, "before_flush", _before_flush)
|
|
event.listen(factory, "after_commit", _after_commit)
|
|
event.listen(factory, "after_transaction_end", _after_transaction_end)
|