connection: clean up failed heartbeat sends

Author: dkropachev

cassandra/connection.py

@@ -1833,7 +1833,15 @@ def __init__(self, connection, owner):
         with connection.lock:
             if connection.in_flight < connection.max_request_id:
                 connection.in_flight += 1
-                connection.send_msg(OptionsMessage(), connection.get_request_id(), self._options_callback)
+                request_id = connection.get_request_id()
+                try:
+                    connection.send_msg(OptionsMessage(), request_id, self._options_callback)
+                except Exception as exc:
+                    connection.in_flight -= 1
+                    if request_id not in connection._requests and request_id not in connection.request_ids:
+                        connection.request_ids.append(request_id)
+                    self._exception = exc
+                    self._event.set()
             else:
                 self._exception = Exception("Failed to send heartbeat because connection 'in_flight' exceeds threshold")
                 self._event.set()

tests/unit/test_connection.py

@@ -21,7 +21,7 @@
 from cassandra import ConsistencyLevel, OperationTimedOut
 from cassandra.cluster import Cluster
 from cassandra.connection import (Connection, HEADER_DIRECTION_TO_CLIENT, ProtocolError,
-                                  locally_supported_compressions, ConnectionHeartbeat, _Frame, Timer, TimerManager,
+                                  locally_supported_compressions, ConnectionHeartbeat, HeartbeatFuture, _Frame, Timer, TimerManager,
                                   ConnectionBusy, ConnectionException, ConnectionShutdown, DefaultEndPoint,
                                   ShardAwarePortGenerator)
 from cassandra.marshal import uint8_pack, uint32_pack, int32_pack
@@ -510,6 +510,21 @@ def test_no_req_ids(self, *args):
         holder.return_connection.assert_has_calls(
             [call(max_connection)] * get_holders.call_count)
 
+    def test_heartbeat_future_releases_request_id_when_send_fails(self, *args):
+        connection = Connection(DefaultEndPoint('1.2.3.4'))
+        connection.push = Mock(side_effect=ConnectionException("write failed"))
+        initial_in_flight = connection.in_flight
+        initial_request_ids = len(connection.request_ids)
+
+        future = HeartbeatFuture(connection, Mock())
+
+        with pytest.raises(ConnectionException):
+            future.wait(0)
+
+        assert connection.in_flight == initial_in_flight
+        assert len(connection.request_ids) == initial_request_ids
+        assert not connection._requests
+
     def test_unexpected_response(self, *args):
         request_id = 999