diff --git a/include/cassandra.h b/include/cassandra.h index c43d5b2e..257ad70e 100644 --- a/include/cassandra.h +++ b/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. + * + * Note: 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. * @@ -2176,13 +2220,6 @@ cass_cluster_set_load_balance_round_robin(CassCluster* cluster); * For each query, all live nodes in a primary 'local' DC are tried first, * followed by any node from other DCs. * - * Note: This is the default, and does not need to be called unless - * switching an existing from another policy or changing settings. - * Without further configuration, a default local_dc is chosen from the - * first connected contact point, and no remote hosts are considered in - * query plans. If relying on this mechanism, be sure to use only contact - * points from the local DC. - * * @deprecated The remote DC settings for DC-aware are not suitable for most * scenarios that require DC failover. There is also unhandled gap between * replication factor number of nodes failing and the full cluster failing. Only @@ -2233,6 +2270,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 9b2cd75a..c01d172c 100644 --- a/scylla-rust-wrapper/src/cluster.rs +++ b/scylla-rust-wrapper/src/cluster.rs @@ -75,9 +75,22 @@ impl LoadBalancingConfig { builder = builder.enable_shuffling_replicas(self.token_aware_shuffling_replicas_enabled); } - if let LoadBalancingKind::DcAware { local_dc } = load_balancing_kind { - builder = builder.prefer_datacenter(local_dc).permit_dc_failover(true) + + match load_balancing_kind { + LoadBalancingKind::DcAware { local_dc } => { + builder = builder.prefer_datacenter(local_dc).permit_dc_failover(true) + } + LoadBalancingKind::RackAware { + local_dc, + local_rack, + } => { + builder = builder + .prefer_datacenter_and_rack(local_dc, local_rack) + .permit_dc_failover(true) + } + LoadBalancingKind::RoundRobin => {} } + if self.latency_awareness_enabled { builder = builder.latency_awareness(self.latency_awareness_builder); } @@ -99,7 +112,13 @@ impl Default for LoadBalancingConfig { #[derive(Clone, Debug)] pub(crate) enum LoadBalancingKind { RoundRobin, - DcAware { local_dc: String }, + DcAware { + local_dc: String, + }, + RackAware { + local_dc: String, + local_rack: String, + }, } #[derive(Clone)] @@ -553,6 +572,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, + ) +} + +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, + 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 either is a null pointer, is an empty string or is not a proper utf-8. + _ => return CassError::CASS_ERROR_LIB_BAD_PARAMS, + }; + + load_balancing_config.load_balancing_kind = Some(LoadBalancingKind::RackAware { + 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, @@ -891,12 +972,7 @@ mod tests { { cass_cluster_set_token_aware_routing(cluster_raw, 0); assert_cass_error_eq!( - cass_cluster_set_load_balance_dc_aware( - cluster_raw, - "eu\0".as_ptr() as *const i8, - 0, - 0 - ), + cass_cluster_set_load_balance_dc_aware(cluster_raw, c"eu".as_ptr(), 0, 0), CassError::CASS_OK ); cass_cluster_set_latency_aware_routing(cluster_raw, 1); @@ -920,25 +996,98 @@ mod tests { } assert!(!cluster.load_balancing_config.token_awareness_enabled); assert!(cluster.load_balancing_config.latency_awareness_enabled); + + // set preferred rack+dc + assert_cass_error_eq!( + cass_cluster_set_load_balance_rack_aware( + cluster_raw, + c"eu-east".as_ptr(), + c"rack1".as_ptr(), + ), + CassError::CASS_OK + ); + + let node_location_preference = + &cluster.load_balancing_config.load_balancing_kind; + match node_location_preference { + Some(LoadBalancingKind::RackAware { + local_dc, + local_rack, + }) => { + assert_eq!(local_dc, "eu-east"); + assert_eq!(local_rack, "rack1"); + } + _ => panic!("Expected preferred dc and rack"), + } + + // set back to preferred dc + assert_cass_error_eq!( + cass_cluster_set_load_balance_dc_aware(cluster_raw, c"eu".as_ptr(), 0, 0), + CassError::CASS_OK + ); + + let node_location_preference = + &cluster.load_balancing_config.load_balancing_kind; + match node_location_preference { + Some(LoadBalancingKind::DcAware { local_dc }) => { + assert_eq!(local_dc, "eu") + } + _ => panic!("Expected preferred dc"), + } } /* Test invalid configurations */ { // Nonzero deprecated parameters assert_cass_error_eq!( - cass_cluster_set_load_balance_dc_aware( + cass_cluster_set_load_balance_dc_aware(cluster_raw, c"eu".as_ptr(), 1, 0), + CassError::CASS_ERROR_LIB_BAD_PARAMS + ); + assert_cass_error_eq!( + cass_cluster_set_load_balance_dc_aware(cluster_raw, c"eu".as_ptr(), 0, 1), + CassError::CASS_ERROR_LIB_BAD_PARAMS + ); + + // null pointers + assert_cass_error_eq!( + cass_cluster_set_load_balance_dc_aware(cluster_raw, std::ptr::null(), 0, 0), + CassError::CASS_ERROR_LIB_BAD_PARAMS + ); + assert_cass_error_eq!( + cass_cluster_set_load_balance_rack_aware( + cluster_raw, + c"eu".as_ptr(), + std::ptr::null(), + ), + CassError::CASS_ERROR_LIB_BAD_PARAMS + ); + assert_cass_error_eq!( + cass_cluster_set_load_balance_rack_aware( + cluster_raw, + std::ptr::null(), + c"rack".as_ptr(), + ), + CassError::CASS_ERROR_LIB_BAD_PARAMS + ); + + // empty strings + let empty_str = "\0".as_ptr() as *const i8; + assert_cass_error_eq!( + cass_cluster_set_load_balance_dc_aware(cluster_raw, std::ptr::null(), 0, 0), + CassError::CASS_ERROR_LIB_BAD_PARAMS + ); + assert_cass_error_eq!( + cass_cluster_set_load_balance_rack_aware( cluster_raw, - "eu\0".as_ptr() as *const i8, - 1, - 0 + c"eu".as_ptr(), + empty_str, ), CassError::CASS_ERROR_LIB_BAD_PARAMS ); assert_cass_error_eq!( - cass_cluster_set_load_balance_dc_aware( + cass_cluster_set_load_balance_rack_aware( cluster_raw, - "eu\0".as_ptr() as *const i8, - 0, - 1 + empty_str, + c"rack".as_ptr(), ), CassError::CASS_ERROR_LIB_BAD_PARAMS ); diff --git a/scylla-rust-wrapper/src/exec_profile.rs b/scylla-rust-wrapper/src/exec_profile.rs index cf8e4275..6ccd2056 100644 --- a/scylla-rust-wrapper/src/exec_profile.rs +++ b/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,