NeteaseUrl/core/http_client.py

# -*- coding: utf-8 -*-
"""异步 HTTP 客户端（aiohttp）：连接池 + 有界重试 + 响应体上限。

替代旧的长连接 `requests.Session`/裸 `requests.get/post`：
- 单一 `ClientSession` 复用 TCP 连接（keep-alive），连接池限制防上游打爆；
- 重试口径与旧实现一致：网络异常/5xx 退避重试一次，4xx（429/460 风控）**不**重试；
- 响应体超过 `max_bytes` 直接拒绝（防止被劫持/异常上游的超大 body 进内存）。

异常层次（调用方按需映射为业务异常，例如 music_api.APIException）：
    UpstreamError
      ├── UpstreamNetworkError  网络/超时（可重试）
      ├── UpstreamHTTPError     非 2xx（5xx 可重试，4xx 不可）
      └── UpstreamTooLarge      响应体超限
"""

import asyncio
import json
from collections.abc import AsyncIterator, Mapping
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from typing import Any

import aiohttp

DEFAULT_TIMEOUT = 30
DEFAULT_MAX_BYTES = 8 * 1024 * 1024

# UA 与 Referer 由调用方传入（music_api.APIConstants）
_EMPTY: Mapping[str, str] = {}


class UpstreamError(Exception):
    """上游请求错误基类。"""


class UpstreamNetworkError(UpstreamError):
    """网络异常/超时（可重试）。"""


class UpstreamHTTPError(UpstreamError):
    """上游返回非 2xx 状态码。"""

    def __init__(self, status: int, message: str = ''):
        super().__init__(f"HTTP {status} {message}".strip())
        self.status = status


class UpstreamTooLarge(UpstreamError):
    """响应体超过大小上限。"""


@dataclass
class HTTPResult:
    """一次请求的完整结果（供需要状态码/头/cookie 的调用点使用）。"""

    status: int
    url: str
    text: str
    headers: Mapping[str, str] = field(default_factory=dict)
    cookies: Mapping[str, str] = field(default_factory=dict)

    def json(self) -> Any:
        return json.loads(self.text)


class AsyncHTTPClient:
    """带连接池的异步 HTTP 客户端。"""

    def __init__(
        self,
        *,
        total_timeout: float = DEFAULT_TIMEOUT,
        max_bytes: int = DEFAULT_MAX_BYTES,
        pool_limit: int = 100,
        pool_per_host: int = 32,
        user_agent: str | None = None,
        referer: str | None = None,
    ):
        self._total_timeout = total_timeout
        self._max_bytes = max_bytes
        self._pool_limit = pool_limit
        self._pool_per_host = pool_per_host
        self._default_headers: dict[str, str] = {}
        if user_agent:
            self._default_headers['User-Agent'] = user_agent
        if referer:
            self._default_headers['Referer'] = referer
        self._session: aiohttp.ClientSession | None = None
        self._loop_sessions: dict[int, aiohttp.ClientSession] = {}
        self._lock = asyncio.Lock()

    async def _session_get(self) -> aiohttp.ClientSession:
        if self._session is not None and not self._session.closed:
            return self._session
        async with self._lock:
            if self._session is None or self._session.closed:
                connector = aiohttp.TCPConnector(
                    limit=self._pool_limit,
                    limit_per_host=self._pool_per_host,
                    ttl_dns_cache=300,
                )
                self._session = aiohttp.ClientSession(
                    connector=connector,
                    timeout=aiohttp.ClientTimeout(total=self._total_timeout),
                    headers=dict(self._default_headers),
                )
        return self._session

    async def _session_for_loop(self) -> aiohttp.ClientSession:
        """按事件循环隔离 session。

        同一进程内可能有两个 loop（应用主 loop + 同步桥的后台 loop）。
        aiohttp 的 ClientSession 不能跨 loop 复用，故按 loop 分别持有。
        """
        key = id(asyncio.get_running_loop())
        sess = self._loop_sessions.get(key)
        if sess is not None and not sess.closed:
            return sess
        async with self._lock:
            sess = self._loop_sessions.get(key)
            if sess is None or sess.closed:
                connector = aiohttp.TCPConnector(
                    limit=self._pool_limit,
                    limit_per_host=self._pool_per_host,
                    ttl_dns_cache=300,
                )
                sess = aiohttp.ClientSession(
                    connector=connector,
                    timeout=aiohttp.ClientTimeout(total=self._total_timeout),
                    headers=dict(self._default_headers),
                )
                self._loop_sessions[key] = sess
        return sess

    async def close(self) -> None:
        if self._session is not None and not self._session.closed:
            await self._session.close()
        self._session = None
        for sess in self._loop_sessions.values():
            if not sess.closed:
                await sess.close()
        self._loop_sessions.clear()

    # ---------- 底层：单次请求（不重试） ----------

    async def _once(
        self,
        method: str,
        url: str,
        *,
        headers: Mapping[str, str] | None = None,
        cookies: Mapping[str, str] | None = None,
        params: Mapping[str, Any] | None = None,
        data: Any = None,
        json_body: Any = None,
        timeout: float | None = None,
        allow_redirects: bool = True,
    ) -> HTTPResult:
        session = await self._session_for_loop()
        req_timeout = aiohttp.ClientTimeout(total=timeout) if timeout else None
        try:
            async with session.request(
                method,
                url,
                headers=dict(headers) if headers else None,
                cookies=dict(cookies) if cookies else None,
                params=params,
                data=data,
                json=json_body,
                timeout=req_timeout,
                allow_redirects=allow_redirects,
            ) as resp:
                body = await resp.read()
                if len(body) > self._max_bytes:
                    raise UpstreamTooLarge(f"上游响应体超过上限 {self._max_bytes} 字节")
                if resp.status >= 400:
                    raise UpstreamHTTPError(resp.status, resp.reason or '')
                return HTTPResult(
                    status=resp.status,
                    url=str(resp.url),
                    text=body.decode(resp.charset or 'utf-8', errors='replace'),
                    headers=dict(resp.headers),
                    cookies={k: v.value for k, v in resp.cookies.items()},
                )
        except aiohttp.ClientError as e:
            raise UpstreamNetworkError(str(e)) from e
        except TimeoutError as e:
            raise UpstreamNetworkError('请求超时') from e

    async def _with_retry(self, fn, *, attempts: int = 2, backoff: float = 1.0) -> HTTPResult:
        """网络异常/5xx 退避重试；4xx 与超限立即抛出不重试。"""
        last: UpstreamError | None = None
        for i in range(attempts):
            try:
                return await fn()
            except UpstreamHTTPError as e:
                if e.status >= 500 and i + 1 < attempts:
                    last = e
                    await asyncio.sleep(backoff)
                    continue
                raise
            except UpstreamNetworkError as e:
                last = e
                if i + 1 < attempts:
                    await asyncio.sleep(backoff)
                    continue
                raise
        # 理论上不可达
        raise last or UpstreamError('上游请求失败')

    # ---------- 高层 API ----------

    async def request(self, method: str, url: str, **kw) -> HTTPResult:
        return await self._with_retry(lambda: self._once(method, url, **kw))

    async def get_text(self, url: str, **kw) -> str:
        return (await self.request('GET', url, **kw)).text

    async def post_text(self, url: str, **kw) -> str:
        return (await self.request('POST', url, **kw)).text

    async def get_json(self, url: str, **kw) -> Any:
        return (await self.request('GET', url, **kw)).json()

    async def post_json(self, url: str, **kw) -> Any:
        return (await self.request('POST', url, **kw)).json()

    @asynccontextmanager
    async def stream(
        self,
        method: str,
        url: str,
        *,
        headers: Mapping[str, str] | None = None,
        cookies: Mapping[str, str] | None = None,
        timeout: float | None = None,
        allow_redirects: bool = True,
    ) -> AsyncIterator[aiohttp.ClientResponse]:
        """流式请求（下载/音视频代理用）：yield 原始响应，边读边写不整段入内存。"""
        session = await self._session_for_loop()
        req_timeout = aiohttp.ClientTimeout(total=timeout) if timeout else None
        try:
            async with session.request(
                method,
                url,
                headers=dict(headers) if headers else None,
                cookies=dict(cookies) if cookies else None,
                timeout=req_timeout,
                allow_redirects=allow_redirects,
            ) as resp:
                if resp.status >= 400:
                    raise UpstreamHTTPError(resp.status, resp.reason or '')
                yield resp
        except aiohttp.ClientError as e:
            raise UpstreamNetworkError(str(e)) from e

    async def open(
        self,
        method: str,
        url: str,
        *,
        headers: Mapping[str, str] | None = None,
        cookies: Mapping[str, str] | None = None,
        timeout: float | None = None,
        allow_redirects: bool = True,
    ) -> aiohttp.ClientResponse:
        """打开流式响应并返回**未关闭**的 ClientResponse（调用方负责 close）。

        与 stream() 的区别：不把响应交给 `async with`，便于「先读状态码/头决定是否
        转发，再流式读取」的代理场景（如音频 Range 代理）。
        """
        session = await self._session_for_loop()
        req_timeout = aiohttp.ClientTimeout(total=timeout) if timeout else None
        try:
            return await session.request(
                method,
                url,
                headers=dict(headers) if headers else None,
                cookies=dict(cookies) if cookies else None,
                timeout=req_timeout,
                allow_redirects=allow_redirects,
            )
        except aiohttp.ClientError as e:
            raise UpstreamNetworkError(str(e)) from e


# 进程级共享客户端（UA/Referer 由 music_api 在初始化后注入）
client = AsyncHTTPClient()