NeteaseUrl/services/_zipcore.py

# -*- coding: utf-8 -*-
"""批量 ZIP 打包引擎：两个打包 worker 的共用骨架。

音乐歌单（``services.batch``）与播客整台（``services.podcast``）的打包 worker
此前是两份约 300 行的平行实现——任务状态机、断线暂停、全局批量名额、
「逐条失败不计整批」、取消与异常的四路终态收敛、账本结算全部各写一遍，
连注释都逐字重复（上一轮修掉的 ``zip_valid`` 未绑定缺陷正是同型出现在这两处）。
真正不同的只有五件事：清单来源、单条怎么下、包内条目怎么拼、并发模型、
交付包名与文案。

因此本模块保留共用骨架，把上述差异收敛到 :class:`ZipJob`；引擎内任何一处
修订（账本口径、取消检查点、结算幂等）从此只需改一次。

术语：下文「条目」指一个下载单元——音乐歌单里是一首歌，播客整台里是一集。

线程模型（本轮最容易踩坏的地方，改动前请先读）：
- 下载 I/O 在 ``ThreadPoolExecutor`` 里跑，窗口大小 = ``job.concurrency``；
- **ZIP 写入与包内名去重只发生在调用线程**（``as_completed`` 消费结果时），
  :class:`zipfile.ZipFile` 不是线程安全的，且去重读改写 ``used_arcnames``
  需要串行——故下载与落包被拆成 ``download_one`` / ``write_entry`` 两个钩子，
  只有前者进线程池。
"""
import os
import tempfile
import threading
import time
import traceback
import zipfile
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from core.state import (
    BATCH_DURATIONS,
    BATCH_DURATIONS_LOCK,
    DOWNLOAD_TASKS,
    DOWNLOAD_TASKS_LOCK,
)
from music_downloader import (
    DownloadCanceled,
    DownloadException,
)

from ._shared import (
    BATCH_ZIP_SEM,
)
from ._taskcore import (
    _mark_task_terminal,
    _settle_batch_billing,
    _task_work_dir,
    _wait_if_paused,
)
from .api_service import (
    api_service,
)


class ItemFailed(Exception):
    """单个条目失败（不影响整批）。

    ``label`` 已是最终展示文案（引擎不再拼后缀）。与
    :class:`~music_downloader.DownloadCanceled` 必须区分——取消要中止整批，
    本异常只记一次失败。
    """

    def __init__(self, label: str) -> None:
        super().__init__(label)
        self.label = label


@dataclass
class ZipContext:
    """单次打包运行中，引擎交给 ZipJob 的上下文（引擎持有，任务侧只读）。"""

    task_id: str
    params: dict[str, Any]
    cancel_ev: threading.Event
    work_dir: Path
    zf: zipfile.ZipFile
    progress: Callable[..., None]
    total: int
    used_arcnames: set[str] = field(default_factory=set)


@dataclass
class _Tally:
    """成功/失败计数。

    两个 worker 此前把 ``ok_count`` 当外层局部变量，异常处理器读到的是「异常
    发生那一刻」的值（``zip_valid`` 为假时该值不影响退款，但日志里的「成功N」
    会失真）。收成一个对象后，异常路径读到的恒是真实值。
    """

    ok: int = 0
    failed: list[str] = field(default_factory=list)


class ZipJob:
    """一次批量打包任务的「因任务而异」部分。

    子类至少实现 :meth:`prepare` / :meth:`download_one` / :meth:`write_entry` /
    :meth:`zip_filename`，其余钩子按需覆写。**实例每次运行新建**——下载器信息
    预取池等运行期状态就挂在实例上。
    """

    tag = '批量打包'            # 日志前缀
    noun = '项'                 # 计数单位（文案里的「N 首/N 集」）
    item_noun = '条目'          # 全称（文案里的「部分xx已打包完成」）
    tmp_prefix = 'batch_zip_'   # 临时 ZIP 文件名前缀
    start_msg = '打包中...'      # 进入 running 时的进度文案
    concurrency = 1             # 每窗口条目数（1 = 顺序）
    record_duration = False     # 是否把耗时记入 BATCH_DURATIONS（供前端估算等待）

    # ---- 必须实现 ----

    def prepare(self, params: dict[str, Any], task_id: str, failed: list[str]) -> list[Any]:
        """组装条目清单。可抛 ``DownloadException`` 让任务按业务失败收敛。

        ``failed`` 可被直接追加——清单阶段就有条目已知失败（如按 ID 列表时个别
        详情拉取失败），这些失败要与下载阶段的一起计入。
        """
        raise NotImplementedError

    def download_one(self, item: Any, idx: int, ctx: ZipContext) -> Any:
        """下载单个条目（**在下载线程池内执行**）。

        返回的 ``payload`` 会原样交给 :meth:`write_entry`（放临时文件名、
        已取到的元信息等）。失败请抛异常：``ItemFailed`` 用其自带文案，
        其他异常由引擎按 ``fail_label(item)(异常)`` 记录。
        """
        raise NotImplementedError

    def write_entry(self, item: Any, idx: int, payload: Any, ctx: ZipContext) -> None:
        """把下载结果写入 ZIP（**在引擎调用线程内执行，天然串行**）。

        包内名去重在这里做——需要读写 ``ctx.used_arcnames``，必须串行。
        """
        raise NotImplementedError

    def zip_filename(self, ok_count: int) -> str:
        """交付 ZIP 的文件名。"""
        raise NotImplementedError

    # ---- 可选钩子 ----

    def prefetch(self, window: list[Any], ctx: ZipContext) -> None:
        """窗口级预热（每窗口一次、提交下载前调用）。默认无操作。"""

    def fail_label(self, item: Any) -> str:
        """单条失败文案前缀（引擎拼成 ``f'{label}({异常})'``）。"""
        return str(item)

    def done_msg(self, item: Any, idx: int, ctx: ZipContext) -> str:
        """条目成功后的进度文案（空串 = 沿用当前文案）。默认不追加文案。"""
        return ''

    def log_label(self) -> str:
        """日志里区隔本次任务的名字（如电台名），默认空串。"""
        return ''

    def all_failed_msg(self, failed: list[str], refunded: int) -> str:
        """全部条目失败时的用户可见文案（含退费说明）。"""
        return f"全部失败（{len(failed)} {self.noun}）" + (
            f"，已退回 {refunded} 次" if refunded else '')

    def partial_msg(self, refunded: int) -> str:
        """已打包部分交付时的用户可见文案（仅异常路径会用到）。"""
        return f"部分{self.item_noun}已打包完成，可领取" + (
            f"，退回 {refunded} 次" if refunded else '')

    def supplement(self, task_id: str, params: dict[str, Any], count: int) -> None:
        """清单确定后的计费补扣（默认无操作）。"""

    def on_delivered(self, items: list[Any], task_id: str, params: dict[str, Any]) -> None:
        """交付落定后的操作记录（默认无操作）。"""

    def cleanup(self) -> None:
        """释放运行期资源（如信息预取池）。默认无操作。"""


def _discard(path: str | None) -> None:
    """删除未发布的临时 ZIP（坏包/空包/取消包）——失败静默（清理非关键路径）。"""
    if not path:
        return
    try:
        os.remove(path)
    except OSError:
        pass


def run_zip_job(job: ZipJob, task_id: str, *,
                settle: Callable[[str, int, str], int] = _settle_batch_billing) -> None:
    """执行一次批量打包任务（在 worker 线程内调用，自身不向外抛异常）。

    ``settle`` 由调用方注入，默认走 :func:`services._taskcore._settle_batch_billing`。
    做成参数是为了让测试能在**模块级**替换结算函数（包装层在调用时读取模块全局，
    替换即生效）。
    """
    zip_tmp_path: str | None = None
    # zip_valid 必须在函数头初始化：ZIP 构建前的业务错误（参数缺失/无条目/
    # 补扣失败）会 early raise，异常处理器若引用未绑定名会 NameError 二次崩线程
    zip_valid = False
    tally = _Tally()
    items: list[Any] = []
    try:
        with DOWNLOAD_TASKS_LOCK:
            task = DOWNLOAD_TASKS[task_id]
            task['status'] = 'running'
            task['msg'] = job.start_msg
        params = task['params']
        cancel_ev = task['cancel']
        pause_ev = task['pause']
        # 计耗时只算「下载+打包」这一段（不含排队），故不用 task['start']
        t0 = time.time()

        items = job.prepare(params, task_id, tally.failed)
        with DOWNLOAD_TASKS_LOCK:
            DOWNLOAD_TASKS[task_id]['total'] = len(items)
        job.supplement(task_id, params, len(items))

        def _progress(msg: str = '') -> None:
            with DOWNLOAD_TASKS_LOCK:
                t = DOWNLOAD_TASKS[task_id]
                t['done'] = tally.ok + len(tally.failed)
                t['failed'] = len(tally.failed)
                if msg:
                    t['msg'] = msg

        # ZIP 构建放 downloads 卷（原实现落系统 %TEMP%——500 首无损可达 GB 级，
        # 小系统盘一次大单即塞爆），且建在任务专属目录（见 _task_work_dir）：
        # downloads 根目录会被缓存清扫按 mtime 回收
        work_dir = _task_work_dir(task_id)
        with tempfile.NamedTemporaryFile(prefix=job.tmp_prefix, suffix='.zip',
                                         dir=str(work_dir), delete=False) as zip_tmp:
            zip_tmp_path = zip_tmp.name
            # 音频本身已压缩，条目按 ZIP_STORED（免压缩）大幅提速；文本条目仍压缩
            with zipfile.ZipFile(zip_tmp, 'w', zipfile.ZIP_DEFLATED) as zf:
                ctx = ZipContext(task_id=task_id, params=params, cancel_ev=cancel_ev,
                                 work_dir=work_dir, zf=zf, progress=_progress,
                                 total=len(items))
                _fill_zip(job, items, ctx, tally, task, pause_ev)
        # zip_valid 只在 ZIP with 块正常退出（中央目录完整写出）后才置 True。
        # 此前两个 worker 在这之后还各留了一段 `if not zip_valid:` 失败处理——
        # 紧跟赋值、中间无任何可改zip_valid的语句，是**不可达**的复审残留，已删。
        zip_valid = True

        if job.record_duration and tally.ok > 0:
            with BATCH_DURATIONS_LOCK:
                BATCH_DURATIONS.append(time.time() - t0)

        # 全部失败（含清单非空但一条都没成）：不发布空 ZIP、不按交付结算，
        # 转 error 并全额退费。原音乐批量用 `ok==0 and failed`、播客用早抛
        # DownloadException——这里统一为显式收敛，两边的账本终态一致
        if tally.ok == 0:
            _discard(zip_tmp_path)
            zip_tmp_path = None
            refunded = settle(task_id, 0, 'error')
            _mark_task_terminal(task_id, 'error', job.all_failed_msg(tally.failed, refunded))
            with DOWNLOAD_TASKS_LOCK:
                t = DOWNLOAD_TASKS.get(task_id)
                if t is not None:
                    t['failed'] = len(tally.failed)
                    t['refunded'] = refunded
            api_service.logger.info(
                f"{job.tag}全部失败: {task_id} 失败{len(tally.failed)} 退回{refunded}次")
            return

        # 正常交付：按入包成功数结算（'settled' 是 _settle_batch_billing 的默认值，
        # 此处显式写出，让「成功/取消/失败」三种账本终态在本文件里一目了然）
        refunded = settle(task_id, tally.ok, 'settled')
        with DOWNLOAD_TASKS_LOCK:
            t = DOWNLOAD_TASKS.get(task_id)
            # 发布前锁内复查取消/失效状态——取消请求落在「检查后、拿锁前」的窗口
            # 也不得覆盖终态
            if t is None or t.get('status') in ('canceled', 'error'):
                raise DownloadCanceled("任务已取消")
            if cancel_ev.is_set():
                raise DownloadCanceled("任务已取消")
            t['status'] = 'done'
            t['completed_at'] = time.time()
            t['done'] = tally.ok + len(tally.failed)
            t['failed'] = len(tally.failed)
            t['msg'] = f"成功{tally.ok} 失败{len(tally.failed)}" + (
                f" 退回{refunded}次" if refunded else '')
            t['refunded'] = refunded
            t['result'] = {
                'zip_path': zip_tmp_path,
                'filename': job.zip_filename(tally.ok),
                'audio_files': [],   # 缓存文件不归任务所有，不进入清理清单
            }
        job.on_delivered(items, task_id, params)
        api_service.logger.info(
            f"{job.tag}完成: {job.log_label()}成功{tally.ok} 失败{len(tally.failed)} "
            f"退回{refunded}次 ({task_id})")
    except DownloadCanceled:
        # 部分成功后取消——已入包部分照常交付（成果守恒），未完成部分按账本退费。
        # ZIP 未正常完成（写坏）时不按 ok_count 交付——按全额退费收敛
        delivered = tally.ok if zip_valid else 0
        refunded = settle(task_id, delivered, 'canceled')
        if delivered > 0 and zip_tmp_path:
            _mark_task_terminal(task_id, 'canceled',
                                f"已取消（已打包 {tally.ok} {job.noun}可领取）" +
                                (f"，退回 {refunded} 次" if refunded else ''))
            with DOWNLOAD_TASKS_LOCK:
                t = DOWNLOAD_TASKS.get(task_id)
                if t is not None:
                    t['refunded'] = refunded
                    t['result'] = {
                        'zip_path': zip_tmp_path,
                        'filename': job.zip_filename(tally.ok),
                        'audio_files': [],
                    }
        else:
            _discard(zip_tmp_path)
            _mark_task_terminal(task_id, 'canceled',
                                '任务已取消' + (
                                    f"（已退回 {refunded} 次下载次数）" if refunded else ''))
            with DOWNLOAD_TASKS_LOCK:
                t = DOWNLOAD_TASKS.get(task_id)
                if t is not None:
                    t['refunded'] = refunded
        api_service.logger.info(
            f"{job.tag}已取消: {task_id}（成功{tally.ok} 退回{refunded}次）")
    except DownloadException as e:
        # 业务失败（清单阶段：无可下载条目 / 补扣下载次数失败 / 上游不可用）。
        # 文案是我们自己拼的用户可读文本，保留原文（同播客单集下载）
        delivered = tally.ok if zip_valid else 0
        refunded = settle(task_id, delivered, 'error')
        _apply_error(task_id, str(e), refunded)
        _discard(zip_tmp_path)
        api_service.logger.error(
            f"{job.tag}失败: {task_id} {e}（成功{tally.ok} 退回{refunded}次）")
    except Exception as e:
        # 非预期异常原文 logger 先行（结算本身可能再抛，先落下 traceback），
        # 用户可见文案泛化；带上退款结果再补一条，避免排障时只看到「异常」
        api_service.logger.error(f"{job.tag}异常: {task_id} {e}\n{traceback.format_exc()}")
        delivered = tally.ok if zip_valid else 0
        refunded = settle(task_id, delivered, 'error')
        if delivered > 0 and zip_tmp_path:
            _apply_error(task_id, job.partial_msg(refunded), refunded)
            with DOWNLOAD_TASKS_LOCK:
                t = DOWNLOAD_TASKS.get(task_id)
                if t is not None:
                    t['result'] = {
                        'zip_path': zip_tmp_path,
                        'filename': job.zip_filename(tally.ok),
                        'audio_files': [],
                    }
        else:
            _apply_error(task_id, '打包失败，请稍后重试' + (
                f"（已退回 {refunded} 次下载次数）" if refunded else ''), refunded)
            _discard(zip_tmp_path)
        api_service.logger.error(
            f"{job.tag}异常: {task_id}（成功{tally.ok} 退回{refunded}次）")
    finally:
        # 回收信息预取线程池，避免线程泄漏（未启动的预取任务一并取消）
        job.cleanup()


def _apply_error(task_id: str, msg: str, refunded: int) -> None:
    """异常/业务失败的终态收敛：置 error（终态先到先得）、记账退款额。

    走 :func:`_mark_task_terminal` 而非就地写任务字典——后者是播客侧的旧写法，
    除了重复一遍「已终态不覆盖」判断外，还漏掉了 worker 线程已死时按账本补结算
    的兜底（``/F-E``）。
    """
    _mark_task_terminal(task_id, 'error', msg)
    with DOWNLOAD_TASKS_LOCK:
        t = DOWNLOAD_TASKS.get(task_id)
        if t is not None:
            t['refunded'] = refunded


def _fill_zip(job: ZipJob, items: list[Any], ctx: ZipContext, tally: _Tally,
              task: dict[str, Any], pause_ev: threading.Event) -> None:
    """逐窗口下载并写入 ZIP（返回时中央目录尚未写出，由调用方的 with 退出时落盘）。

    窗口 = ``job.concurrency`` 个条目：下载并行、落包串行。窗口大小取 1 时
    退化为顺序下载（播客整台即如此）。
    """
    concurrency = max(1, int(job.concurrency))
    pool: ThreadPoolExecutor | None = None
    try:
        # 不用 with 管理下载池——取消时 with 退出会 wait=True 等在途 future
        # （其内部阻塞在预取 HTTP 上最长 30s），改为不等待 + 取消排队任务
        pool = ThreadPoolExecutor(max_workers=concurrency)
        offset = 0
        while offset < len(items):
            if ctx.cancel_ev.is_set():
                raise DownloadCanceled("任务已取消")
            # 断线暂停：释放名额等待恢复/取消，恢复后重新进入排队
            _wait_if_paused(task, pause_ev, ctx.cancel_ev, tally.ok + len(tally.failed))
            # 每个窗口持有 BATCH_ZIP_SEM（同一时刻全局最多 N 个批量任务在下载）
            # 等待名额期间可被取消，不再无限阻塞线程与退款
            if not BATCH_ZIP_SEM.acquire_cancellable(ctx.cancel_ev):
                raise DownloadCanceled("任务已取消")
            try:
                # 暂停/取消落在「等待恢复返回」与「拿到名额」之间：立刻放名额回到
                # 循环顶部重新处理暂停（原播客侧缺这一步，会多下完这一集才停）
                if (ctx.cancel_ev.is_set() or pause_ev.is_set()
                        or task['status'] == 'paused'):
                    continue
                window = items[offset:offset + concurrency]
                start_idx = offset + 1     # 条目序号从 1 起（进度文案用）
                offset += len(window)
                job.prefetch(window, ctx)   # 只预取本窗口
                futures = {
                    pool.submit(job.download_one, item, start_idx + i, ctx): (item, start_idx + i)
                    for i, item in enumerate(window)
                }
                for future in as_completed(futures):
                    if ctx.cancel_ev.is_set():
                        raise DownloadCanceled("任务已取消")
                    item, idx = futures[future]
                    try:
                        payload = future.result()
                        # 落包在主线程：ZipFile 非线程安全 + 包内名去重要串行
                        job.write_entry(item, idx, payload, ctx)
                    except DownloadCanceled:
                        raise
                    except ItemFailed as f:
                        tally.failed.append(f.label)
                    except Exception as e:
                        tally.failed.append(f"{job.fail_label(item)}({e})")
                    else:
                        tally.ok += 1
                    ctx.progress(job.done_msg(item, idx, ctx))
            finally:
                BATCH_ZIP_SEM.release()
    finally:
        if pool is not None:
            try:
                # 取消/异常时不等待在途下载线程（其可能阻塞在预取 HTTP 上），
                # 排队任务一并取消；在途线程自行收尾（.part 残片由续传/清扫兜底）
                pool.shutdown(wait=False, cancel_futures=True)
            except TypeError:
                pool.shutdown(wait=False)