# -*- 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 # 第 95 条（P2-6）/第 98 条口径：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（FastAPI 主 loop + 遗留 Flask 的同步桥 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()