Skip to content

Commit

Permalink
Fix aio-libs#7306 - Set ClientWebSocketResponse.close_code correctly …
Browse files Browse the repository at this point in the history
…in concurrent closing scenario (aio-libs#7680)

(cherry picked from commit 30850ba)
  • Loading branch information
ttsia committed Oct 10, 2023
1 parent 492e8ee commit 70b3536
Show file tree
Hide file tree
Showing 3 changed files with 31 additions and 2 deletions.
1 change: 1 addition & 0 deletions CHANGES/7306.bugfix
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Fixed ``ClientWebSocketResponse.close_code`` being erroneously set to ``None`` when there are concurrent async tasks receiving data and closing the connection.
5 changes: 3 additions & 2 deletions aiohttp/client_ws.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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

Expand Down
27 changes: 27 additions & 0 deletions tests/test_client_ws_functional.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down

0 comments on commit 70b3536

Please sign in to comment.