[PR #7684/30850ba backport][3.9] Fix #7306 - Set ClientWebSocketRespo…

…nse.close_code correctly in concu… (#7690)

This is a backport of PR #7680
as merged into master
(30850ba).

Fixes the issue when If one asyncio task is waiting on receiving data
and another asyncio task is closing the connection. The
ClientWebSocketResponse.close_code will be None after connection closed.

CHANGES/7306.bugfix

@@ -0,0 +1 @@
+Fixed ``ClientWebSocketResponse.close_code`` being erroneously set to ``None`` when there are concurrent async tasks receiving data and closing the connection.

aiohttp/client_ws.py

@@ -181,7 +181,8 @@ async def send_json(
     async def close(self, *, code: int = WSCloseCode.OK, message: bytes = b"") -> bool:
         # we need to break `receive()` cycle first,
         # `close()` may be called from different task
-        if self._waiting is not None and not self._closed:
+        if self._waiting is not None and not self._closing:
+            self._closing = True
             self._reader.feed_data(WS_CLOSING_MESSAGE, 0)
             await self._waiting
 
@@ -200,7 +201,7 @@ async def close(self, *, code: int = WSCloseCode.OK, message: bytes = b"") -> bo
                 self._response.close()
                 return True
 
-            if self._closing:
+            if self._close_code:
                 self._response.close()
                 return True
 

tests/test_client_ws_functional.py

@@ -262,6 +262,33 @@ async def handler(request):
     assert msg.type == aiohttp.WSMsgType.CLOSED
 
 
+async def test_concurrent_task_close(aiohttp_client) -> None:
+    async def handler(request):
+        ws = web.WebSocketResponse()
+        await ws.prepare(request)
+        await ws.receive()
+        return ws
+
+    app = web.Application()
+    app.router.add_route("GET", "/", handler)
+
+    client = await aiohttp_client(app)
+    async with client.ws_connect("/") as resp:
+        # wait for the message in a separate task
+        task = asyncio.create_task(resp.receive())
+
+        # Make sure we start to wait on receiving message before closing the connection
+        await asyncio.sleep(0.1)
+
+        closed = await resp.close()
+
+        await task
+
+        assert closed
+        assert resp.closed
+        assert resp.close_code == 1000
+
+
 async def test_close_from_server(aiohttp_client) -> None:
     loop = asyncio.get_event_loop()
     closed = loop.create_future()