Merge pull request #170 from muzarski/shuffled-replicas

exec_profile: token aware routing with shuffling setting

Author: dkropachev

scylla-rust-wrapper/src/cluster.rs

@@ -40,6 +40,7 @@ const DRIVER_VERSION: &str = env!("CARGO_PKG_VERSION");
 #[derive(Clone, Debug)]
 pub(crate) struct LoadBalancingConfig {
     pub(crate) token_awareness_enabled: bool,
+    pub(crate) token_aware_shuffling_replicas_enabled: bool,
     pub(crate) dc_awareness: Option<DcAwareness>,
     pub(crate) latency_awareness_enabled: bool,
     pub(crate) latency_awareness_builder: LatencyAwarenessBuilder,
@@ -49,6 +50,11 @@ impl LoadBalancingConfig {
     // as it results in panic due to DefaultPolicyBuilder::build() spawning a tokio task.
     pub(crate) async fn build(self) -> Arc<dyn LoadBalancingPolicy> {
         let mut builder = DefaultPolicyBuilder::new().token_aware(self.token_awareness_enabled);
+        if self.token_awareness_enabled {
+            // Cpp-driver enables shuffling replicas only if token aware routing is enabled.
+            builder =
+                builder.enable_shuffling_replicas(self.token_aware_shuffling_replicas_enabled);
+        }
         if let Some(dc_awareness) = self.dc_awareness.as_ref() {
             builder = builder
                 .prefer_datacenter(dc_awareness.local_dc.clone())
@@ -64,6 +70,7 @@ impl Default for LoadBalancingConfig {
     fn default() -> Self {
         Self {
             token_awareness_enabled: true,
+            token_aware_shuffling_replicas_enabled: true,
             dc_awareness: None,
             latency_awareness_enabled: false,
             latency_awareness_builder: Default::default(),
@@ -619,6 +626,18 @@ pub unsafe extern "C" fn cass_cluster_set_token_aware_routing(
     cluster.load_balancing_config.token_awareness_enabled = enabled != 0;
 }
 
+#[no_mangle]
+pub unsafe extern "C" fn cass_cluster_set_token_aware_routing_shuffle_replicas(
+    cluster_raw: *mut CassCluster,
+    enabled: cass_bool_t,
+) {
+    let cluster = ptr_to_ref_mut(cluster_raw);
+
+    cluster
+        .load_balancing_config
+        .token_aware_shuffling_replicas_enabled = enabled != 0;
+}
+
 #[no_mangle]
 pub unsafe extern "C" fn cass_cluster_set_retry_policy(
     cluster_raw: *mut CassCluster,

scylla-rust-wrapper/src/exec_profile.rs

@@ -421,6 +421,19 @@ pub unsafe extern "C" fn cass_execution_profile_set_token_aware_routing(
     CassError::CASS_OK
 }
 
+#[no_mangle]
+pub unsafe extern "C" fn cass_execution_profile_set_token_aware_routing_shuffle_replicas(
+    profile: *mut CassExecProfile,
+    enabled: cass_bool_t,
+) -> CassError {
+    let profile_builder = ptr_to_ref_mut(profile);
+    profile_builder
+        .load_balancing_config
+        .token_aware_shuffling_replicas_enabled = enabled != 0;
+
+    CassError::CASS_OK
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;