Skip to content

Commit

Permalink
lbp_config: remove DcAwareness and introduce NodeLocationPreference enum
Browse files Browse the repository at this point in the history 
This is because in the following commits we will be introducing rack
awareness.
  • Loading branch information
@muzarski
muzarski committed Oct 17, 2024
1 parent 3379c0d commit c95b1c0
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 18 deletions.
35 changes: 22 additions & 13 deletions scylla-rust-wrapper/src/cluster.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,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) node_location_preference: NodeLocationPreference,
pub(crate) latency_awareness_enabled: bool,
pub(crate) latency_awareness_builder: LatencyAwarenessBuilder,
}
Expand All @@ -55,10 +55,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 NodeLocationPreference::Datacenter { local_dc } = self.node_location_preference {
builder = builder.prefer_datacenter(local_dc).permit_dc_failover(true)
}
if self.latency_awareness_enabled {
builder = builder.latency_awareness(self.latency_awareness_builder);
Expand All @@ -71,16 +69,17 @@ impl Default for LoadBalancingConfig {
Self {
token_awareness_enabled: true,
token_aware_shuffling_replicas_enabled: true,
dc_awareness: None,
node_location_preference: NodeLocationPreference::Any,
latency_awareness_enabled: false,
latency_awareness_builder: Default::default(),
}
}
}

#[derive(Clone, Debug)]
pub(crate) struct DcAwareness {
pub(crate) local_dc: String,
pub(crate) enum NodeLocationPreference {
Any,
Datacenter { local_dc: String },
}

#[derive(Clone)]
Expand Down Expand Up @@ -414,7 +413,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.node_location_preference = NodeLocationPreference::Any;
}

#[no_mangle]
Expand Down Expand Up @@ -453,7 +452,8 @@ 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.node_location_preference =
NodeLocationPreference::Datacenter { local_dc };

CassError::CASS_OK
}
Expand Down Expand Up @@ -808,7 +808,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.node_location_preference,
NodeLocationPreference::Any
);
assert!(cluster.load_balancing_config.token_awareness_enabled);
assert!(!cluster.load_balancing_config.latency_awareness_enabled);
}
Expand All @@ -835,8 +838,14 @@ mod tests {
40,
);

let dc_awareness = cluster.load_balancing_config.dc_awareness.as_ref().unwrap();
assert_eq!(dc_awareness.local_dc, "eu");
let node_location_preference =
&cluster.load_balancing_config.node_location_preference;
match node_location_preference {
NodeLocationPreference::Datacenter { 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);
}
Expand Down
21 changes: 16 additions & 5 deletions scylla-rust-wrapper/src/exec_profile.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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, NodeLocationPreference};
use crate::retry_policy::CassRetryPolicy;
use crate::retry_policy::RetryPolicy::{
DefaultRetryPolicy, DowngradingConsistencyRetryPolicy, FallthroughRetryPolicy,
Expand Down Expand Up @@ -353,7 +353,9 @@ 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
.node_location_preference = NodeLocationPreference::Any;

CassError::CASS_OK
}
Expand Down Expand Up @@ -473,7 +475,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.node_location_preference,
NodeLocationPreference::Any
);
assert!(profile.load_balancing_config.token_awareness_enabled);
assert!(!profile.load_balancing_config.latency_awareness_enabled);
}
Expand All @@ -500,8 +505,14 @@ mod tests {
40,
);

let dc_awareness = profile.load_balancing_config.dc_awareness.as_ref().unwrap();
assert_eq!(dc_awareness.local_dc, "eu");
let node_location_preference =
&profile.load_balancing_config.node_location_preference;
match node_location_preference {
NodeLocationPreference::Datacenter { 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);
}
Expand Down

0 comments on commit c95b1c0

Please sign in to comment.