Leave websocket transport open if receive times out or is cancelled (#8251)

Author: bdraco

CHANGES/8251.bugfix.rst

@@ -0,0 +1,4 @@
+Leave websocket transport open if receive times out or is cancelled
+-- by :user:`bdraco`.
+
+This restores the behavior prior to the change in #7978.

aiohttp/web_ws.py

@@ -488,8 +488,7 @@ async def receive(self, timeout: Optional[float] = None) -> WSMessage:
                     waiter = self._waiting
                     set_result(waiter, True)
                     self._waiting = None
-            except (asyncio.CancelledError, asyncio.TimeoutError):
-                self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE)
+            except asyncio.TimeoutError:
                 raise
             except EofStream:
                 self._close_code = WSCloseCode.OK

tests/test_web_websocket.py

@@ -445,7 +445,8 @@ async def test_receive_timeouterror(make_request: Any, loop: Any) -> None:
     with pytest.raises(asyncio.TimeoutError):
         await ws.receive()
 
-    assert len(ws._req.transport.close.mock_calls) == 1
+    # Should not close the connection on timeout
+    assert len(ws._req.transport.close.mock_calls) == 0
 
 
 async def test_multiple_receive_on_close_connection(make_request: Any) -> None:

tests/test_web_websocket_functional.py

@@ -2,6 +2,8 @@
 # HTTP websocket server functional tests
 
 import asyncio
+import contextlib
+import sys
 from typing import Any, Optional
 
 import pytest
@@ -828,3 +830,94 @@ async def ws_handler(request):
     resp = await client.get("/api/null", timeout=aiohttp.ClientTimeout(total=1))
     assert (await resp.json()) == {"err": None}
     resp.close()
+
+
+async def test_receive_being_cancelled_keeps_connection_open(
+    loop: Any, aiohttp_client: Any
+) -> None:
+    closed = loop.create_future()
+
+    async def handler(request):
+        ws = web.WebSocketResponse(autoping=False)
+        await ws.prepare(request)
+
+        task = asyncio.create_task(ws.receive())
+        await asyncio.sleep(0)
+        task.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await task
+
+        msg = await ws.receive()
+        assert msg.type == WSMsgType.PING
+        await asyncio.sleep(0)
+        await ws.pong("data")
+
+        msg = await ws.receive()
+        assert msg.type == WSMsgType.CLOSE
+        assert msg.data == WSCloseCode.OK
+        assert msg.extra == "exit message"
+        closed.set_result(None)
+        return ws
+
+    app = web.Application()
+    app.router.add_get("/", handler)
+    client = await aiohttp_client(app)
+
+    ws = await client.ws_connect("/", autoping=False)
+
+    await asyncio.sleep(0)
+    await ws.ping("data")
+
+    msg = await ws.receive()
+    assert msg.type == WSMsgType.PONG
+    assert msg.data == b"data"
+
+    await ws.close(code=WSCloseCode.OK, message="exit message")
+
+    await closed
+
+
+async def test_receive_timeout_keeps_connection_open(
+    loop: Any, aiohttp_client: Any
+) -> None:
+    closed = loop.create_future()
+    timed_out = loop.create_future()
+
+    async def handler(request):
+        ws = web.WebSocketResponse(autoping=False)
+        await ws.prepare(request)
+
+        task = asyncio.create_task(ws.receive(sys.float_info.min))
+        with contextlib.suppress(asyncio.TimeoutError):
+            await task
+
+        timed_out.set_result(None)
+
+        msg = await ws.receive()
+        assert msg.type == WSMsgType.PING
+        await asyncio.sleep(0)
+        await ws.pong("data")
+
+        msg = await ws.receive()
+        assert msg.type == WSMsgType.CLOSE
+        assert msg.data == WSCloseCode.OK
+        assert msg.extra == "exit message"
+        closed.set_result(None)
+        return ws
+
+    app = web.Application()
+    app.router.add_get("/", handler)
+    client = await aiohttp_client(app)
+
+    ws = await client.ws_connect("/", autoping=False)
+
+    await timed_out
+    await ws.ping("data")
+
+    msg = await ws.receive()
+    assert msg.type == WSMsgType.PONG
+    assert msg.data == b"data"
+
+    await ws.close(code=WSCloseCode.OK, message="exit message")
+
+    await closed