[redis-rs][core] Move connection refresh to the background

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

@@ -5,11 +5,16 @@ use crate::cluster_topology::TopologyHash;
 use dashmap::DashMap;
 use futures::FutureExt;
 use rand::seq::IteratorRandom;
+use std::collections::{HashMap, HashSet};
 use std::net::IpAddr;
 use std::sync::atomic::Ordering;
 use std::sync::Arc;
 use telemetrylib::Telemetry;
 
+use tracing::debug;
+
+use tokio::task::JoinHandle;
+
 /// Count the number of connections in a connections_map object
 macro_rules! count_connections {
     ($conn_map:expr) => {{
@@ -134,11 +139,59 @@ impl<Connection> std::fmt::Display for ConnectionsMap<Connection> {
     }
 }
 
+// This struct is used to track the status of each address refresh
+pub(crate) struct RefreshConnectionStates<Connection> {
+    // Holds all the failed addresses that started a refresh task.
+    pub(crate) refresh_addresses_started: HashSet<String>,
+    // Follow the refresh ops on the connections
+    pub(crate) refresh_address_in_progress: HashMap<String, JoinHandle<()>>,
+    // Holds all the refreshed addresses that are ready to be inserted into the connection_map
+    pub(crate) refresh_addresses_done: HashMap<String, Option<ClusterNode<Connection>>>,
+}
+
+impl<Connection> RefreshConnectionStates<Connection> {
+    pub(crate) fn clear_refresh_state(&mut self) {
+        let addresses: Vec<String> = self
+            .refresh_address_in_progress
+            .iter()
+            .map(|entry| entry.0.clone())
+            .collect();
+
+        debug!(
+            "clear_refresh_state: removing all refresh data and tasks for addresses: {:?}",
+            addresses
+        );
+
+        for address in addresses {
+            if let Some(refresh_task) = self.refresh_address_in_progress.remove(&address) {
+                // Check if handle exists before calling abort
+                if !refresh_task.is_finished() {
+                    refresh_task.abort();
+                }
+            }
+        }
+
+        self.refresh_addresses_started.clear();
+        self.refresh_addresses_done.clear();
+    }
+}
+
+impl<Connection> Default for RefreshConnectionStates<Connection> {
+    fn default() -> Self {
+        Self {
+            refresh_addresses_started: HashSet::new(),
+            refresh_address_in_progress: HashMap::new(),
+            refresh_addresses_done: HashMap::new(),
+        }
+    }
+}
+
 pub(crate) struct ConnectionsContainer<Connection> {
     connection_map: DashMap<String, ClusterNode<Connection>>,
     pub(crate) slot_map: SlotMap,
     read_from_replica_strategy: ReadFromReplicaStrategy,
     topology_hash: TopologyHash,
+    pub(crate) refresh_conn_state: RefreshConnectionStates<Connection>,
 }
 
 impl<Connection> Drop for ConnectionsContainer<Connection> {
@@ -155,6 +208,7 @@ impl<Connection> Default for ConnectionsContainer<Connection> {
             slot_map: Default::default(),
             read_from_replica_strategy: ReadFromReplicaStrategy::AlwaysFromPrimary,
             topology_hash: 0,
+            refresh_conn_state: Default::default(),
         }
     }
 }
@@ -182,6 +236,7 @@ where
             slot_map,
             read_from_replica_strategy,
             topology_hash,
+            refresh_conn_state: Default::default(),
         }
     }
 
@@ -572,6 +627,7 @@ mod tests {
             connection_map,
             read_from_replica_strategy: ReadFromReplicaStrategy::AZAffinity("use-1a".to_string()),
             topology_hash: 0,
+            refresh_conn_state: Default::default(),
         }
     }
 
@@ -628,6 +684,7 @@ mod tests {
             connection_map,
             read_from_replica_strategy: strategy,
             topology_hash: 0,
+            refresh_conn_state: Default::default(),
         }
     }