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 实现体积守卫与取消响应。

异步口径：RPC 由本模块直接打 aria2 的 JSON-RPC（HTTP POST + JSON-RPC 2.0），
全程 await——原先依赖的 aria2p 内部是 requests + websocket 的同步实现，每次
RPC 都要经 asyncio.to_thread 卸载，是本项目最后一处「把同步 HTTP 藏在 to_thread
里」的地方。轮询间隔用 asyncio.sleep 让出事件循环。剩下的 to_thread 只有两处，
都是真正的阻塞卸载：进程创建（subprocess.Popen）与强杀回收（kill + wait）。
"""

from __future__ import annotations

import asyncio
import atexit
import contextlib
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 core.http_client import client as _http

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


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

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

    returncode: int | None

    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'

# 单次 RPC 超时：本机回环，正常在毫秒级；给足余量以免大库时误判网络失败
_RPC_TIMEOUT = 10.0

# tellStatus 返回的**终态**：任务已不在下载队列，只能清结果表
_TERMINAL_STATUSES = frozenset({'complete', 'error', 'removed'})

_backend: Aria2Backend | None = None
_unavailable = False   # 启动失败后置位，不再反复重试


# 启动单飞锁按事件循环隔离：asyncio.Lock 绑定创建它的 loop，应用主 loop 与测试
# 各自的 loop 并存，跨 loop 复用会抛 "bound to a different event loop"。
_START_LOCKS: dict[int, asyncio.Lock] = {}


def _start_lock() -> asyncio.Lock:
    key = id(asyncio.get_running_loop())
    lock = _START_LOCKS.get(key)
    if lock is None:
        lock = asyncio.Lock()
        _START_LOCKS[key] = lock
    return lock


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


class Aria2RpcError(RuntimeError):
    """aria2 在 JSON-RPC 响应体里回报的 error 对象（协议层失败，非 HTTP 层）。"""


class _Aria2Rpc:
    """aria2 JSON-RPC 极简客户端（只覆盖本模块用到的 5 个方法）。

    替代 aria2p：后者内部是 requests + websocket 的同步实现，每个 RPC 都得
    to_thread 卸载；而 aria2 的 RPC 就是朴素的 HTTP POST + JSON-RPC 2.0，
    直接复用 core.http_client 的进程级连接池即可，且连接池本就为异步而生。

    一律走 ``request_once``（**不重试**），与 aria2p 口径一致：addUri /
    forceRemove 这类变更调用若被重发可能重复提交，RPC 抖动由上层 ``fetch``
    收敛成 UpstreamNetworkError、交给下载器的退避重试统一处理。
    """

    def __init__(self, port: int, secret: str) -> None:
        self._url = f'http://127.0.0.1:{port}/jsonrpc'
        self._token = f'token:{secret}'
        self._seq = 0

    async def _call(self, method: str, *params: Any) -> Any:
        self._seq += 1
        body = {
            'jsonrpc': '2.0',
            'id': str(self._seq),
            'method': f'aria2.{method}',
            'params': [self._token, *params],
        }
        resp = await _http.request_once(
            'POST', self._url, json_body=body, timeout=_RPC_TIMEOUT)
        data = resp.json()
        if not isinstance(data, dict):
            raise Aria2RpcError(f'aria2 RPC 响应不是 JSON 对象: {resp.text[:200]!r}')
        err = data.get('error')
        if err:
            raise Aria2RpcError(f'aria2 RPC {method} 失败: {err}')
        return data.get('result')

    async def get_version(self) -> str:
        """探活（启动时就绪判定）。"""
        return str(await self._call('getVersion'))

    async def add_uri(self, url: str, options: dict[str, str]) -> str:
        """提交下载，返回 GID。"""
        return str(await self._call('addUri', [url], options))

    async def tell_status(self, gid: str) -> dict[str, Any]:
        """查询任务状态（不传 keys，取全量字段）。"""
        res = await self._call('tellStatus', gid)
        return res if isinstance(res, dict) else {}

    async def force_remove(self, gid: str) -> str:
        """停掉在途任务（不做联系 tracker 等耗时动作）。"""
        return str(await self._call('forceRemove', gid))

    async def remove_download_result(self, gid: str) -> str:
        """清掉已终结任务在结果表里的条目（否则结果表只增不减）。"""
        return str(await self._call('removeDownloadResult', gid))


class _Rpc(Protocol):
    """aria2 RPC 的最小接口。

    与 _ProcessHandle 同理：只声明本模块真正调用的方法，既不绑死具体实现，
    也让测试能用轻量替身驱动整套状态机（不必起 HTTP 假服务器）。
    """

    async def get_version(self) -> str: ...
    async def add_uri(self, url: str, options: dict[str, str]) -> str: ...
    async def tell_status(self, gid: str) -> dict[str, Any]: ...
    async def force_remove(self, gid: str) -> str: ...
    async def remove_download_result(self, gid: str) -> str: ...


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

    def __init__(self, rpc: _Rpc, proc: _ProcessHandle, port: int) -> None:
        self._rpc = rpc
        self._proc = proc
        self.port = port
        self._closed = False

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

    def close(self) -> None:
        """终止 aria2c 进程（幂等）。

        **同步**：由 atexit 调用，那时已经没有事件循环，且进程回收本就是阻塞操作。
        这里不关 RPC 连接——它走 core.http_client 的**进程级共享连接池**，关掉会
        连带影响内置流式通道（原先 aria2p 关的是它自己私有的 requests 会话）。
        """
        if self._closed:
            return
        self._closed = True
        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:
            gid = await self._rpc.add_uri(url, options)
        except Exception as e:
            raise UpstreamNetworkError(f'aria2 提交任务失败: {e}') from e

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

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

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

    async def _abort(
        self, gid: str, target: Path, *, drop_files: bool, status: str,
    ) -> None:
        """移除任务。与 ``aria2p.remove(force=True, files=drop_files, clean=True)``
        逐条对齐（该库的默认行为就是还要清控制文件）：

        - 已终结（complete/error/removed）→ 只能清结果表（removeDownloadResult）；
          在途 → forceRemove 先停掉，再顺手清结果表；
        - 控制文件（target.aria2）**总是**删除——这是 aria2p 的 clean=True 默认值，
          也是「取消 / 体积超限」后下轮必须重新下载的依据；
        - 只有 drop_files 才删已落地的数据文件（超时场景保留残片供续传）。
        任一步失败只告警、不影响主流程（与 aria2p 的 logger.exception 同口径），
        且与 aria2p 一致：主移除动作失败时**不**删数据文件。
        """
        removed = True
        try:
            if status in _TERMINAL_STATUSES:
                await self._rpc.remove_download_result(gid)
            else:
                await self._rpc.force_remove(gid)
                with contextlib.suppress(Exception):
                    await self._rpc.remove_download_result(gid)
        except Exception as e:
            removed = False
            logger.warning(f'终止 aria2 任务失败（不影响主流程）: {e}')
        self._drop_control_file(target)
        if drop_files and removed:
            try:
                target.unlink(missing_ok=True)
            except OSError 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


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

    首次调用尝试拉起常驻进程并等 RPC 就绪（最长 ``_RPC_READY_TIMEOUT`` 秒，
    等待用 ``await asyncio.sleep`` 让出事件循环，故本函数是**原生异步**、
    不需要调用方再包 to_thread）；之后命中缓存。失败结果进程内记忆，
    不反复重试。并发首调由启动单飞锁收敛，只拉起一个进程。
    """
    global _backend, _unavailable
    if _unavailable:
        return None
    if _backend is not None:
        return _backend
    async with _start_lock():
        if _unavailable:
            return None
        if _backend is None:
            _backend = await _try_start()
            if _backend is None:
                _unavailable = True
    return _backend


def _reap(proc: _ProcessHandle) -> None:
    """强杀并回收子进程（阻塞）。只应由 asyncio.to_thread 调用。"""
    try:
        proc.kill()
    except OSError:
        return
    with contextlib.suppress(Exception):
        proc.wait(timeout=5)


async def _try_start() -> Aria2Backend | None:
    """拉起 aria2c --enable-rpc 并等待 RPC 就绪；任何失败都返回 None。"""
    exe = shutil.which(ARIA2_BIN)
    if not exe:
        logger.info('未检测到 aria2c，下载走内置流式通道')
        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:
        # 进程创建是真阻塞调用（Windows CreateProcess 数十毫秒），
        # 且本函数在事件循环内被 await，故经 to_thread 卸载（真正的阻塞卸载）。
        # exe 来自 shutil.which，参数全部由本模块构造，不经任何用户输入
        proc = await asyncio.to_thread(
            subprocess.Popen, cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    except OSError as e:
        logger.warning(f'aria2c 启动失败，降级内置下载: {e}')
        return None

    rpc = _Aria2Rpc(port, 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:
            await rpc.get_version()
            break
        except Exception:
            if time.monotonic() > deadline:
                logger.warning(f'aria2c RPC {_RPC_READY_TIMEOUT:.0f}s 内未就绪，降级内置下载')
                await asyncio.to_thread(_reap, proc)
                return None
            await asyncio.sleep(0.2)

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


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