Repository: aaugustin/websockets (python)

Issue:
Pong's CancelledError races with ConnectionClosed error
I upgraded to the current version of this library and I find it difficult to handle `asyncio.CancelledError` in case the pong is being cancelled due to a disconnect.

What happens in my server is that I'm waiting for pongs in one task while sending and receiving in another task. Both tasks are being waited for. The result is that in some occasions, `asyncio.CancelledError` is raised before `websockets.ConnectionClosed` and vice versa. When the `asyncio.CancelledError` is raised first, there is no context on why it has been raised.

Another problem with `asyncio.CancelledError` is that there's no way (IIRC) of finding out whether the current task has been cancelled or the task I'm waiting for has been cancelled. But this mainly is an awkward concept of asyncio and not this libraries fault. It seems one has to be extremely careful on when to cancel a task and when not to cancel a task and instead use other ways to signal cancellation (which highly depends on the task). This may clash with the design choice of this library.

TL;DR: I would prefer it if this library would raise `websockets.ConnectionClosed` in the pong's future instead of cancelling the pong.

Pull request:
Abort pending pings with ConnectionClosed.
Fix #464.

Patch:
diff --git a/docs/changelog.rst b/docs/changelog.rst
index c895baa..630ee87 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -3,11 +3,21 @@ Changelog
 
 .. currentmodule:: websockets
 
-6.1
+7.0
 ...
 
 *In development*
 
+.. warning::
+
+  **Version 7.0 changes how a :meth:`~protocol.WebSocketCommonProtocol.ping`
+  that hasn't received a pong yet behaves when the connection is closed.**
+
+  The ping — as in ``ping = await websocket.ping()`` — used to be canceled
+  when the connection is closed, so that ``await ping`` raised
+  :exc:`~concurrent.futures.CancelledError`. Now ``await ping`` raises
+  :exc:`~exceptions.ConnectionClosed` like other public APIs.
+
 * websockets sends Ping frames at regular intervals and closes the connection
   if it doesn't receive a matching Pong frame. See :class:`~protocol.WebSocketCommonProtocol` for details.
 
diff --git a/src/websockets/protocol.py b/src/websockets/protocol.py
index b9bbda3..34fa488 100644
--- a/src/websockets/protocol.py
+++ b/src/websockets/protocol.py
@@ -919,19 +919,6 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
             if self.keepalive_ping_task is not None:
                 self.keepalive_ping_task.cancel()
 
-            # Cancel all pending pings because they'll never receive a pong.
-            for ping in self.pings.values():
-                ping.cancel()
-            if self.pings:
-                pings_hex = ', '.join(
-                    binascii.hexlify(ping_id).decode() or '[empty]'
-                    for ping_id in self.pings
-                )
-                plural = 's' if len(self.pings) > 1 else ''
-                logger.debug(
-                    "%s - canceled pending ping%s: %s", self.side, plural, pings_hex
-                )
-
             # A client should wait for a TCP close from the server.
             if self.is_client and self.transfer_data_task is not None:
                 if (yield from self.wait_for_connection_lost()):
@@ -1059,6 +1046,30 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
 
         return self.close_connection_task
 
+    def abort_keepalive_pings(self):
+        """
+        Raise ConnectionClosed in pending keepalive pings.
+
+        They'll never receive a pong once the connection is closed.
+
+        """
+        assert self.state is State.CLOSED
+        exc = ConnectionClosed(self.close_code, self.close_reason)
+        exc.__cause__ = self.transfer_data_exc  # emulate raise ... from ...
+
+        for ping in self.pings.values():
+            ping.set_exception(exc)
+
+        if self.pings:
+            pings_hex = ', '.join(
+                binascii.hexlify(ping_id).decode() or '[empty]'
+                for ping_id in self.pings
+            )
+            plural = 's' if len(self.pings) > 1 else ''
+            logger.debug(
+                "%s - aborted pending ping%s: %s", self.side, plural, pings_hex
+            )
+
     # asyncio.StreamReaderProtocol methods
 
     def connection_made(self, transport):
@@ -1122,6 +1133,7 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
             self.close_code,
             self.close_reason or '[empty]',
         )
+        self.abort_keepalive_pings()
         # 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.


Test patch:
diff --git a/tests/test_protocol.py b/tests/test_protocol.py
index 133a015..063960a 100644
--- a/tests/test_protocol.py
+++ b/tests/test_protocol.py
@@ -605,13 +605,14 @@ class CommonTests:
         self.run_loop_once()
         self.assertTrue(ping.done())
 
-    def test_cancel_ping(self):
+    def test_abort_ping(self):
         ping = self.loop.run_until_complete(self.protocol.ping())
         # Remove the frame from the buffer, else close_connection() complains.
         self.last_sent_frame()
-        self.assertFalse(ping.cancelled())
+        self.assertFalse(ping.done())
         self.close_connection()
-        self.assertTrue(ping.cancelled())
+        self.assertTrue(ping.done())
+        self.assertIsInstance(ping.exception(), ConnectionClosed)
 
     def test_acknowledge_previous_pings(self):
         pings = [