# -*- coding: utf-8 -*- """批量 ZIP 打包引擎：两个打包 worker 的共用骨架。 音乐歌单（``services.batch``）与播客整台（``services.podcast``）的打包 worker 此前是两份约 300 行的平行实现——任务状态机、断线暂停、全局批量名额、 「逐条失败不计整批」、取消与异常的四路终态收敛、账本结算全部各写一遍， 连注释都逐字重复（上一轮修掉的 ``zip_valid`` 未绑定缺陷正是同型出现在这两处）。 真正不同的只有五件事：清单来源、单条怎么下、包内条目怎么拼、并发模型、 交付包名与文案。 因此本模块保留共用骨架，把上述差异收敛到 :class:`ZipJob`；引擎内任何一处 修订（账本口径、取消检查点、结算幂等）从此只需改一次。 术语：下文「条目」指一个下载单元——音乐歌单里是一首歌，播客整台里是一集。 并发模型（本轮最容易踩坏的地方，改动前请先读）： - 一个窗口 = ``job.concurrency`` 个条目，窗口内 ``asyncio.gather`` 并发下载； - **ZIP 写入与包内名去重只在引擎协程里串行发生**（按窗口顺序消费结果）—— :class:`zipfile.ZipFile` 不是并发安全的，且去重读改写 ``used_arcnames`` 需要串行；因此 ``download_one``（并发）与 ``write_entry``（串行）是两个钩子， 只有前者进 gather。``write_entry`` 里的重型落盘（几百 MB 的音频条目） 由实现方自行经 ``asyncio.to_thread`` 卸载，避免卡住事件循环。 ``run_zip_job`` 的四个终态（成功 / 全失败 / 取消 / 异常）原先各写一段 8~20 行的 收尾内联在 try/except 里。此处把「一次运行」的状态收进 :class:`_ZipRun`， 四个终态各自成函数——改账本口径或补一种终态时，不必再在 167 行里找落点。 """ import asyncio import os import tempfile import time import traceback import zipfile from collections.abc import Awaitable, Callable 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: Any 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（供前端估算等待） # ---- 必须实现 ---- async def prepare(self, params: dict[str, Any], task_id: str, failed: list[str]) -> list[Any]: """组装条目清单。可抛 ``DownloadException`` 让任务按业务失败收敛。 ``failed`` 可被直接追加——清单阶段就有条目已知失败（如按 ID 列表时个别 详情拉取失败），这些失败要与下载阶段的一起计入。 """ raise NotImplementedError async def download_one(self, item: Any, idx: int, ctx: ZipContext) -> Any: """下载单个条目（**在窗口内并发执行**）。 返回的 ``payload`` 会原样交给 :meth:`write_entry`（放临时文件名、 已取到的元信息等）。失败请抛异常：``ItemFailed`` 用其自带文案， 其他异常由引擎按 ``fail_label(item)(异常)`` 记录。 """ raise NotImplementedError async def write_entry(self, item: Any, idx: int, payload: Any, ctx: ZipContext) -> None: """把下载结果写入 ZIP（**由引擎按窗口顺序串行调用**）。 包内名去重在这里做——需要读写 ``ctx.used_arcnames``，必须串行。 重型落盘请自行 ``await asyncio.to_thread(...)``：直接在协程里写 几百 MB 会把事件循环卡住。 """ raise NotImplementedError def zip_filename(self, ok_count: int) -> str: """交付 ZIP 的文件名。""" raise NotImplementedError # ---- 可选钩子 ---- async 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 '') async def supplement(self, task_id: str, params: dict[str, Any], count: int) -> None: """清单确定后的计费补扣（默认无操作）。""" async def on_delivered(self, items: list[Any], task_id: str, params: dict[str, Any]) -> None: """交付落定后的操作记录（默认无操作）。""" async def cleanup(self) -> None: """释放运行期资源（如信息预取任务）。默认无操作。""" def _discard(path: str | None) -> None: """删除未发布的临时 ZIP（坏包/空包/取消包）——失败静默（清理非关键路径）。""" if not path: return try: os.remove(path) except OSError: pass @dataclass class _ZipRun: """一次 ``run_zip_job`` 的共享状态。 四个终态收尾函数都要读同一批东西（任务 id、结算入口、计数、临时 ZIP 路径、 ZIP 是否完整）。原先它们是散在 167 行函数体里的局部变量，被各 except 分支 分别访问——收成一个对象后，收尾函数不必逐个透传参数，也让「谁真的用到了 哪些状态」一眼可见。 """ job: ZipJob task_id: str settle: Callable[[str, int, str], Awaitable[int]] tally: _Tally = field(default_factory=_Tally) task: dict[str, Any] = field(default_factory=dict) params: dict[str, Any] = field(default_factory=dict) cancel_ev: Any = None pause_ev: Any = None items: list[Any] = field(default_factory=list) zip_tmp_path: str | None = None zip_valid: bool = False t0: float = 0.0 def _progress_of(run: _ZipRun) -> Callable[..., None]: """构造写任务进度的回调（窗口内每完成一条调用一次）。""" def _progress(msg: str = '') -> None: with DOWNLOAD_TASKS_LOCK: t = DOWNLOAD_TASKS[run.task_id] t['done'] = run.tally.ok + len(run.tally.failed) t['failed'] = len(run.tally.failed) if msg: t['msg'] = msg return _progress async def _start_zip_run(run: _ZipRun) -> None: """置 running、组装清单、写回 total、执行计费补扣。""" with DOWNLOAD_TASKS_LOCK: task = DOWNLOAD_TASKS[run.task_id] task['status'] = 'running' task['msg'] = run.job.start_msg run.task = task run.params = task['params'] run.cancel_ev = task['cancel'] run.pause_ev = task['pause'] # 计耗时只算「下载+打包」这一段（不含排队），故不用 task['start'] run.t0 = time.time() run.items = await run.job.prepare(run.params, run.task_id, run.tally.failed) with DOWNLOAD_TASKS_LOCK: DOWNLOAD_TASKS[run.task_id]['total'] = len(run.items) await run.job.supplement(run.task_id, run.params, len(run.items)) async def _build_zip_archive(run: _ZipRun) -> None: """建临时 ZIP 并逐窗口填充。返回时 ``zip_valid`` 才置位。 ZIP 构建放 downloads 卷（原实现落系统 %TEMP%——500 首无损可达 GB 级， 小系统盘一次大单即塞爆），且建在任务专属目录（见 ``_task_work_dir``）： downloads 根目录会被缓存清扫按 mtime 回收。 """ job = run.job assert run.cancel_ev is not None and run.pause_ev is not None work_dir = _task_work_dir(run.task_id) with tempfile.NamedTemporaryFile(prefix=job.tmp_prefix, suffix='.zip', dir=str(work_dir), delete=False) as zip_tmp: run.zip_tmp_path = zip_tmp.name # 音频本身已压缩，条目按 ZIP_STORED（免压缩）大幅提速；文本条目仍压缩 with zipfile.ZipFile(zip_tmp, 'w', zipfile.ZIP_DEFLATED) as zf: ctx = ZipContext(task_id=run.task_id, params=run.params, cancel_ev=run.cancel_ev, work_dir=work_dir, zf=zf, progress=_progress_of(run), total=len(run.items)) await _fill_zip(job, run.items, ctx, run.tally, run.task, run.pause_ev) # zip_valid 只在 ZIP with 块正常退出（中央目录完整写出）后才置 True。 # 此前两个 worker 在这之后还各留了一段 `if not zip_valid:` 失败处理—— # 紧跟赋值、中间无任何可改zip_valid的语句，是**不可达**的复审残留，已删。 run.zip_valid = True if job.record_duration and run.tally.ok > 0: with BATCH_DURATIONS_LOCK: BATCH_DURATIONS.append(time.time() - run.t0) async def _finish_all_failed(run: _ZipRun) -> None: """全部失败（含清单非空但一条都没成）：不发布空 ZIP、不按交付结算，转 error 并全额退费。 原音乐批量用 ``ok==0 and failed``、播客用早抛 ``DownloadException``——这里 统一为显式收敛，两边的账本终态一致。 """ job, tally = run.job, run.tally _discard(run.zip_tmp_path) run.zip_tmp_path = None refunded = await run.settle(run.task_id, 0, 'error') await _mark_task_terminal(run.task_id, 'error', job.all_failed_msg(tally.failed, refunded)) with DOWNLOAD_TASKS_LOCK: t = DOWNLOAD_TASKS.get(run.task_id) if t is not None: t['failed'] = len(tally.failed) t['refunded'] = refunded api_service.logger.info( f"{job.tag}全部失败: {run.task_id} 失败{len(tally.failed)} 退回{refunded}次") async def _publish_delivery(run: _ZipRun) -> None: """正常交付：按入包成功数结算，锁内复查取消/终态后发布结果。 ``'settled'`` 是 ``_settle_batch_billing`` 的默认值，此处显式写出，让 「成功/取消/失败」三种账本终态在本文件里一目了然。 """ job, tally = run.job, run.tally assert run.cancel_ev is not None refunded = await run.settle(run.task_id, tally.ok, 'settled') with DOWNLOAD_TASKS_LOCK: t = DOWNLOAD_TASKS.get(run.task_id) # 发布前锁内复查取消/失效状态——取消请求落在「检查后、拿锁前」的窗口 # 也不得覆盖终态 if t is None or t.get('status') in ('canceled', 'error'): raise DownloadCanceled("任务已取消") if run.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': run.zip_tmp_path, 'filename': job.zip_filename(tally.ok), 'audio_files': [], # 缓存文件不归任务所有，不进入清理清单 } await job.on_delivered(run.items, run.task_id, run.params) api_service.logger.info( f"{job.tag}完成: {job.log_label()}成功{tally.ok} 失败{len(tally.failed)} " f"退回{refunded}次 ({run.task_id})") async def _finish_canceled(run: _ZipRun) -> None: """部分成功后取消——已入包部分照常交付（成果守恒），未完成部分按账本退费。 ZIP 未正常完成（写坏）时不按 ``ok_count`` 交付——按全额退费收敛。 """ job, tally = run.job, run.tally delivered = tally.ok if run.zip_valid else 0 refunded = await run.settle(run.task_id, delivered, 'canceled') if delivered > 0 and run.zip_tmp_path: await _mark_task_terminal(run.task_id, 'canceled', f"已取消（已打包 {tally.ok} {job.noun}可领取）" + (f"，退回 {refunded} 次" if refunded else '')) with DOWNLOAD_TASKS_LOCK: t = DOWNLOAD_TASKS.get(run.task_id) if t is not None: t['refunded'] = refunded t['result'] = { 'zip_path': run.zip_tmp_path, 'filename': job.zip_filename(tally.ok), 'audio_files': [], } else: _discard(run.zip_tmp_path) await _mark_task_terminal(run.task_id, 'canceled', '任务已取消' + ( f"（已退回 {refunded} 次下载次数）" if refunded else '')) with DOWNLOAD_TASKS_LOCK: t = DOWNLOAD_TASKS.get(run.task_id) if t is not None: t['refunded'] = refunded api_service.logger.info( f"{job.tag}已取消: {run.task_id}（成功{tally.ok} 退回{refunded}次）") async def _finish_business_error(run: _ZipRun, exc: DownloadException) -> None: """业务失败（清单阶段：无可下载条目 / 补扣下载次数失败 / 上游不可用）。 文案是我们自己拼的用户可读文本，保留原文（同播客单集下载）。 """ job, tally = run.job, run.tally delivered = tally.ok if run.zip_valid else 0 refunded = await run.settle(run.task_id, delivered, 'error') await _apply_error(run.task_id, str(exc), refunded) _discard(run.zip_tmp_path) api_service.logger.error( f"{job.tag}失败: {run.task_id} {exc}（成功{tally.ok} 退回{refunded}次）") async def _finish_unexpected(run: _ZipRun, exc: Exception) -> None: """非预期异常：原文 logger 先行（结算本身可能再抛，先落下 traceback）， 用户可见文案泛化；带上退款结果再补一条，避免排障时只看到「异常」。""" job, tally = run.job, run.tally api_service.logger.error(f"{job.tag}异常: {run.task_id} {exc}\n{traceback.format_exc()}") delivered = tally.ok if run.zip_valid else 0 refunded = await run.settle(run.task_id, delivered, 'error') if delivered > 0 and run.zip_tmp_path: await _apply_error(run.task_id, job.partial_msg(refunded), refunded) with DOWNLOAD_TASKS_LOCK: t = DOWNLOAD_TASKS.get(run.task_id) if t is not None: t['result'] = { 'zip_path': run.zip_tmp_path, 'filename': job.zip_filename(tally.ok), 'audio_files': [], } else: await _apply_error(run.task_id, '打包失败，请稍后重试' + ( f"（已退回 {refunded} 次下载次数）" if refunded else ''), refunded) _discard(run.zip_tmp_path) api_service.logger.error( f"{job.tag}异常: {run.task_id}（成功{tally.ok} 退回{refunded}次）") async def run_zip_job(job: ZipJob, task_id: str, *, settle: Callable[[str, int, str], Awaitable[int]] = _settle_batch_billing) -> None: """执行一次批量打包任务（在 worker 协程内调用，自身不向外抛异常）。 ``settle`` 由调用方注入，默认走 :func:`services._taskcore._settle_batch_billing`。 做成参数是为了让测试能在**模块级**替换结算函数（包装层在调用时读取模块全局， 替换即生效）。 ``zip_valid`` 在这里的等价物是 ``run.zip_valid``，它**在构建前必为假**—— ZIP 构建前的业务错误（参数缺失/无条目/补扣失败）会 early raise，收尾函数 据此不按 ``ok_count`` 交付。原先它是必须手工在函数头初始化的局部变量， 漏掉就让异常处理器 NameError 二次崩线程；收进 dataclass 后这个约束由构造 函数承担。 """ run = _ZipRun(job=job, task_id=task_id, settle=settle) try: await _start_zip_run(run) await _build_zip_archive(run) if run.tally.ok == 0: await _finish_all_failed(run) return await _publish_delivery(run) except DownloadCanceled: await _finish_canceled(run) except DownloadException as e: await _finish_business_error(run, e) except Exception as e: await _finish_unexpected(run, e) finally: # 回收信息预取任务，避免任务泄漏（未启动的预取任务一并取消） await job.cleanup() async def _apply_error(task_id: str, msg: str, refunded: int) -> None: """异常/业务失败的终态收敛：置 error（终态先到先得）、记账退款额。 走 :func:`_mark_task_terminal` 而非就地写任务字典——后者是播客侧的旧写法， 除了重复一遍「已终态不覆盖」判断外，还漏掉了 worker 已结束未结算时按账本 补结算的兜底（``/F-E``）。 """ await _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 async def _fill_zip(job: ZipJob, items: list[Any], ctx: ZipContext, tally: _Tally, task: dict[str, Any], pause_ev) -> None: """逐窗口下载并写入 ZIP（返回时中央目录尚未写出，由调用方的 with 退出时落盘）。 窗口 = ``job.concurrency`` 个条目：下载并发、落包串行。窗口大小取 1 时 退化为顺序下载（播客整台即如此）。 取消语义与旧的线程池版本对齐：等待名额期间可被取消；窗口内某条抛 DownloadCanceled 时整批中止（在途下载因每次分块都查取消位而很快退出）。 """ concurrency = max(1, int(job.concurrency)) offset = 0 while offset < len(items): if ctx.cancel_ev.is_set(): raise DownloadCanceled("任务已取消") # 断线暂停：释放名额等待恢复/取消，恢复后重新进入排队 await _wait_if_paused(task, pause_ev, ctx.cancel_ev, tally.ok + len(tally.failed)) # 每个窗口持有 BATCH_ZIP_SEM（同一时刻全局最多 N 个批量任务在下载） # 等待名额期间可被取消，不再无限等待与延迟退款 if not await 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) await job.prefetch(window, ctx) # 只预取本窗口 results = await asyncio.gather( *(job.download_one(item, start_idx + i, ctx) for i, item in enumerate(window)), return_exceptions=True, ) for i, result in enumerate(results): if ctx.cancel_ev.is_set(): raise DownloadCanceled("任务已取消") item, idx = window[i], start_idx + i try: if isinstance(result, BaseException): raise result # 落包串行：ZipFile 非并发安全 + 包内名去重要按序 await job.write_entry(item, idx, result, 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()