Merge pull request #405 from havaker/fallback

piodul · web-flow · commit 4ad601a77303 · 2022-03-03T13:31:54.000+01:00
Load balancing fallback
diff --git a/scylla/src/transport/load_balancing/dc_aware_round_robin.rs b/scylla/src/transport/load_balancing/dc_aware_round_robin.rs
@@ -1,4 +1,4 @@
-use super::{ChildLoadBalancingPolicy, LoadBalancingPolicy, Statement};
+use super::{ChildLoadBalancingPolicy, LoadBalancingPolicy, Plan, Statement};
 use crate::transport::{cluster::ClusterData, node::Node};
 use std::sync::{
     atomic::{AtomicUsize, Ordering},
@@ -51,11 +51,7 @@ const EMPTY_NODE_LIST: &Vec<Arc<Node>> = &vec![];
 const ORDER_TYPE: Ordering = Ordering::Relaxed;
 
 impl LoadBalancingPolicy for DcAwareRoundRobinPolicy {
-    fn plan<'a>(
-        &self,
-        _statement: &Statement,
-        cluster: &'a ClusterData,
-    ) -> Box<dyn Iterator<Item = Arc<Node>> + Send + Sync + 'a> {
+    fn plan<'a>(&self, _statement: &Statement, cluster: &'a ClusterData) -> Plan<'a> {
         let index = self.index.fetch_add(1, ORDER_TYPE);
 
         let local_nodes = self.retrieve_local_nodes(cluster);
@@ -123,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() {
@@ -131,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);
     }
 }
diff --git a/scylla/src/transport/load_balancing/mod.rs b/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;
@@ -23,14 +23,21 @@ pub struct Statement<'a> {
     pub keyspace: Option<&'a str>,
 }
 
+impl<'a> Statement<'a> {
+    fn empty() -> Self {
+        Self {
+            token: None,
+            keyspace: None,
+        }
+    }
+}
+
+pub type Plan<'a> = Box<dyn Iterator<Item = Arc<Node>> + Send + Sync + 'a>;
+
 /// Policy that decides which nodes to contact for each query
 pub trait LoadBalancingPolicy: Send + Sync {
     /// It is used for each query to find which nodes to query first
-    fn plan<'a>(
-        &self,
-        statement: &Statement,
-        cluster: &'a ClusterData,
-    ) -> Box<dyn Iterator<Item = Arc<Node>> + Send + Sync + 'a>;
+    fn plan<'a>(&self, statement: &Statement, cluster: &'a ClusterData) -> Plan<'a>;
 
     /// Returns name of load balancing policy
     fn name(&self) -> String;
@@ -46,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
     }
diff --git a/scylla/src/transport/load_balancing/round_robin.rs b/scylla/src/transport/load_balancing/round_robin.rs
@@ -1,4 +1,4 @@
-use super::{ChildLoadBalancingPolicy, LoadBalancingPolicy, Statement};
+use super::{ChildLoadBalancingPolicy, LoadBalancingPolicy, Plan, Statement};
 use crate::transport::{cluster::ClusterData, node::Node};
 use std::sync::{
     atomic::{AtomicUsize, Ordering},
@@ -28,11 +28,7 @@ impl Default for RoundRobinPolicy {
 const ORDER_TYPE: Ordering = Ordering::Relaxed;
 
 impl LoadBalancingPolicy for RoundRobinPolicy {
-    fn plan<'a>(
-        &self,
-        _statement: &Statement,
-        cluster: &'a ClusterData,
-    ) -> Box<dyn Iterator<Item = Arc<Node>> + Send + Sync + 'a> {
+    fn plan<'a>(&self, _statement: &Statement, cluster: &'a ClusterData) -> Plan<'a> {
         let index = self.index.fetch_add(1, ORDER_TYPE);
 
         let nodes_count = cluster.all_nodes.len();
@@ -75,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]
@@ -83,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);
     }
 }
diff --git a/scylla/src/transport/load_balancing/token_aware.rs b/scylla/src/transport/load_balancing/token_aware.rs
@@ -1,8 +1,10 @@
-use super::{ChildLoadBalancingPolicy, LoadBalancingPolicy, Statement};
+use super::{ChildLoadBalancingPolicy, LoadBalancingPolicy, Plan, Statement};
 use crate::routing::Token;
 use crate::transport::topology::Strategy;
 use crate::transport::{cluster::ClusterData, node::Node};
 use itertools::Itertools;
+use std::collections::HashSet;
+use std::net::SocketAddr;
 use std::{collections::HashMap, sync::Arc};
 use tracing::trace;
 
@@ -103,11 +105,7 @@ impl TokenAwarePolicy {
 }
 
 impl LoadBalancingPolicy for TokenAwarePolicy {
-    fn plan<'a>(
-        &self,
-        statement: &Statement,
-        cluster: &'a ClusterData,
-    ) -> Box<dyn Iterator<Item = Arc<Node>> + Send + Sync + 'a> {
+    fn plan<'a>(&self, statement: &Statement, cluster: &'a ClusterData) -> Plan<'a> {
         match statement.token {
             Some(token) => {
                 let keyspace = statement.keyspace.and_then(|k| cluster.keyspaces.get(k));
@@ -142,7 +140,20 @@ impl LoadBalancingPolicy for TokenAwarePolicy {
                     "TokenAware"
                 );
 
-                self.child_policy.apply_child_policy(replicas)
+                let fallback_plan = {
+                    let replicas_set: HashSet<SocketAddr> =
+                        replicas.iter().map(|node| node.address).collect();
+
+                    self.child_policy
+                        .plan(&Statement::empty(), cluster)
+                        .filter(move |node| !replicas_set.contains(&node.address))
+                };
+
+                let plan = self
+                    .child_policy
+                    .apply_child_policy(replicas)
+                    .chain(fallback_plan);
+                Box::new(plan)
             }
             // fallback to child policy
             None => {
@@ -422,11 +433,7 @@ mod tests {
     struct DumbPolicy {}
 
     impl LoadBalancingPolicy for DumbPolicy {
-        fn plan<'a>(
-            &self,
-            _: &Statement,
-            _: &'a ClusterData,
-        ) -> Box<dyn Iterator<Item = Arc<Node>> + Send + Sync + 'a> {
+        fn plan<'a>(&self, _: &Statement, _: &'a ClusterData) -> Plan<'a> {
             let empty_node_list: Vec<Arc<Node>> = Vec::new();
 
             Box::new(empty_node_list.into_iter())
