Fix timer handle churn in websocket heartbeat

aiohttp/client_ws.py

@@ -7,7 +7,7 @@
 
 from .client_exceptions import ClientError, ServerTimeoutError
 from .client_reqrep import ClientResponse
-from .helpers import call_later, set_result
+from .helpers import calculate_timeout_when, set_result
 from .http import (
     WS_CLOSED_MESSAGE,
     WS_CLOSING_MESSAGE,
@@ -72,6 +72,7 @@
         self._autoping = autoping
         self._heartbeat = heartbeat
         self._heartbeat_cb: Optional[asyncio.TimerHandle] = None
+        self._heartbeat_when: float = 0.0
         if heartbeat is not None:
             self._pong_heartbeat = heartbeat / 2.0
         self._pong_response_cb: Optional[asyncio.TimerHandle] = None
@@ -85,48 +86,62 @@
         self._reset_heartbeat()
 
     def _cancel_heartbeat(self) -> None:
-        if self._pong_response_cb is not None:
-            self._pong_response_cb.cancel()
-            self._pong_response_cb = None
-
+        self._cancel_pong_response_cb()
         if self._heartbeat_cb is not None:
             self._heartbeat_cb.cancel()
             self._heartbeat_cb = None
 
+    def _cancel_pong_response_cb(self) -> None:
+        if self._pong_response_cb is not None:
+            self._pong_response_cb.cancel()
+            self._pong_response_cb = None
+
     def _reset_heartbeat(self) -> None:
-        self._cancel_heartbeat()
-
-        if self._heartbeat is not None:
-            self._heartbeat_cb = call_later(
-                self._send_heartbeat,
-                self._heartbeat,
-                self._loop,
-                timeout_ceil_threshold=(
-                    self._conn._connector._timeout_ceil_threshold
-                    if self._conn is not None
-                    else 5
-                ),
-            )
+        if self._heartbeat is None:
+            return
+        self._cancel_pong_response_cb()
+        assert self._loop is not None
+        loop = self._loop
+        conn = self._conn
+        timeout_ceil_threshold = (
+            conn._connector._timeout_ceil_threshold if conn is not None else 5
+        )
+        now = loop.time()
+        when = calculate_timeout_when(now, self._heartbeat, timeout_ceil_threshold)
+        self._heartbeat_when = when
+        if self._heartbeat_cb is None:
+            # We do not cancel the previous heartbeat_cb here because
+            # it generates a significant amount of TimerHandle churn
+            # which causes asyncio to rebuild the heap frequently.
+            # Instead _send_heartbeat() will reschedule the next
+            # heartbeat if it fires too early.
+            self._heartbeat_cb = loop.call_at(when, self._send_heartbeat)
 
     def _send_heartbeat(self) -> None:
-        if self._heartbeat is not None and not self._closed:
-            # 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.
-            self._loop.create_task(self._writer.ping())  # type: ignore[unused-awaitable]
-
-            if self._pong_response_cb is not None:
-                self._pong_response_cb.cancel()
-            self._pong_response_cb = call_later(
-                self._pong_not_received,
-                self._pong_heartbeat,
-                self._loop,
-                timeout_ceil_threshold=(
-                    self._conn._connector._timeout_ceil_threshold
-                    if self._conn is not None
-                    else 5
-                ),
+        self._heartbeat_cb = None
+        if self._heartbeat is None or self._closed:
+            return
+        loop = self._loop
+        now = loop.time()
+        if now < self._heartbeat_when:
+            # Heartbeat fired too early, reschedule
+            self._heartbeat_cb = loop.call_at(
+                self._heartbeat_when, self._send_heartbeat
             )
+            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
+        )
+        when = calculate_timeout_when(now, self._pong_heartbeat, timeout_ceil_threshold)
+        self._cancel_pong_response_cb()
+        self._pong_response_cb = loop.call_at(when, self._pong_not_received)
 
     def _pong_not_received(self) -> None:
         if not self._closed:

aiohttp/helpers.py

@@ -598,12 +598,23 @@ def call_later(
     loop: asyncio.AbstractEventLoop,
     timeout_ceil_threshold: float = 5,
 ) -> Optional[asyncio.TimerHandle]:
-    if timeout is not None and timeout > 0:
-        when = loop.time() + timeout
-        if timeout > timeout_ceil_threshold:
-            when = ceil(when)
-        return loop.call_at(when, cb)
-    return None
+    if timeout is None or timeout <= 0:
+        return None
+    now = loop.time()
+    when = calculate_timeout_when(now, timeout, timeout_ceil_threshold)
+    return loop.call_at(when, cb)
+
+
+def calculate_timeout_when(
+    loop_time: float,
+    timeout: float,
+    timeout_ceiling_threshold: float,
+) -> float:
+    """Calculate when to execute a timeout."""
+    when = loop_time + timeout
+    if timeout > timeout_ceiling_threshold:
+        when = ceil(when)
+    return when
 
 
 class TimeoutHandle:

aiohttp/web_ws.py

@@ -11,7 +11,7 @@
 
 from . import hdrs
 from .abc import AbstractStreamWriter
-from .helpers import call_later, set_exception, set_result
+from .helpers import calculate_timeout_when, set_exception, set_result
 from .http import (
     WS_CLOSED_MESSAGE,
     WS_CLOSING_MESSAGE,
@@ -74,6 +74,7 @@
         "_autoclose",
         "_autoping",
         "_heartbeat",
+        "_heartbeat_when",
         "_heartbeat_cb",
         "_pong_heartbeat",
         "_pong_response_cb",
@@ -112,6 +113,7 @@
         self._autoclose = autoclose
         self._autoping = autoping
         self._heartbeat = heartbeat
+        self._heartbeat_when = 0.0
         self._heartbeat_cb: Optional[asyncio.TimerHandle] = None
         if heartbeat is not None:
             self._pong_heartbeat = heartbeat / 2.0
@@ -120,50 +122,63 @@
         self._max_msg_size = max_msg_size
 
     def _cancel_heartbeat(self) -> None:
-        if self._pong_response_cb is not None:
-            self._pong_response_cb.cancel()
-            self._pong_response_cb = None
-
+        self._cancel_pong_response_cb()
         if self._heartbeat_cb is not None:
             self._heartbeat_cb.cancel()
             self._heartbeat_cb = None
 
-    def _reset_heartbeat(self) -> None:
-        self._cancel_heartbeat()
+    def _cancel_pong_response_cb(self) -> None:
+        if self._pong_response_cb is not None:
+            self._pong_response_cb.cancel()
+            self._pong_response_cb = None
 
-        if self._heartbeat is not None:
-            assert self._loop is not None
-            self._heartbeat_cb = call_later(
-                self._send_heartbeat,
-                self._heartbeat,
-                self._loop,
-                timeout_ceil_threshold=(
-                    self._req._protocol._timeout_ceil_threshold
-                    if self._req is not None
-                    else 5
-                ),
-            )
+    def _reset_heartbeat(self) -> None:
+        if self._heartbeat is None:
+            return
+        self._cancel_pong_response_cb()
+        req = self._req
+        timeout_ceil_threshold = (
+            req._protocol._timeout_ceil_threshold if req is not None else 5
+        )
+        assert self._loop is not None
+        loop = self._loop
+        now = loop.time()
+        when = calculate_timeout_when(now, self._heartbeat, timeout_ceil_threshold)
+        self._heartbeat_when = when
+        if self._heartbeat_cb is None:
+            # We do not cancel the previous heartbeat_cb here because
+            # it generates a significant amount of TimerHandle churn
+            # which causes asyncio to rebuild the heap frequently.
+            # Instead _send_heartbeat() will reschedule the next
+            # heartbeat if it fires too early.
+            self._heartbeat_cb = loop.call_at(when, self._send_heartbeat)
 
     def _send_heartbeat(self) -> None:
-        if self._heartbeat is not None and not self._closed:
-            assert self._loop is not None and self._writer is not None
-            # 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.
-            self._loop.create_task(self._writer.ping())  # type: ignore[unused-awaitable]
-
-            if self._pong_response_cb is not None:
-                self._pong_response_cb.cancel()
-            self._pong_response_cb = call_later(
-                self._pong_not_received,
-                self._pong_heartbeat,
-                self._loop,
-                timeout_ceil_threshold=(
-                    self._req._protocol._timeout_ceil_threshold
-                    if self._req is not None
-                    else 5
-                ),
+        self._heartbeat_cb = None
+        if self._heartbeat is None or self._closed:
+            return
+        assert self._loop is not None and self._writer is not None
+        loop = self._loop
+        now = loop.time()
+        if now < self._heartbeat_when:
+            # Heartbeat fired too early, reschedule
+            self._heartbeat_cb = loop.call_at(
+                self._heartbeat_when, self._send_heartbeat
             )
+            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
+        )
+        when = calculate_timeout_when(now, self._pong_heartbeat, timeout_ceil_threshold)
+        self._cancel_pong_response_cb()
+        self._pong_response_cb = loop.call_at(when, self._pong_not_received)
 
     def _pong_not_received(self) -> None:
         if self._req is not None and self._req.transport is not None: