policies: hide in-flight hosts from query plans

Author: dkropachev

cassandra/policies.py

@@ -143,6 +143,17 @@ def _is_ignored_zero_token_host(self, host):
     def _filter_zero_token_hosts(self, hosts):
         return tuple(h for h in hosts if not self._is_ignored_zero_token_host(h))
 
+    def _is_in_flight_host(self, host):
+        return ((getattr(host, '_currently_handling_node_up', False) is True or
+                 getattr(host, '_currently_handling_node_addition', False) is True) and
+                getattr(host, 'is_up', None) is not True)
+
+    def _is_ignored_query_plan_host(self, host):
+        return self._is_ignored_zero_token_host(host) or self._is_in_flight_host(host)
+
+    def _filter_query_plan_hosts(self, hosts):
+        return tuple(h for h in hosts if not self._is_ignored_query_plan_host(h))
+
     def distance(self, host):
         """
         Returns a measure of how remote a :class:`~.pool.Host` is in
@@ -211,7 +222,7 @@ def make_query_plan(self, working_keyspace=None, query=None):
         pos = self._position
         self._position += 1
 
-        hosts = self._filter_zero_token_hosts(self._live_hosts)
+        hosts = self._filter_query_plan_hosts(self._live_hosts)
         length = len(hosts)
         if length:
             pos %= length
@@ -303,15 +314,15 @@ def make_query_plan(self, working_keyspace=None, query=None):
         pos = self._position
         self._position += 1
 
-        local_live = self._filter_zero_token_hosts(self._dc_live_hosts.get(self.local_dc, ()))
+        local_live = self._filter_query_plan_hosts(self._dc_live_hosts.get(self.local_dc, ()))
         pos = (pos % len(local_live)) if local_live else 0
         for host in islice(cycle(local_live), pos, pos + len(local_live)):
             yield host
 
         # the dict can change, so get candidate DCs iterating over keys of a copy
         other_dcs = [dc for dc in self._dc_live_hosts.copy().keys() if dc != self.local_dc]
         for dc in other_dcs:
-            remote_live = self._filter_zero_token_hosts(self._dc_live_hosts.get(dc, ()))
+            remote_live = self._filter_query_plan_hosts(self._dc_live_hosts.get(dc, ()))
             for host in remote_live[:self.used_hosts_per_remote_dc]:
                 yield host
 
@@ -422,14 +433,14 @@ def make_query_plan(self, working_keyspace=None, query=None):
         pos = self._position
         self._position += 1
 
-        local_rack_live = self._filter_zero_token_hosts(self._live_hosts.get((self.local_dc, self.local_rack), ()))
+        local_rack_live = self._filter_query_plan_hosts(self._live_hosts.get((self.local_dc, self.local_rack), ()))
         pos = (pos % len(local_rack_live)) if local_rack_live else 0
         # Slice the cyclic iterator to start from pos and include the next len(local_live) elements
         # This ensures we get exactly one full cycle starting from pos
         for host in islice(cycle(local_rack_live), pos, pos + len(local_rack_live)):
             yield host
 
-        local_live = [host for host in self._filter_zero_token_hosts(self._dc_live_hosts.get(self.local_dc, ()))
+        local_live = [host for host in self._filter_query_plan_hosts(self._dc_live_hosts.get(self.local_dc, ()))
                       if host.rack != self.local_rack]
         pos = (pos % len(local_live)) if local_live else 0
         for host in islice(cycle(local_live), pos, pos + len(local_live)):
@@ -438,7 +449,7 @@ def make_query_plan(self, working_keyspace=None, query=None):
         # the dict can change, so get candidate DCs iterating over keys of a copy
         for dc, remote_live in self._dc_live_hosts.copy().items():
             if dc != self.local_dc:
-                remote_live = self._filter_zero_token_hosts(remote_live)
+                remote_live = self._filter_query_plan_hosts(remote_live)
                 for host in remote_live[:self.used_hosts_per_remote_dc]:
                     yield host
 
@@ -531,7 +542,7 @@ def make_query_plan(self, working_keyspace=None, query=None):
         child = self._child_policy
         if query is None or query.routing_key is None or keyspace is None:
             for host in child.make_query_plan(keyspace, query):
-                if not self._is_ignored_zero_token_host(host):
+                if not self._is_ignored_query_plan_host(host):
                     yield host
             return
 
@@ -553,15 +564,15 @@ def make_query_plan(self, working_keyspace=None, query=None):
         def yield_in_order(hosts):
             for distance in [HostDistance.LOCAL_RACK, HostDistance.LOCAL, HostDistance.REMOTE]:
                 for replica in hosts:
-                    if (not self._is_ignored_zero_token_host(replica) and
+                    if (not self._is_ignored_query_plan_host(replica) and
                             replica.is_up and child.distance(replica) == distance):
                         yield replica
 
         # yield replicas: local_rack, local, remote
         yield from yield_in_order(replicas)
         # yield rest of the cluster: local_rack, local, remote
         yield from yield_in_order([host for host in child.make_query_plan(keyspace, query)
-                                   if host not in replicas and not self._is_ignored_zero_token_host(host)])
+                                   if host not in replicas and not self._is_ignored_query_plan_host(host)])
 
     def on_up(self, *args, **kwargs):
         return self._child_policy.on_up(*args, **kwargs)
@@ -739,7 +750,7 @@ def make_query_plan(self, working_keyspace=None, query=None):
             working_keyspace=working_keyspace, query=query
         )
         for host in child_qp:
-            if not self._is_ignored_zero_token_host(host) and self.predicate(host):
+            if not self._is_ignored_query_plan_host(host) and self.predicate(host):
                 yield host
 
     def check_supported(self):
@@ -1390,14 +1401,14 @@ def make_query_plan(self, working_keyspace=None, query=None):
         target_host = self._cluster_metadata.get_host(addr)
 
         child = self._child_policy
-        if target_host and target_host.is_up and not self._is_ignored_zero_token_host(target_host):
+        if target_host and target_host.is_up and not self._is_ignored_query_plan_host(target_host):
             yield target_host
             for h in child.make_query_plan(keyspace, query):
-                if h != target_host and not self._is_ignored_zero_token_host(h):
+                if h != target_host and not self._is_ignored_query_plan_host(h):
                     yield h
         else:
             for h in child.make_query_plan(keyspace, query):
-                if not self._is_ignored_zero_token_host(h):
+                if not self._is_ignored_query_plan_host(h):
                     yield h
 
 

tests/unit/test_cluster.py

@@ -155,6 +155,27 @@ def test_on_add_waits_for_all_session_pool_futures_before_marking_host_up(self):
         finally:
             cluster.shutdown()
 
+    def test_on_add_excludes_host_from_query_plan_until_pool_futures_complete(self):
+        cluster = Cluster(protocol_version=4)
+        host = Host("127.0.0.1", SimpleConvictionPolicy, datacenter="dc1", rack="rack1", host_id=uuid.uuid4())
+        pending_future = Future()
+        session = Mock()
+        session.add_or_renew_pool.return_value = pending_future
+        session.update_created_pools.return_value = set()
+        cluster.sessions = [session]
+
+        try:
+            cluster.on_add(host, refresh_nodes=False)
+
+            load_balancer = cluster.profile_manager.default.load_balancing_policy
+            assert host not in list(load_balancer.make_query_plan())
+
+            pending_future.set_result(True)
+
+            assert list(load_balancer.make_query_plan()) == [host]
+        finally:
+            cluster.shutdown()
+
     def test_on_up_waits_for_all_session_pool_futures_before_marking_host_up(self):
         cluster = Cluster(protocol_version=4)
         host = Host("127.0.0.1", SimpleConvictionPolicy, host_id=uuid.uuid4())
@@ -186,6 +207,28 @@ def test_on_up_waits_for_all_session_pool_futures_before_marking_host_up(self):
         finally:
             cluster.shutdown()
 
+    def test_on_up_excludes_host_from_query_plan_until_pool_futures_complete(self):
+        cluster = Cluster(protocol_version=4)
+        host = Host("127.0.0.1", SimpleConvictionPolicy, datacenter="dc1", rack="rack1", host_id=uuid.uuid4())
+        host.set_down()
+        pending_future = Future()
+        session = Mock()
+        session.add_or_renew_pool.return_value = pending_future
+        session.update_created_pools.return_value = set()
+        cluster.sessions = [session]
+
+        try:
+            cluster.on_up(host)
+
+            load_balancer = cluster.profile_manager.default.load_balancing_policy
+            assert host not in list(load_balancer.make_query_plan())
+
+            pending_future.set_result(True)
+
+            assert list(load_balancer.make_query_plan()) == [host]
+        finally:
+            cluster.shutdown()
+
     def test_invalid_contact_point_types(self):
         with pytest.raises(ValueError):
             Cluster(contact_points=[None], protocol_version=4, connect_timeout=1)