{lb,plan}: call fallback() even if pick() is None

We want to keep regarding pick() as a cheap, happy path function,
which does not allocate. In LWT optimisation, however, only the primary
replica can be computed without allocations (due to specifics of Network
Topology Strategy, replicas are not assigned greedily, but their
distribution across racks is being balanced). If a policy (in our case,
the default policy) recognizes that a picked replica is down, it would
try to pick another, and this computation would be expensive for
ReplicasOrdered. Instead, having recognized during LWT optimised case
that a picked replica is down, the policy returns None from pick()
to hint that further computation will be expensive.
The plan logic is hence altered to call fallback() even if pick()
returns None.
In the non-LWT case, as subsequent calls to pick() are still cheap,
pick() will still try to find next replicas if one is recognized
to be down.

A test is added. It asserts that `fallback()` is called if `pick()`
returned None.
For the test to be possible to be written, a convenience constructor is
added for `Node` under `cfg(test)`.

Author: wprzytula

scylla/src/transport/load_balancing/mod.rs

@@ -53,9 +53,8 @@ pub type FallbackPlan<'a> = Box<dyn Iterator<Item = NodeRef<'a>> + Send + Sync +
 /// `pick` and `fallback`. `pick` returns a first node to contact for a given query, `fallback`
 /// returns the rest of the load balancing plan.
 ///
-/// `fallback` is called only after a failed send to `pick`ed node (or when executing
-/// speculatively).
-/// If a `pick` returns `None`, `fallback` will not be called.
+/// `fallback` is called not only if a send to `pick`ed node failed (or when executing
+/// speculatively), but also if `pick` returns `None`.
 ///
 /// Usually the driver needs only the first node from load balancing plan (most queries are send
 /// successfully, and there is no need to retry).

scylla/src/transport/load_balancing/plan.rs

@@ -54,9 +54,24 @@ impl<'a> Iterator for Plan<'a> {
                     self.state = PlanState::Picked(picked);
                     Some(picked)
                 } else {
-                    error!("Load balancing policy returned an empty plan! The query cannot be executed. Routing info: {:?}", self.routing_info);
-                    self.state = PlanState::PickedNone;
-                    None
+                    // `pick()` returned None, which semantically means that a first node cannot be computed _cheaply_.
+                    // This, however, does not imply that fallback would return an empty plan, too.
+                    // For instance, as a side effect of LWT optimisation in Default Policy, pick() may return None
+                    // when the primary replica is down. `fallback()` will nevertheless return the remaining replicas,
+                    // if there are such.
+                    let mut iter = self.policy.fallback(self.routing_info, self.cluster);
+                    let first_fallback_node = iter.next();
+                    if let Some(node) = first_fallback_node {
+                        self.state = PlanState::Fallback {
+                            iter,
+                            node_to_filter_out: node,
+                        };
+                        Some(node)
+                    } else {
+                        error!("Load balancing policy returned an empty plan! The query cannot be executed. Routing info: {:?}", self.routing_info);
+                        self.state = PlanState::PickedNone;
+                        None
+                    }
                 }
             }
             PlanState::Picked(node) => {
@@ -85,3 +100,65 @@ impl<'a> Iterator for Plan<'a> {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::{net::SocketAddr, str::FromStr, sync::Arc};
+
+    use crate::transport::{
+        locator::test::{create_locator, mock_metadata_for_token_aware_tests},
+        Node, NodeAddr,
+    };
+
+    use super::*;
+
+    fn expected_nodes() -> Vec<Arc<Node>> {
+        vec![Arc::new(Node::new_for_test(
+            NodeAddr::Translatable(SocketAddr::from_str("127.0.0.1:9042").unwrap()),
+            None,
+            None,
+        ))]
+    }
+
+    #[derive(Debug)]
+    struct PickingNonePolicy {
+        expected_nodes: Vec<Arc<Node>>,
+    }
+    impl LoadBalancingPolicy for PickingNonePolicy {
+        fn pick<'a>(
+            &'a self,
+            _query: &'a RoutingInfo,
+            _cluster: &'a ClusterData,
+        ) -> Option<NodeRef<'a>> {
+            None
+        }
+
+        fn fallback<'a>(
+            &'a self,
+            _query: &'a RoutingInfo,
+            _cluster: &'a ClusterData,
+        ) -> FallbackPlan<'a> {
+            Box::new(self.expected_nodes.iter())
+        }
+
+        fn name(&self) -> String {
+            "PickingNone".into()
+        }
+    }
+
+    #[tokio::test]
+    async fn plan_calls_fallback_even_if_pick_returned_none() {
+        let policy = PickingNonePolicy {
+            expected_nodes: expected_nodes(),
+        };
+        let locator = create_locator(&mock_metadata_for_token_aware_tests());
+        let cluster_data = ClusterData {
+            known_peers: Default::default(),
+            keyspaces: Default::default(),
+            locator,
+        };
+        let routing_info = RoutingInfo::default();
+        let plan = Plan::new(&policy, &routing_info, &cluster_data);
+        assert_eq!(Vec::from_iter(plan.cloned()), policy.expected_nodes);
+    }
+}

scylla/src/transport/node.rs

@@ -214,3 +214,25 @@ impl Hash for Node {
         self.host_id.hash(state);
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    impl Node {
+        pub(crate) fn new_for_test(
+            address: NodeAddr,
+            datacenter: Option<String>,
+            rack: Option<String>,
+        ) -> Self {
+            Self {
+                host_id: Uuid::new_v4(),
+                address,
+                datacenter,
+                rack,
+                pool: None,
+                down_marker: false.into(),
+            }
+        }
+    }
+}