Filter that node is alive

scylladb · Jul 2, 2023 · 533fac4 · 533fac4
1 parent 146655e
commit 533fac4
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/scylla/src/transport/load_balancing/default.rs b/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()

diff --git a/scylla/src/transport/load_balancing/enforce_node.rs b/scylla/src/transport/load_balancing/enforce_node.rs
@@ -1,7 +1,9 @@
-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.
 #[derive(Debug)]
 pub struct EnforceTargetNodePolicy {
     target_node: uuid::Uuid,
@@ -21,6 +23,7 @@ impl LoadBalancingPolicy for EnforceTargetNodePolicy {
         cluster
             .known_peers
             .get(&self.target_node)
+            .filter(DefaultPolicy::is_alive)
             .or_else(|| self.fallback.pick(query, cluster))
     }