diff --git a/scylla-rust-wrapper/src/cluster.rs b/scylla-rust-wrapper/src/cluster.rs index e14e2f59..55af70a3 100644 --- a/scylla-rust-wrapper/src/cluster.rs +++ b/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, 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 { 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, diff --git a/scylla-rust-wrapper/src/exec_profile.rs b/scylla-rust-wrapper/src/exec_profile.rs index 58c11167..bbe8f062 100644 --- a/scylla-rust-wrapper/src/exec_profile.rs +++ b/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::*;