control-connection: reconnect after late endpoint swap

Author: dkropachev

cassandra/cluster.py

@@ -2550,7 +2550,8 @@ def on_down_potentially_blocking(
             self, host: Host, is_host_addition: bool,
             down_epoch: Optional[int] = None,
             expected_endpoint: Optional[EndPoint] = None,
-            profile_manager_already_notified: bool = False) -> Any:
+            profile_manager_already_notified: bool = False,
+            control_connection_already_notified: bool = False) -> Any:
         pending_up_epoch = None
         try:
             down_endpoint = None
@@ -2580,7 +2581,8 @@ def on_down_potentially_blocking(
                     else:
                         if not profile_manager_already_notified:
                             self.profile_manager.on_down(host)
-                        self.control_connection.on_down(host)
+                        if not control_connection_already_notified:
+                            self.control_connection.on_down(host)
             else:
                 log.debug("Not signalling down for stale down handling on node %s; endpoint changed from %s",
                           host, expected_endpoint)
@@ -2640,7 +2642,8 @@ def on_down_potentially_blocking(
                 self._handle_pending_node_up(host, pending_up_epoch)
 
     def on_down(self, host, is_host_addition, expect_host_to_be_down=False,
-                expected_endpoint=None, profile_manager_already_notified=False):
+                expected_endpoint=None, profile_manager_already_notified=False,
+                control_connection_already_notified=False):
         """
         Intended for internal use only.
         """
@@ -2725,7 +2728,8 @@ def on_down(self, host, is_host_addition, expect_host_to_be_down=False,
 
         future = self.on_down_potentially_blocking(
             host, is_host_addition, down_epoch, expected_endpoint,
-            profile_manager_already_notified)
+            profile_manager_already_notified,
+            control_connection_already_notified)
         if future is None:
             pending_down = None
             pending_up_epoch = None
@@ -5024,18 +5028,20 @@ def _refresh_node_list_and_token_map(self, connection, preloaded_results=None,
                     reconnector.cancel()
                 with host.lock:
                     old_endpoint = host.endpoint
+                    self._cluster.profile_manager.on_down(host)
+                    self.on_down(host)
                 self._cluster.on_down(
                     host, is_host_addition=False, expect_host_to_be_down=True,
                     expected_endpoint=old_endpoint,
-                    profile_manager_already_notified=True)
+                    profile_manager_already_notified=True,
+                    control_connection_already_notified=True)
 
                 with host.lock:
                     if not self._cluster._endpoints_match(host.endpoint, old_endpoint):
                         log.debug("[control connection] Not updating host ip from %s to %s for (%s); "
                                   "endpoint changed to %s",
                                   old_endpoint, endpoint, host_id, host.endpoint)
                         continue
-                    self._cluster.profile_manager.on_down(host)
                     host.endpoint = endpoint
                     self._cluster.metadata.update_host(host, old_endpoint)
                 self._cluster.on_up(host, expected_endpoint=endpoint)

tests/unit/test_cluster.py

@@ -1199,7 +1199,7 @@ def test_forced_down_is_not_discounted_by_connected_pool(self):
 
         assert not host.is_up
         cluster.on_down_potentially_blocking.assert_called_once_with(
-            host, False, ANY, endpoint)
+            host, False, ANY, endpoint, False, False)
 
     @staticmethod
     def _state(cluster, host):

tests/unit/test_control_connection.py

@@ -134,7 +134,8 @@ def on_up(self, host, expected_endpoint=None):
         pass
 
     def on_down(self, host, is_host_addition, expect_host_to_be_down=False,
-                expected_endpoint=None, profile_manager_already_notified=False):
+                expected_endpoint=None, profile_manager_already_notified=False,
+                control_connection_already_notified=False):
         self.down_host = host
         self.down_expected_endpoint = expected_endpoint
 
@@ -463,7 +464,8 @@ def test_change_client_route_endpoint_when_only_port_changes(self):
         self.cluster.on_down.assert_called_once_with(
             host, is_host_addition=False, expect_host_to_be_down=True,
             expected_endpoint=old_endpoint,
-            profile_manager_already_notified=True)
+            profile_manager_already_notified=True,
+            control_connection_already_notified=True)
         self.cluster.on_up.assert_called_once_with(
             host, expected_endpoint=new_endpoint)
 
@@ -504,6 +506,44 @@ def test_endpoint_change_preserves_live_policy_hosts_when_down_handler_runs_late
             cluster.executor = original_executor
             cluster.shutdown()
 
+    def test_endpoint_change_reconnects_control_connection_when_down_handler_runs_late(self):
+        old_endpoint = DefaultEndPoint("127.0.0.1")
+        new_endpoint = DefaultEndPoint("127.0.0.2")
+        host_id = uuid.uuid4()
+        cluster = Cluster(contact_points=[])
+        original_executor = cluster.executor
+        cluster.executor = RunOnResultExecutor()
+        cluster._start_reconnector = Mock()
+        cluster.control_connection._connection = Mock(endpoint=old_endpoint)
+        cluster.control_connection.reconnect = Mock()
+
+        try:
+            host = Host(
+                old_endpoint, SimpleConvictionPolicy,
+                datacenter="dc1", rack="rack1", host_id=host_id)
+            host.set_up()
+            cluster.metadata.add_or_return_host(host)
+            cluster.profile_manager.populate(cluster, [host])
+            cluster.endpoint_factory = Mock()
+            cluster.endpoint_factory.create.return_value = new_endpoint
+            cluster.control_connection._token_meta_enabled = False
+
+            preloaded_results = _node_meta_results(
+                local_results=([], []),
+                peer_results=(
+                    ["rpc_address", "peer", "data_center", "rack", "host_id"],
+                    [["127.0.0.2", "127.0.0.2", "dc1", "rack1", host_id]]))
+
+            cluster.control_connection._refresh_node_list_and_token_map(
+                Mock(), preloaded_results=preloaded_results)
+            cluster.executor.run_all()
+
+            cluster.control_connection.reconnect.assert_called_once_with()
+        finally:
+            cluster.control_connection._connection = None
+            cluster.executor = original_executor
+            cluster.shutdown()
+
     def test_stale_control_connection_failure_is_endpoint_fenced(self):
         host_id = uuid.uuid4()
         old_endpoint = ClientRoutesEndPoint(