load_balancing: hash round robin's index

When RoundRobinPolicy is used as a child policy, calling its plan method
twice will cause the round robin's index to increase by 2. This could
led to some nodes not receiving any requests.

Let's consider the following scenario:
- Round robin's index is even
- User sends a query with a token equal to X
- TokenAwarePolicy computes a set of replicas for this query - {A, B}
- TokenAwarePolicy runs a RoundRobin child policy on that set and gets
  {A, B}. Round robin's index is odd
- TokenAwarePolicy runs a RoundRobin child policy to generate fallback.
  Round robin's index is even

Because the list of steps starts and ends with the same parity of the
round robin index, the beginning of load balancing plan will be the same
for every query with token equal to X.

Hashing round robin index before rotating is a mitigation to this
problem (@piodul's suggestion).

Author: havaker

scylla/src/transport/load_balancing/dc_aware_round_robin.rs

@@ -119,6 +119,7 @@ mod tests {
     use super::*;
 
     use crate::transport::load_balancing::tests;
+    use std::collections::HashSet;
 
     #[tokio::test]
     async fn test_dc_aware_round_robin_policy() {
@@ -127,23 +128,27 @@ mod tests {
         let local_dc = "eu".to_string();
         let policy = DcAwareRoundRobinPolicy::new(local_dc);
 
-        let plans = (0..4)
+        let plans = (0..32)
             .map(|_| {
                 tests::get_plan_and_collect_node_identifiers(
                     &policy,
                     &tests::EMPTY_STATEMENT,
                     &cluster,
                 )
             })
-            .collect::<Vec<_>>();
+            .collect::<HashSet<_>>();
 
         let expected_plans = vec![
             vec![1, 2, 3, 4, 5],
+            vec![1, 2, 3, 5, 4],
             vec![2, 3, 1, 5, 4],
+            vec![2, 3, 1, 4, 5],
             vec![3, 1, 2, 4, 5],
-            vec![1, 2, 3, 5, 4],
-        ];
+            vec![3, 1, 2, 5, 4],
+        ]
+        .into_iter()
+        .collect::<HashSet<_>>();
 
-        assert_eq!(plans, expected_plans);
+        assert_eq!(expected_plans, plans);
     }
 }

scylla/src/transport/load_balancing/mod.rs

@@ -6,7 +6,7 @@
 use super::{cluster::ClusterData, node::Node};
 use crate::routing::Token;
 
-use std::sync::Arc;
+use std::{collections::hash_map::DefaultHasher, hash::Hasher, sync::Arc};
 
 mod dc_aware_round_robin;
 mod round_robin;
@@ -53,10 +53,21 @@ pub trait ChildLoadBalancingPolicy: LoadBalancingPolicy {
     ) -> Box<dyn Iterator<Item = Arc<Node>> + Send + Sync>;
 }
 
-// Does safe modulo
-fn compute_rotation(index: usize, count: usize) -> usize {
-    if count != 0 {
-        index % count
+// Hashing round robin's index is a mitigation to problems that occur when a
+// `RoundRobin::apply_child_policy()` is called twice by a parent policy.
+fn round_robin_index_hash(index: usize) -> u64 {
+    let mut hasher = DefaultHasher::new();
+    hasher.write_usize(index);
+
+    hasher.finish()
+}
+
+// Does safe modulo and additionally hashes the index
+fn compute_rotation(round_robin_index: usize, sequence_length: usize) -> usize {
+    if sequence_length > 1 {
+        let hash = round_robin_index_hash(round_robin_index);
+
+        (hash % sequence_length as u64) as usize
     } else {
         0
     }

scylla/src/transport/load_balancing/round_robin.rs

@@ -71,6 +71,7 @@ mod tests {
     use super::*;
 
     use crate::transport::load_balancing::tests;
+    use std::collections::HashSet;
 
     // ConnectionKeeper (which lives in Node) requires context of Tokio runtime
     #[tokio::test]
@@ -79,25 +80,27 @@ mod tests {
 
         let policy = RoundRobinPolicy::new();
 
-        let plans = (0..6)
+        let plans = (0..16)
             .map(|_| {
                 tests::get_plan_and_collect_node_identifiers(
                     &policy,
                     &tests::EMPTY_STATEMENT,
                     &cluster,
                 )
             })
-            .collect::<Vec<_>>();
+            .collect::<HashSet<_>>();
 
+        // Check if `plans` contains all possible round robin plans
         let expected_plans = vec![
             vec![1, 2, 3, 4, 5],
             vec![2, 3, 4, 5, 1],
             vec![3, 4, 5, 1, 2],
             vec![4, 5, 1, 2, 3],
             vec![5, 1, 2, 3, 4],
-            vec![1, 2, 3, 4, 5],
-        ];
+        ]
+        .into_iter()
+        .collect::<HashSet<Vec<_>>>();
 
-        assert_eq!(plans, expected_plans);
+        assert_eq!(expected_plans, plans);
     }
 }