Make LoadBalancingPolicy shard-aware

Co-authored-by: Wojciech Przytuła <[email protected]>

examples/Cargo.toml

@@ -20,6 +20,7 @@ uuid = "1.0"
 tower = "0.4"
 stats_alloc = "0.1"
 clap = { version = "3.2.4", features = ["derive"] }
+rand = "0.8.5"
 
 [[example]]
 name = "auth"

examples/custom_load_balancing_policy.rs

@@ -1,20 +1,37 @@
 use anyhow::Result;
+use rand::thread_rng;
+use rand::Rng;
 use scylla::transport::NodeRef;
 use scylla::{
     load_balancing::{LoadBalancingPolicy, RoutingInfo},
+    routing::Shard,
     transport::{ClusterData, ExecutionProfile},
     Session, SessionBuilder,
 };
 use std::{env, sync::Arc};
 
 /// Example load balancing policy that prefers nodes from favorite datacenter
+/// This is, of course, very naive, as it is completely non token-aware.
+/// For more realistic implementation, see [`DefaultPolicy`](scylla::load_balancing::DefaultPolicy).
 #[derive(Debug)]
 struct CustomLoadBalancingPolicy {
     fav_datacenter_name: String,
 }
 
+fn with_random_shard(node: NodeRef) -> (NodeRef, 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)
+}
+
 impl LoadBalancingPolicy for CustomLoadBalancingPolicy {
-    fn pick<'a>(&'a self, _info: &'a RoutingInfo, cluster: &'a ClusterData) -> Option<NodeRef<'a>> {
+    fn pick<'a>(
+        &'a self,
+        _info: &'a RoutingInfo,
+        cluster: &'a ClusterData,
+    ) -> Option<(NodeRef<'a>, Shard)> {
         self.fallback(_info, cluster).next()
     }
 
@@ -28,9 +45,9 @@ impl LoadBalancingPolicy for CustomLoadBalancingPolicy {
             .unique_nodes_in_datacenter_ring(&self.fav_datacenter_name);
 
         match fav_dc_nodes {
-            Some(nodes) => Box::new(nodes.iter()),
+            Some(nodes) => Box::new(nodes.iter().map(with_random_shard)),
             // If there is no dc with provided name, fallback to other datacenters
-            None => Box::new(cluster.get_nodes_info().iter()),
+            None => Box::new(cluster.get_nodes_info().iter().map(with_random_shard)),
         }
     }