schema-agreement: fix retry deadline cleanup

Author: dkropachev

cassandra/cluster.py

@@ -4269,6 +4269,7 @@ def wait_for_schema_agreement(self, connection=None, preloaded_results=None, wai
                 if schema_mismatches is None:
                     return True
 
+                elapsed = self._time.time() - start
                 log.debug("[control connection] Schemas mismatched, trying again")
                 remaining = total_timeout - elapsed
                 if remaining > 0:

cassandra/connection.py

@@ -642,16 +642,25 @@ def maybe_request_more(self):
         if space_in_queue >= max_queue_size / 2:
             self.update_next_pages(space_in_queue)
 
+    def _send_revise_request(self, request, callback):
+        with self.connection.lock:
+            request_id = self.connection.get_request_id()
+            try:
+                self.connection.send_msg(request, request_id, callback)
+            except Exception:
+                if request_id not in self.connection._requests and request_id not in self.connection.request_ids:
+                    self.connection.request_ids.append(request_id)
+                raise
+
     def update_next_pages(self, num_next_pages):
         try:
             self._state.num_pages_requested += num_next_pages
             log.debug("Updating backpressure for session %s from %s", self.stream_id, self.connection.host)
-            with self.connection.lock:
-                self.connection.send_msg(ReviseRequestMessage(ReviseRequestMessage.RevisionType.PAGING_BACKPRESSURE,
-                                                              self.stream_id,
-                                                              next_pages=num_next_pages),
-                                         self.connection.get_request_id(),
-                                         self._on_backpressure_response)
+            self._send_revise_request(
+                ReviseRequestMessage(ReviseRequestMessage.RevisionType.PAGING_BACKPRESSURE,
+                                     self.stream_id,
+                                     next_pages=num_next_pages),
+                self._on_backpressure_response)
         except ConnectionShutdown as ex:
             log.debug("Failed to update backpressure for session %s from %s, connection is shutdown",
                       self.stream_id, self.connection.host)
@@ -668,11 +677,10 @@ def _on_backpressure_response(self, response):
     def cancel(self):
         try:
             log.debug("Canceling paging session %s from %s", self.stream_id, self.connection.host)
-            with self.connection.lock:
-                self.connection.send_msg(ReviseRequestMessage(ReviseRequestMessage.RevisionType.PAGING_CANCEL,
-                                                              self.stream_id),
-                                         self.connection.get_request_id(),
-                                         self._on_cancel_response)
+            self._send_revise_request(
+                ReviseRequestMessage(ReviseRequestMessage.RevisionType.PAGING_CANCEL,
+                                     self.stream_id),
+                self._on_cancel_response)
         except ConnectionShutdown:
             log.debug("Failed to cancel session %s from %s, connection is shutdown",
                       self.stream_id, self.connection.host)

tests/unit/test_connection.py

@@ -410,6 +410,19 @@ def callback(conn, error):
         assert len(c.request_ids) == initial_request_ids
         assert not c._requests
 
+    def test_continuous_paging_cancel_releases_request_id_when_send_fails(self):
+        c = self.make_connection()
+        c.push = Mock(side_effect=ConnectionException("write failed"))
+        state = Mock(max_queue_size=100, num_pages_requested=0, num_pages_received=0)
+        session = c.new_continuous_paging_session(1, Mock(), Mock(), state)
+        initial_request_ids = len(c.request_ids)
+
+        with pytest.raises(ConnectionException):
+            session.cancel()
+
+        assert len(c.request_ids) == initial_request_ids
+        assert not c._requests
+
 
 @patch('cassandra.connection.ConnectionHeartbeat._raise_if_stopped')
 class ConnectionHeartbeatTest(unittest.TestCase):

tests/unit/test_control_connection.py

@@ -377,6 +377,21 @@ def test_wait_for_schema_agreement_does_not_sleep_past_deadline_after_mismatch(s
         assert not self.control_connection.wait_for_schema_agreement()
         assert self.time.clock == self.cluster.max_schema_agreement_wait
 
+    def test_wait_for_schema_agreement_counts_query_time_before_mismatch_retry_sleep(self):
+        self.cluster.max_schema_agreement_wait = 0.1
+        peer_columns = self.connection.peer_results[0]
+        mismatching_peer_rows = [list(row) for row in self.connection.peer_results[1]]
+        mismatching_peer_rows[1][2] = 'b'
+
+        def wait_for_responses(*args, **kwargs):
+            self.time.sleep(0.09)
+            return _node_meta_results(self.connection.local_results, (peer_columns, mismatching_peer_rows))
+
+        self.connection.wait_for_responses.side_effect = wait_for_responses
+
+        assert not self.control_connection.wait_for_schema_agreement()
+        self.assertAlmostEqual(self.time.clock, self.cluster.max_schema_agreement_wait)
+
     def test_wait_for_schema_agreement_skipping(self):
         """
         If rpc_address or schema_version isn't set, the host should be skipped