diff --git a/include/cassandra.h b/include/cassandra.h index 3eaaded2..a795aba6 100644 --- a/include/cassandra.h +++ b/include/cassandra.h @@ -2219,6 +2219,46 @@ cass_cluster_set_load_balance_dc_aware_n(CassCluster* cluster, unsigned used_hosts_per_remote_dc, cass_bool_t allow_remote_dcs_for_local_cl); + +/** + * Configures the cluster 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. + * + * @public @memberof CassCluster + * + * @param[in] cluster + * @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_cluster_set_load_balance_rack_aware(CassCluster* cluster, + const char* local_dc, + const char* local_rack); + + +/** + * Same as cass_cluster_set_load_balance_rack_aware(), but with lengths for string + * parameters. + * + * @public @memberof CassCluster + * + * @param[in] cluster + * @param[in] local_dc + * @param[in] local_dc_length + * @return same as cass_cluster_set_load_balance_dc_aware() + * + * @see cass_cluster_set_load_balance_dc_aware() + */ +CASS_EXPORT CassError +cass_cluster_set_load_balance_rack_aware_n(CassCluster* cluster, + const char* local_dc, + size_t local_dc_length, + const char* local_rack, + size_t local_rack_length); + + /** * Configures the cluster to use token-aware request routing or not. * diff --git a/scylla-rust-wrapper/extern/cassandra.h b/scylla-rust-wrapper/extern/cassandra.h index 3eaaded2..41c370c9 100644 --- a/scylla-rust-wrapper/extern/cassandra.h +++ b/scylla-rust-wrapper/extern/cassandra.h @@ -1087,6 +1087,7 @@ 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 token-aware request routing or not. * @@ -2219,6 +2220,46 @@ cass_cluster_set_load_balance_dc_aware_n(CassCluster* cluster, unsigned used_hosts_per_remote_dc, cass_bool_t allow_remote_dcs_for_local_cl); + +/** + * Configures the cluster 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. + * + * @public @memberof CassCluster + * + * @param[in] cluster + * @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_cluster_set_load_balance_rack_aware(CassCluster* cluster, + const char* local_dc, + const char* local_rack); + + +/** + * Same as cass_cluster_set_load_balance_rack_aware(), but with lengths for string + * parameters. + * + * @public @memberof CassCluster + * + * @param[in] cluster + * @param[in] local_dc + * @param[in] local_dc_length + * @return same as cass_cluster_set_load_balance_dc_aware() + * + * @see cass_cluster_set_load_balance_dc_aware() + */ +CASS_EXPORT CassError +cass_cluster_set_load_balance_rack_aware_n(CassCluster* cluster, + const char* local_dc, + size_t local_dc_length, + const char* local_rack, + size_t local_rack_length); + + /** * Configures the cluster to use token-aware request routing or not. * diff --git a/scylla-rust-wrapper/src/cluster.rs b/scylla-rust-wrapper/src/cluster.rs index 3f70f6b6..e7c5b136 100644 --- a/scylla-rust-wrapper/src/cluster.rs +++ b/scylla-rust-wrapper/src/cluster.rs @@ -55,9 +55,22 @@ impl LoadBalancingConfig { builder = builder.enable_shuffling_replicas(self.token_aware_shuffling_replicas_enabled); } - if let NodeLocationPreference::Datacenter { local_dc } = self.node_location_preference { - builder = builder.prefer_datacenter(local_dc).permit_dc_failover(true) + + match self.node_location_preference { + NodeLocationPreference::Datacenter { local_dc } => { + builder = builder.prefer_datacenter(local_dc).permit_dc_failover(true) + } + NodeLocationPreference::DatacenterAndRack { + local_dc, + local_rack, + } => { + builder = builder + .prefer_datacenter_and_rack(local_dc, local_rack) + .permit_dc_failover(true) + } + NodeLocationPreference::Any => {} } + if self.latency_awareness_enabled { builder = builder.latency_awareness(self.latency_awareness_builder); } @@ -79,7 +92,13 @@ impl Default for LoadBalancingConfig { #[derive(Clone, Debug)] pub(crate) enum NodeLocationPreference { Any, - Datacenter { local_dc: String }, + Datacenter { + local_dc: String, + }, + DatacenterAndRack { + local_dc: String, + local_rack: String, + }, } #[derive(Clone)] @@ -477,6 +496,68 @@ pub unsafe extern "C" fn cass_cluster_set_load_balance_dc_aware_n( ) } +#[no_mangle] +pub unsafe extern "C" fn cass_cluster_set_load_balance_rack_aware( + cluster_raw: *mut CassCluster, + local_dc_raw: *const c_char, + local_rack_raw: *const c_char, +) -> CassError { + cass_cluster_set_load_balance_rack_aware_n( + cluster_raw, + local_dc_raw, + strlen(local_dc_raw), + local_rack_raw, + strlen(local_rack_raw), + ) +} + +#[no_mangle] +pub unsafe extern "C" fn cass_cluster_set_load_balance_rack_aware_n( + cluster_raw: *mut CassCluster, + local_dc_raw: *const c_char, + local_dc_length: size_t, + local_rack_raw: *const c_char, + local_rack_length: size_t, +) -> CassError { + let cluster = ptr_to_ref_mut(cluster_raw); + + set_load_balance_rack_aware_n( + &mut cluster.load_balancing_config, + local_dc_raw, + local_dc_length, + local_rack_raw, + local_rack_length, + ) +} + +unsafe fn set_load_balance_rack_aware_n( + load_balancing_config: &mut LoadBalancingConfig, + local_dc_raw: *const c_char, + local_dc_length: size_t, + local_rack_raw: *const c_char, + local_rack_length: size_t, +) -> CassError { + let (local_dc, local_rack) = match ( + ptr_to_cstr_n(local_dc_raw, local_dc_length), + ptr_to_cstr_n(local_rack_raw, local_rack_length), + ) { + (Some(local_dc_str), Some(local_rack_str)) + if local_dc_length > 0 && local_rack_length > 0 => + { + (local_dc_str.to_owned(), local_rack_str.to_owned()) + } + // One of them is either null pointer, is an empty string or is not a proper utf-8. + _ => return CassError::CASS_ERROR_LIB_BAD_PARAMS, + }; + + load_balancing_config.node_location_preference = NodeLocationPreference::DatacenterAndRack { + local_dc, + local_rack, + }; + + CassError::CASS_OK +} + #[no_mangle] pub unsafe extern "C" fn cass_cluster_set_cloud_secure_connection_bundle_n( _cluster_raw: *mut CassCluster,