Skip to content

Commit

Permalink
[PR #8546/a561fa99 backport][3.10] Fix WebSocket server heartbeat tim…
Browse files Browse the repository at this point in the history
…eout logic (#8573)

Co-authored-by: J. Nick Koston <nick@koston.org>
Co-authored-by: Arcadiy Ivanov <arcadiy@ivanov.biz>
  • Loading branch information
bdraco and arcivanov authored Aug 1, 2024
1 parent bf5a66f commit 0e6bbbc
Show file tree
Hide file tree
Showing 3 changed files with 43 additions and 4 deletions.
7 changes: 7 additions & 0 deletions CHANGES/8540.bugfix.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
Fixed WebSocket server heartbeat timeout logic to terminate `receive` and return :py:class:`~aiohttp.ServerTimeoutError` -- by :user:`arcivanov`.

When a WebSocket pong message was not received, the
:py:meth:`~aiohttp.ClientWebSocketResponse.receive` operation did not terminate.
This change causes `_pong_not_received` to feed the `reader` an error message, causing
pending `receive` to terminate and return the error message. The error message contains
the exception :py:class:`~aiohttp.ServerTimeoutError`.
8 changes: 6 additions & 2 deletions aiohttp/client_ws.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
import sys
from typing import Any, Optional, cast

from .client_exceptions import ClientError
from .client_exceptions import ClientError, ServerTimeoutError
from .client_reqrep import ClientResponse
from .helpers import call_later, set_result
from .http import (
Expand Down Expand Up @@ -122,8 +122,12 @@ def _pong_not_received(self) -> None:
if not self._closed:
self._closed = True
self._close_code = WSCloseCode.ABNORMAL_CLOSURE
self._exception = asyncio.TimeoutError()
self._exception = ServerTimeoutError()
self._response.close()
if self._waiting and not self._closing:
self._reader.feed_data(
WSMessage(WSMsgType.ERROR, self._exception, None)
)

@property
def closed(self) -> bool:
Expand Down
32 changes: 30 additions & 2 deletions tests/test_client_ws_functional.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
import pytest

import aiohttp
from aiohttp import hdrs, web
from aiohttp import ServerTimeoutError, WSMsgType, hdrs, web
from aiohttp.http import WSCloseCode
from aiohttp.pytest_plugin import AiohttpClient

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


async def test_send_recv_compress(aiohttp_client) -> None:
async def test_heartbeat_no_pong_concurrent_receive(aiohttp_client: Any) -> None:
ping_received = False

async def handler(request):
nonlocal ping_received
ws = web.WebSocketResponse(autoping=False)
await ws.prepare(request)
msg = await ws.receive()
ping_received = msg.type is aiohttp.WSMsgType.PING
ws._reader.feed_eof = lambda: None
await asyncio.sleep(10.0)

app = web.Application()
app.router.add_route("GET", "/", handler)

client = await aiohttp_client(app)
resp = await client.ws_connect("/", heartbeat=0.1)
resp._reader.feed_eof = lambda: None

# Connection should be closed roughly after 1.5x heartbeat.
msg = await resp.receive(5.0)
assert ping_received
assert resp.close_code is WSCloseCode.ABNORMAL_CLOSURE
assert msg
assert msg.type is WSMsgType.ERROR
assert isinstance(msg.data, ServerTimeoutError)


async def test_send_recv_compress(aiohttp_client: Any) -> None:
async def handler(request):
ws = web.WebSocketResponse()
await ws.prepare(request)
Expand Down

0 comments on commit 0e6bbbc

Please sign in to comment.