connection: avoid double heartbeat in-flight release

Author: dkropachev

cassandra/connection.py

@@ -1820,8 +1820,12 @@ def __init__(self, connection, owner):
                 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:
+                    if connection.is_control_connection:
+                        connection.in_flight -= 1
+                    # send_msg() registers the callback before writing to the socket,
+                    # so a write failure must unwind that registration here.
+                    connection._requests.pop(request_id, None)
+                    if request_id not in connection.request_ids:
                         connection.request_ids.append(request_id)
                     self._exception = exc
                     self._event.set()

tests/unit/test_connection.py

@@ -466,14 +466,24 @@ def test_no_req_ids(self, *args):
     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"))
+        owner = Mock()
         initial_in_flight = connection.in_flight
         initial_request_ids = len(connection.request_ids)
 
-        future = HeartbeatFuture(connection, Mock())
+        # HostConnection.return_connection releases the heartbeat's in-flight slot.
+        def return_connection(conn):
+            with conn.lock:
+                conn.in_flight -= 1
+
+        owner.return_connection.side_effect = return_connection
+
+        future = HeartbeatFuture(connection, owner)
 
         with pytest.raises(ConnectionException):
             future.wait(0)
 
+        owner.return_connection(connection)
+
         assert connection.in_flight == initial_in_flight
         assert len(connection.request_ids) == initial_request_ids
         assert not connection._requests