lbp_config: remove DcAwareness and introduce LoadBalancingKind enum

Introduced for multiple reasons:
- rack awareness will be introduced soon
- preparation for another refactor in next commit.

Author: muzarski

scylla-rust-wrapper/src/cluster.rs

@@ -51,7 +51,7 @@ const DRIVER_VERSION: &str = env!("CARGO_PKG_VERSION");
 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) load_balancing_kind: LoadBalancingKind,
     pub(crate) latency_awareness_enabled: bool,
     pub(crate) latency_awareness_builder: LatencyAwarenessBuilder,
 }
@@ -65,10 +65,8 @@ impl LoadBalancingConfig {
             builder =
                 builder.enable_shuffling_replicas(self.token_aware_shuffling_replicas_enabled);
         }
-        if let Some(dc_awareness) = self.dc_awareness {
-            builder = builder
-                .prefer_datacenter(dc_awareness.local_dc)
-                .permit_dc_failover(true)
+        if let LoadBalancingKind::DcAware { local_dc } = self.load_balancing_kind {
+            builder = builder.prefer_datacenter(local_dc).permit_dc_failover(true)
         }
         if self.latency_awareness_enabled {
             builder = builder.latency_awareness(self.latency_awareness_builder);
@@ -81,16 +79,17 @@ impl Default for LoadBalancingConfig {
         Self {
             token_awareness_enabled: true,
             token_aware_shuffling_replicas_enabled: true,
-            dc_awareness: None,
+            load_balancing_kind: LoadBalancingKind::RoundRobin,
             latency_awareness_enabled: false,
             latency_awareness_builder: Default::default(),
         }
     }
 }
 
 #[derive(Clone, Debug)]
-pub(crate) struct DcAwareness {
-    pub(crate) local_dc: String,
+pub(crate) enum LoadBalancingKind {
+    RoundRobin,
+    DcAware { local_dc: String },
 }
 
 #[derive(Clone)]
@@ -457,7 +456,7 @@ pub unsafe extern "C" fn cass_cluster_set_credentials_n(
 #[no_mangle]
 pub unsafe extern "C" fn cass_cluster_set_load_balance_round_robin(cluster_raw: *mut CassCluster) {
     let cluster = ptr_to_ref_mut(cluster_raw);
-    cluster.load_balancing_config.dc_awareness = None;
+    cluster.load_balancing_config.load_balancing_kind = LoadBalancingKind::RoundRobin;
 }
 
 #[no_mangle]
@@ -496,7 +495,7 @@ pub(crate) unsafe fn set_load_balance_dc_aware_n(
         .unwrap()
         .to_string();
 
-    load_balancing_config.dc_awareness = Some(DcAwareness { local_dc });
+    load_balancing_config.load_balancing_kind = LoadBalancingKind::DcAware { local_dc };
 
     CassError::CASS_OK
 }
@@ -851,7 +850,10 @@ mod tests {
                 /* Test valid configurations */
                 let cluster = ptr_to_ref(cluster_raw);
                 {
-                    assert_matches!(cluster.load_balancing_config.dc_awareness, None);
+                    assert_matches!(
+                        cluster.load_balancing_config.load_balancing_kind,
+                        LoadBalancingKind::RoundRobin
+                    );
                     assert!(cluster.load_balancing_config.token_awareness_enabled);
                     assert!(!cluster.load_balancing_config.latency_awareness_enabled);
                 }
@@ -878,8 +880,13 @@ mod tests {
                         40,
                     );
 
-                    let dc_awareness = cluster.load_balancing_config.dc_awareness.as_ref().unwrap();
-                    assert_eq!(dc_awareness.local_dc, "eu");
+                    let load_balancing_kind = &cluster.load_balancing_config.load_balancing_kind;
+                    match load_balancing_kind {
+                        LoadBalancingKind::DcAware { local_dc } => {
+                            assert_eq!(local_dc, "eu")
+                        }
+                        _ => panic!("Expected preferred dc"),
+                    }
                     assert!(!cluster.load_balancing_config.token_awareness_enabled);
                     assert!(cluster.load_balancing_config.latency_awareness_enabled);
                 }

scylla-rust-wrapper/src/exec_profile.rs

@@ -17,7 +17,7 @@ use crate::argconv::{free_boxed, ptr_to_cstr_n, ptr_to_ref, ptr_to_ref_mut, strl
 use crate::batch::CassBatch;
 use crate::cass_error::CassError;
 use crate::cass_types::CassConsistency;
-use crate::cluster::{set_load_balance_dc_aware_n, LoadBalancingConfig};
+use crate::cluster::{set_load_balance_dc_aware_n, LoadBalancingConfig, LoadBalancingKind};
 use crate::retry_policy::CassRetryPolicy;
 use crate::retry_policy::RetryPolicy::{
     DefaultRetryPolicy, DowngradingConsistencyRetryPolicy, FallthroughRetryPolicy,
@@ -353,7 +353,7 @@ pub unsafe extern "C" fn cass_execution_profile_set_load_balance_round_robin(
     profile: *mut CassExecProfile,
 ) -> CassError {
     let profile_builder = ptr_to_ref_mut(profile);
-    profile_builder.load_balancing_config.dc_awareness = None;
+    profile_builder.load_balancing_config.load_balancing_kind = LoadBalancingKind::RoundRobin;
 
     CassError::CASS_OK
 }
@@ -473,7 +473,10 @@ mod tests {
                 /* Test valid configurations */
                 let profile = ptr_to_ref(profile_raw);
                 {
-                    assert_matches!(profile.load_balancing_config.dc_awareness, None);
+                    assert_matches!(
+                        profile.load_balancing_config.load_balancing_kind,
+                        LoadBalancingKind::RoundRobin
+                    );
                     assert!(profile.load_balancing_config.token_awareness_enabled);
                     assert!(!profile.load_balancing_config.latency_awareness_enabled);
                 }
@@ -500,8 +503,13 @@ mod tests {
                         40,
                     );
 
-                    let dc_awareness = profile.load_balancing_config.dc_awareness.as_ref().unwrap();
-                    assert_eq!(dc_awareness.local_dc, "eu");
+                    let load_balancing_kind = &profile.load_balancing_config.load_balancing_kind;
+                    match load_balancing_kind {
+                        LoadBalancingKind::DcAware { local_dc } => {
+                            assert_eq!(local_dc, "eu")
+                        }
+                        _ => panic!("Expected preferred dc"),
+                    }
                     assert!(!profile.load_balancing_config.token_awareness_enabled);
                     assert!(profile.load_balancing_config.latency_awareness_enabled);
                 }