Restore timeout diagnostics and libev cleanup

Author: dkropachev

cassandra/__init__.py

@@ -687,10 +687,29 @@ class OperationTimedOut(DriverException):
     The last :class:`~.Host` this operation was attempted against.
     """
 
-    def __init__(self, errors=None, last_host=None):
+    timeout = None
+    """
+    The timeout value (in seconds) that was in effect when the operation
+    timed out, or ``None`` if not applicable.
+    """
+
+    in_flight = None
+    """
+    The number of in-flight requests on the connection at the time of
+    the timeout (includes orphaned requests), or ``None`` if not applicable.
+    """
+
+    def __init__(self, errors=None, last_host=None, timeout=None, in_flight=None):
         self.errors = errors
         self.last_host = last_host
+        self.timeout = timeout
+        self.in_flight = in_flight
         message = "errors=%s, last_host=%s" % (self.errors, self.last_host)
+        if self.timeout is not None:
+            message += " (timeout=%ss" % self.timeout
+            if self.in_flight is not None:
+                message += ", in_flight=%d" % self.in_flight
+            message += ")"
         Exception.__init__(self, message)
 
 

cassandra/cluster.py

@@ -1694,7 +1694,9 @@ def add_execution_profile(self, name, profile, pool_wait_timeout=5):
             futures.update(session.update_created_pools())
         _, not_done = wait_futures(futures, pool_wait_timeout)
         if not_done:
-            raise OperationTimedOut("Failed to create all new connection pools in the %ss timeout.")
+            raise OperationTimedOut(
+                "Failed to create all new connection pools in the %ss timeout." % pool_wait_timeout,
+                timeout=pool_wait_timeout)
 
     def connection_factory(self, endpoint, host_conn = None, *args, **kwargs):
         """
@@ -4794,6 +4796,7 @@ def _on_timeout(self, _attempts=0):
             )
             return
 
+        conn_in_flight = None
         if self._connection is not None:
             try:
                 self._connection._requests.pop(self._req_id)
@@ -4804,9 +4807,13 @@ def _on_timeout(self, _attempts=0):
             except KeyError:
                 key = "Connection defunct by heartbeat"
                 errors = {key: "Client request timeout. See Session.execute[_async](timeout)"}
-                self._set_final_exception(OperationTimedOut(errors, self._current_host))
+                self._set_final_exception(OperationTimedOut(errors, self._current_host,
+                                                            timeout=self.timeout,
+                                                            in_flight=self._connection.in_flight))
                 return
 
+            # Capture connection stats before pool.return_connection() can alter state.
+            conn_in_flight = self._connection.in_flight
             pool = self.session._pools.get(self._current_host)
             if pool and not pool.is_shutdown:
                 # Do not return the stream ID to the pool yet. We cannot reuse it
@@ -4831,7 +4838,9 @@ def _on_timeout(self, _attempts=0):
                 host = str(connection.endpoint) if connection else 'unknown'
                 errors = {host: "Request timed out while waiting for schema agreement. See Session.execute[_async](timeout) and Cluster.max_schema_agreement_wait."}
 
-        self._set_final_exception(OperationTimedOut(errors, self._current_host))
+        self._set_final_exception(OperationTimedOut(errors, self._current_host,
+                                                    timeout=self.timeout,
+                                                    in_flight=conn_in_flight))
 
     def _on_speculative_execute(self):
         self._timer = None

cassandra/connection.py

@@ -1827,7 +1827,10 @@ def wait(self, timeout):
             if self._exception:
                 raise self._exception
         else:
-            raise OperationTimedOut("Connection heartbeat timeout after %s seconds" % (timeout,), self.connection.endpoint)
+            raise OperationTimedOut("Connection heartbeat timeout after %s seconds" % (timeout,),
+                                    self.connection.endpoint,
+                                    timeout=timeout,
+                                    in_flight=self.connection.in_flight)
 
     def _options_callback(self, response):
         if isinstance(response, SupportedMessage):

cassandra/io/libevreactor.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 import atexit
 from collections import deque
-from functools import partial
 import logging
 import os
 import socket
@@ -45,6 +44,12 @@ def _cleanup(loop):
         loop._cleanup()
 
 
+def _atexit_cleanup():
+    global _global_loop
+    if _global_loop is not None:
+        _cleanup(_global_loop)
+
+
 class LibevLoop(object):
 
     def __init__(self):
@@ -233,7 +238,7 @@ def _loop_will_run(self, prepare):
 
 
 _global_loop = None
-atexit.register(partial(_cleanup, _global_loop))
+atexit.register(_atexit_cleanup)
 
 
 class LibevConnection(Connection):

tests/unit/io/test_libevreactor.py

@@ -88,6 +88,19 @@ def test_watchers_are_finished(self):
 
         _global_loop._shutdown = False
 
+    def test_atexit_cleanup_uses_current_global_loop(self):
+        from cassandra.io import libevreactor
+
+        old_loop = libevreactor._global_loop
+        current_loop = Mock()
+        libevreactor._global_loop = current_loop
+        self.addCleanup(setattr, libevreactor, '_global_loop', old_loop)
+
+        with patch.object(libevreactor, '_cleanup') as cleanup:
+            libevreactor._atexit_cleanup()
+
+        cleanup.assert_called_once_with(current_loop)
+
 
 class LibevTimerPatcher(unittest.TestCase):
 

tests/unit/test_connection.py

@@ -520,6 +520,8 @@ def send_msg(msg, req_id, msg_callback):
         assert isinstance(exc, OperationTimedOut)
         assert exc.errors == 'Connection heartbeat timeout after 0.05 seconds'
         assert exc.last_host == DefaultEndPoint('localhost')
+        assert exc.timeout == 0.05
+        assert isinstance(exc.in_flight, int)
         holder.return_connection.assert_has_calls(
             [call(connection)] * get_holders.call_count)
 

tests/unit/test_response_future.py

@@ -142,6 +142,8 @@ def test_heartbeat_defunct_deadlock(self):
 
         connection = MagicMock(spec=Connection)
         connection._requests = {}
+        connection.in_flight = 5
+        connection.orphaned_request_ids = set()
 
         pool = Mock()
         pool.is_shutdown = False
@@ -162,8 +164,10 @@ def test_heartbeat_defunct_deadlock(self):
 
         # Simulate ResponseFuture timing out
         rf._on_timeout()
-        with pytest.raises(OperationTimedOut, match="Connection defunct by heartbeat"):
+        with pytest.raises(OperationTimedOut, match="Connection defunct by heartbeat") as exc_info:
             rf.result()
+        assert exc_info.value.timeout == 1
+        assert exc_info.value.in_flight == 5
 
     def test_read_timeout_error_message(self):
         session = self.make_session()
@@ -653,7 +657,7 @@ def test_timeout_does_not_release_stream_id(self):
         pool = self.make_pool()
         session._pools.get.return_value = pool
         connection = Mock(spec=Connection, lock=RLock(), _requests={}, request_ids=deque(),
-                orphaned_request_ids=set(), orphaned_threshold=256)
+                orphaned_request_ids=set(), orphaned_threshold=256, in_flight=3)
         pool.borrow_connection.return_value = (connection, 1)
 
         rf = self.make_response_future(session)
@@ -663,8 +667,10 @@ def test_timeout_does_not_release_stream_id(self):
 
         rf._on_timeout()
         pool.return_connection.assert_called_once_with(connection, stream_was_orphaned=True)
-        with pytest.raises(OperationTimedOut, match="Client request timeout"):
+        with pytest.raises(OperationTimedOut, match="Client request timeout") as exc_info:
             rf.result()
+        assert exc_info.value.timeout == 1
+        assert exc_info.value.in_flight == 3
 
         assert len(connection.request_ids) == 0, \
             "Request IDs should be empty but it's not: {}".format(connection.request_ids)