Fix aio-libs#7306 - Set ClientWebSocketResponse.close_code correctly in concurrent closing scenario (aio-libs#7680)

ttsia · web-flow · commit 30850babb43a · 2023-10-09T19:31:23.000+01:00
diff --git a/CHANGES/7306.bugfix b/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.
diff --git a/aiohttp/client_ws.py b/aiohttp/client_ws.py
@@ -191,7 +191,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
 
@@ -210,7 +211,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
 
diff --git a/tests/test_client_ws_functional.py b/tests/test_client_ws_functional.py
@@ -233,6 +233,33 @@ async def handler(request):
     assert msg.type == aiohttp.WSMsgType.CLOSED
 
 
+async def test_concurrent_task_close(aiohttp_client: Any) -> 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_concurrent_close(aiohttp_client: Any) -> None:
     client_ws = None
 

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