Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix WebSocket server heartbeat timeout logic #8546

Merged
merged 7 commits into from
Aug 1, 2024
Merged
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions CHANGES/8540.bugfix.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
WSS: Fix heartbeat timeout logic -- by :user:`arcivanov`.
bdraco marked this conversation as resolved.
Show resolved Hide resolved

The :py:meth:`~aiohttp.ClientWebSocketResponse._pong_not_received` did not
terminate a :py:meth:`~aiohttp.ClientWebSocketResponse.receive` operation.
bdraco marked this conversation as resolved.
Show resolved Hide resolved
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 @@ -5,7 +5,7 @@
import sys
from typing import Any, Final, 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 @@ -132,8 +132,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
30 changes: 29 additions & 1 deletion tests/test_client_ws_functional.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
import pytest

import aiohttp
from aiohttp import hdrs, web
from aiohttp import ServerTimeoutError, WSMsgType, hdrs, web
from aiohttp.client_ws import ClientWSTimeout
from aiohttp.http import WSCloseCode
from aiohttp.pytest_plugin import AiohttpClient
Expand Down Expand Up @@ -685,6 +685,34 @@ async def handler(request):
assert resp.close_code is WSCloseCode.ABNORMAL_CLOSURE


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()
Expand Down
Loading