Update TokenAwarePolicy.make_query_plan to schedule to replicas first

End result should be like this:
1. Replicas from the same RACK (if rack is specified)
2. Replicas from the same DC, but Remote RACK (if rack is specified)
3. Replicas from the same DC (if rack is not specified)
3. Replicas from the remote DC
5. Non-replicas in the same order

Author: dkropachev

cassandra/policies.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import random
+import threading
 
 from collections import namedtuple
 from itertools import islice, cycle, groupby, repeat
@@ -526,14 +527,23 @@ def make_query_plan(self, working_keyspace=None, query=None):
         if self.shuffle_replicas:
             shuffle(replicas)
 
-        for replica in replicas:
-            if replica.is_up and child.distance(replica) in [HostDistance.LOCAL, HostDistance.LOCAL_RACK]:
-                yield replica
+        def yield_in_order(hosts):
+            for replica in hosts:
+                if replica.is_up and child.distance(replica) == HostDistance.LOCAL_RACK:
+                    yield replica
 
-        for host in child.make_query_plan(keyspace, query):
-            # skip if we've already listed this host
-            if host not in replicas or child.distance(host) == HostDistance.REMOTE:
-                yield host
+            for replica in hosts:
+                if replica.is_up and child.distance(replica) == HostDistance.LOCAL:
+                    yield replica
+
+            for replica in hosts:
+                if replica.is_up and child.distance(replica) == HostDistance.REMOTE:
+                    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])
 
     def on_up(self, *args, **kwargs):
         return self._child_policy.on_up(*args, **kwargs)

cassandra/pool.py

@@ -250,7 +250,8 @@ def __str__(self):
 
     def __repr__(self):
         dc = (" %s" % (self._datacenter,)) if self._datacenter else ""
-        return "<%s: %s%s>" % (self.__class__.__name__, self.endpoint, dc)
+        rack = (" %s" % (self._rack,)) if self._rack else ""
+        return "<%s: %s%s%s>" % (self.__class__.__name__, self.endpoint, dc, rack)
 
 
 class _ReconnectionHandler(object):

tests/unit/test_policies.py

@@ -636,7 +636,7 @@ def get_replicas(keyspace, packed_key):
 
         cluster.metadata.get_replicas.side_effect = get_replicas
 
-        policy = TokenAwarePolicy(DCAwareRoundRobinPolicy("dc1", used_hosts_per_remote_dc=1))
+        policy = TokenAwarePolicy(DCAwareRoundRobinPolicy("dc1", used_hosts_per_remote_dc=2))
         policy.populate(cluster, hosts)
 
         for i in range(4):
@@ -648,14 +648,74 @@ def get_replicas(keyspace, packed_key):
             assert qplan[0] in replicas
             assert qplan[0].datacenter == "dc1"
 
-            # then the local non-replica
-            assert qplan[1] not in replicas
+            # then the replica from remote DC
+            assert qplan[1] in replicas
+            assert qplan[1].datacenter == "dc2"
+
+            # then non-replica from local DC
+            assert qplan[2] not in replicas
+            assert qplan[2].datacenter == "dc1"
+
+            # and only then non-replica from remote DC
+            assert qplan[3] not in replicas
+            assert qplan[3].datacenter == "dc2"
+
+            assert 4 == len(qplan)
+
+    def test_wrap_rack_aware(self):
+        cluster = Mock(spec=Cluster)
+        cluster.metadata = Mock(spec=Metadata)
+        cluster.control_connection._tablets_routing_v1 = False
+        hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy) for i in range(8)]
+        for host in hosts:
+            host.set_up()
+        hosts[0].set_location_info("dc1", "rack1")
+        hosts[1].set_location_info("dc1", "rack2")
+        hosts[2].set_location_info("dc2", "rack1")
+        hosts[3].set_location_info("dc2", "rack2")
+        hosts[4].set_location_info("dc1", "rack1")
+        hosts[5].set_location_info("dc1", "rack2")
+        hosts[6].set_location_info("dc2", "rack1")
+        hosts[7].set_location_info("dc2", "rack2")
+
+        def get_replicas(keyspace, packed_key):
+            index = struct.unpack('>i', packed_key)[0]
+            # return one node from each DC
+            if index % 2 == 0:
+                return [hosts[0], hosts[1], hosts[2], hosts[3]]
+            else:
+                return [hosts[4], hosts[5], hosts[6], hosts[7]]
+
+        cluster.metadata.get_replicas.side_effect = get_replicas
+
+        policy = TokenAwarePolicy(RackAwareRoundRobinPolicy("dc1", "rack1", used_hosts_per_remote_dc=4))
+        policy.populate(cluster, hosts)
+
+        for i in range(4):
+            query = Statement(routing_key=struct.pack('>i', i), keyspace='keyspace_name')
+            qplan = list(policy.make_query_plan(None, query))
+            replicas = get_replicas(None, struct.pack('>i', i))
+
+            print(qplan)
+            print(replicas)
+
+            # first should be replica from local rack local dc
+            assert qplan[0] in replicas
+            assert qplan[0].datacenter == "dc1"
+            assert qplan[0].rack == "rack1"
+
+            # second should be replica from remote rack local dc
+            assert qplan[1] in replicas
             assert qplan[1].datacenter == "dc1"
+            assert qplan[1].rack == "rack2"
 
-            # then one of the remotes (used_hosts_per_remote_dc is 1, so we
-            # shouldn't see two remotes)
+            # third and forth should be replica from the remote dcs
+            assert qplan[2] in replicas
             assert qplan[2].datacenter == "dc2"
-            assert 3 == len(qplan)
+            assert qplan[3] in replicas
+            assert qplan[3].datacenter == "dc2"
+
+            assert 8 == len(qplan)
 
     class FakeCluster:
         def __init__(self):
@@ -1403,7 +1463,7 @@ def setUp(self):
 
     def _check_init(self, hfp):
         assert hfp._child_policy is self.child_policy
-        assert isinstance(hfp._hosts_lock, LockType)
+        assert isinstance(hfp._known_hosts_lock, LockType)
 
         # we can't use a simple assertIs because we wrap the function
         arg0, arg1 = Mock(name='arg0'), Mock(name='arg1')