mirror of
https://github.com/nonebot/nonebot2.git
synced 2026-04-17 06:07:55 +00:00
🐛 Fix: 修正 http/websocket client timeout (#3923)
Co-authored-by: Ju4tCode <42488585+yanyongyu@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
@@ -38,6 +38,7 @@ from nonebot.drivers import WebSocket as BaseWebSocket
|
||||
from nonebot.drivers.none import Driver as NoneDriver
|
||||
from nonebot.exception import WebSocketClosed
|
||||
from nonebot.internal.driver import (
|
||||
DEFAULT_TIMEOUT,
|
||||
Cookies,
|
||||
CookieTypes,
|
||||
HeaderTypes,
|
||||
@@ -45,6 +46,7 @@ from nonebot.internal.driver import (
|
||||
Timeout,
|
||||
TimeoutTypes,
|
||||
)
|
||||
from nonebot.utils import UNSET, UnsetType, exclude_unset
|
||||
|
||||
try:
|
||||
import aiohttp
|
||||
@@ -63,7 +65,7 @@ class Session(HTTPClientSession):
|
||||
headers: HeaderTypes = None,
|
||||
cookies: CookieTypes = None,
|
||||
version: str | HTTPVersion = HTTPVersion.H11,
|
||||
timeout: TimeoutTypes = None,
|
||||
timeout: TimeoutTypes | UnsetType = UNSET,
|
||||
proxy: str | None = None,
|
||||
):
|
||||
self._client: aiohttp.ClientSession | None = None
|
||||
@@ -85,15 +87,32 @@ class Session(HTTPClientSession):
|
||||
else:
|
||||
raise RuntimeError(f"Unsupported HTTP version: {version}")
|
||||
|
||||
_timeout = None
|
||||
if isinstance(timeout, Timeout):
|
||||
self._timeout = aiohttp.ClientTimeout(
|
||||
total=timeout.total,
|
||||
connect=timeout.connect,
|
||||
sock_read=timeout.read,
|
||||
timeout_kwargs: dict[str, float | None] = exclude_unset(
|
||||
{
|
||||
"total": timeout.total,
|
||||
"connect": timeout.connect,
|
||||
"sock_read": timeout.read,
|
||||
}
|
||||
)
|
||||
else:
|
||||
self._timeout = aiohttp.ClientTimeout(timeout)
|
||||
if timeout_kwargs:
|
||||
_timeout = aiohttp.ClientTimeout(**timeout_kwargs) # type: ignore
|
||||
elif timeout is not UNSET:
|
||||
_timeout = aiohttp.ClientTimeout(connect=timeout, sock_read=timeout)
|
||||
|
||||
if _timeout is None:
|
||||
_timeout = aiohttp.ClientTimeout(
|
||||
**exclude_unset(
|
||||
{
|
||||
"total": DEFAULT_TIMEOUT.total,
|
||||
"connect": DEFAULT_TIMEOUT.connect,
|
||||
"sock_read": DEFAULT_TIMEOUT.read,
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
self._timeout = _timeout
|
||||
self._proxy = proxy
|
||||
|
||||
@property
|
||||
@@ -102,6 +121,25 @@ class Session(HTTPClientSession):
|
||||
raise RuntimeError("Session is not initialized")
|
||||
return self._client
|
||||
|
||||
def _get_timeout(self, timeout: TimeoutTypes | UnsetType) -> aiohttp.ClientTimeout:
|
||||
_timeout = None
|
||||
if isinstance(timeout, Timeout):
|
||||
timeout_kwargs: dict[str, float | None] = exclude_unset(
|
||||
{
|
||||
"total": timeout.total,
|
||||
"connect": timeout.connect,
|
||||
"sock_read": timeout.read,
|
||||
}
|
||||
)
|
||||
if timeout_kwargs:
|
||||
_timeout = aiohttp.ClientTimeout(**timeout_kwargs) # type: ignore
|
||||
elif timeout is not UNSET:
|
||||
_timeout = aiohttp.ClientTimeout(connect=timeout, sock_read=timeout)
|
||||
|
||||
if _timeout is None:
|
||||
return self._timeout
|
||||
return _timeout
|
||||
|
||||
@override
|
||||
async def request(self, setup: Request) -> Response:
|
||||
if self._params:
|
||||
@@ -121,15 +159,6 @@ class Session(HTTPClientSession):
|
||||
if cookie.value is not None
|
||||
)
|
||||
|
||||
if isinstance(setup.timeout, Timeout):
|
||||
timeout = aiohttp.ClientTimeout(
|
||||
total=setup.timeout.total,
|
||||
connect=setup.timeout.connect,
|
||||
sock_read=setup.timeout.read,
|
||||
)
|
||||
else:
|
||||
timeout = aiohttp.ClientTimeout(setup.timeout)
|
||||
|
||||
async with await self.client.request(
|
||||
setup.method,
|
||||
url,
|
||||
@@ -138,7 +167,7 @@ class Session(HTTPClientSession):
|
||||
cookies=cookies,
|
||||
headers=setup.headers,
|
||||
proxy=setup.proxy or self._proxy,
|
||||
timeout=timeout,
|
||||
timeout=self._get_timeout(setup.timeout),
|
||||
) as response:
|
||||
return Response(
|
||||
response.status,
|
||||
@@ -171,15 +200,6 @@ class Session(HTTPClientSession):
|
||||
if cookie.value is not None
|
||||
)
|
||||
|
||||
if isinstance(setup.timeout, Timeout):
|
||||
timeout = aiohttp.ClientTimeout(
|
||||
total=setup.timeout.total,
|
||||
connect=setup.timeout.connect,
|
||||
sock_read=setup.timeout.read,
|
||||
)
|
||||
else:
|
||||
timeout = aiohttp.ClientTimeout(setup.timeout)
|
||||
|
||||
async with self.client.request(
|
||||
setup.method,
|
||||
url,
|
||||
@@ -188,7 +208,7 @@ class Session(HTTPClientSession):
|
||||
cookies=cookies,
|
||||
headers=setup.headers,
|
||||
proxy=setup.proxy or self._proxy,
|
||||
timeout=timeout,
|
||||
timeout=self._get_timeout(setup.timeout),
|
||||
) as response:
|
||||
response_headers = response.headers.copy()
|
||||
# aiohttp does not guarantee fixed-size chunks; re-chunk to exact size
|
||||
@@ -270,13 +290,39 @@ class Mixin(HTTPClientMixin, WebSocketClientMixin):
|
||||
else:
|
||||
raise RuntimeError(f"Unsupported HTTP version: {setup.version}")
|
||||
|
||||
timeout = None
|
||||
if isinstance(setup.timeout, Timeout):
|
||||
timeout = aiohttp.ClientWSTimeout(
|
||||
ws_receive=setup.timeout.read, # type: ignore
|
||||
ws_close=setup.timeout.total, # type: ignore
|
||||
timeout_kwargs: dict[str, float | None] = exclude_unset(
|
||||
{
|
||||
"ws_receive": setup.timeout.read,
|
||||
"ws_close": (
|
||||
setup.timeout.total
|
||||
if setup.timeout.close is UNSET
|
||||
else setup.timeout.close
|
||||
),
|
||||
}
|
||||
)
|
||||
if timeout_kwargs:
|
||||
timeout = aiohttp.ClientWSTimeout(**timeout_kwargs)
|
||||
elif setup.timeout is not UNSET:
|
||||
timeout = aiohttp.ClientWSTimeout(
|
||||
ws_receive=setup.timeout, # type: ignore
|
||||
ws_close=setup.timeout, # type: ignore
|
||||
)
|
||||
|
||||
if timeout is None:
|
||||
timeout = aiohttp.ClientWSTimeout(
|
||||
**exclude_unset(
|
||||
{
|
||||
"ws_receive": DEFAULT_TIMEOUT.read,
|
||||
"ws_close": (
|
||||
DEFAULT_TIMEOUT.total
|
||||
if DEFAULT_TIMEOUT.close is UNSET
|
||||
else DEFAULT_TIMEOUT.close
|
||||
),
|
||||
}
|
||||
)
|
||||
)
|
||||
else:
|
||||
timeout = aiohttp.ClientWSTimeout(ws_close=setup.timeout or 10.0) # type: ignore
|
||||
|
||||
async with aiohttp.ClientSession(version=version, trust_env=True) as session:
|
||||
async with session.ws_connect(
|
||||
@@ -295,7 +341,7 @@ class Mixin(HTTPClientMixin, WebSocketClientMixin):
|
||||
headers: HeaderTypes = None,
|
||||
cookies: CookieTypes = None,
|
||||
version: str | HTTPVersion = HTTPVersion.H11,
|
||||
timeout: TimeoutTypes = None,
|
||||
timeout: TimeoutTypes | UnsetType = UNSET,
|
||||
proxy: str | None = None,
|
||||
) -> Session:
|
||||
return Session(
|
||||
|
||||
Reference in New Issue
Block a user