Filter that node is alive

Author: Ten0

scylla/src/transport/load_balancing/default.rs

@@ -594,7 +594,7 @@ impl DefaultPolicy {
         vec.into_iter()
     }
 
-    fn is_alive(node: &NodeRef<'_>) -> bool {
+    pub(crate) fn is_alive(node: &NodeRef<'_>) -> bool {
         // For now, we leave this as stub, until we have time to improve node events.
         // node.is_enabled() && !node.is_down()
         node.is_enabled()

scylla/src/transport/load_balancing/enforce_node.rs

@@ -1,7 +1,11 @@
-use super::{FallbackPlan, LoadBalancingPolicy, NodeRef, RoutingInfo};
+use super::{DefaultPolicy, FallbackPlan, LoadBalancingPolicy, NodeRef, RoutingInfo};
 use crate::transport::{cluster::ClusterData, Node};
 use std::sync::Arc;
 
+/// This policy will always return the same node, unless it is not available anymore, in which case it will
+/// fallback to the provided policy.
+///
+/// This is meant to be used for shard-aware batching.
 #[derive(Debug)]
 pub struct EnforceTargetNodePolicy {
     target_node: uuid::Uuid,
@@ -21,6 +25,7 @@ impl LoadBalancingPolicy for EnforceTargetNodePolicy {
         cluster
             .known_peers
             .get(&self.target_node)
+            .filter(DefaultPolicy::is_alive)
             .or_else(|| self.fallback.pick(query, cluster))
     }