schema-agreement: retry busy control connection

Author: dkropachev

cassandra/cluster.py

@@ -4312,7 +4312,23 @@ def _wait_for_schema_agreement(self, connection=None, preloaded_results=None, wa
                               "response during schema agreement check: %s", timeout)
                     elapsed = self._time.time() - start
                     continue
-                except (ConnectionException, ConnectionBusy) as exc:
+                except ConnectionBusy as exc:
+                    elapsed = self._time.time() - start
+                    if schema_mismatches is None:
+                        fallback_wait = total_timeout - elapsed
+                        fallback = self._wait_for_schema_agreement_through_session(fallback_wait)
+                        if fallback is not None:
+                            return fallback
+                        raise
+
+                    log.debug("[control connection] Connection busy during schema agreement check: %s",
+                              exc)
+                    remaining = total_timeout - elapsed
+                    if remaining > 0:
+                        self._time.sleep(min(0.2, remaining))
+                        elapsed = self._time.time() - start
+                    continue
+                except ConnectionException as exc:
                     if isinstance(exc, ConnectionShutdown) and self._is_shutdown:
                         log.debug("[control connection] Aborting wait for schema match due to shutdown")
                         return None

tests/unit/test_control_connection.py

@@ -323,17 +323,23 @@ def test_wait_for_schema_agreement_does_not_accept_session_fallback_after_known_
         assert not self.control_connection.wait_for_schema_agreement()
         session.wait_for_schema_agreement.assert_not_called()
 
-    def test_wait_for_schema_agreement_does_not_accept_session_fallback_after_mismatch_then_busy(self):
+    def test_wait_for_schema_agreement_retries_control_connection_after_mismatch_then_busy(self):
         session = Mock(is_shutdown=False)
         session.wait_for_schema_agreement.return_value = True
         self.cluster.sessions = [session]
-        self.connection.peer_results[1][1][2] = 'b'
+
+        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'
+        matching_peer_rows = [list(row) for row in self.connection.peer_results[1]]
         self.connection.wait_for_responses.side_effect = [
-            _node_meta_results(self.connection.local_results, self.connection.peer_results),
-            ConnectionBusy("overloaded")]
+            _node_meta_results(self.connection.local_results, (peer_columns, mismatching_peer_rows)),
+            ConnectionBusy("overloaded"),
+            _node_meta_results(self.connection.local_results, (peer_columns, matching_peer_rows))]
 
-        assert not self.control_connection.wait_for_schema_agreement()
+        assert self.control_connection.wait_for_schema_agreement()
         session.wait_for_schema_agreement.assert_not_called()
+        assert self.connection.wait_for_responses.call_count == 3
 
     def test_wait_for_schema_agreement_does_not_exceed_configured_wait_with_session_fallback(self):
         session = Mock(is_shutdown=False)