default_policy: Prefer local rack nodes

Until now, the default policy respected rack preference only when it
came to replicas. Considering non-token-aware queries (e.g. unprepared
statements), they would never prioritize the preferred rack.
Now, local rack nodes are always preferred before other nodes, both in
`pick()` and in `fallback()`. Tests are adjusted for this behaviour.

scylla/src/transport/load_balancing/default.rs

@@ -170,8 +170,22 @@ or refrain from preferring datacenters (which may ban all other datacenters, if
         // some alive local node.
         // If there was no preferred datacenter specified, all nodes are treated as local.
         let nodes = self.preferred_node_set(cluster);
-        let picked = Self::pick_node(nodes, &self.pick_predicate);
-        if let Some(alive_local) = picked {
+
+        if let ReplicaLocationPreference::DatacenterAndRack(dc, rack) = &self.preferences {
+            // Try to pick some alive local rack random node.
+            let rack_predicate = Self::make_rack_predicate(
+                &self.pick_predicate,
+                ReplicaLocationCriteria::DatacenterAndRack(dc, rack),
+            );
+            let local_rack_picked = Self::pick_node(nodes, rack_predicate);
+
+            if let Some(alive_local_rack) = local_rack_picked {
+                return Some(alive_local_rack);
+            }
+        }
+
+        // Try to pick some alive local random node.
+        if let Some(alive_local) = Self::pick_node(nodes, &self.pick_predicate) {
             return Some(alive_local);
         }
 
@@ -276,6 +290,17 @@ or refrain from preferring datacenters (which may ban all other datacenters, if
 
         // Get a list of all local alive nodes, and apply a round robin to it
         let local_nodes = self.preferred_node_set(cluster);
+
+        let maybe_local_rack_nodes =
+            if let ReplicaLocationPreference::DatacenterAndRack(dc, rack) = &self.preferences {
+                let rack_predicate = Self::make_rack_predicate(
+                    &self.pick_predicate,
+                    ReplicaLocationCriteria::DatacenterAndRack(dc, rack),
+                );
+                Either::Left(Self::round_robin_nodes(local_nodes, rack_predicate))
+            } else {
+                Either::Right(std::iter::empty::<NodeRef<'a>>())
+            };
         let robined_local_nodes = Self::round_robin_nodes(local_nodes, Self::is_alive);
 
         let all_nodes = cluster.replica_locator().unique_nodes_in_global_ring();
@@ -301,6 +326,7 @@ or refrain from preferring datacenters (which may ban all other datacenters, if
 
         // Construct a fallback plan as a composition of replicas, local nodes and remote nodes.
         let plan = maybe_replicas
+            .chain(maybe_local_rack_nodes)
             .chain(robined_local_nodes)
             .chain(maybe_remote_nodes)
             .chain(maybe_down_local_nodes)
@@ -1575,7 +1601,8 @@ mod tests {
                 expected_groups: ExpectedGroupsBuilder::new()
                     .deterministic([B]) // pick local rack replicas
                     .deterministic([C]) // fallback replicas
-                    .group([A, G]) // local nodes
+                    .group([A]) // local rack nodes
+                    .group([G]) // local DC nodes
                     .build(),
             },
             // Keyspace SS with RF=2 with enabled rack-awareness and no local-rack replica
@@ -1600,7 +1627,8 @@ mod tests {
                 //                                      r2  r1  r1  r1  r2  r1  r1
                 expected_groups: ExpectedGroupsBuilder::new()
                     .group([A]) // pick local DC
-                    .group([C, G, B]) // local nodes
+                    .group([G]) // local rack nodes
+                    .group([C, B]) // local DC nodes
                     .build(),
             },
         ];