diff --git a/utils/request.py b/utils/request.py index 97c9a16..ecfb64a 100644 --- a/utils/request.py +++ b/utils/request.py @@ -1,4 +1,5 @@ # -*- coding: utf-8 -*- +import asyncio from typing import * import aiohttp @@ -16,9 +17,21 @@ http_session: Optional[aiohttp.ClientSession] = None def init(): global http_session - http_session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=10)) + http_session = aiohttp.ClientSession( + response_class=CustomClientResponse, + timeout=aiohttp.ClientTimeout(total=10), + ) async def shut_down(): if http_session is not None: await http_session.close() + + +class CustomClientResponse(aiohttp.ClientResponse): + # 因为aiohttp的BUG,当底层连接断开时,_wait_released可能会抛出CancelledError,导致上层协程结束。这里改个错误类型 + async def _wait_released(self): + try: + return await super()._wait_released() + except asyncio.CancelledError as e: + raise aiohttp.ClientConnectionError('Connection released') from e