exec_profile: cass_execution_profile_set_load_balance_rack_aware[_n]

This is an extension to the extension.
cpp-driver does not implement it for some reason.

Author: muzarski

include/cassandra.h

@@ -1087,6 +1087,50 @@ cass_execution_profile_set_load_balance_dc_aware_n(CassExecProfile* profile,
                                                    unsigned used_hosts_per_remote_dc,
                                                    cass_bool_t allow_remote_dcs_for_local_cl);
 
+
+/**
+ * Configures the execution profile to use Rack-aware load balancing.
+ * For each query, all live nodes in a primary 'local' rack are tried first,
+ * followed by nodes from local DC and then nodes from other DCs.
+ *
+ * <b>Note:</b> Profile-based load balancing policy is disabled by default.
+ * cluster load balancing policy is used when profile does not contain a policy.
+ *
+ * @public @memberof CassExecProfile
+ *
+ * @param[in] profile
+ * @param[in] local_dc The primary data center to try first
+ * @param[in] local_rack The primary rack to try first
+ * @return CASS_OK if successful, otherwise an error occurred
+ */
+CASS_EXPORT CassError
+cass_execution_profile_set_load_balance_rack_aware(CassExecProfile* profile,
+                                                   const char* local_dc,
+                                                   const char* local_rack);
+
+
+/**
+ * Same as cass_execution_profile_set_load_balance_rack_aware(), but with lengths for string
+ * parameters.
+ *
+ * @public @memberof CassExecProfile
+ *
+ * @param[in] profile
+ * @param[in] local_dc
+ * @param[in] local_dc_length
+ * @return same cass_execution_profile_set_load_balance_rack_aware()
+ *
+ * @see cass_execution_profile_set_load_balance_rack_aware()
+ * @see cass_cluster_set_load_balance_rack_aware_n()
+ */
+CASS_EXPORT CassError
+cass_execution_profile_set_load_balance_rack_aware_n(CassExecProfile* profile,
+                                                     const char* local_dc,
+                                                     size_t local_dc_length,
+                                                     const char* local_rack,
+                                                     size_t local_rack_length);
+
+
 /**
  * Configures the execution profile to use token-aware request routing or not.
  *

scylla-rust-wrapper/src/cluster.rs

@@ -606,7 +606,7 @@ pub unsafe extern "C" fn cass_cluster_set_load_balance_rack_aware_n(
     )
 }
 
-unsafe fn set_load_balance_rack_aware_n(
+pub(crate) unsafe fn set_load_balance_rack_aware_n(
     load_balancing_config: &mut LoadBalancingConfig,
     local_dc_raw: *const c_char,
     local_dc_length: size_t,

scylla-rust-wrapper/src/exec_profile.rs

@@ -17,7 +17,10 @@ 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, LoadBalancingKind};
+use crate::cluster::{
+    set_load_balance_dc_aware_n, set_load_balance_rack_aware_n, LoadBalancingConfig,
+    LoadBalancingKind,
+};
 use crate::retry_policy::CassRetryPolicy;
 use crate::retry_policy::RetryPolicy::{
     DefaultRetryPolicy, DowngradingConsistencyRetryPolicy, FallthroughRetryPolicy,
@@ -357,6 +360,40 @@ pub unsafe extern "C" fn cass_execution_profile_set_load_balance_dc_aware_n(
     )
 }
 
+#[no_mangle]
+pub unsafe extern "C" fn cass_execution_profile_set_load_balance_rack_aware(
+    profile: *mut CassExecProfile,
+    local_dc_raw: *const c_char,
+    local_rack_raw: *const c_char,
+) -> CassError {
+    cass_execution_profile_set_load_balance_rack_aware_n(
+        profile,
+        local_dc_raw,
+        strlen(local_dc_raw),
+        local_rack_raw,
+        strlen(local_rack_raw),
+    )
+}
+
+#[no_mangle]
+pub unsafe extern "C" fn cass_execution_profile_set_load_balance_rack_aware_n(
+    profile: *mut CassExecProfile,
+    local_dc_raw: *const c_char,
+    local_dc_length: size_t,
+    local_rack_raw: *const c_char,
+    local_rack_length: size_t,
+) -> CassError {
+    let profile_builder = ptr_to_ref_mut(profile);
+
+    set_load_balance_rack_aware_n(
+        &mut profile_builder.load_balancing_config,
+        local_dc_raw,
+        local_dc_length,
+        local_rack_raw,
+        local_rack_length,
+    )
+}
+
 #[no_mangle]
 pub unsafe extern "C" fn cass_execution_profile_set_load_balance_round_robin(
     profile: *mut CassExecProfile,