Improve robustness of functions for getting connections in Node
Node::random_connection and Node::connection_for_token were failing
immediately if the chosen connection was broken. Now if the ideal
connection is not available they try getting any other working
connection.
/// Node represents a cluster node along with it's data and connections
-use crate::routing::{Shard, ShardInfo, Token};
+use crate::routing::{ShardInfo, Token};
use crate::transport::connection::VerifiedKeyspaceName;
use crate::transport::connection::{Connection, ConnectionConfig};
use crate::transport::connection_keeper::{ConnectionKeeper, ShardInfoSender};
@@ -7,7 +7,8 @@ use crate::transport::errors::QueryError;
use futures::future::join_all;
use futures::{future::RemoteHandle, FutureExt};
-use rand::Rng;
+use rand::{seq::SliceRandom, Rng};
+use std::convert::TryInto;
use std::hash::{Hash, Hasher};
use std::net::SocketAddr;
use std::sync::{Arc, RwLock};
@@ -113,6 +114,7 @@ impl Node {
}
/// Get connection which should be used to connect using given token
+ /// If this connection is broken get any random connection to this Node
pub async fn connection_for_token(&self, token: Token) -> Result<Arc<Connection>, QueryError> {
let connections: Arc<NodeConnections> = self.connections.read().unwrap().clone();
@@ -122,10 +124,8 @@ impl Node {
shard_info,
shard_conns,
} => {
- let shard: Shard = shard_info.shard_of(token);
- // It's guaranteed by NodeWorker that for each shard < shard_info.nr_shards
- // there is an entry for this shard in connections map
- shard_conns[shard as usize].get_connection().await
+ let shard: usize = shard_info.shard_of(token).try_into().unwrap();
+ Self::connection_for_shard(shard, shard_conns).await
}
}
}
@@ -140,13 +140,39 @@ impl Node {
shard_info,
shard_conns,
} => {
- let shard = rand::thread_rng().gen_range(0..shard_info.nr_shards as u32);
- // It's guaranteed by NodeWorker that shard_conns.len() == shard_info.nr_shards
- shard_conns[shard as usize].get_connection().await
+ let shard: usize = rand::thread_rng().gen_range(0..shard_info.nr_shards).into();
+ Self::connection_for_shard(shard, shard_conns).await
}
}
}
+ // Tries to get a connection to given shard, if it's broken returns any working connection
+ async fn connection_for_shard(
+ shard: usize,
+ shard_conns: &[ConnectionKeeper],
+ ) -> Result<Arc<Connection>, QueryError> {
+ // Try getting the desired connection
+ let mut last_error: QueryError = match shard_conns[shard].get_connection().await {
+ Ok(connection) => return Ok(connection),
+ Err(e) => e,
+ };
+
+ // If this fails try getting any other in random order
+ let mut shards_to_try: Vec<usize> =
+ (shard..shard_conns.len()).chain(0..shard).skip(1).collect();
+
+ shards_to_try.shuffle(&mut rand::thread_rng());
+
+ for shard in shards_to_try {
+ match shard_conns[shard].get_connection().await {
+ Ok(conn) => return Ok(conn),
+ Err(e) => last_error = e,
+ }
+ }
+
+ return Err(last_error);
+ }
+
pub async fn use_keyspace(
&self,
keyspace_name: VerifiedKeyspaceName,
