# -*- 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`` 两个钩子， 只有前者进线程池。 ``run_zip_job`` 的四个终态（成功 / 全失败 / 取消 / 异常）原先各写一段 8~20 行的 收尾内联在 try/except 里。此处把「一次运行」的状态收进 :class:`_ZipRun`， 四个终态各自成函数——改账本口径或补一种终态时，不必再在 167 行里找落点。 """ 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 @dataclass class _ZipRun: """一次 ``run_zip_job`` 的共享状态。 四个终态收尾函数都要读同一批东西（任务 id、结算入口、计数、临时 ZIP 路径、 ZIP 是否完整）。原先它们是散在 167 行函数体里的局部变量，被各 except 分支 分别访问——收成一个对象后，收尾函数不必逐个透传参数，也让「谁真的用到了 哪些状态」一眼可见。 """ job: ZipJob task_id: str settle: Callable[[str, int, str], 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: threading.Event | None = None pause_ev: threading.Event | None = 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 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 = run.job.prepare(run.params, run.task_id, run.tally.failed) with DOWNLOAD_TASKS_LOCK: DOWNLOAD_TASKS[run.task_id]['total'] = len(run.items) run.job.supplement(run.task_id, run.params, len(run.items)) 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)) _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) 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 = run.settle(run.task_id, 0, 'error') _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}次") def _publish_delivery(run: _ZipRun) -> None: """正常交付：按入包成功数结算，锁内复查取消/终态后发布结果。 ``'settled'`` 是 ``_settle_batch_billing`` 的默认值，此处显式写出，让 「成功/取消/失败」三种账本终态在本文件里一目了然。 """ job, tally = run.job, run.tally assert run.cancel_ev is not None refunded = 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': [], # 缓存文件不归任务所有，不进入清理清单 } 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})") 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 = run.settle(run.task_id, delivered, 'canceled') if delivered > 0 and run.zip_tmp_path: _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) _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}次）") 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 = run.settle(run.task_id, delivered, 'error') _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}次）") 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 = run.settle(run.task_id, delivered, 'error') if delivered > 0 and run.zip_tmp_path: _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: _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}次）") 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_valid`` 在这里的等价物是 ``run.zip_valid``，它**在构建前必为假**—— ZIP 构建前的业务错误（参数缺失/无条目/补扣失败）会 early raise，收尾函数 据此不按 ``ok_count`` 交付。原先它是必须手工在函数头初始化的局部变量， 漏掉就让异常处理器 NameError 二次崩线程；收进 dataclass 后这个约束由构造 函数承担。 """ run = _ZipRun(job=job, task_id=task_id, settle=settle) try: _start_zip_run(run) _build_zip_archive(run) if run.tally.ok == 0: _finish_all_failed(run) return _publish_delivery(run) except DownloadCanceled: _finish_canceled(run) except DownloadException as e: _finish_business_error(run, e) except Exception as e: _finish_unexpected(run, e) 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)