Skip to content

Commit 0e6bbbc

Browse files
bdracoarcivanov
andauthored
[PR #8546/a561fa99 backport][3.10] Fix WebSocket server heartbeat timeout logic (#8573)
Co-authored-by: J. Nick Koston <[email protected]> Co-authored-by: Arcadiy Ivanov <[email protected]>
1 parent bf5a66f commit 0e6bbbc

File tree

3 files changed

+43
-4
lines changed

3 files changed

+43
-4
lines changed

CHANGES/8540.bugfix.rst

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
Fixed WebSocket server heartbeat timeout logic to terminate `receive` and return :py:class:`~aiohttp.ServerTimeoutError` -- by :user:`arcivanov`.
2+
3+
When a WebSocket pong message was not received, the
4+
:py:meth:`~aiohttp.ClientWebSocketResponse.receive` operation did not terminate.
5+
This change causes `_pong_not_received` to feed the `reader` an error message, causing
6+
pending `receive` to terminate and return the error message. The error message contains
7+
the exception :py:class:`~aiohttp.ServerTimeoutError`.

aiohttp/client_ws.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
import sys
55
from typing import Any, Optional, cast
66

7-
from .client_exceptions import ClientError
7+
from .client_exceptions import ClientError, ServerTimeoutError
88
from .client_reqrep import ClientResponse
99
from .helpers import call_later, set_result
1010
from .http import (
@@ -122,8 +122,12 @@ def _pong_not_received(self) -> None:
122122
if not self._closed:
123123
self._closed = True
124124
self._close_code = WSCloseCode.ABNORMAL_CLOSURE
125-
self._exception = asyncio.TimeoutError()
125+
self._exception = ServerTimeoutError()
126126
self._response.close()
127+
if self._waiting and not self._closing:
128+
self._reader.feed_data(
129+
WSMessage(WSMsgType.ERROR, self._exception, None)
130+
)
127131

128132
@property
129133
def closed(self) -> bool:

tests/test_client_ws_functional.py

Lines changed: 30 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import pytest
66

77
import aiohttp
8-
from aiohttp import hdrs, web
8+
from aiohttp import ServerTimeoutError, WSMsgType, hdrs, web
99
from aiohttp.http import WSCloseCode
1010
from aiohttp.pytest_plugin import AiohttpClient
1111

@@ -624,7 +624,35 @@ async def handler(request):
624624
assert resp.close_code is WSCloseCode.ABNORMAL_CLOSURE
625625

626626

627-
async def test_send_recv_compress(aiohttp_client) -> None:
627+
async def test_heartbeat_no_pong_concurrent_receive(aiohttp_client: Any) -> None:
628+
ping_received = False
629+
630+
async def handler(request):
631+
nonlocal ping_received
632+
ws = web.WebSocketResponse(autoping=False)
633+
await ws.prepare(request)
634+
msg = await ws.receive()
635+
ping_received = msg.type is aiohttp.WSMsgType.PING
636+
ws._reader.feed_eof = lambda: None
637+
await asyncio.sleep(10.0)
638+
639+
app = web.Application()
640+
app.router.add_route("GET", "/", handler)
641+
642+
client = await aiohttp_client(app)
643+
resp = await client.ws_connect("/", heartbeat=0.1)
644+
resp._reader.feed_eof = lambda: None
645+
646+
# Connection should be closed roughly after 1.5x heartbeat.
647+
msg = await resp.receive(5.0)
648+
assert ping_received
649+
assert resp.close_code is WSCloseCode.ABNORMAL_CLOSURE
650+
assert msg
651+
assert msg.type is WSMsgType.ERROR
652+
assert isinstance(msg.data, ServerTimeoutError)
653+
654+
655+
async def test_send_recv_compress(aiohttp_client: Any) -> None:
628656
async def handler(request):
629657
ws = web.WebSocketResponse()
630658
await ws.prepare(request)

0 commit comments

Comments
 (0)