transport: load_balancing: Add preferred_rack field to Default LBP

Updated documentation in the docs book.

Author: Gor Stepanyan

docs/source/load-balancing/default-policy.md

@@ -11,6 +11,7 @@ for queries with non-local consistency mode is also supported.
 `DefaultPolicyBuilder` with the following default values:
 
 - `preferred_datacenter`: `None`
+- `preferred_rack`: `None`
 - `is_token_aware`: `true`
 - `permit_dc_failover`: `false`
 - `latency_awareness`: `None`
@@ -25,6 +26,7 @@ use scylla::load_balancing::DefaultPolicy;
 
 let default_policy = DefaultPolicy::builder()
         .prefer_datacenter("dc1".to_string())
+        .prefer_rack("rack1".to_string())
         .token_aware(true)
         .permit_dc_failover(true)
         .build();
@@ -48,6 +50,14 @@ false), the default policy will only include local nodes in load balancing
 plans. Remote nodes will be excluded, even if they are alive and available to
 serve requests.
 
+#### Preferred Rack
+
+The `preferred_rack` field in `DefaultPolicy` allows the load balancing policy to
+prioritize nodes based on their availability zones in the preferred datacenter.
+When a preferred rack is set, the policy will first return replicas in the local rack
+in the preferred datacenter, and then the other replicas in the datacenter.
+When a preferred datacenter is not set, setting preferred rack will not have any effect.
+
 #### Datacenter Failover
 
 In the event of a datacenter outage or network failure, the nodes in that

scylla/src/transport/load_balancing/default.rs

@@ -21,6 +21,7 @@ use tracing::warn;
 /// Latency awareness is available, although not recommended.
 pub struct DefaultPolicy {
     preferred_datacenter: Option<String>,
+    preferred_rack: Option<String>,
     is_token_aware: bool,
     permit_dc_failover: bool,
     pick_predicate: Box<dyn Fn(&NodeRef) -> bool + Send + Sync>,
@@ -32,6 +33,7 @@ impl fmt::Debug for DefaultPolicy {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("DefaultPolicy")
             .field("preferred_datacenter", &self.preferred_datacenter)
+            .field("preferred_rack", &self.preferred_rack)
             .field("is_token_aware", &self.is_token_aware)
             .field("permit_dc_failover", &self.permit_dc_failover)
             .field("latency_awareness", &self.latency_awareness)
@@ -364,6 +366,7 @@ impl Default for DefaultPolicy {
     fn default() -> Self {
         Self {
             preferred_datacenter: None,
+            preferred_rack: None,
             is_token_aware: true,
             permit_dc_failover: false,
             pick_predicate: Box::new(Self::is_alive),
@@ -390,6 +393,7 @@ impl Default for DefaultPolicy {
 #[derive(Clone, Debug)]
 pub struct DefaultPolicyBuilder {
     preferred_datacenter: Option<String>,
+    preferred_rack: Option<String>,
     is_token_aware: bool,
     permit_dc_failover: bool,
     latency_awareness: Option<LatencyAwarenessBuilder>,
@@ -401,6 +405,7 @@ impl DefaultPolicyBuilder {
     pub fn new() -> Self {
         Self {
             preferred_datacenter: None,
+            preferred_rack: None,
             is_token_aware: true,
             permit_dc_failover: false,
             latency_awareness: None,
@@ -421,6 +426,7 @@ impl DefaultPolicyBuilder {
 
         Arc::new(DefaultPolicy {
             preferred_datacenter: self.preferred_datacenter,
+            preferred_rack: self.preferred_rack,
             is_token_aware: self.is_token_aware,
             permit_dc_failover: self.permit_dc_failover,
             pick_predicate,
@@ -429,6 +435,19 @@ impl DefaultPolicyBuilder {
         })
     }
 
+    /// Sets the rack to be preferred by this policy
+    ///
+    /// Allows the load balancing policy to prioritize nodes based on their availability zones
+    /// in the preferred datacenter.
+    /// When a preferred rack is set, the policy will first return replicas in the local rack
+    /// in the preferred datacenter, and then the other replicas in the datacenter.
+    ///
+    /// When a preferred datacenter is not set, setting preferred rack will not have any effect.
+    pub fn prefer_rack(mut self, rack_name: String) -> Self {
+        self.preferred_rack = Some(rack_name);
+        self
+    }
+
     /// Sets the datacenter to be preferred by this policy.
     ///
     /// Allows the load balancing policy to prioritize nodes based on their location.
@@ -1786,6 +1805,7 @@ mod latency_awareness {
 
             DefaultPolicy {
                 preferred_datacenter: Some("eu".to_owned()),
+                preferred_rack: None,
                 permit_dc_failover: true,
                 is_token_aware: true,
                 pick_predicate,