[PR #8641/0a88bab backport][3.10] Fix WebSocket ping tasks being prematurely garbage collected (#8646)

Author: bdraco

CHANGES/8641.bugfix.rst

@@ -0,0 +1,3 @@
+Fixed WebSocket ping tasks being prematurely garbage collected -- by :user:`bdraco`.
+
+There was a small risk that WebSocket ping tasks would be prematurely garbage collected because the event loop only holds a weak reference to the task. The garbage collection risk has been fixed by holding a strong reference to the task. Additionally, the task is now scheduled eagerly with Python 3.12+ to increase the chance it can be completed immediately and avoid having to hold any references to the task.

aiohttp/client_ws.py

@@ -72,6 +72,7 @@ def __init__(
  self._exception: Optional[BaseException] = None
  self._compress = compress
  self._client_notakeover = client_notakeover
+ self._ping_task: Optional[asyncio.Task[None]] = None
 
  self._reset_heartbeat()
 
@@ -80,6 +81,9 @@ def _cancel_heartbeat(self) -> None:
  if self._heartbeat_cb is not None:
  self._heartbeat_cb.cancel()
  self._heartbeat_cb = None
+ if self._ping_task is not None:
+ self._ping_task.cancel()
+ self._ping_task = None
 
  def _cancel_pong_response_cb(self) -> None:
  if self._pong_response_cb is not None:
@@ -118,11 +122,6 @@ def _send_heartbeat(self) -> None:
  )
  return
 
- # fire-and-forget a task is not perfect but maybe ok for
- # sending ping. Otherwise we need a long-living heartbeat
- # task in the class.
- loop.create_task(self._writer.ping()) # type: ignore[unused-awaitable]
-
  conn = self._conn
  timeout_ceil_threshold = (
  conn._connector._timeout_ceil_threshold if conn is not None else 5
@@ -131,6 +130,22 @@ def _send_heartbeat(self) -> None:
  self._cancel_pong_response_cb()
  self._pong_response_cb = loop.call_at(when, self._pong_not_received)
 
+ if sys.version_info >= (3, 12):
+ # Optimization for Python 3.12, try to send the ping
+ # immediately to avoid having to schedule
+ # the task on the event loop.
+ ping_task = asyncio.Task(self._writer.ping(), loop=loop, eager_start=True)
+ else:
+ ping_task = loop.create_task(self._writer.ping())
+
+ if not ping_task.done():
+ self._ping_task = ping_task
+ ping_task.add_done_callback(self._ping_task_done)
+
+ def _ping_task_done(self, task: "asyncio.Task[None]") -> None:
+ """Callback for when the ping task completes."""
+ self._ping_task = None
+
  def _pong_not_received(self) -> None:
  if not self._closed:
  self._set_closed()

aiohttp/web_ws.py

@@ -96,12 +96,16 @@ def __init__(
  self._pong_response_cb: Optional[asyncio.TimerHandle] = None
  self._compress = compress
  self._max_msg_size = max_msg_size
+ self._ping_task: Optional[asyncio.Task[None]] = None
 
  def _cancel_heartbeat(self) -> None:
  self._cancel_pong_response_cb()
  if self._heartbeat_cb is not None:
  self._heartbeat_cb.cancel()
  self._heartbeat_cb = None
+ if self._ping_task is not None:
+ self._ping_task.cancel()
+ self._ping_task = None
 
  def _cancel_pong_response_cb(self) -> None:
  if self._pong_response_cb is not None:
@@ -141,11 +145,6 @@ def _send_heartbeat(self) -> None:
  )
  return
 
- # fire-and-forget a task is not perfect but maybe ok for
- # sending ping. Otherwise we need a long-living heartbeat
- # task in the class.
- loop.create_task(self._writer.ping()) # type: ignore[unused-awaitable]
-
  req = self._req
  timeout_ceil_threshold = (
  req._protocol._timeout_ceil_threshold if req is not None else 5
@@ -154,6 +153,22 @@ def _send_heartbeat(self) -> None:
  self._cancel_pong_response_cb()
  self._pong_response_cb = loop.call_at(when, self._pong_not_received)
 
+ if sys.version_info >= (3, 12):
+ # Optimization for Python 3.12, try to send the ping
+ # immediately to avoid having to schedule
+ # the task on the event loop.
+ ping_task = asyncio.Task(self._writer.ping(), loop=loop, eager_start=True)
+ else:
+ ping_task = loop.create_task(self._writer.ping())
+
+ if not ping_task.done():
+ self._ping_task = ping_task
+ ping_task.add_done_callback(self._ping_task_done)
+
+ def _ping_task_done(self, task: "asyncio.Task[None]") -> None:
+ """Callback for when the ping task completes."""
+ self._ping_task = None
+
  def _pong_not_received(self) -> None:
  if self._req is not None and self._req.transport is not None:
  self._set_closed()

tests/test_client_ws_functional.py

@@ -1,6 +1,7 @@
 import asyncio
 import sys
 from typing import Any, NoReturn
+from unittest import mock
 
 import pytest
 
@@ -727,8 +728,53 @@ async def handler(request):
  assert isinstance(msg.data, ServerTimeoutError)
 
 
-async def test_send_recv_compress(aiohttp_client: Any) -> None:
- async def handler(request):
+async def test_close_websocket_while_ping_inflight(
+ aiohttp_client: AiohttpClient,
+) -> None:
+ """Test closing the websocket while a ping is in-flight."""
+ ping_received = False
+
+ async def handler(request: web.Request) -> NoReturn:
+ nonlocal ping_received
+ ws = web.WebSocketResponse(autoping=False)
+ await ws.prepare(request)
+ msg = await ws.receive()
+ assert msg.type is aiohttp.WSMsgType.BINARY
+ msg = await ws.receive()
+ ping_received = msg.type is aiohttp.WSMsgType.PING
+ await ws.receive()
+ assert False
+
+ app = web.Application()
+ app.router.add_route("GET", "/", handler)
+
+ client = await aiohttp_client(app)
+ resp = await client.ws_connect("/", heartbeat=0.1)
+ await resp.send_bytes(b"ask")
+
+ cancelled = False
+ ping_stated = False
+
+ async def delayed_ping() -> None:
+ nonlocal cancelled, ping_stated
+ ping_stated = True
+ try:
+ await asyncio.sleep(1)
+ except asyncio.CancelledError:
+ cancelled = True
+ raise
+
+ with mock.patch.object(resp._writer, "ping", delayed_ping):
+ await asyncio.sleep(0.1)
+
+ await resp.close()
+ await asyncio.sleep(0)
+ assert ping_stated is True
+ assert cancelled is True
+
+
+async def test_send_recv_compress(aiohttp_client: AiohttpClient) -> None:
+ async def handler(request: web.Request) -> web.WebSocketResponse:
  ws = web.WebSocketResponse()
  await ws.prepare(request)