Add option to set pong waiters on connection close.

Author: aaugustin

src/websockets/sync/connection.py

@@ -111,7 +111,7 @@ def __init__(
         self.recv_events_thread.start()
 
         # Mapping of ping IDs to pong waiters, in chronological order.
-        self.pong_waiters: dict[bytes, tuple[threading.Event, float]] = {}
+        self.pong_waiters: dict[bytes, tuple[threading.Event, float, bool]] = {}
 
         self.latency: float = 0
         """
@@ -554,7 +554,11 @@ def close(self, code: int = CloseCode.NORMAL_CLOSURE, reason: str = "") -> None:
             # They mean that the connection is closed, which was the goal.
             pass
 
-    def ping(self, data: Data | None = None) -> threading.Event:
+    def ping(
+        self,
+        data: Data | None = None,
+        ack_on_close: bool = False,
+    ) -> threading.Event:
         """
         Send a Ping_.
 
@@ -566,6 +570,12 @@ def ping(self, data: Data | None = None) -> threading.Event:
         Args:
             data: Payload of the ping. A :class:`str` will be encoded to UTF-8.
                 If ``data`` is :obj:`None`, the payload is four random bytes.
+            ack_on_close: when this option is :obj:`True`, the event will also
+                be set when the connection is closed. While this avoids getting
+                stuck waiting for a pong that will never arrive, it requires
+                checking that the state of the connection is still ``OPEN`` to
+                confirm that a pong was received, rather than the connection
+                being closed.
 
         Returns:
             An event that will be set when the corresponding pong is received.
@@ -599,7 +609,7 @@ def ping(self, data: Data | None = None) -> threading.Event:
                 data = struct.pack("!I", random.getrandbits(32))
 
             pong_waiter = threading.Event()
-            self.pong_waiters[data] = (pong_waiter, time.monotonic())
+            self.pong_waiters[data] = (pong_waiter, time.monotonic(), ack_on_close)
             self.protocol.send_ping(data)
             return pong_waiter
 
@@ -660,7 +670,11 @@ def acknowledge_pings(self, data: bytes) -> None:
             # Acknowledge all previous pings too in that case.
             ping_id = None
             ping_ids = []
-            for ping_id, (pong_waiter, ping_timestamp) in self.pong_waiters.items():
+            for ping_id, (
+                pong_waiter,
+                ping_timestamp,
+                ack_on_close,
+            ) in self.pong_waiters.items():
                 ping_ids.append(ping_id)
                 pong_waiter.set()
                 if ping_id == data:
@@ -673,6 +687,19 @@ def acknowledge_pings(self, data: bytes) -> None:
             for ping_id in ping_ids:
                 del self.pong_waiters[ping_id]
 
+    def acknowledge_pending_pings(self) -> None:
+        """
+        Acknowledge pending pings when the connection is closed.
+
+        """
+        assert self.protocol.state is CLOSED
+
+        for pong_waiter, _ping_timestamp, ack_on_close in self.pong_waiters.values():
+            if ack_on_close:
+                pong_waiter.set()
+
+        self.pong_waiters.clear()
+
     def recv_events(self) -> None:
         """
         Read incoming data from the socket and process events.
@@ -944,3 +971,6 @@ def close_socket(self) -> None:
 
         # Abort recv() with a ConnectionClosed exception.
         self.recv_messages.close()
+
+        # Acknowledge pings sent with the ack_on_close option.
+        self.acknowledge_pending_pings()

tests/sync/test_connection.py

@@ -685,6 +685,15 @@ def test_acknowledge_previous_ping(self):
         self.remote_connection.pong("that")
         self.assertTrue(pong_waiter.wait(MS))
 
+    def test_acknowledge_ping_on_close(self):
+        """ping with ack_on_close is acknowledged when the connection is closed."""
+        with self.drop_frames_rcvd():  # drop automatic response to ping
+            pong_waiter_ack_on_close = self.connection.ping("this", ack_on_close=True)
+            pong_waiter = self.connection.ping("that")
+        self.connection.close()
+        self.assertTrue(pong_waiter_ack_on_close.wait(MS))
+        self.assertFalse(pong_waiter.wait(MS))
+
     def test_ping_duplicate_payload(self):
         """ping rejects the same payload until receiving the pong."""
         with self.drop_frames_rcvd():  # drop automatic response to ping