connection: clean up failed async keyspace sends

Author: dkropachev

cassandra/connection.py

@@ -1745,7 +1745,13 @@ def process_result(result):
         # acquire a new request id
         request_id = self.get_request_id()
 
-        self.send_msg(query, request_id, process_result)
+        try:
+            self.send_msg(query, request_id, process_result)
+        except Exception as exc:
+            with self.lock:
+                if request_id not in self._requests and request_id not in self.request_ids:
+                    self.request_ids.append(request_id)
+            callback(self, exc)
 
     @property
     def is_idle(self):

tests/unit/test_connection.py

@@ -22,7 +22,7 @@
 from cassandra.cluster import Cluster
 from cassandra.connection import (Connection, HEADER_DIRECTION_TO_CLIENT, ProtocolError,
                                   locally_supported_compressions, ConnectionHeartbeat, _Frame, Timer, TimerManager,
-                                  ConnectionException, ConnectionShutdown, DefaultEndPoint, ShardAwarePortGenerator)
+                                  ConnectionException, ConnectionShutdown, ConnectionBusy, DefaultEndPoint, ShardAwarePortGenerator)
 from cassandra.marshal import uint8_pack, uint32_pack, int32_pack
 from cassandra.protocol import (write_stringmultimap, write_int, write_string,
                                 SupportedMessage, ProtocolHandler, ResultMessage, QueryMessage,
@@ -389,6 +389,26 @@ def test_wait_for_responses_releases_request_id_when_send_raises_after_registrat
         assert len(c.request_ids) == initial_request_ids
         assert not c._requests
 
+    def test_set_keyspace_async_reports_send_failure_and_releases_request_id(self):
+        c = self.make_connection()
+        c.push = Mock(side_effect=ConnectionException("write failed"))
+        initial_in_flight = c.in_flight
+        initial_request_ids = len(c.request_ids)
+        callback_errors = []
+
+        def callback(conn, error):
+            callback_errors.append(error)
+            with conn.lock:
+                conn.in_flight -= 1
+
+        c.set_keyspace_async("ks", callback)
+
+        assert len(callback_errors) == 1
+        assert isinstance(callback_errors[0], ConnectionException)
+        assert c.in_flight == initial_in_flight
+        assert len(c.request_ids) == initial_request_ids
+        assert not c._requests
+
 
 @patch('cassandra.connection.ConnectionHeartbeat._raise_if_stopped')
 class ConnectionHeartbeatTest(unittest.TestCase):