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.

Sign up for GitHub

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

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

Merged
merged 2 commits into from
Aug 28, 2024
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9165165
SlotMap refactor: Added new NodesMap, changed shard addresses to be s…
barshaul Aug 27, 2024
4ae711b
Addressing comments
barshaul Aug 27, 2024
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
34 changes: 16 additions & 18 deletions redis/src/cluster.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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));
}
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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),
Expand All @@ -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)
}

Expand All @@ -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,
{
Expand All @@ -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());
Expand Down
39 changes: 21 additions & 18 deletions redis/src/cluster_async/connections_container.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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)
Expand Down Expand Up @@ -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>> {
Expand Down
40 changes: 11 additions & 29 deletions redis/src/cluster_async/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -396,7 +396,7 @@ where
&self,
slot: u16,
slot_addr: SlotAddr,
) -> Option<String> {
) -> Option<Arc<String>> {
self.conn_lock
.read()
.await
Expand Down Expand Up @@ -444,7 +444,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
Expand Down Expand Up @@ -1020,7 +1020,6 @@ where
Self::refresh_slots_and_subscriptions_with_retries(
connection.inner.clone(),
&RefreshPolicy::NotThrottable,
None,
)
.await?;

Expand Down Expand Up @@ -1164,7 +1163,6 @@ where
if let Err(err) = Self::refresh_slots_and_subscriptions_with_retries(
inner.clone(),
&RefreshPolicy::Throttable,
None,
)
.await
{
Expand Down Expand Up @@ -1336,7 +1334,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,
Expand Down Expand Up @@ -1388,10 +1385,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?;
Comment on lines -1391 to -1394
Copy link
Author

@barshaul barshaul Aug 27, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removed as per the design decision, updates based on MOVED errors will be done synchronously and not as part of refresh_slots.

}
in_progress.store(false, Ordering::Relaxed);

Expand All @@ -1400,15 +1393,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.
Expand All @@ -1418,7 +1402,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)
}
Expand Down Expand Up @@ -1629,21 +1613,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));
Expand All @@ -1669,18 +1652,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.push_sender.clone(),
)
.await;
if let Ok(node) = node {
connections.0.insert(addr.into(), node);
connections.0.insert(addr, node);
}
connections
},
Expand Down Expand Up @@ -2024,7 +2007,6 @@ where
*future = Box::pin(Self::refresh_slots_and_subscriptions_with_retries(
self.inner.clone(),
&RefreshPolicy::Throttable,
None,
));
Poll::Ready(Err(err))
}
Expand Down Expand Up @@ -2271,12 +2253,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,
),
)));
}
Expand Down
Loading