rakuten-api/app/gateway/task_queue.py

"""任务队列业务逻辑：状态机、租约、并发度、幂等

所有写操作通过 `self._lock` 串行化。该锁同时保护「检查 + 写入」的原子序列
（如 lease 时「检查是否已有 active 任务 → 取 queued → 改为 leased」），避免两个
lease 请求并发取到同一个任务。

长轮询通过 `self._cond`（基于同一把锁的 Condition）唤醒：submit / report（terminal）
/ reclaim 失败等任何可能让「队列前进」的事件都 notify_all。lease 在无可领任务时
wait_for，超时返回 None。

绝不自动重投（docs/order-gateway.md §5）：sweep 把过期任务置 stale 后**不**回到
queued，普通 lease 取不到它；恢复只能走 reclaim。
"""
from __future__ import annotations

import asyncio
import logging
import secrets
from datetime import datetime, timedelta, timezone

from app.shared.task_state import (
    ACTIVE_STATUSES,
    LEASABLE_STATUSES,
    RECLAIMABLE_STATUSES,
    TERMINAL_STATUSES,
    TaskStatus,
)
from app.gateway.callback import CallbackNotifier, stale_payload, terminal_payload
from app.gateway.db import GatewayDB, ReportRow, TaskRow
from app.gateway.models import (
    GatewayHealthData,
    LeaseData,
    QueuedAlertEntry,
    ReportEntry,
    RenewData,
    ReclaimData,
    ReportData,
    SubmitOrderData,
    TaskDetail,
    TaskListData,
    WorkerHealthEntry,
)
from app.shared.errors import (
    InvalidTaskStateError,
    LeaseInvalidError,
    TaskNotFoundError,
)

logger = logging.getLogger(__name__)


def utcnow() -> datetime:
    return datetime.now(timezone.utc)


def to_iso(dt: datetime) -> str:
    """统一时间戳格式：ISO8601 UTC，秒精度，Z 结尾"""
    return dt.astimezone(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")


def parse_iso(s: str) -> datetime:
    """解析 to_iso 产出的字符串（也兼容带 +00:00 的变体）"""
    if s.endswith("Z"):
        s = s[:-1] + "+00:00"
    return datetime.fromisoformat(s)


def _generate_task_id() -> str:
    """服务端生成的 task_id：日期 + 8 字节随机 hex，避免与上游自带的冲突"""
    return f"po-{to_iso(utcnow())[:10].replace('-', '')}-{secrets.token_hex(4)}"


def _task_to_detail(task: TaskRow, latest_state: str | None, reports: list[ReportEntry]) -> TaskDetail:
    return TaskDetail(
        task_id=task.task_id,
        site=task.site,
        intent=task.intent,
        callback_url=task.callback_url,
        status=TaskStatus(task.status),
        lease_owner=task.lease_owner,
        lease_expires_at=task.lease_expires_at,
        lease_count=task.lease_count,
        created_at=task.created_at,
        updated_at=task.updated_at,
        latest_state=latest_state,
        reports=reports,
    )


class TaskQueue:
    """任务队列业务逻辑

    持有一个 GatewayDB 实例。所有写操作与「检查 + 写入」原子序列都在 `self._lock`
    保护下；长轮询等待通过 `self._cond` 唤醒。
    """

    def __init__(
        self,
        db: GatewayDB,
        *,
        lease_ttl_seconds: int,
        worker_offline_alert_seconds: int,
        notifier: CallbackNotifier | None = None,
    ):
        self._db = db
        self._lease_ttl = lease_ttl_seconds
        self._worker_offline_alert_seconds = worker_offline_alert_seconds
        # 终结类事件通知（§4.8）。None 表示完全关闭回调；测试可替换为记录型假实现
        self.notifier = notifier
        self._lock = asyncio.Lock()
        self._cond = asyncio.Condition(self._lock)

    # ---- 提交 ----

    async def submit(
        self,
        *,
        task_id: str | None,
        site: str,
        intent: dict,
        callback_url: str | None = None,
    ) -> SubmitOrderData:
        """上游提交下单意图。task_id 缺省时服务端生成；重复提交幂等

        callback_url 只在首次创建时写入；幂等重发（created=false）返回既有任务，
        不更新任何字段。
        """
        import json

        tid = task_id or _generate_task_id()
        now = to_iso(utcnow())
        row = TaskRow(
            task_id=tid,
            site=site,
            intent_json=json.dumps(intent, ensure_ascii=False),
            callback_url=callback_url,
            status=TaskStatus.QUEUED.value,
            lease_owner=None,
            lease_expires_at=None,
            lease_count=0,
            created_at=now,
            updated_at=now,
        )
        async with self._lock:
            created = await self._db.insert_task(row)
            if created:
                self._cond.notify_all()
            else:
                existing = await self._db.get_task(tid)
                # 既不可能创建失败又拿不到既有行：除非被并发删除，按罕见错误处理
                if existing is None:
                    raise RuntimeError(f"任务 {tid} 既未新建也无法读取，状态异常")
                row = existing
        return SubmitOrderData(task_id=row.task_id, status=TaskStatus(row.status), created=created)

    # ---- 领取 ----

    async def lease(
        self,
        *,
        worker_id: str,
        wait: int,
        site: str | None,
        max_wait: int,
    ) -> LeaseData | None:
        """本地长轮询领取

        - 全局并发度 1：已有 leased/running 任务时立即返回 None
        - 无可领任务时挂起最多 `min(wait, max_wait)` 秒，超时返回 None
        - 拿到任务时 queued → leased，写入 lease_owner / lease_expires_at / lease_count
        - 顺带刷新 worker 心跳（last_seen_at），这就是 lease 兼任心跳的设计
        """
        effective_wait = max(0, min(wait, max_wait))
        deadline = utcnow() + timedelta(seconds=effective_wait)

        async with self._lock:
            # 先把过期租约扫到 stale，避免「卡死的任务」堵住新任务
            self._notify_stale(await self._sweep_locked())

            await self._db.upsert_worker(worker_id, to_iso(utcnow()))

            while True:
                result = await self._try_pick_locked(worker_id, site)
                if result is not None:
                    return result

                remaining = (deadline - utcnow()).total_seconds()
                if remaining <= 0:
                    return None
                try:
                    await asyncio.wait_for(self._cond.wait(), timeout=remaining)
                except asyncio.TimeoutError:
                    return None

    async def _try_pick_locked(self, worker_id: str, site: str | None) -> LeaseData | None:
        """锁内尝试领取一次。返回 None 表示此刻无可领任务"""
        # 全局并发度 1：已有 leased/running 立即返回空
        active = await self._db.list_tasks_in_statuses(tuple(s.value for s in ACTIVE_STATUSES))
        if active:
            return None

        task = await self._db.pick_queued_task(site)
        if task is None:
            return None

        return await self._lease_to_locked(task, worker_id, known_state=None)

    async def _lease_to_locked(
        self,
        task: TaskRow,
        worker_id: str,
        *,
        known_state: str | None,
    ) -> LeaseData:
        """把一个 queued 或 stale 任务租给 worker。调用方必须持有 self._lock"""
        expires = to_iso(utcnow() + timedelta(seconds=self._lease_ttl))
        new_count = task.lease_count + 1
        await self._db.update_task(
            task.task_id,
            status=TaskStatus.LEASED.value,
            lease_owner=worker_id,
            lease_expires_at=expires,
            lease_count=new_count,
            updated_at=to_iso(utcnow()),
        )
        # 已知状态：取最近一次 report 的 state（首次领取为 None）
        if known_state is None:
            known_state = await self._db.get_latest_state(task.task_id)
        return LeaseData(
            task_id=task.task_id,
            site=task.site,
            intent=task.intent,
            lease_expires_at=expires,
            lease_count=new_count,
            known_state=known_state,  # type: ignore[arg-type]
        )

    # ---- 续租 ----

    async def renew(self, task_id: str, worker_id: str) -> RenewData:
        async with self._lock:
            task = await self._db.get_task(task_id)
            if task is None:
                raise TaskNotFoundError(task_id)
            if task.lease_owner != worker_id:
                raise LeaseInvalidError(f"worker {worker_id} 不是任务 {task_id} 的租约持有者")
            if TaskStatus(task.status) in TERMINAL_STATUSES:
                raise LeaseInvalidError(f"任务 {task_id} 已终结，无法续租")
            if TaskStatus(task.status) not in ACTIVE_STATUSES:
                raise LeaseInvalidError(f"任务 {task_id} 状态 {task.status}，无可续租约")

            expires = to_iso(utcnow() + timedelta(seconds=self._lease_ttl))
            await self._db.update_task(
                task_id,
                lease_expires_at=expires,
                updated_at=to_iso(utcnow()),
            )
            return RenewData(task_id=task_id, lease_expires_at=expires, lease_count=task.lease_count)

    # ---- 回报 ----

    async def report(
        self,
        task_id: str,
        *,
        worker_id: str,
        state: str,
        payable_yen: int | None,
        pay_deadline: str | None,
        site_order_id: str | None,
        evidence_ref: str | None,
        detail: str,
        terminal: bool,
        terminal_status: str | None,
    ) -> ReportData:
        async with self._lock:
            task = await self._db.get_task(task_id)
            if task is None:
                raise TaskNotFoundError(task_id)
            if task.lease_owner != worker_id:
                raise LeaseInvalidError(f"worker {worker_id} 不是任务 {task_id} 的租约持有者")
            status = TaskStatus(task.status)

            if status in TERMINAL_STATUSES:
                # 任务已终结：仍然追加 report（规格 §6 末尾「任务已 terminal 的仍可上报，
                # gateway 追加到 task_reports」），但不允许改任务状态
                if terminal and terminal_status is not None and terminal_status != status.value:
                    raise InvalidTaskStateError(
                        f"任务 {task_id} 已是终态 {status.value}，不能改为 {terminal_status}"
                    )
            elif status == TaskStatus.LEASED:
                # 首次 report：leased → running
                await self._db.update_task(
                    task_id, status=TaskStatus.RUNNING.value, updated_at=to_iso(utcnow())
                )

            now = to_iso(utcnow())
            report_row = ReportRow(
                task_id=task_id,
                state=state,
                payable_yen=payable_yen,
                pay_deadline=pay_deadline,
                site_order_id=site_order_id,
                evidence_ref=evidence_ref,
                detail=detail,
                reported_at=now,
            )
            inserted = await self._db.upsert_report(report_row)

            final_status = status.value
            if terminal:
                final_status = self._resolve_terminal_status(
                    current=status, explicit=terminal_status, state=state
                )
                await self._db.update_task(
                    task_id,
                    status=final_status,
                    clear_lease=True,
                    updated_at=now,
                )
                # 任务终结可能让并发度 1 的闸门放开，唤醒等着的 lease
                self._cond.notify_all()
                # 终结类事件回调（§4.8）：仅「本次真正把任务推入终态」才通知——
                # 已终结任务的后续 report（如付款后监控的 shipped/delivered）不重复通知
                if status not in TERMINAL_STATUSES:
                    self._notify_terminal(task, report_row, final_status)

            return ReportData(task_id=task_id, status=TaskStatus(final_status), recorded=inserted)

    @staticmethod
    def _resolve_terminal_status(
        *, current: TaskStatus, explicit: str | None, state: str
    ) -> str:
        """terminal=true 时决定任务的终态。

        显式指定优先；否则按订单 state 推断：paid → succeeded，cancelled → failed，
        其余（如 awaiting_payment 永久搁置）→ needs_human。

        注意 cancelled 既可能来自明确失败，也可能是站点侧自动取消（如付款期限过期），
        统一记 failed 让上游看到需要处理。
        """
        if explicit is not None:
            if explicit not in {s.value for s in TERMINAL_STATUSES}:
                raise InvalidTaskStateError(f"terminal_status={explicit} 不是终态")
            return explicit
        if state == "paid":
            return TaskStatus.SUCCEEDED.value
        if state == "cancelled":
            return TaskStatus.FAILED.value
        # awaiting_payment 等无明确结局的状态被标 terminal 时，多半是 worker 检测到
        # 3DS 之类人工环节，按 needs_human 处理
        return TaskStatus.NEEDS_HUMAN.value

    # ---- 恢复（绝不自动重投，见规格 §5）----

    async def reclaim(self, task_id: str, worker_id: str) -> ReclaimData:
        """把 stale 任务重新租给 worker。lease_count 必然 > 1"""
        async with self._lock:
            task = await self._db.get_task(task_id)
            if task is None:
                raise TaskNotFoundError(task_id)
            if TaskStatus(task.status) not in RECLAIMABLE_STATUSES:
                raise InvalidTaskStateError(
                    f"任务 {task_id} 状态 {task.status}，不能 reclaim（只能 reclaim stale 任务）"
                )

            known_state = await self._db.get_latest_state(task_id)
            lease = await self._lease_to_locked(task, worker_id, known_state=known_state)
            # 复用 LeaseData 的字段构造 ReclaimData（结构一致）
            return ReclaimData(
                task_id=lease.task_id,
                site=lease.site,
                intent=lease.intent,
                lease_expires_at=lease.lease_expires_at,
                lease_count=lease.lease_count,
                known_state=lease.known_state,
            )

    # ---- 过期清扫 ----

    async def sweep(self) -> int:
        """把过期的 leased/running 任务置 stale。返回清扫条数"""
        async with self._lock:
            swept = await self._sweep_locked()
            self._notify_stale(swept)
            return len(swept)

    async def _sweep_locked(self) -> list[TaskRow]:
        """锁内执行 sweep。调用方必须持有 self._lock。返回本次被置 stale 的任务"""
        active = await self._db.list_tasks_in_statuses(tuple(s.value for s in ACTIVE_STATUSES))
        now = utcnow()
        swept: list[TaskRow] = []
        for task in active:
            if task.lease_expires_at and parse_iso(task.lease_expires_at) < now:
                await self._db.update_task(
                    task.task_id,
                    status=TaskStatus.STALE.value,
                    clear_lease=True,
                    updated_at=to_iso(now),
                )
                swept.append(task)
                logger.warning(
                    "租约过期，任务置 stale（不自动重投）：task_id=%s site=%s",
                    task.task_id,
                    task.site,
                )
        if swept:
            # stale 后腾出了「全局并发度 1」的槽位，唤醒等待者
            # （虽然 stale 任务不会被普通 lease 取到，但队列里可能还有 queued）
            self._cond.notify_all()
        return swept

    # ---- 终结类事件回调（§4.8）----
    # notify 只是调度后台发送任务，持锁调用是安全的：真正的 HTTP 发送不碰这把锁

    def _notify_terminal(self, task: TaskRow, report_row: ReportRow, final_status: str) -> None:
        if self.notifier is not None and task.callback_url:
            self.notifier.notify(
                task.callback_url, terminal_payload(task, report_row, final_status)
            )

    def _notify_stale(self, swept: list[TaskRow]) -> None:
        if self.notifier is None:
            return
        for task in swept:
            if task.callback_url:
                self.notifier.notify(task.callback_url, stale_payload(task))

    # ---- 查询 ----

    async def get_task_detail(self, task_id: str) -> TaskDetail:
        task = await self._db.get_task(task_id)
        if task is None:
            raise TaskNotFoundError(task_id)
        reports = await self._db.get_reports(task_id)
        latest = reports[-1].state if reports else None
        return _task_to_detail(
            task,
            latest_state=latest,
            reports=[
                ReportEntry(
                    state=r.state,
                    payable_yen=r.payable_yen,
                    pay_deadline=r.pay_deadline,
                    site_order_id=r.site_order_id,
                    evidence_ref=r.evidence_ref,
                    detail=r.detail,
                    reported_at=r.reported_at,
                )
                for r in reports
            ],
        )

    async def list_tasks(
        self,
        *,
        status: str | None,
        site: str | None,
        limit: int,
        offset: int,
    ) -> TaskListData:
        rows, total = await self._db.list_tasks(
            status=status, site=site, limit=limit, offset=offset
        )
        items: list[TaskDetail] = []
        for row in rows:
            reports = await self._db.get_reports(row.task_id)
            latest = reports[-1].state if reports else None
            items.append(
                _task_to_detail(
                    row,
                    latest_state=latest,
                    reports=[
                        ReportEntry(
                            state=r.state,
                            payable_yen=r.payable_yen,
                            pay_deadline=r.pay_deadline,
                            site_order_id=r.site_order_id,
                            evidence_ref=r.evidence_ref,
                            detail=r.detail,
                            reported_at=r.reported_at,
                        )
                        for r in reports
                    ],
                )
            )
        return TaskListData(items=items, total=total, limit=limit, offset=offset)

    # ---- 健康检查 ----

    async def health_snapshot(self) -> GatewayHealthData:
        """规格 §4.7：只做 worker 失联 + 任务长时间无人领两条兜底告警"""
        async with self._lock:
            queued_count = await self._db.count_by_status(TaskStatus.QUEUED.value)
            active_rows = await self._db.list_tasks_in_statuses(
                tuple(s.value for s in ACTIVE_STATUSES)
            )
            queued_rows = await self._db.list_tasks_in_statuses(
                (TaskStatus.QUEUED.value,)
            )
            workers = await self._db.list_workers()

        now = utcnow()
        offline_workers: list[WorkerHealthEntry] = []
        worker_entries: list[WorkerHealthEntry] = []
        for w in workers:
            last_seen = parse_iso(w.last_seen_at)
            age = int((now - last_seen).total_seconds())
            entry = WorkerHealthEntry(
                worker_id=w.worker_id,
                last_seen_at=w.last_seen_at,
                last_seen_seconds=age,
            )
            worker_entries.append(entry)
            if age > self._worker_offline_alert_seconds:
                offline_workers.append(entry)

        stale_queued_threshold = utcnow() - timedelta(seconds=self._worker_offline_alert_seconds)
        stale_queued: list[QueuedAlertEntry] = []
        for t in queued_rows:
            created = parse_iso(t.created_at)
            if created < stale_queued_threshold:
                stale_queued.append(
                    QueuedAlertEntry(
                        task_id=t.task_id,
                        site=t.site,
                        created_at=t.created_at,
                        age_seconds=int((now - created).total_seconds()),
                    )
                )

        active_details = [
            _task_to_detail(
                t,
                latest_state=await self._db.get_latest_state(t.task_id),
                reports=[],
            )
            for t in active_rows
        ]

        degraded = bool(offline_workers or stale_queued)
        return GatewayHealthData(
            status="degraded" if degraded else "ok",
            queued_count=queued_count,
            active_tasks=active_details,
            workers=worker_entries,
            offline_workers=offline_workers,
            stale_queued_tasks=stale_queued,
        )