NeteaseUrl/services/_taskcore.py

# -*- coding: utf-8 -*-
"""任务生命周期原语（启动 / 终态 / 文件收尾 / 结算）。

worker 自本轮起是协程：``_start_task`` 把 ``worker(*args)`` 包成 asyncio 任务挂到
应用事件循环上（原先起一条线程）。任务字典里的句柄随之由 ``worker_thread`` 换成
``worker_task``，监控循环对「worker 已结束但没人走终态结算」的兜底判断改为
``task.done()``——语义不变（settle 幂等，无双退）。

取消/暂停位仍是 ``threading.Event``：置位方可能是监控线程或请求线程，读取方是
worker 协程，只有"置位"与"轮询读取"两种操作，既不需要跨循环唤醒也不会被事件
循环卡住，比 asyncio.Event 更省心。
"""

import asyncio
import threading
import time
import traceback
from pathlib import Path
from typing import Any

import auth_system
from core.state import (
    DOWNLOAD_TASKS,
    DOWNLOAD_TASKS_LOCK,
)
from music_downloader import (
    DownloadCanceled,
)

from .api_service import (
    api_service,
)

# 存活期的后台任务强引用——asyncio 只持弱引用，不留引用的话任务可能中途被 GC
_BACKGROUND_TASKS: set[asyncio.Task] = set()


def _log_worker_crash(handle: asyncio.Task) -> None:
    """worker 的异常逃逸到任务层时留痕。

    正常情况下 worker 自带四路终态收敛，不会抛到这里；真抛出来说明收尾路径本身
    出了问题，必须留下 traceback，否则任务永远停在 running 且无人知晓。
    """
    if handle.cancelled():
        return
    exc = handle.exception()
    if exc is not None:
        api_service.logger.error(f"任务 worker 异常退出: {exc}", exc_info=exc)


async def _start_task(task_id: str, worker, args: tuple = ()) -> str | None:
    """启动任务 worker（启动失败时收敛任务终态并按账本退费，不留僵尸 queued 记录）。
    退款只走账本（本处 settle 后装饰器看到账已关，不会重复退）。
    任务句柄存入任务字典——监控循环对 4h 超时任务若发现 worker 已结束
    （无法自行走到终态结算），可代为按账本结算退款（settle 幂等，无双退）。"""
    try:
        loop = asyncio.get_running_loop()
        handle = loop.create_task(worker(*args), name=f'task:{task_id}')
    except Exception as e:
        with DOWNLOAD_TASKS_LOCK:
            t = DOWNLOAD_TASKS.pop(task_id, None)
        if t:
            try:
                await auth_system.ledger_settle(task_id, 0, 'submit_failed')
            except Exception:
                pass
        api_service.logger.error(f"任务启动失败: {task_id} {e}")
        return f'任务启动失败：{e}'
    _BACKGROUND_TASKS.add(handle)
    handle.add_done_callback(_BACKGROUND_TASKS.discard)
    handle.add_done_callback(_log_worker_crash)
    with DOWNLOAD_TASKS_LOCK:
        t = DOWNLOAD_TASKS.get(task_id)
        if t is not None:
            t['worker_task'] = handle
    return None


async def _mark_task_terminal(task_id: str, status: str, msg: str = '') -> None:
    """任务进入终态的统一入口（记录 completed_at，保留期从完成时刻起算，
    而不是从提交时刻起算）。
    /F-E：置 error/canceled 终态时若 worker 已结束（无人再走正常结算），
    代为按账本结算退款——ledger_settle 幂等，worker 已自行结算过则此处为 0。
    终态先到先得——已终态（done/error/canceled）的任务
    不再被覆盖（堵住 canceled→error 的非法迁移边：取消后解析阶段的自身失败
    不得把"已取消"改写成"失败"，completed_at 也不再刷新）。"""
    need_settle = False
    with DOWNLOAD_TASKS_LOCK:
        t = DOWNLOAD_TASKS.get(task_id)
        if t and t.get('status') in ('done', 'error', 'canceled'):
            return
        if t:
            t['status'] = status
            t['msg'] = msg or t.get('msg', '')
            t['completed_at'] = time.time()
            if status in ('canceled', 'error'):
                wt = t.get('worker_task')
                if wt is not None and wt.done():
                    need_settle = True
    # settle 移到任务锁外——ledger_settle 内部要抢 SQLite 写锁
    # （busy 最长 15s），持任务锁调用会让全站进度/提交/取消/领取随库停摆；
    # settle 幂等，锁外调用无双退。落库已迁移为原生异步，直接 await。
    if need_settle:
        try:
            await auth_system.ledger_settle(task_id, 0, status)
        except Exception:
            pass


def _task_files(task: dict[str, Any]) -> list[str]:
    """任务当前关联的可清理文件路径（交付 ZIP/交付副本；缓存文件不在此列）"""
    result = task.get('result') or {}
    files = []
    for f in (result.get('zip_path'), result.get('file_path')):
        if f:
            files.append(f)
    return files


async def _settle_batch_billing(task_id: str, ok_count: int, final_status: str = 'settled') -> int:
    """批量任务结算（走持久账本，按实际预扣-实际交付幂等退款；
    重复调用/重启后调用都只会退一次）"""
    try:
        refund = await auth_system.ledger_settle(task_id, ok_count, final_status)
    except Exception:
        # 本函数多在异常处理器内被调用——结算再抛错（库忙超时）
        # 会杀死 worker 收尾路径，跳过终态标记；失败留周期对账兜底（幂等）
        api_service.logger.error(
            f"批量结算异常（留对账兜底）: {task_id}\n{traceback.format_exc()}")
        return 0
    if refund:
        try:
            with DOWNLOAD_TASKS_LOCK:
                t = DOWNLOAD_TASKS.get(task_id) or {}
            uid = int(t.get('billing_user_id') or 0)
            u = await auth_system.get_user_by_id(uid) if uid else None
            uname = u['username'] if u else 'system'
            remaining = u['duration_seconds'] if u else '?'
            await auth_system.add_log('success', uname, 'system', '批量退费',
                f'失败部分 {refund} 次已退回，剩余 {remaining} 次')
        except Exception:
            pass
    return refund


def _task_work_dir(task_id: str) -> Path:
    """任务专属工作目录（交付 ZIP 与中间文件都落这里）。

    原实现把临时包建在 downloads 根目录，会被缓存清扫线程按 mtime 回收
    （Linux 上已打开文件同样能被 unlink），排队超时的任务于是出现
    「扣费后领取 404」的悬空交付；任务目录从此对根目录清扫免疫。
    音乐歌单打包与播客整台打包共用本目录约定。
    """
    path = api_service.downloads_path / 'tasks' / task_id
    path.mkdir(parents=True, exist_ok=True)
    return path


async def _wait_if_paused(task: dict[str, Any], pause_ev: threading.Event,
                          cancel_ev: threading.Event, done: int) -> None:
    """断线暂停闸门：置 paused 态并等待恢复或取消，恢复后回到 running。

    两个打包 worker（音乐歌单 / 播客整台）共用同一套暂停语义——此前各自持有
    一份逐字相同的副本（连注释都只差一行），是「同型缺陷同时出现在两处」的温床。

    调用点位于窗口循环顶部，等待期间**不持有** BATCH_ZIP_SEM：名额交还队列
    给下一位，这正是「断线不占坑」的实现方式。

    ``done`` 由调用方按当前进度传入（ok_count + len(failed)）——写成参数而非
    捕获闭包，是为了让本函数脱离两个 worker 各自的局部状态。
    """
    if not (pause_ev.is_set() or task.get('status') == 'paused'):
        return
    with DOWNLOAD_TASKS_LOCK:
        task['status'] = 'paused'
        task['done'] = done
        task['msg'] = '已暂停（重新打开页面自动恢复）'
    while not cancel_ev.is_set() and (pause_ev.is_set() or task.get('status') == 'paused'):
        await asyncio.sleep(1)
    if cancel_ev.is_set():
        raise DownloadCanceled("任务已取消")
    pause_ev.clear()
    with DOWNLOAD_TASKS_LOCK:
        task['status'] = 'running'
        task['msg'] = '继续下载...'