NeteaseUrl/music_downloader/_aria2.py

# -*- coding: utf-8 -*-
"""aria2c 常驻 RPC 下载后端（可选加速通道；不可用时调用方降级）。

为什么用常驻 RPC 而不是每个任务起一个 aria2c 进程：
- 批量任务最多 500 首，逐首歌都付一次进程启动 + RPC 握手开销不划算；
- 取消/查询只需一次 RPC，不必在进程表里追踪 N 个子进程。

安全与稳健：
- RPC 只监听 127.0.0.1 并强制随机 secret —— 无鉴权的 aria2 RPC 允许本机任意
  进程操纵下载队列（读写任意文件），暴露到容器外等于开了一个后门；
- aria2c 缺失 / 进程早退 / RPC 未就绪，一律返回 None 静默降级，
  由调用方回落到内置的 core.http_client 流式通道；
- aria2 没有「单文件体积上限」选项，故用轮询 tellStatus 实现体积守卫与取消响应。

异步口径：aria2p 客户端本身是同步阻塞实现（内部走 urllib），故每次 RPC 调用
（提交/查询/移除）都经 ``asyncio.to_thread`` 卸载；轮询间隔改 ``asyncio.sleep``
让出事件循环。进程启动（``_try_start``）最长阻塞 8s，调用方同样经 to_thread 触发。
"""

from __future__ import annotations

import asyncio
import atexit
import secrets
import shutil
import socket
import subprocess
import threading
import time
from collections.abc import Callable
from pathlib import Path
from typing import Any, Protocol

from core.http_client import UpstreamNetworkError

from ._http import DEFAULT_UA
from ._log import logger
from .models import DownloadCanceled, DownloadException


class _ProcessHandle(Protocol):
    """aria2c 子进程的最小接口。

    只依赖本模块真正调用的方法，既避免绑死 subprocess.Popen 具体类型，
    也让测试可以用轻量替身驱动生命周期逻辑。
    """

    def poll(self) -> int | None: ...
    def terminate(self) -> None: ...
    def kill(self) -> None: ...
    def wait(self, timeout: float | None = None) -> int: ...

# aria2 可执行文件名。Docker 镜像内由 apk 提供；宿主未装则整条通道降级
ARIA2_BIN = 'aria2c'

# RPC 就绪等待上限：超过即判定启动失败并降级（不拖慢首次下载）
_RPC_READY_TIMEOUT = 8.0
# 状态轮询间隔：取消与体积守卫的响应粒度（0.3s 足够，RPC 开销可忽略）
_POLL_INTERVAL = 0.3
# 单文件体积守卫默认值，与 downloader 内置通道口径一致
_DEFAULT_MAX_BYTES = 500 * 1024 * 1024
# 分片数与单服务器连接数：网易云直链支持 Range，多连接可提速；
# 数值保守（aria2 默认即 5），过高易触发 CDN 风控
_SPLIT = 5
_MIN_SPLIT_SIZE = '1M'

_lock = threading.Lock()
_backend: Aria2Backend | None = None
_unavailable = False   # 启动失败后置位，不再反复重试


def _free_port() -> int:
    """向内核要一个空闲端口（bind 后立即释放，交给 aria2c 监听）。"""
    with socket.socket() as s:
        s.bind(('127.0.0.1', 0))
        return int(s.getsockname()[1])


class Aria2Backend:
    """aria2c RPC 后端。请通过 get_backend() 获取单例，不要直接构造。"""

    def __init__(self, api: Any, proc: _ProcessHandle, port: int) -> None:
        self._api = api
        self._proc = proc
        self.port = port
        self._closed = False

    # ---------- 生命周期 ----------

    def close(self) -> None:
        """关闭 RPC 客户端并终止 aria2c 进程（幂等）。"""
        if self._closed:
            return
        self._closed = True
        try:
            self._api.client.close()
        except Exception:
            pass   # 客户端关闭失败不影响后续进程回收
        if self._proc.poll() is None:
            self._proc.terminate()
            try:
                self._proc.wait(timeout=5)
            except subprocess.TimeoutExpired:
                self._proc.kill()
                self._proc.wait(timeout=5)

    # ---------- 下载 ----------

    async def fetch(
        self,
        url: str,
        target: Path,
        *,
        ua: str = DEFAULT_UA,
        referer: str | None = None,
        cancel_event: threading.Event | None = None,
        progress_cb: Callable[[int, int], None] | None = None,
        max_bytes: int = _DEFAULT_MAX_BYTES,
        timeout: float = 600.0,
    ) -> None:
        """把 url 下载到 target（落地即最终字节，含断点续传）。

        Raises:
            DownloadCanceled: cancel_event 置位（残片一并删除）
            DownloadException: 体积超限（残片一并删除）、aria2 报告任务失败、总超时
            UpstreamNetworkError: RPC 抖动等可重试的网络类失败（保留残片续传）
        """
        options = {
            'dir': str(target.parent),
            'out': target.name,
            'continue': 'true',             # 断点续传（控制文件 target.aria2 与残片同放）
            'auto-file-renaming': 'false',  # 禁止另存为 target.1，否则与 .part 语义冲突
            'allow-overwrite': 'true',
            'split': str(_SPLIT),
            'max-connection-per-server': str(_SPLIT),
            'min-split-size': _MIN_SPLIT_SIZE,
            'user-agent': ua,
            'summary-interval': '0',
            'quiet': 'true',
        }
        if referer:
            options['referer'] = referer

        try:
            download = await asyncio.to_thread(self._api.add_uris, [url], options=options)
        except Exception as e:
            raise UpstreamNetworkError(f'aria2 提交任务失败: {e}') from e

        deadline = time.monotonic() + timeout
        while True:
            if cancel_event is not None and cancel_event.is_set():
                await asyncio.to_thread(self._abort, download, drop_files=True)
                raise DownloadCanceled('下载已取消')
            try:
                await asyncio.to_thread(download.update)
            except Exception as e:
                # RPC 抖动按可重试的网络失败上报，交给外层退避重试
                raise UpstreamNetworkError(f'aria2 状态查询失败: {e}') from e

            done = int(download.completed_length or 0)
            total = int(download.total_length or 0)
            if progress_cb is not None:
                try:
                    progress_cb(done, total)
                except Exception:
                    pass   # 进度回调绝不影响下载主流程
            if done > max_bytes:
                await asyncio.to_thread(self._abort, download, drop_files=True)
                raise DownloadException(
                    f'文件超过大小上限 {max_bytes // (1024 * 1024)}MB，已中止')
            if download.is_complete:
                self._drop_control_file(target)
                return
            if download.has_failed:
                # 保留残片与控制文件，等待外层重试时续传
                raise UpstreamNetworkError(
                    f'aria2 下载失败: {download.error_message or download.status}')
            if time.monotonic() > deadline:
                await asyncio.to_thread(self._abort, download, drop_files=False)
                raise UpstreamNetworkError('aria2 下载超时')
            await asyncio.sleep(_POLL_INTERVAL)

    # ---------- 内部 ----------

    def _abort(self, download: Any, *, drop_files: bool) -> None:
        """强制移除任务。drop_files=True 时连同已下载字节一起删除。"""
        try:
            self._api.remove([download], force=True, files=drop_files)
        except Exception as e:
            logger.warning(f'终止 aria2 任务失败（不影响主流程）: {e}')

    @staticmethod
    def _drop_control_file(target: Path) -> None:
        """清理 aria2 控制文件（正常情况下 aria2 已自行删除，这里兜底）。"""
        ctrl = target.with_name(target.name + '.aria2')
        try:
            ctrl.unlink(missing_ok=True)
        except OSError:
            pass


def get_backend() -> Aria2Backend | None:
    """返回 aria2 后端单例；aria2c 不可用/启动失败时返回 None（调用方降级）。

    首次调用尝试拉起常驻进程（最长阻塞 _RPC_READY_TIMEOUT 秒），之后命中缓存；
    失败结果进程内记忆，不反复重试。**同步阻塞**——异步调用方须经
    ``asyncio.to_thread(get_backend)`` 触发，避免首次启动期间冻住事件循环。
    """
    global _backend, _unavailable
    if _unavailable:
        return None
    if _backend is not None:
        return _backend
    with _lock:
        if _unavailable:
            return None
        if _backend is None:
            _backend = _try_start()
            if _backend is None:
                _unavailable = True
    return _backend


def _try_start() -> Aria2Backend | None:
    """拉起 aria2c --enable-rpc 并等待 RPC 就绪；任何失败都返回 None。"""
    exe = shutil.which(ARIA2_BIN)
    if not exe:
        logger.info('未检测到 aria2c，下载走内置流式通道')
        return None
    try:
        import aria2p
    except ImportError as e:
        logger.info(f'aria2p 不可用，下载走内置流式通道: {e}')
        return None

    port = _free_port()
    secret = secrets.token_urlsafe(24)
    cmd = [
        exe,
        '--enable-rpc',
        f'--rpc-listen-port={port}',
        f'--rpc-secret={secret}',
        '--rpc-listen-all=false',
        '--rpc-allow-origin-all=false',
        '--no-conf',
        '--quiet=true',
    ]
    try:
        # exe 来自 shutil.which，参数全部由本模块构造，不经任何用户输入
        proc = subprocess.Popen(
            cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    except OSError as e:
        logger.warning(f'aria2c 启动失败，降级内置下载: {e}')
        return None

    client = aria2p.Client(host='http://127.0.0.1', port=port, secret=secret)
    deadline = time.monotonic() + _RPC_READY_TIMEOUT
    while True:
        if proc.poll() is not None:
            logger.warning(f'aria2c 进程提前退出（code={proc.returncode}），降级内置下载')
            return None
        try:
            client.get_version()
            break
        except Exception:
            if time.monotonic() > deadline:
                logger.warning(f'aria2c RPC {_RPC_READY_TIMEOUT:.0f}s 内未就绪，降级内置下载')
                proc.kill()
                proc.wait(timeout=5)
                return None
            time.sleep(0.2)

    atexit.register(_shutdown)
    logger.info(f'aria2c 常驻 RPC 已就绪（127.0.0.1:{port}，{_SPLIT} 连接/文件）')
    return Aria2Backend(aria2p.API(client), proc, port)


def _shutdown() -> None:
    """进程退出时回收 aria2c（atexit 只注册一次）。"""
    global _backend
    if _backend is not None:
        _backend.close()
        _backend = None