load_balancing: add a fallback to token-aware policy

Fallback load balancing plan is chained to the primary plan. Generating
this fallback is realized with a `plan` method of token-aware's child
policy (instead of `apply_child_policy`). It was done to avoid
unnecessary `Arc` cloning required to construct `plan` parameter for
`apply_child_policy()`.

Author: havaker

scylla/src/transport/load_balancing/mod.rs

@@ -23,6 +23,15 @@ 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

scylla/src/transport/load_balancing/token_aware.rs

@@ -3,6 +3,8 @@ 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;
 
@@ -138,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 => {