LBP: Return Option<Shard> instead of Shard

Lorak-mmk · Lorak-mmk · commit 6ecc4449d759 · 2024-03-27T11:35:05.000+01:00
This was already documented as such, but due to an oversight the code
was in disagreement with documentation. Approach from the documentation
is better, because the currently implemented approach prevented
deduplication in Plan from working correctly.
diff --git a/examples/custom_load_balancing_policy.rs b/examples/custom_load_balancing_policy.rs
@@ -18,20 +18,20 @@ struct CustomLoadBalancingPolicy {
     fav_datacenter_name: String,
 }
 
-fn with_random_shard(node: NodeRef) -> (NodeRef, Shard) {
+fn with_random_shard(node: NodeRef) -> (NodeRef, Option<Shard>) {
     let nr_shards = node
         .sharder()
         .map(|sharder| sharder.nr_shards.get())
         .unwrap_or(1);
-    (node, thread_rng().gen_range(0..nr_shards) as Shard)
+    (node, Some(thread_rng().gen_range(0..nr_shards) as Shard))
 }
 
 impl LoadBalancingPolicy for CustomLoadBalancingPolicy {
     fn pick<'a>(
         &'a self,
         _info: &'a RoutingInfo,
         cluster: &'a ClusterData,
-    ) -> Option<(NodeRef<'a>, Shard)> {
+    ) -> Option<(NodeRef<'a>, Option<Shard>)> {
         self.fallback(_info, cluster).next()
     }
 
diff --git a/scylla/src/transport/load_balancing/default.rs b/scylla/src/transport/load_balancing/default.rs
diff --git a/scylla/src/transport/load_balancing/mod.rs b/scylla/src/transport/load_balancing/mod.rs
@@ -39,7 +39,8 @@ pub struct RoutingInfo<'a> {
 ///
 /// It is computed on-demand, only if querying the most preferred node fails
 /// (or when speculative execution is triggered).
-pub type FallbackPlan<'a> = Box<dyn Iterator<Item = (NodeRef<'a>, Shard)> + Send + Sync + 'a>;
+pub type FallbackPlan<'a> =
+    Box<dyn Iterator<Item = (NodeRef<'a>, Option<Shard>)> + Send + Sync + 'a>;
 
 /// Policy that decides which nodes and shards to contact for each query.
 ///
@@ -67,7 +68,7 @@ pub trait LoadBalancingPolicy: Send + Sync + std::fmt::Debug {
         &'a self,
         query: &'a RoutingInfo,
         cluster: &'a ClusterData,
-    ) -> Option<(NodeRef<'a>, Shard)>;
+    ) -> Option<(NodeRef<'a>, Option<Shard>)>;
 
     /// Returns all contact-appropriate nodes for a given query.
     fn fallback<'a>(&'a self, query: &'a RoutingInfo, cluster: &'a ClusterData)
diff --git a/scylla/src/transport/load_balancing/plan.rs b/scylla/src/transport/load_balancing/plan.rs
@@ -1,3 +1,4 @@
+use rand::{thread_rng, Rng};
 use tracing::error;
 
 use super::{FallbackPlan, LoadBalancingPolicy, NodeRef, RoutingInfo};
@@ -6,20 +7,65 @@ use crate::{routing::Shard, transport::ClusterData};
 enum PlanState<'a> {
     Created,
     PickedNone, // This always means an abnormal situation: it means that no nodes satisfied locality/node filter requirements.
-    Picked((NodeRef<'a>, Shard)),
+    Picked((NodeRef<'a>, Option<Shard>)),
     Fallback {
         iter: FallbackPlan<'a>,
-        node_to_filter_out: (NodeRef<'a>, Shard),
+        target_to_filter_out: (NodeRef<'a>, Option<Shard>),
     },
 }
 
-/// The list of nodes constituting the query plan.
+/// The list of targets constituting the query plan. Target here is a pair `(NodeRef<'a>, Shard)`.
 ///
-/// The plan is partly lazily computed, with the first node computed
-/// eagerly in the first place and the remaining nodes computed on-demand
+/// The plan is partly lazily computed, with the first target computed
+/// eagerly in the first place and the remaining targets computed on-demand
 /// (all at once).
 /// This significantly reduces the allocation overhead on "the happy path"
-/// (when the first node successfully handles the request),
+/// (when the first target successfully handles the request).
+///
+/// `Plan` implements `Iterator<Item=(NodeRef<'a>, Shard)>` but LoadBalancingPolicy
+/// returns `Option<Shard>` instead of `Shard` both in `pick` and in `fallback`.
+/// `Plan` handles the `None` case by using random shard for a given node.
+/// There is currently no way to configure RNG used by `Plan`.
+/// If you don't want `Plan` to do randomize shards or you want to control the RNG,
+/// use custom LBP that will always return non-`None` shards.
+/// Example of LBP that always uses shard 0, preventing `Plan` from using random numbers:
+///
+/// ```
+/// # use std::sync::Arc;
+/// # use scylla::load_balancing::LoadBalancingPolicy;
+/// # use scylla::load_balancing::RoutingInfo;
+/// # use scylla::transport::ClusterData;
+/// # use scylla::transport::NodeRef;
+/// # use scylla::routing::Shard;
+/// # use scylla::load_balancing::FallbackPlan;
+///
+/// #[derive(Debug)]
+/// struct NonRandomLBP {
+///     inner: Arc<dyn LoadBalancingPolicy>,
+/// }
+/// impl LoadBalancingPolicy for NonRandomLBP {
+///     fn pick<'a>(
+///         &'a self,
+///         info: &'a RoutingInfo,
+///         cluster: &'a ClusterData,
+///     ) -> Option<(NodeRef<'a>, Option<Shard>)> {
+///         self.inner
+///             .pick(info, cluster)
+///             .map(|(node, shard)| (node, shard.or(Some(0))))
+///     }
+///
+///     fn fallback<'a>(&'a self, info: &'a RoutingInfo, cluster: &'a ClusterData) -> FallbackPlan<'a> {
+///         Box::new(self.inner
+///             .fallback(info, cluster)
+///             .map(|(node, shard)| (node, shard.or(Some(0)))))
+///     }
+///
+///     fn name(&self) -> String {
+///         "NonRandomLBP".to_string()
+///     }
+/// }
+/// ```
+
 pub struct Plan<'a> {
     policy: &'a dyn LoadBalancingPolicy,
     routing_info: &'a RoutingInfo<'a>,
@@ -41,6 +87,21 @@ impl<'a> Plan<'a> {
             state: PlanState::Created,
         }
     }
+
+    fn with_random_shard_if_unknown(
+        (node, shard): (NodeRef<'_>, Option<Shard>),
+    ) -> (NodeRef<'_>, Shard) {
+        (
+            node,
+            shard.unwrap_or_else(|| {
+                let nr_shards = node
+                    .sharder()
+                    .map(|sharder| sharder.nr_shards.get())
+                    .unwrap_or(1);
+                thread_rng().gen_range(0..nr_shards).into()
+            }),
+        )
+    }
 }
 
 impl<'a> Iterator for Plan<'a> {
@@ -52,7 +113,7 @@ impl<'a> Iterator for Plan<'a> {
                 let picked = self.policy.pick(self.routing_info, self.cluster);
                 if let Some(picked) = picked {
                     self.state = PlanState::Picked(picked);
-                    Some(picked)
+                    Some(Self::with_random_shard_if_unknown(picked))
                 } else {
                     // `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.
@@ -64,9 +125,9 @@ impl<'a> Iterator for Plan<'a> {
                     if let Some(node) = first_fallback_node {
                         self.state = PlanState::Fallback {
                             iter,
-                            node_to_filter_out: node,
+                            target_to_filter_out: node,
                         };
-                        Some(node)
+                        Some(Self::with_random_shard_if_unknown(node))
                     } else {
                         error!("Load balancing policy returned an empty plan! The query cannot be executed. Routing info: {:?}", self.routing_info);
                         self.state = PlanState::PickedNone;
@@ -77,20 +138,20 @@ impl<'a> Iterator for Plan<'a> {
             PlanState::Picked(node) => {
                 self.state = PlanState::Fallback {
                     iter: self.policy.fallback(self.routing_info, self.cluster),
-                    node_to_filter_out: *node,
+                    target_to_filter_out: *node,
                 };
 
                 self.next()
             }
             PlanState::Fallback {
                 iter,
-                node_to_filter_out,
+                target_to_filter_out: node_to_filter_out,
             } => {
                 for node in iter {
                     if node == *node_to_filter_out {
                         continue;
                     } else {
-                        return Some(node);
+                        return Some(Self::with_random_shard_if_unknown(node));
                     }
                 }
 
@@ -135,7 +196,7 @@ mod tests {
             &'a self,
             _query: &'a RoutingInfo,
             _cluster: &'a ClusterData,
-        ) -> Option<(NodeRef<'a>, Shard)> {
+        ) -> Option<(NodeRef<'a>, Option<Shard>)> {
             None
         }
 
@@ -147,7 +208,7 @@ mod tests {
             Box::new(
                 self.expected_nodes
                     .iter()
-                    .map(|(node_ref, shard)| (node_ref, *shard)),
+                    .map(|(node_ref, shard)| (node_ref, Some(*shard))),
             )
         }
 
diff --git a/scylla/tests/integration/consistency.rs b/scylla/tests/integration/consistency.rs
@@ -379,7 +379,7 @@ impl LoadBalancingPolicy for RoutingInfoReportingWrapper {
         &'a self,
         query: &'a RoutingInfo,
         cluster: &'a scylla::transport::ClusterData,
-    ) -> Option<(NodeRef<'a>, Shard)> {
+    ) -> Option<(NodeRef<'a>, Option<Shard>)> {
         self.routing_info_tx
             .send(OwnedRoutingInfo::from(query.clone()))
             .unwrap();
diff --git a/scylla/tests/integration/execution_profiles.rs b/scylla/tests/integration/execution_profiles.rs
@@ -51,9 +51,13 @@ impl<const NODE: u8> LoadBalancingPolicy for BoundToPredefinedNodePolicy<NODE> {
         &'a self,
         _info: &'a RoutingInfo,
         cluster: &'a ClusterData,
-    ) -> Option<(NodeRef<'a>, Shard)> {
+    ) -> Option<(NodeRef<'a>, Option<Shard>)> {
         self.report_node(Report::LoadBalancing);
-        cluster.get_nodes_info().iter().next().map(|node| (node, 0))
+        cluster
+            .get_nodes_info()
+            .iter()
+            .next()
+            .map(|node| (node, None))
     }
 
     fn fallback<'a>(
diff --git a/scylla/tests/integration/utils.rs b/scylla/tests/integration/utils.rs
@@ -19,12 +19,12 @@ pub(crate) fn setup_tracing() {
         .try_init();
 }
 
-fn with_pseudorandom_shard(node: NodeRef) -> (NodeRef, Shard) {
+fn with_pseudorandom_shard(node: NodeRef) -> (NodeRef, Option<Shard>) {
     let nr_shards = node
         .sharder()
         .map(|sharder| sharder.nr_shards.get())
         .unwrap_or(1);
-    (node, ((nr_shards - 1) % 42) as Shard)
+    (node, Some(((nr_shards - 1) % 42) as Shard))
 }
 
 #[derive(Debug)]
@@ -34,7 +34,7 @@ impl LoadBalancingPolicy for FixedOrderLoadBalancer {
         &'a self,
         _info: &'a scylla::load_balancing::RoutingInfo,
         cluster: &'a scylla::transport::ClusterData,
-    ) -> Option<(NodeRef<'a>, Shard)> {
+    ) -> Option<(NodeRef<'a>, Option<Shard>)> {
         cluster
             .get_nodes_info()
             .iter()
