redis-core: add az awareness to read strategy

Signed-off-by: Adar Ovadia <adarov@amazon.com>

Author: Adar Ovadia

glide-core/Cargo.toml

@@ -9,6 +9,7 @@ authors = ["Valkey GLIDE Maintainers"]
 
 [dependencies]
 bytes = "1"
+copstr = "0"
 futures = "^0.3"
 redis = { path = "./redis-rs/redis", features = [
     "aio",

glide-core/redis-rs/redis/Cargo.toml

@@ -22,6 +22,7 @@ bench = false
 # These two are generally really common simple dependencies so it does not seem
 # much of a point to optimize these, but these could in theory be removed for
 # an indirection through std::Formatter.
+copstr = "0"
 ryu = "1.0"
 itoa = "1.0"
 

glide-core/redis-rs/redis/src/cluster.rs

@@ -229,7 +229,10 @@ where
     ) -> RedisResult<Self> {
         let connection = Self {
             connections: RefCell::new(HashMap::new()),
-            slots: RefCell::new(SlotMap::new(vec![], cluster_params.read_from_replicas)),
+            slots: RefCell::new(SlotMap::new(
+                vec![],
+                cluster_params.read_from_replicas.clone(),
+            )),
             auto_reconnect: RefCell::new(true),
             cluster_params,
             read_timeout: RefCell::new(None),
@@ -387,7 +390,7 @@ where
                 "can't parse node address",
             )))?;
             match parse_and_count_slots(&value, self.cluster_params.tls, addr).map(|slots_data| {
-                SlotMap::new(slots_data.1, self.cluster_params.read_from_replicas)
+                SlotMap::new(slots_data.1, self.cluster_params.read_from_replicas.clone())
             }) {
                 Ok(new_slots) => {
                     result = Ok(new_slots);

glide-core/redis-rs/redis/src/cluster_async/connections_container.rs

@@ -1,4 +1,4 @@
-use crate::cluster_async::{ClusterParams, ConnectionFuture};
+use crate::cluster_async::ConnectionFuture;
 use crate::cluster_routing::{Route, SlotAddr};
 use crate::cluster_slotmap::{ReadFromReplicaStrategy, SlotMap, SlotMapValue};
 use crate::cluster_topology::TopologyHash;
@@ -233,10 +233,10 @@ where
         }
     }
 
-    fn round_robin_read_from_az_awareness_replica(
+    pub(crate) fn round_robin_read_from_az_awareness_replica(
         &self,
         slot_map_value: &SlotMapValue,
-        user_az: String,
+        client_az: String,
     ) -> Option<ConnectionAndAddress<Connection>> {
         let addrs = &slot_map_value.addrs;
         let initial_index = slot_map_value
@@ -266,7 +266,7 @@ where
             // Check if this replica’s AZ matches the user’s AZ.
             if let Some(connection_details) = self.connection_details_for_address(replica.as_str())
             {
-                if connection_details.1.az.as_deref() == Some(&user_az) {
+                if connection_details.1.az.as_deref() == Some(&client_az) {
                     // Attempt to update `latest_used_replica` with the index of this replica.
                     let _ = slot_map_value.latest_used_replica.compare_exchange_weak(
                         initial_index,
@@ -280,42 +280,33 @@ where
         }
     }
 
-    fn lookup_route(
-        &self,
-        route: &Route,
-        cluster_params: &Option<ClusterParams>,
-    ) -> Option<ConnectionAndAddress<Connection>> {
+    fn lookup_route(&self, route: &Route) -> Option<ConnectionAndAddress<Connection>> {
         let slot_map_value = self.slot_map.slot_value_for_route(route)?;
         let addrs = &slot_map_value.addrs;
         if addrs.replicas.is_empty() {
             return self.connection_for_address(addrs.primary.as_str());
         }
 
-        //
         match route.slot_addr() {
             // Master strategy will be in use when the command is not read_only
             SlotAddr::Master => self.connection_for_address(addrs.primary.as_str()),
             // ReplicaOptional strategy will be in use when the command is read_only
-            SlotAddr::ReplicaOptional => match self.read_from_replica_strategy {
+            SlotAddr::ReplicaOptional => match &self.read_from_replica_strategy {
                 ReadFromReplicaStrategy::AlwaysFromPrimary => {
                     self.connection_for_address(addrs.primary.as_str())
                 }
                 ReadFromReplicaStrategy::RoundRobin => {
                     self.round_robin_read_from_replica(slot_map_value)
                 }
-                ReadFromReplicaStrategy::AZAffinity => self
-                    .round_robin_read_from_az_awareness_replica(
-                        slot_map_value,
-                        cluster_params.as_ref().unwrap().client_az.clone()?,
-                    ),
+                ReadFromReplicaStrategy::AZAffinity(az) => {
+                    self.round_robin_read_from_az_awareness_replica(slot_map_value, az.to_string())
+                }
             },
             // when the user strategy per command is replica_preffered
-            SlotAddr::ReplicaRequired => match self.read_from_replica_strategy {
-                ReadFromReplicaStrategy::AZAffinity => self
-                    .round_robin_read_from_az_awareness_replica(
-                        slot_map_value,
-                        cluster_params.as_ref().unwrap().client_az.clone()?,
-                    ),
+            SlotAddr::ReplicaRequired => match &self.read_from_replica_strategy {
+                ReadFromReplicaStrategy::AZAffinity(az) => {
+                    self.round_robin_read_from_az_awareness_replica(slot_map_value, az.to_string())
+                }
                 _ => self.round_robin_read_from_replica(slot_map_value),
             },
         }
@@ -325,17 +316,9 @@ where
         &self,
         route: &Route,
     ) -> Option<ConnectionAndAddress<Connection>> {
-        self.connection_for_route_with_params(route, None)
-    }
-
-    pub(crate) fn connection_for_route_with_params(
-        &self,
-        route: &Route,
-        cluster_params: Option<ClusterParams>,
-    ) -> Option<ConnectionAndAddress<Connection>> {
-        self.lookup_route(route, &cluster_params).or_else(|| {
+        self.lookup_route(route).or_else(|| {
             if route.slot_addr() != SlotAddr::Master {
-                self.lookup_route(&Route::new(route.slot(), SlotAddr::Master), &cluster_params)
+                self.lookup_route(&Route::new(route.slot(), SlotAddr::Master))
             } else {
                 None
             }
@@ -513,7 +496,7 @@ mod tests {
     }
 
     fn create_container_with_az_strategy(
-        strategy: ReadFromReplicaStrategy,
+        // strategy: ReadFromReplicaStrategy,
         use_management_connections: bool,
     ) -> ConnectionsContainer<usize> {
         let slot_map = SlotMap::new(
@@ -571,7 +554,7 @@ mod tests {
         ConnectionsContainer {
             slot_map,
             connection_map,
-            read_from_replica_strategy: strategy,
+            read_from_replica_strategy: ReadFromReplicaStrategy::AZAffinity("use-1a".to_string()),
             topology_hash: 0,
         }
     }
@@ -802,28 +785,17 @@ mod tests {
 
     #[test]
     fn get_connection_for_az_affinity_route() {
-        let container =
-            create_container_with_az_strategy(ReadFromReplicaStrategy::AZAffinity, false);
-        let mut cluster_params = ClusterParams::default();
-
-        cluster_params.client_az = Some("use-1a".to_string());
+        let container = create_container_with_az_strategy(false);
 
         // slot number is not exits
         assert!(container
-            .connection_for_route_with_params(
-                &Route::new(1001, SlotAddr::ReplicaOptional),
-                Some(cluster_params.clone())
-            )
+            .connection_for_route(&Route::new(1001, SlotAddr::ReplicaOptional))
             .is_none());
-
         // Get the replica that holds the slot 1002
         assert_eq!(
             21,
             container
-                .connection_for_route_with_params(
-                    &Route::new(1002, SlotAddr::ReplicaOptional),
-                    Some(cluster_params.clone())
-                )
+                .connection_for_route(&Route::new(1002, SlotAddr::ReplicaOptional))
                 .unwrap()
                 .1
         );
@@ -832,20 +804,14 @@ mod tests {
         assert_eq!(
             2,
             container
-                .connection_for_route_with_params(
-                    &Route::new(1500, SlotAddr::Master),
-                    Some(cluster_params.clone())
-                )
+                .connection_for_route(&Route::new(1500, SlotAddr::Master))
                 .unwrap()
                 .1
         );
 
         // receive one of the replicas that holds the slot 2001 and is in the availability zone of the client ("use-1a")
         assert!(one_of(
-            container.connection_for_route_with_params(
-                &Route::new(2001, SlotAddr::ReplicaRequired),
-                Some(cluster_params.clone())
-            ),
+            container.connection_for_route(&Route::new(2001, SlotAddr::ReplicaRequired)),
             &[31, 33],
         ));
 
@@ -854,10 +820,7 @@ mod tests {
         assert_eq!(
             31,
             container
-                .connection_for_route_with_params(
-                    &Route::new(2001, SlotAddr::ReplicaOptional),
-                    Some(cluster_params.clone())
-                )
+                .connection_for_route(&Route::new(2001, SlotAddr::ReplicaOptional))
                 .unwrap()
                 .1
         );
@@ -867,10 +830,7 @@ mod tests {
         assert_eq!(
             32,
             container
-                .connection_for_route_with_params(
-                    &Route::new(2001, SlotAddr::ReplicaOptional),
-                    Some(cluster_params.clone())
-                )
+                .connection_for_route(&Route::new(2001, SlotAddr::ReplicaOptional))
                 .unwrap()
                 .1
         );
@@ -880,15 +840,38 @@ mod tests {
         assert_eq!(
             3,
             container
-                .connection_for_route_with_params(
-                    &Route::new(2001, SlotAddr::ReplicaOptional),
-                    Some(cluster_params.clone())
-                )
+                .connection_for_route(&Route::new(2001, SlotAddr::ReplicaOptional))
                 .unwrap()
                 .1
         );
     }
 
+    #[test]
+    fn get_connection_for_az_affinity_route_round_robin() {
+        let container = create_container_with_az_strategy(false);
+
+        let mut addresses = vec![
+            container
+                .connection_for_route(&Route::new(2001, SlotAddr::ReplicaOptional))
+                .unwrap()
+                .1,
+            container
+                .connection_for_route(&Route::new(2001, SlotAddr::ReplicaOptional))
+                .unwrap()
+                .1,
+            container
+                .connection_for_route(&Route::new(2001, SlotAddr::ReplicaOptional))
+                .unwrap()
+                .1,
+            container
+                .connection_for_route(&Route::new(2001, SlotAddr::ReplicaOptional))
+                .unwrap()
+                .1,
+        ];
+        addresses.sort();
+        assert_eq!(addresses, vec![31, 31, 33, 33]);
+    }
+
     #[test]
     fn get_connection_by_address() {
         let container = create_container();

glide-core/redis-rs/redis/src/cluster_async/connections_logic.rs

@@ -14,7 +14,6 @@ use crate::{
 };
 use std::net::SocketAddr;
 
-use future::ok;
 use futures::prelude::*;
 use futures_util::{future::BoxFuture, join};
 use tracing::warn;
@@ -374,22 +373,8 @@ where
     match info_res {
         Ok(value) => {
             let info_dict: Result<InfoDict, RedisError> = FromRedisValue::from_redis_value(&value);
-            if let Ok(info_dict) = info_dict {
-                if let Some(az) = info_dict.get::<String>("availability_zone") {
-                    conn_details.az = Some(az);
-                    Ok(())
-                } else {
-                    Err(RedisError::from((
-                        ErrorKind::ResponseError,
-                        "Failed to get availability_zone from info",
-                    )))
-                }
-            } else {
-                Err(RedisError::from((
-                    ErrorKind::ResponseError,
-                    "Failed to parse info command",
-                )))
-            }
+            conn_details.az = info_dict?.get::<String>("availability_zone");
+            Ok(())
         }
         Err(_) => {
             // Handle the error case for the INFO command

glide-core/redis-rs/redis/src/cluster_async/mod.rs

@@ -1043,7 +1043,7 @@ where
             conn_lock: RwLock::new(ConnectionsContainer::new(
                 Default::default(),
                 connections,
-                cluster_params.read_from_replicas,
+                cluster_params.read_from_replicas.clone(),
                 0,
             )),
             cluster_params: cluster_params.clone(),
@@ -1807,7 +1807,7 @@ where
         *write_guard = ConnectionsContainer::new(
             new_slots,
             new_connections,
-            inner.cluster_params.read_from_replicas,
+            inner.cluster_params.read_from_replicas.clone(),
             topology_hash,
         );
         Ok(())
@@ -2026,9 +2026,7 @@ where
                 )
             }
             InternalSingleNodeRouting::SpecificNode(route) => {
-                match read_guard
-                    .connection_for_route_with_params(&route, Some(core.cluster_params.clone()))
-                {
+                match read_guard.connection_for_route(&route) {
                     Some((conn, address)) => ConnectionCheck::Found((conn, address)),
                     None => {
                         // No connection is found for the given route:
@@ -2498,7 +2496,7 @@ where
             curr_retry,
             inner.cluster_params.tls,
             num_of_nodes_to_query,
-            inner.cluster_params.read_from_replicas,
+            inner.cluster_params.read_from_replicas.clone(),
         ),
         failed_addresses,
     )

glide-core/redis-rs/redis/src/cluster_client.rs

@@ -33,7 +33,6 @@ use tokio::sync::mpsc;
 struct BuilderParams {
     password: Option<String>,
     username: Option<String>,
-    client_az: Option<String>,
     read_from_replicas: ReadFromReplicaStrategy,
     tls: Option<TlsMode>,
     #[cfg(feature = "tls-rustls")]
@@ -131,7 +130,6 @@ impl SlotsRefreshRateLimit {
 pub struct ClusterParams {
     pub(crate) password: Option<String>,
     pub(crate) username: Option<String>,
-    pub(crate) client_az: Option<String>,
     pub(crate) read_from_replicas: ReadFromReplicaStrategy,
     /// tls indicates tls behavior of connections.
     /// When Some(TlsMode), connections use tls and verify certification depends on TlsMode.
@@ -167,7 +165,6 @@ impl ClusterParams {
         Ok(Self {
             password: value.password,
             username: value.username,
-            client_az: value.client_az,
             read_from_replicas: value.read_from_replicas,
             tls: value.tls,
             retry_params: value.retries_configuration,
@@ -390,6 +387,22 @@ impl ClusterClientBuilder {
         self
     }
 
+    /// Sets the read strategy on all new connections, based on the specified policy.
+    ///
+    /// Using the specified `read_strategy`, this function configures whether read queries will be
+    /// routed to replica nodes or primary nodes. If `ReadFromReplicaStrategy::AZAffinity` is set,
+    /// read requests will first attempt to access replicas in the same availability zone, falling
+    /// back to other replicas or the primary if needed. If `ReadFromReplicaStrategy::RoundRobin` is chosen, reads are distributed
+    /// across replicas for load balancing, while `ReadFromReplicaStrategy::AlwaysFromPrimary` ensures all read and write queries
+    /// are directed to the primary node.
+    ///
+    /// # Parameters
+    /// - `read_strategy`: Defines the replica routing strategy.
+    pub fn read_from(mut self, read_strategy: ReadFromReplicaStrategy) -> ClusterClientBuilder {
+        self.builder_params.read_from_replicas = read_strategy;
+        self
+    }
+
     /// Enables periodic topology checks for this client.
     ///
     /// If enabled, periodic topology checks will be executed at the configured intervals to examine whether there