Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactored the connections map from HashMap to DashMap #180

Merged
merged 3 commits into from
Aug 7, 2024
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

5 changes: 4 additions & 1 deletion redis/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,9 @@ crc16 = { version = "0.4", optional = true }
rand = { version = "0.8", optional = true }
derivative = { version = "2.2.0", optional = true }

# Only needed for async cluster
dashmap = { version = "6.0.1", optional = true }
eifrah-aws marked this conversation as resolved.
Show resolved Hide resolved

# Only needed for async_std support
async-std = { version = "1.8.0", optional = true }
async-trait = { version = "0.1.24", optional = true }
Expand Down Expand Up @@ -124,7 +127,7 @@ tokio-native-tls-comp = ["tokio-comp", "tls-native-tls", "tokio-native-tls"]
tokio-rustls-comp = ["tokio-comp", "tls-rustls", "tokio-rustls"]
connection-manager = ["futures", "aio", "tokio-retry"]
streams = []
cluster-async = ["cluster", "futures", "futures-util"]
cluster-async = ["cluster", "futures", "futures-util", "dashmap"]
keep-alive = ["socket2"]
sentinel = ["rand"]
tcp_nodelay = []
Expand Down
85 changes: 46 additions & 39 deletions redis/src/cluster_async/connections_container.rs
Original file line number Diff line number Diff line change
@@ -1,14 +1,13 @@
use crate::cluster_async::ConnectionFuture;
use crate::cluster_routing::{Route, SlotAddr};
use crate::cluster_slotmap::{ReadFromReplicaStrategy, SlotMap, SlotMapValue};
use crate::cluster_topology::TopologyHash;
use arcstr::ArcStr;
use dashmap::DashMap;
use futures::FutureExt;
use rand::seq::IteratorRandom;
use std::collections::HashMap;
use std::net::IpAddr;

use crate::cluster_routing::{Route, SlotAddr};
use crate::cluster_slotmap::{ReadFromReplicaStrategy, SlotMap, SlotMapValue};
use crate::cluster_topology::TopologyHash;

/// A struct that encapsulates a network connection along with its associated IP address.
#[derive(Clone, Eq, PartialEq, Debug)]
pub struct ConnectionWithIp<Connection> {
Expand Down Expand Up @@ -86,11 +85,12 @@ pub(crate) enum ConnectionType {
PreferManagement,
}

pub(crate) struct ConnectionsMap<Connection>(pub(crate) HashMap<ArcStr, ClusterNode<Connection>>);
pub(crate) struct ConnectionsMap<Connection>(pub(crate) DashMap<ArcStr, ClusterNode<Connection>>);
eifrah-aws marked this conversation as resolved.
Show resolved Hide resolved

impl<Connection> std::fmt::Display for ConnectionsMap<Connection> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
for (address, node) in self.0.iter() {
for item in self.0.iter() {
let (address, node) = (item.key(), item.value());
match node.user_connection.ip {
Some(ip) => writeln!(f, "{address} - {ip}")?,
None => writeln!(f, "{address}")?,
Expand All @@ -101,7 +101,7 @@ impl<Connection> std::fmt::Display for ConnectionsMap<Connection> {
}

pub(crate) struct ConnectionsContainer<Connection> {
connection_map: HashMap<ArcStr, ClusterNode<Connection>>,
connection_map: DashMap<ArcStr, ClusterNode<Connection>>,
pub(crate) slot_map: SlotMap,
read_from_replica_strategy: ReadFromReplicaStrategy,
topology_hash: TopologyHash,
Expand Down Expand Up @@ -213,9 +213,10 @@ where
pub(crate) fn all_node_connections(
&self,
) -> impl Iterator<Item = ConnectionAndAddress<Connection>> + '_ {
self.connection_map
.iter()
.map(move |(address, node)| (address.clone(), node.user_connection.conn.clone()))
self.connection_map.iter().map(move |item| {
let (node, address) = (item.key(), item.value());
(node.clone(), address.user_connection.conn.clone())
})
}

pub(crate) fn all_primary_connections(
Expand All @@ -228,16 +229,19 @@ where
}

pub(crate) fn node_for_address(&self, address: &str) -> Option<ClusterNode<Connection>> {
self.connection_map.get(address).cloned()
self.connection_map
.get(address)
.map(|item| item.value().clone())
}

pub(crate) fn connection_for_address(
&self,
address: &str,
) -> Option<ConnectionAndAddress<Connection>> {
self.connection_map
.get_key_value(address)
.map(|(address, conn)| (address.clone(), conn.user_connection.conn.clone()))
self.connection_map.get(address).map(|item| {
let (address, conn) = (item.key(), item.value());
(address.clone(), conn.user_connection.conn.clone())
})
}

pub(crate) fn random_connections(
Expand All @@ -249,14 +253,15 @@ where
.iter()
.choose_multiple(&mut rand::thread_rng(), amount)
.into_iter()
.map(move |(address, node)| {
.map(move |item| {
let (address, node) = (item.key(), item.value());
let conn = node.get_connection(&conn_type);
(address.clone(), conn)
})
}

pub(crate) fn replace_or_add_connection_for_address(
&mut self,
&self,
address: impl Into<ArcStr>,
node: ClusterNode<Connection>,
) -> ArcStr {
Expand All @@ -265,8 +270,10 @@ where
address
}

pub(crate) fn remove_node(&mut self, address: &ArcStr) -> Option<ClusterNode<Connection>> {
self.connection_map.remove(address)
pub(crate) fn remove_node(&self, address: &ArcStr) -> Option<ClusterNode<Connection>> {
self.connection_map
.remove(address)
.map(|(_key, value)| value)
}

pub(crate) fn len(&self) -> usize {
Expand Down Expand Up @@ -302,13 +309,13 @@ mod tests {
}
}
}
fn remove_nodes(container: &mut ConnectionsContainer<usize>, addresss: &[&str]) {
fn remove_nodes(container: &ConnectionsContainer<usize>, addresss: &[&str]) {
barshaul marked this conversation as resolved.
Show resolved Hide resolved
for address in addresss {
container.remove_node(&(*address).into());
}
}

fn remove_all_connections(container: &mut ConnectionsContainer<usize>) {
fn remove_all_connections(container: &ConnectionsContainer<usize>) {
remove_nodes(
container,
&[
Expand Down Expand Up @@ -366,7 +373,7 @@ mod tests {
],
ReadFromReplicaStrategy::AlwaysFromPrimary, // this argument shouldn't matter, since we overload the RFR strategy.
);
let mut connection_map = HashMap::new();
let connection_map = DashMap::new();
connection_map.insert(
"primary1".into(),
create_cluster_node(1, use_management_connections),
Expand Down Expand Up @@ -514,7 +521,7 @@ mod tests {

#[test]
fn get_replica_connection_for_replica_route_if_some_but_not_all_replicas_were_removed() {
let mut container = create_container();
let container = create_container();
container.remove_node(&"replica3-2".into());

assert_eq!(
Expand All @@ -540,8 +547,8 @@ mod tests {

#[test]
fn get_primary_connection_for_replica_route_if_all_replicas_were_removed() {
let mut container = create_container();
remove_nodes(&mut container, &["replica2-1", "replica3-1", "replica3-2"]);
let container = create_container();
remove_nodes(&container, &["replica2-1", "replica3-1", "replica3-2"]);

assert_eq!(
2,
Expand Down Expand Up @@ -593,15 +600,15 @@ mod tests {

#[test]
fn get_connection_by_address_returns_none_if_connection_was_removed() {
let mut container = create_container();
let container = create_container();
container.remove_node(&"primary1".into());

assert!(container.connection_for_address("primary1").is_none());
}

#[test]
fn get_connection_by_address_returns_added_connection() {
let mut container = create_container();
let container = create_container();
let address = container.replace_or_add_connection_for_address(
"foobar",
ClusterNode::new_only_with_user_conn(4),
Expand Down Expand Up @@ -630,8 +637,8 @@ mod tests {

#[test]
fn get_random_connections_returns_none_if_all_connections_were_removed() {
let mut container = create_container();
remove_all_connections(&mut container);
let container = create_container();
remove_all_connections(&container);

assert_eq!(
0,
Expand All @@ -643,8 +650,8 @@ mod tests {

#[test]
fn get_random_connections_returns_added_connection() {
let mut container = create_container();
remove_all_connections(&mut container);
let container = create_container();
remove_all_connections(&container);
let address = container.replace_or_add_connection_for_address(
"foobar",
ClusterNode::new_only_with_user_conn(4),
Expand Down Expand Up @@ -694,7 +701,7 @@ mod tests {

#[test]
fn get_all_user_connections_returns_added_connection() {
let mut container = create_container();
let container = create_container();
container.replace_or_add_connection_for_address(
"foobar",
ClusterNode::new_only_with_user_conn(4),
Expand All @@ -711,7 +718,7 @@ mod tests {

#[test]
fn get_all_user_connections_does_not_return_removed_connection() {
let mut container = create_container();
let container = create_container();
container.remove_node(&"primary1".into());

let mut connections: Vec<_> = container
Expand All @@ -738,7 +745,7 @@ mod tests {

#[test]
fn get_all_primaries_does_not_return_removed_connection() {
let mut container = create_container();
let container = create_container();
container.remove_node(&"primary1".into());

let mut connections: Vec<_> = container
Expand All @@ -752,7 +759,7 @@ mod tests {

#[test]
fn len_is_adjusted_on_removals_and_additions() {
let mut container = create_container();
let container = create_container();

assert_eq!(container.len(), 6);

Expand All @@ -769,7 +776,7 @@ mod tests {
#[test]
fn len_is_not_adjusted_on_removals_of_nonexisting_connections_or_additions_of_existing_connections(
) {
let mut container = create_container();
let container = create_container();

assert_eq!(container.len(), 6);

Expand All @@ -785,7 +792,7 @@ mod tests {

#[test]
fn remove_node_returns_connection_if_it_exists() {
let mut container = create_container();
let container = create_container();

let connection = container.remove_node(&"primary1".into());
assert_eq!(connection, Some(ClusterNode::new_only_with_user_conn(1)));
Expand All @@ -796,7 +803,7 @@ mod tests {

#[test]
fn test_is_empty() {
let mut container = create_container();
let container = create_container();

assert!(!container.is_empty());
container.remove_node(&"primary1".into());
Expand Down Expand Up @@ -829,7 +836,7 @@ mod tests {

#[test]
fn is_primary_returns_false_for_removed_node() {
let mut container = create_container();
let container = create_container();
let address = "primary1".into();
container.remove_node(&address);

Expand Down
23 changes: 12 additions & 11 deletions redis/src/cluster_async/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@ use crate::{
};
#[cfg(all(not(feature = "tokio-comp"), feature = "async-std-comp"))]
use async_std::task::{spawn, JoinHandle};
use dashmap::DashMap;
#[cfg(all(not(feature = "tokio-comp"), feature = "async-std-comp"))]
use futures::executor::block_on;
use std::{
Expand Down Expand Up @@ -1090,10 +1091,10 @@ where
.buffer_unordered(initial_nodes.len())
.fold(
(
ConnectionsMap(HashMap::with_capacity(initial_nodes.len())),
ConnectionsMap(DashMap::with_capacity(initial_nodes.len())),
None,
),
|mut connections: (ConnectionMap<C>, Option<String>), addr_conn_res| async move {
|connections: (ConnectionMap<C>, Option<String>), addr_conn_res| async move {
match addr_conn_res {
Ok((addr, node)) => {
connections.0 .0.insert(addr.into(), node);
Expand Down Expand Up @@ -1156,14 +1157,14 @@ where
conn_type: RefreshConnectionType,
) {
info!("Started refreshing connections to {:?}", addresses);
let mut connections_container = inner.conn_lock.write().await;
let connections_container = inner.conn_lock.read().await;
let cluster_params = &inner.cluster_params;
let subscriptions_by_address = &inner.subscriptions_by_address;
let push_sender = &inner.push_sender;

stream::iter(addresses.into_iter())
.fold(
&mut *connections_container,
&*connections_container,
|connections_container, address| async move {
let node_option = connections_container.remove_node(&address);

Expand Down Expand Up @@ -1476,12 +1477,12 @@ where
drop(subs_by_address_guard);

if !addrs_to_refresh.is_empty() {
let mut conns_write_guard = inner.conn_lock.write().await;
let conns_read_guard = inner.conn_lock.read().await;
// have to remove or otherwise the refresh_connection wont trigger node recreation
for addr_to_refresh in addrs_to_refresh.iter() {
conns_write_guard.remove_node(addr_to_refresh);
conns_read_guard.remove_node(addr_to_refresh);
}
drop(conns_write_guard);
drop(conns_read_guard);
// immediately trigger connection reestablishment
Self::refresh_connections(
inner.clone(),
Expand Down Expand Up @@ -1609,8 +1610,8 @@ where
.await;
let new_connections: ConnectionMap<C> = stream::iter(addresses_and_connections_iter)
.fold(
ConnectionsMap(HashMap::with_capacity(nodes_len)),
|mut connections, (addr, node)| async {
ConnectionsMap(DashMap::with_capacity(nodes_len)),
|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 =
Expand Down Expand Up @@ -1923,7 +1924,7 @@ where
{
Ok(node) => {
let connection_clone = node.user_connection.conn.clone().await;
let mut connections = core.conn_lock.write().await;
let connections = core.conn_lock.read().await;
let address = connections.replace_or_add_connection_for_address(addr, node);
drop(connections);
(address, connection_clone)
Expand Down Expand Up @@ -1993,7 +1994,7 @@ where
if !is_primary {
// If the connection is a replica, remove the connection and retry.
// The connection will be established again on the next call to refresh slots once the replica is no longer in loading state.
core.conn_lock.write().await.remove_node(&address);
core.conn_lock.read().await.remove_node(&address);
} else {
// If the connection is primary, just sleep and retry
let sleep_duration = core.cluster_params.retry_params.wait_time_for_retry(retry);
Expand Down