Add keepalive to the new asyncio implementation.

Author: aaugustin

docs/faq/common.rst

@@ -97,7 +97,7 @@ There are two main reasons why latency may increase:
 * Poor network connectivity.
 * More traffic than the recipient can handle.
 
-See the discussion of :doc:`timeouts <../topics/timeouts>` for details.
+See the discussion of :doc:`keepalive <../topics/keepalive>` for details.
 
 If websockets' default timeout of 20 seconds is too short for your use case,
 you can adjust it with the ``ping_timeout`` argument.
@@ -146,7 +146,7 @@ It closes the connection if it doesn't get a pong within 20 seconds.
 
 You can adjust this behavior with ``ping_interval`` and ``ping_timeout``.
 
-See :doc:`../topics/timeouts` for details.
+See :doc:`../topics/keepalive` for details.
 
 How do I respond to pings?
 --------------------------

docs/howto/upgrade.rst

@@ -70,17 +70,6 @@ Missing features
     If your application relies on one of them, you should stick to the original
     implementation until the new implementation supports it in a future release.
 
-Keepalive
-.........
-
-The new implementation doesn't provide a :ref:`keepalive mechanism <keepalive>`
-yet.
-
-As a consequence, :func:`~asyncio.client.connect` and
-:func:`~asyncio.server.serve` don't accept the ``ping_interval`` and
-``ping_timeout`` arguments and the
-:attr:`~legacy.protocol.WebSocketCommonProtocol.latency` property doesn't exist.
-
 HTTP Basic Authentication
 .........................
 

docs/reference/features.rst

@@ -53,9 +53,9 @@ Both sides
     +------------------------------------+--------+--------+--------+--------+
     | Send a pong                        | ✅     | ✅     | ✅     | ✅     |
     +------------------------------------+--------+--------+--------+--------+
-    | Keepalive                          | ❌     | ❌     | —      | ✅     |
+    | Keepalive                          | ✅     | ❌     | —      | ✅     |
     +------------------------------------+--------+--------+--------+--------+
-    | Heartbeat                          | ❌     | ❌     | —      | ✅     |
+    | Heartbeat                          | ✅     | ❌     | —      | ✅     |
     +------------------------------------+--------+--------+--------+--------+
     | Perform the closing handshake      | ✅     | ✅     | ✅     | ✅     |
     +------------------------------------+--------+--------+--------+--------+

docs/reference/new-asyncio/client.rst

@@ -43,6 +43,8 @@ Using a connection
 
     .. autoproperty:: remote_address
 
+    .. autoattribute:: latency
+
     .. autoproperty:: state
 
     The following attributes are available after the opening handshake,

docs/reference/new-asyncio/common.rst

@@ -33,6 +33,8 @@ Both sides (new :mod:`asyncio`)
 
     .. autoproperty:: remote_address
 
+    .. autoattribute:: latency
+
     .. autoproperty:: state
 
     The following attributes are available after the opening handshake,

docs/reference/new-asyncio/server.rst

@@ -64,6 +64,8 @@ Using a connection
 
     .. autoproperty:: remote_address
 
+    .. autoattribute:: latency
+
     .. autoproperty:: state
 
     The following attributes are available after the opening handshake,

docs/topics/broadcast.rst

@@ -206,8 +206,8 @@ How do we avoid running out of memory when slow clients can't keep up with the
 broadcast rate, then? The most straightforward option is to disconnect them.
 
 If a client gets too far behind, eventually it reaches the limit defined by
-``ping_timeout`` and websockets terminates the connection. You can refer to
-the discussion of :doc:`keepalive and timeouts <timeouts>` for details.
+``ping_timeout`` and websockets terminates the connection. You can refer to the
+discussion of :doc:`keepalive <keepalive>` for details.
 
 How :func:`~asyncio.connection.broadcast` works
 -----------------------------------------------

docs/topics/index.rst

@@ -11,7 +11,7 @@ Get a deeper understanding of how websockets is built and why.
    authentication
    broadcast
    compression
-   timeouts
+   keepalive
    design
    memory
    security

docs/topics/timeouts.rst renamed to docs/topics/keepalive.rst

@@ -1,5 +1,5 @@
-Timeouts
-========
+Keepalive and latency
+=====================
 
 .. currentmodule:: websockets
 
@@ -49,9 +49,9 @@ This mechanism serves two purposes:
    application gets a :exc:`~exceptions.ConnectionClosed` exception.
 
 Timings are configurable with the ``ping_interval`` and ``ping_timeout``
-arguments of :func:`~client.connect` and :func:`~server.serve`. Shorter values
-will detect connection drops faster but they will increase network traffic and
-they will be more sensitive to latency.
+arguments of :func:`~asyncio.client.connect` and :func:`~asyncio.server.serve`.
+Shorter values will detect connection drops faster but they will increase
+network traffic and they will be more sensitive to latency.
 
 Setting ``ping_interval`` to :obj:`None` disables the whole keepalive and
 heartbeat mechanism.
@@ -111,6 +111,6 @@ Latency between a client and a server may increase for two reasons:
   than the client can accept.
 
 The latency measured during the last exchange of Ping and Pong frames is
-available in the :attr:`~legacy.protocol.WebSocketCommonProtocol.latency`
-attribute. Alternatively, you can measure the latency at any time with the
-:attr:`~legacy.protocol.WebSocketCommonProtocol.ping` method.
+available in the :attr:`~asyncio.connection.Connection.latency` attribute.
+Alternatively, you can measure the latency at any time with the
+:attr:`~asyncio.connection.Connection.ping` method.

src/websockets/asyncio/client.py

@@ -37,24 +37,29 @@ class ClientConnection(Connection):
     :exc:`~websockets.exceptions.ConnectionClosedError` when the connection is
     closed with any other code.
 
+    The ``ping_interval``, ``ping_timeout``, ``close_timeout``, ``max_queue``,
+    and ``write_limit`` arguments the same meaning as in :func:`connect`.
+
     Args:
         protocol: Sans-I/O connection.
-        close_timeout: Timeout for closing the connection in seconds.
-            :obj:`None` disables the timeout.
 
     """
 
     def __init__(
         self,
         protocol: ClientProtocol,
         *,
+        ping_interval: float | None = 20,
+        ping_timeout: float | None = 20,
         close_timeout: float | None = 10,
         max_queue: int | tuple[int, int | None] = 16,
         write_limit: int | tuple[int, int | None] = 2**15,
     ) -> None:
         self.protocol: ClientProtocol
         super().__init__(
             protocol,
+            ping_interval=ping_interval,
+            ping_timeout=ping_timeout,
             close_timeout=close_timeout,
             max_queue=max_queue,
             write_limit=write_limit,
@@ -84,7 +89,9 @@ async def handshake(
         if self.response is None:
             raise ConnectionError("connection closed during handshake")
 
-        if self.protocol.handshake_exc is not None:
+        if self.protocol.handshake_exc is None:
+            self.start_keepalive()
+        else:
             try:
                 async with asyncio_timeout(self.close_timeout):
                     await self.connection_lost_waiter
@@ -146,6 +153,10 @@ class connect:
             :doc:`compression guide <../../topics/compression>` for details.
         open_timeout: Timeout for opening the connection in seconds.
             :obj:`None` disables the timeout.
+        ping_interval: Interval between keepalive pings in seconds.
+            :obj:`None` disables keepalive.
+        ping_timeout: Timeout for keepalive pings in seconds.
+            :obj:`None` disables timeouts.
         close_timeout: Timeout for closing the connection in seconds.
             :obj:`None` disables the timeout.
         max_size: Maximum size of incoming messages in bytes.
@@ -208,6 +219,8 @@ def __init__(
         compression: str | None = "deflate",
         # Timeouts
         open_timeout: float | None = 10,
+        ping_interval: float | None = 20,
+        ping_timeout: float | None = 20,
         close_timeout: float | None = 10,
         # Limits
         max_size: int | None = 2**20,
@@ -256,6 +269,8 @@ def factory() -> ClientConnection:
             # This is a connection in websockets and a protocol in asyncio.
             connection = create_connection(
                 protocol,
+                ping_interval=ping_interval,
+                ping_timeout=ping_timeout,
                 close_timeout=close_timeout,
                 max_queue=max_queue,
                 write_limit=write_limit,

src/websockets/asyncio/connection.py

@@ -24,7 +24,13 @@
 from ..http11 import Request, Response
 from ..protocol import CLOSED, OPEN, Event, Protocol, State
 from ..typing import Data, LoggerLike, Subprotocol
-from .compatibility import TimeoutError, aiter, anext, asyncio_timeout_at
+from .compatibility import (
+    TimeoutError,
+    aiter,
+    anext,
+    asyncio_timeout,
+    asyncio_timeout_at,
+)
 from .messages import Assembler
 
 
@@ -48,11 +54,15 @@ def __init__(
         self,
         protocol: Protocol,
         *,
+        ping_interval: float | None = 20,
+        ping_timeout: float | None = 20,
         close_timeout: float | None = 10,
         max_queue: int | tuple[int, int | None] = 16,
         write_limit: int | tuple[int, int | None] = 2**15,
     ) -> None:
         self.protocol = protocol
+        self.ping_interval = ping_interval
+        self.ping_timeout = ping_timeout
         self.close_timeout = close_timeout
         if isinstance(max_queue, int):
             max_queue = (max_queue, None)
@@ -95,6 +105,21 @@ def __init__(
         # Mapping of ping IDs to pong waiters, in chronological order.
         self.pong_waiters: dict[bytes, tuple[asyncio.Future[float], float]] = {}
 
+        self.latency: float = 0
+        """
+        Latency of the connection, in seconds.
+
+        This value is updated after sending a ping frame and receiving a
+        matching pong frame. Before the first ping, :attr:`latency` is ``0``.
+
+        By default, websockets enables a :ref:`keepalive <keepalive>` mechanism
+        that sends ping frames automatically at regular intervals. You can also
+        send ping frames and measure latency with :meth:`ping`.
+        """
+
+        # Task that sends keepalive pings. None when ping_interval is None.
+        self.keepalive_task: asyncio.Task[None] | None = None
+
         # Exception raised while reading from the connection, to be chained to
         # ConnectionClosed in order to show why the TCP connection dropped.
         self.recv_exc: BaseException | None = None
@@ -144,7 +169,8 @@ def state(self) -> State:
 
         This attribute is provided for completeness. Typical applications
         shouldn't check its value. Instead, they should call :meth:`~recv` or
-        :meth:`send` and handle :exc:`~exceptions.ConnectionClosed` exceptions.
+        :meth:`send` and handle :exc:`~websockets.exceptions.ConnectionClosed`
+        exceptions.
 
         """
         return self.protocol.state
@@ -540,7 +566,7 @@ async def wait_closed(self) -> None:
         """
         await asyncio.shield(self.connection_lost_waiter)
 
-    async def ping(self, data: Data | None = None) -> Awaitable[None]:
+    async def ping(self, data: Data | None = None) -> Awaitable[float]:
         """
         Send a Ping_.
 
@@ -643,8 +669,10 @@ def acknowledge_pings(self, data: bytes) -> None:
         ping_ids = []
         for ping_id, (pong_waiter, ping_timestamp) in self.pong_waiters.items():
             ping_ids.append(ping_id)
-            pong_waiter.set_result(pong_timestamp - ping_timestamp)
+            latency = pong_timestamp - ping_timestamp
+            pong_waiter.set_result(latency)
             if ping_id == data:
+                self.latency = latency
                 break
         else:
             raise AssertionError("solicited pong not found in pings")
@@ -664,7 +692,8 @@ def abort_pings(self) -> None:
         exc = self.protocol.close_exc
 
         for pong_waiter, _ping_timestamp in self.pong_waiters.values():
-            pong_waiter.set_exception(exc)
+            if not pong_waiter.done():
+                pong_waiter.set_exception(exc)
             # If the exception is never retrieved, it will be logged when ping
             # is garbage-collected. This is confusing for users.
             # Given that ping is done (with an exception), canceling it does
@@ -673,6 +702,50 @@ def abort_pings(self) -> None:
 
         self.pong_waiters.clear()
 
+    async def keepalive(self) -> None:
+        """
+        Send a Ping frame and wait for a Pong frame at regular intervals.
+
+        """
+        assert self.ping_interval is not None
+        latency = 0.0
+        try:
+            while True:
+                # If self.ping_timeout > latency > self.ping_interval, pings
+                # will be sent immediately after receiving pongs. The period
+                # will be longer than self.ping_interval.
+                await asyncio.sleep(self.ping_interval - latency)
+
+                self.logger.debug("% sending keepalive ping")
+                pong_waiter = await self.ping()
+
+                if self.ping_timeout is not None:
+                    try:
+                        async with asyncio_timeout(self.ping_timeout):
+                            latency = await pong_waiter
+                        self.logger.debug("% received keepalive pong")
+                    except asyncio.TimeoutError:
+                        if self.debug:
+                            self.logger.debug("! timed out waiting for keepalive pong")
+                        async with self.send_context():
+                            self.protocol.fail(
+                                CloseCode.INTERNAL_ERROR,
+                                "keepalive ping timeout",
+                            )
+                        break
+        except ConnectionClosed:
+            pass
+        except Exception:
+            self.logger.error("keepalive ping failed", exc_info=True)
+
+    def start_keepalive(self) -> None:
+        """
+        Run :meth:`keepalive` in a task, unless keepalive is disabled.
+
+        """
+        if self.ping_interval is not None:
+            self.keepalive_task = self.loop.create_task(self.keepalive())
+
     @contextlib.asynccontextmanager
     async def send_context(
         self,
@@ -835,11 +908,15 @@ def connection_lost(self, exc: Exception | None) -> None:
         self.protocol.receive_eof()  # receive_eof is idempotent
         self.recv_messages.close()
         self.set_recv_exc(exc)
+        self.abort_pings()
+        # If keepalive() was waiting for a pong, abort_pings() terminated it.
+        # If it was sleeping until the next ping, we need to cancel it now
+        if self.keepalive_task is not None:
+            self.keepalive_task.cancel()
         # If self.connection_lost_waiter isn't pending, that's a bug, because:
         # - it's set only here in connection_lost() which is called only once;
         # - it must never be canceled.
         self.connection_lost_waiter.set_result(None)
-        self.abort_pings()
 
         # Adapted from asyncio.streams.FlowControlMixin
         if self.paused:  # pragma: no cover