SlotMap refactor: Added NodesMap, sharing shard addresses between shard nodes and slot map values. (#185)

* SlotMap refactor: Added new NodesMap, changed shard addresses to be shard between shard nodes and slot map values

Author: barshaul

redis/src/cluster.rs

@@ -35,14 +35,14 @@
 //!     .expire(key, 60).ignore()
 //!     .query(&mut connection).unwrap();
 //! ```
+use rand::{seq::IteratorRandom, thread_rng, Rng};
 use std::cell::RefCell;
 use std::collections::HashSet;
 use std::str::FromStr;
+use std::sync::Arc;
 use std::thread;
 use std::time::Duration;
 
-use rand::{seq::IteratorRandom, thread_rng, Rng};
-
 use crate::cluster_pipeline::UNROUTABLE_ERROR;
 use crate::cluster_routing::{
     MultipleNodeRoutingInfo, ResponsePolicy, Routable, SingleNodeRoutingInfo, SlotAddr,
@@ -343,22 +343,20 @@ where
         let mut slots = self.slots.borrow_mut();
         *slots = self.create_new_slots()?;
 
-        let mut nodes = slots.values().flatten().collect::<Vec<_>>();
-        nodes.sort_unstable();
-        nodes.dedup();
-
+        let nodes = slots.all_node_addresses();
         let mut connections = self.connections.borrow_mut();
         *connections = nodes
             .into_iter()
             .filter_map(|addr| {
-                if connections.contains_key(addr) {
-                    let mut conn = connections.remove(addr).unwrap();
+                let addr = addr.to_string();
+                if connections.contains_key(&addr) {
+                    let mut conn = connections.remove(&addr).unwrap();
                     if conn.check_connection() {
                         return Some((addr.to_string(), conn));
                     }
                 }
 
-                if let Ok(mut conn) = self.connect(addr) {
+                if let Ok(mut conn) = self.connect(&addr) {
                     if conn.check_connection() {
                         return Some((addr.to_string(), conn));
                     }
@@ -424,7 +422,7 @@ where
         if let Some(addr) = slots.slot_addr_for_route(route) {
             Ok((
                 addr.to_string(),
-                self.get_connection_by_addr(connections, addr)?,
+                self.get_connection_by_addr(connections, &addr)?,
             ))
         } else {
             // try a random node next.  This is safe if slots are involved
@@ -495,13 +493,13 @@ where
     fn execute_on_all<'a>(
         &'a self,
         input: Input,
-        addresses: HashSet<&'a str>,
+        addresses: HashSet<Arc<String>>,
         connections: &'a mut HashMap<String, C>,
-    ) -> Vec<RedisResult<(&'a str, Value)>> {
+    ) -> Vec<RedisResult<(Arc<String>, Value)>> {
         addresses
             .into_iter()
             .map(|addr| {
-                let connection = self.get_connection_by_addr(connections, addr)?;
+                let connection = self.get_connection_by_addr(connections, &addr)?;
                 match input {
                     Input::Slice { cmd, routable: _ } => connection.req_packed_command(cmd),
                     Input::Cmd(cmd) => connection.req_command(cmd),
@@ -526,16 +524,16 @@ where
         input: Input,
         slots: &'a mut SlotMap,
         connections: &'a mut HashMap<String, C>,
-    ) -> Vec<RedisResult<(&'a str, Value)>> {
-        self.execute_on_all(input, slots.addresses_for_all_nodes(), connections)
+    ) -> Vec<RedisResult<(Arc<String>, Value)>> {
+        self.execute_on_all(input, slots.all_node_addresses(), connections)
     }
 
     fn execute_on_all_primaries<'a>(
         &'a self,
         input: Input,
         slots: &'a mut SlotMap,
         connections: &'a mut HashMap<String, C>,
-    ) -> Vec<RedisResult<(&'a str, Value)>> {
+    ) -> Vec<RedisResult<(Arc<String>, Value)>> {
         self.execute_on_all(input, slots.addresses_for_all_primaries(), connections)
     }
 
@@ -545,7 +543,7 @@ where
         slots: &'a mut SlotMap,
         connections: &'a mut HashMap<String, C>,
         routes: &'b [(Route, Vec<usize>)],
-    ) -> Vec<RedisResult<(&'a str, Value)>>
+    ) -> Vec<RedisResult<(Arc<String>, Value)>>
     where
         'b: 'a,
     {
@@ -557,7 +555,7 @@ where
                     ErrorKind::IoError,
                     "Couldn't find connection",
                 )))?;
-                let connection = self.get_connection_by_addr(connections, addr)?;
+                let connection = self.get_connection_by_addr(connections, &addr)?;
                 let (_, indices) = routes.get(index).unwrap();
                 let cmd =
                     crate::cluster_routing::command_for_multi_slot_indices(&input, indices.iter());

redis/src/cluster_async/connections_container.rs

@@ -147,32 +147,32 @@ where
 
     /// Returns true if the address represents a known primary node.
     pub(crate) fn is_primary(&self, address: &String) -> bool {
-        self.connection_for_address(address).is_some()
-            && self
-                .slot_map
-                .values()
-                .any(|slot_addrs| slot_addrs.primary.as_str() == address)
+        self.connection_for_address(address).is_some() && self.slot_map.is_primary(address)
     }
 
     fn round_robin_read_from_replica(
         &self,
         slot_map_value: &SlotMapValue,
     ) -> Option<ConnectionAndAddress<Connection>> {
-        let addrs = &slot_map_value.addrs;
+        let addrs = &slot_map_value
+            .addrs
+            .read()
+            .expect("Failed to obtain ShardAddrs's read lock");
         let initial_index = slot_map_value
-            .latest_used_replica
+            .last_used_replica
             .load(std::sync::atomic::Ordering::Relaxed);
         let mut check_count = 0;
         loop {
             check_count += 1;
 
             // Looped through all replicas, no connected replica was found.
-            if check_count > addrs.replicas.len() {
-                return self.connection_for_address(addrs.primary.as_str());
+            if check_count > addrs.replicas().len() {
+                return self.connection_for_address(addrs.primary().as_str());
             }
-            let index = (initial_index + check_count) % addrs.replicas.len();
-            if let Some(connection) = self.connection_for_address(addrs.replicas[index].as_str()) {
-                let _ = slot_map_value.latest_used_replica.compare_exchange_weak(
+            let index = (initial_index + check_count) % addrs.replicas().len();
+            if let Some(connection) = self.connection_for_address(addrs.replicas()[index].as_str())
+            {
+                let _ = slot_map_value.last_used_replica.compare_exchange_weak(
                     initial_index,
                     index,
                     std::sync::atomic::Ordering::Relaxed,
@@ -185,16 +185,19 @@ where
 
     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());
+        let addrs = &slot_map_value
+            .addrs
+            .read()
+            .expect("Failed to obtain ShardAddrs's read lock");
+        if addrs.replicas().is_empty() {
+            return self.connection_for_address(addrs.primary().as_str());
         }
 
         match route.slot_addr() {
-            SlotAddr::Master => self.connection_for_address(addrs.primary.as_str()),
+            SlotAddr::Master => self.connection_for_address(addrs.primary().as_str()),
             SlotAddr::ReplicaOptional => match self.read_from_replica_strategy {
                 ReadFromReplicaStrategy::AlwaysFromPrimary => {
-                    self.connection_for_address(addrs.primary.as_str())
+                    self.connection_for_address(addrs.primary().as_str())
                 }
                 ReadFromReplicaStrategy::RoundRobin => {
                     self.round_robin_read_from_replica(slot_map_value)
@@ -232,7 +235,7 @@ where
         self.slot_map
             .addresses_for_all_primaries()
             .into_iter()
-            .flat_map(|addr| self.connection_for_address(addr))
+            .flat_map(|addr| self.connection_for_address(&addr))
     }
 
     pub(crate) fn node_for_address(&self, address: &str) -> Option<ClusterNode<Connection>> {

redis/src/cluster_async/mod.rs

@@ -438,7 +438,7 @@ where
         &self,
         slot: u16,
         slot_addr: SlotAddr,
-    ) -> Option<String> {
+    ) -> Option<Arc<String>> {
         self.conn_lock
             .read()
             .await
@@ -486,7 +486,7 @@ where
     }
 
     // return slots of node
-    pub(crate) async fn get_slots_of_address(&self, node_address: &str) -> Vec<u16> {
+    pub(crate) async fn get_slots_of_address(&self, node_address: Arc<String>) -> Vec<u16> {
         self.conn_lock
             .read()
             .await
@@ -1090,7 +1090,6 @@ where
         Self::refresh_slots_and_subscriptions_with_retries(
             connection.inner.clone(),
             &RefreshPolicy::NotThrottable,
-            None,
         )
         .await?;
 
@@ -1250,7 +1249,6 @@ where
             if let Err(err) = Self::refresh_slots_and_subscriptions_with_retries(
                 inner.clone(),
                 &RefreshPolicy::Throttable,
-                None,
             )
             .await
             {
@@ -1492,7 +1490,6 @@ where
     async fn refresh_slots_and_subscriptions_with_retries(
         inner: Arc<InnerCore<C>>,
         policy: &RefreshPolicy,
-        moved_redirect: Option<RedirectNode>,
     ) -> RedisResult<()> {
         let SlotRefreshState {
             in_progress,
@@ -1544,10 +1541,6 @@ where
                 Self::refresh_slots(inner.clone(), curr_retry)
             })
             .await;
-        } else if moved_redirect.is_some() {
-            // Update relevant slots in the slots map based on the moved_redirect address,
-            // rather than refreshing all slots by querying the cluster nodes for their topology view.
-            Self::update_slots_for_redirect_change(inner.clone(), moved_redirect).await?;
         }
         in_progress.store(false, Ordering::Relaxed);
 
@@ -1556,15 +1549,6 @@ where
         res
     }
 
-    /// Update relevant slots in the slots map based on the moved_redirect address
-    pub(crate) async fn update_slots_for_redirect_change(
-        _inner: Arc<InnerCore<C>>,
-        _moved_redirect: Option<RedirectNode>,
-    ) -> RedisResult<()> {
-        // TODO: Add implementation
-        Ok(())
-    }
-
     /// Determines if the cluster topology has changed and refreshes slots and subscriptions if needed.
     /// Returns `RedisResult` with `true` if changes were detected and slots were refreshed,
     /// or `false` if no changes were found. Raises an error if refreshing the topology fails.
@@ -1574,7 +1558,7 @@ where
     ) -> RedisResult<bool> {
         let topology_changed = Self::check_for_topology_diff(inner.clone()).await;
         if topology_changed {
-            Self::refresh_slots_and_subscriptions_with_retries(inner.clone(), policy, None).await?;
+            Self::refresh_slots_and_subscriptions_with_retries(inner.clone(), policy).await?;
         }
         Ok(topology_changed)
     }
@@ -1797,21 +1781,20 @@ where
         .0?;
         let connections = &*read_guard;
         // Create a new connection vector of the found nodes
-        let mut nodes = new_slots.values().flatten().collect::<Vec<_>>();
-        nodes.sort_unstable();
-        nodes.dedup();
+        let nodes = new_slots.all_node_addresses();
         let nodes_len = nodes.len();
         let addresses_and_connections_iter = stream::iter(nodes)
             .fold(
                 Vec::with_capacity(nodes_len),
                 |mut addrs_and_conns, addr| async move {
+                    let addr = addr.to_string();
                     if let Some(node) = connections.node_for_address(addr.as_str()) {
                         addrs_and_conns.push((addr, Some(node)));
                         return addrs_and_conns;
                     }
                     // If it's a DNS endpoint, it could have been stored in the existing connections vector using the resolved IP address instead of the DNS endpoint's name.
                     // We shall check if a connection is already exists under the resolved IP name.
-                    let (host, port) = match get_host_and_port_from_addr(addr) {
+                    let (host, port) = match get_host_and_port_from_addr(&addr) {
                         Some((host, port)) => (host, port),
                         None => {
                             addrs_and_conns.push((addr, None));
@@ -1837,18 +1820,18 @@ where
                 |connections, (addr, node)| async {
                     let mut cluster_params = inner.cluster_params.clone();
                     let subs_guard = inner.subscriptions_by_address.read().await;
-                    cluster_params.pubsub_subscriptions = subs_guard.get(addr).cloned();
+                    cluster_params.pubsub_subscriptions = subs_guard.get(&addr).cloned();
                     drop(subs_guard);
                     let node = get_or_create_conn(
-                        addr,
+                        &addr,
                         node,
                         &cluster_params,
                         RefreshConnectionType::AllConnections,
                         inner.glide_connection_options.clone(),
                     )
                     .await;
                     if let Ok(node) = node {
-                        connections.0.insert(addr.into(), node);
+                        connections.0.insert(addr, node);
                     }
                     connections
                 },
@@ -2192,7 +2175,6 @@ where
                     *future = Box::pin(Self::refresh_slots_and_subscriptions_with_retries(
                         self.inner.clone(),
                         &RefreshPolicy::Throttable,
-                        None,
                     ));
                     Poll::Ready(Err(err))
                 }
@@ -2439,12 +2421,12 @@ where
 
             match ready!(self.poll_complete(cx)) {
                 PollFlushAction::None => return Poll::Ready(Ok(())),
-                PollFlushAction::RebuildSlots(moved_redirect) => {
+                PollFlushAction::RebuildSlots(_moved_redirect) => {
+                    // TODO: Add logic to update the slots map based on the MOVED error
                     self.state = ConnectionState::Recover(RecoverFuture::RecoverSlots(Box::pin(
                         ClusterConnInner::refresh_slots_and_subscriptions_with_retries(
                             self.inner.clone(),
                             &RefreshPolicy::Throttable,
-                            moved_redirect,
                         ),
                     )));
                 }