PYTHON-2299 Add the "awaited" field to heartbeat events (#457)

Author: ShaneHarvey

pymongo/monitor.py

@@ -212,11 +212,13 @@ def _check_server(self):
         except ReferenceError:
             raise
         except Exception as error:
-            address = self._server_description.address
+            sd = self._server_description
+            address = sd.address
             duration = _time() - start
             if self._publish:
+                awaited = sd.is_server_type_known and sd.topology_version
                 self._listeners.publish_server_heartbeat_failed(
-                    address, duration, error)
+                    address, duration, error, awaited)
             self._reset_connection()
             if isinstance(error, _OperationCancelled):
                 raise
@@ -231,7 +233,6 @@ def _check_once(self):
         """
         address = self._server_description.address
         if self._publish:
-            # PYTHON-2299: Add the "awaited" field to heartbeat events.
             self._listeners.publish_server_heartbeat_started(address)
 
         if self._cancel_context and self._cancel_context.cancelled:
@@ -246,7 +247,7 @@ def _check_once(self):
                                    self._rtt_monitor.average())
             if self._publish:
                 self._listeners.publish_server_heartbeat_succeeded(
-                    address, round_trip_time, response)
+                    address, round_trip_time, response, response.awaitable)
             return sd
 
     def _check_with_socket(self, conn):

pymongo/monitoring.py

@@ -1099,12 +1099,13 @@ class ServerHeartbeatSucceededEvent(_ServerHeartbeatEvent):
     .. versionadded:: 3.3
     """
 
-    __slots__ = ('__duration', '__reply')
+    __slots__ = ('__duration', '__reply', '__awaited')
 
-    def __init__(self, duration, reply, *args):
-        super(ServerHeartbeatSucceededEvent, self).__init__(*args)
+    def __init__(self, duration, reply, connection_id, awaited=False):
+        super(ServerHeartbeatSucceededEvent, self).__init__(connection_id)
         self.__duration = duration
         self.__reply = reply
+        self.__awaited = awaited
 
     @property
     def duration(self):
@@ -1116,10 +1117,20 @@ def reply(self):
         """An instance of :class:`~pymongo.ismaster.IsMaster`."""
         return self.__reply
 
+    @property
+    def awaited(self):
+        """Whether the heartbeat was awaited.
+
+        If true, then :meth:`duration` reflects the sum of the round trip time
+        to the server and the time that the server waited before sending a
+        response.
+        """
+        return self.__awaited
+
     def __repr__(self):
-        return "<%s %s duration: %s, reply: %s>" % (
+        return "<%s %s duration: %s, awaited: %s, reply: %s>" % (
             self.__class__.__name__, self.connection_id,
-            self.duration, self.reply)
+            self.duration, self.awaited, self.reply)
 
 
 class ServerHeartbeatFailedEvent(_ServerHeartbeatEvent):
@@ -1129,12 +1140,13 @@ class ServerHeartbeatFailedEvent(_ServerHeartbeatEvent):
     .. versionadded:: 3.3
     """
 
-    __slots__ = ('__duration', '__reply')
+    __slots__ = ('__duration', '__reply', '__awaited')
 
-    def __init__(self, duration, reply, *args):
-        super(ServerHeartbeatFailedEvent, self).__init__(*args)
+    def __init__(self, duration, reply, connection_id, awaited=False):
+        super(ServerHeartbeatFailedEvent, self).__init__(connection_id)
         self.__duration = duration
         self.__reply = reply
+        self.__awaited = awaited
 
     @property
     def duration(self):
@@ -1146,10 +1158,20 @@ def reply(self):
         """A subclass of :exc:`Exception`."""
         return self.__reply
 
+    @property
+    def awaited(self):
+        """Whether the heartbeat was awaited.
+
+        If true, then :meth:`duration` reflects the sum of the round trip time
+        to the server and the time that the server waited before sending a
+        response.
+        """
+        return self.__awaited
+
     def __repr__(self):
-        return "<%s %s duration: %s, reply: %r>" % (
+        return "<%s %s duration: %s, awaited: %s, reply: %r>" % (
             self.__class__.__name__, self.connection_id,
-            self.duration, self.reply)
+            self.duration, self.awaited, self.reply)
 
 
 class _EventListeners(object):
@@ -1303,7 +1325,7 @@ def publish_server_heartbeat_started(self, connection_id):
                 _handle_exception()
 
     def publish_server_heartbeat_succeeded(self, connection_id, duration,
-                                           reply):
+                                           reply, awaited):
         """Publish a ServerHeartbeatSucceededEvent to all server heartbeat
         listeners.
 
@@ -1312,15 +1334,18 @@ def publish_server_heartbeat_succeeded(self, connection_id, duration,
          - `duration`: The execution time of the event in the highest possible
             resolution for the platform.
          - `reply`: The command reply.
+         - `awaited`: True if the response was awaited.
          """
-        event = ServerHeartbeatSucceededEvent(duration, reply, connection_id)
+        event = ServerHeartbeatSucceededEvent(duration, reply, connection_id,
+                                              awaited)
         for subscriber in self.__server_heartbeat_listeners:
             try:
                 subscriber.succeeded(event)
             except Exception:
                 _handle_exception()
 
-    def publish_server_heartbeat_failed(self, connection_id, duration, reply):
+    def publish_server_heartbeat_failed(self, connection_id, duration, reply,
+                                        awaited):
         """Publish a ServerHeartbeatFailedEvent to all server heartbeat
         listeners.
 
@@ -1329,8 +1354,10 @@ def publish_server_heartbeat_failed(self, connection_id, duration, reply):
          - `duration`: The execution time of the event in the highest possible
             resolution for the platform.
          - `reply`: The command reply.
+         - `awaited`: True if the response was awaited.
          """
-        event = ServerHeartbeatFailedEvent(duration, reply, connection_id)
+        event = ServerHeartbeatFailedEvent(duration, reply, connection_id,
+                                           awaited)
         for subscriber in self.__server_heartbeat_listeners:
             try:
                 subscriber.failed(event)

test/test_monitoring.py

@@ -1466,13 +1466,13 @@ def test_server_heartbeat_event_repr(self):
         self.assertEqual(
             repr(event),
             "<ServerHeartbeatSucceededEvent ('localhost', 27017) "
-            "duration: 0.1, reply: {'ok': 1}>")
+            "duration: 0.1, awaited: False, reply: {'ok': 1}>")
         event = monitoring.ServerHeartbeatFailedEvent(
             delta, 'ERROR', connection_id)
         self.assertEqual(
             repr(event),
             "<ServerHeartbeatFailedEvent ('localhost', 27017) "
-            "duration: 0.1, reply: 'ERROR'>")
+            "duration: 0.1, awaited: False, reply: 'ERROR'>")
 
     def test_server_event_repr(self):
         server_address = ('localhost', 27017)

test/test_streaming_protocol.py

@@ -160,6 +160,14 @@ def hb_started(event):
             return (isinstance(event, monitoring.ServerHeartbeatStartedEvent)
                     and event.connection_id == address)
 
+        def hb_succeeded(event):
+            return (isinstance(event, monitoring.ServerHeartbeatSucceededEvent)
+                    and event.connection_id == address)
+
+        def hb_failed(event):
+            return (isinstance(event, monitoring.ServerHeartbeatFailedEvent)
+                    and event.connection_id == address)
+
         hb_started_events = hb_listener.matching(hb_started)
         # Explanation of the expected heartbeat events:
         # Time: event
@@ -178,6 +186,14 @@ def hb_started(event):
         # This can be reduced to ~15 after SERVER-49220 is fixed.
         self.assertLess(len(hb_started_events), 40)
 
+        # Check the awaited flag.
+        hb_succeeded_events = hb_listener.matching(hb_succeeded)
+        hb_failed_events = hb_listener.matching(hb_failed)
+        self.assertFalse(hb_succeeded_events[0].awaited)
+        self.assertTrue(hb_succeeded_events[1].awaited)
+        self.assertTrue(hb_failed_events[0].awaited)
+        self.assertFalse(hb_failed_events[1].awaited)
+
 
 if __name__ == "__main__":
     unittest.main()