Code review feedback

Author: mehaase

README.md

@@ -89,9 +89,21 @@ feature.
 ```python
 async def heartbeat(ws, timeout, interval):
     '''
-    Send periodic pings on WebSocket ``ws``. Wait up to ``timeout`` seconds to
-    receive a pong before raising an exception. If a pong is received, then wait
-    ``interval`` seconds before sending the next ping.
+    Send periodic pings on WebSocket ``ws``.
+
+    Wait up to ``timeout`` seconds to send a ping and receive a pong. Raises
+    ``TooSlowError`` if the timeout is exceeded. If a pong is received, then
+    wait ``interval`` seconds before sending the next ping.
+
+    This function runs until cancelled.
+
+    :param ws: A WebSocket to send heartbeat pings on.
+    :param float timeout: Timeout in seconds.
+    :param float interval: Interval between receiving pong and sending next
+        ping, in seconds.
+    :raises: ``ConnectionClosed`` if ``ws`` is closed.
+    :raises: ``TooSlowError`` if the timeout expires.
+    :returns: This function runs until cancelled.
     '''
     while True:
         with trio.fail_after(timeout):

examples/client.py

@@ -79,9 +79,20 @@ async def handle_connection(ws, use_heartbeat):
 async def heartbeat(ws, timeout, interval):
     '''
     Send periodic pings on WebSocket ``ws``.
-    After sending a ping, wait up to ``timeout`` seconds to receive a pong
-    before raising an exception. If a pong is received, then wait ``interval``
-    seconds before sending the next ping.
+
+    Wait up to ``timeout`` seconds to send a ping and receive a pong. Raises
+    ``TooSlowError`` if the timeout is exceeded. If a pong is received, then
+    wait ``interval`` seconds before sending the next ping.
+
+    This function runs until cancelled.
+
+    :param ws: A WebSocket to send heartbeat pings on.
+    :param float timeout: Timeout in seconds.
+    :param float interval: Interval between receiving pong and sending next
+        ping, in seconds.
+    :raises: ``ConnectionClosed`` if ``ws`` is closed.
+    :raises: ``TooSlowError`` if the timeout expires.
+    :returns: This function runs until cancelled.
     '''
     while True:
         with trio.fail_after(timeout):

trio_websocket/__init__.py

@@ -425,10 +425,11 @@ async def ping(self, payload=None):
                 format(payload))
         if payload is None:
             payload = struct.pack('!I', random.getrandbits(32))
-        self._pings[payload] = trio.Event()
+        event = trio.Event()
+        self._pings[payload] = event
         self._wsproto.ping(payload)
         await self._write_pending()
-        await self._pings[payload].wait()
+        await event.wait()
 
     async def pong(self, payload=None):
         '''
@@ -589,13 +590,13 @@ async def _handle_pong_received_event(self, event):
             # We received a pong that doesn't match any in-flight pongs. Nothing
             # we can do with it, so ignore it.
             return
-        key, event = self._pings.popitem(0)
-        while key != payload:
-            logger.debug('conn#%d pong [skipped] %r', self._id, key)
-            event.set()
+        while self._pings:
             key, event = self._pings.popitem(0)
-        logger.debug('conn#%d pong %r', self._id, key)
-        event.set()
+            skipped = ' [skipped] ' if payload != key else ' '
+            logger.debug('conn#%d pong%s%r', self._id, skipped, key)
+            event.set()
+            if payload == key:
+                break
 
     async def _reader_task(self):
         ''' A background task that reads network data and generates events. '''
@@ -624,7 +625,7 @@ async def _reader_task(self):
                         event_type)
                     await handler(event)
                 except KeyError:
-                    logger.warning('Received unknown event type: %s',
+                    logger.warning('Received unknown event type: "%s"',
                         event_type)
 
             # Get network data.