cluster/pool: handle stale up and replacement races

Author: dkropachev

cassandra/cluster.py

@@ -2411,8 +2411,9 @@ def _clear_stale_reconnector():
                     host, is_host_addition=False, allow_retry_after_auth_failure=True)
                 if future is not None:
                     have_future = True
-                    future.add_done_callback(callback)
                     futures.add(future)
+            for future in tuple(futures):
+                future.add_done_callback(callback)
         except Exception:
             log.exception("Unexpected failure handling node %s being marked up:", host)
             for future in futures:

cassandra/pool.py

@@ -705,7 +705,13 @@ def _replace(self, connection):
             return
 
         if keyspace:
-            replacement_connection.set_keyspace_blocking(keyspace)
+            try:
+                replacement_connection.set_keyspace_blocking(keyspace)
+            except Exception:
+                log.warning("Failed reconnecting %s. Retrying." % (expected_endpoint,))
+                replacement_connection.close()
+                self._session.submit(self._replace, connection)
+                return
 
         with self._lock:
             if self.is_shutdown:

tests/unit/test_cluster.py

@@ -1090,6 +1090,34 @@ def test_noop_down_during_up_handling_does_not_supersede_up(self):
         assert host.is_up
         assert state.up_epoch is None
 
+    def test_on_up_waits_for_all_pool_futures_when_one_is_already_done(self):
+        completed_pool_future = Future()
+        completed_pool_future.set_result(True)
+        pending_pool_future = Future()
+        completed_session = Mock()
+        completed_session.add_or_renew_pool.return_value = completed_pool_future
+        pending_session = Mock()
+        pending_session.add_or_renew_pool.return_value = pending_pool_future
+        listener = Mock()
+        cluster = self._make_cluster(listener=listener)
+        cluster.sessions = [completed_session, pending_session]
+        cluster._prepare_all_queries = Mock()
+        cluster.profile_manager.distance.return_value = HostDistance.LOCAL
+        host = self._make_host()
+        host.set_down()
+
+        Cluster.on_up(cluster, host)
+
+        assert not host.is_up
+        listener.on_up.assert_not_called()
+
+        pending_pool_future.set_result(False)
+
+        assert not host.is_up
+        listener.on_up.assert_not_called()
+        cluster._start_reconnector.assert_called_once_with(
+            host, is_host_addition=False, expected_endpoint=host.endpoint)
+
     def test_newer_forced_down_during_up_handling_is_preserved(self):
         pool_future = Future()
         session = Mock()

tests/unit/test_host_connection_pool.py

@@ -350,3 +350,35 @@ def mock_connection_factory(self, *args, **kwargs):
 
             # Cleanup executor with proper wait
             session.cluster.executor.shutdown(wait=True)
+
+    def test_replace_retries_when_replacement_keyspace_set_fails(self):
+        host = Host(DefaultEndPoint('127.0.0.1'), SimpleConvictionPolicy,
+                    host_id=uuid.uuid4())
+        session = NonCallableMagicMock(spec=Session, keyspace='ks')
+        session.cluster = MagicMock()
+        session.cluster.shard_aware_options = ShardAwareOptions()
+        session.cluster._endpoints_match.side_effect = Cluster._endpoints_match
+        initial_connection = HashableMock(
+            spec=Connection, in_flight=0, is_defunct=False, is_closed=False,
+            max_request_id=100, signaled_error=False,
+            orphaned_threshold_reached=False,
+            features=ProtocolFeatures(shard_id=0))
+        replacement_connection = HashableMock(
+            spec=Connection, in_flight=0, is_defunct=False, is_closed=False,
+            max_request_id=100, signaled_error=False,
+            orphaned_threshold_reached=False,
+            features=ProtocolFeatures(shard_id=0))
+        replacement_connection.set_keyspace_blocking.side_effect = RuntimeError(
+            "keyspace failed")
+        session.cluster.connection_factory.side_effect = [
+            initial_connection, replacement_connection]
+
+        pool = HostConnection(host, HostDistance.LOCAL, session)
+        pool._is_replacing = True
+
+        pool._replace(initial_connection)
+
+        assert session.submit.call_count == 1
+        submitted_fn, submitted_connection = session.submit.call_args.args
+        assert submitted_fn == pool._replace
+        assert submitted_connection is initial_connection