forked from scylladb/scylla-rust-driver
-
Notifications
You must be signed in to change notification settings - Fork 0
/
node.rs
135 lines (116 loc) · 3.82 KB
/
node.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
/// Node represents a cluster node along with it's data and connections
use crate::routing::{Sharder, Token};
use crate::transport::connection::Connection;
use crate::transport::connection::VerifiedKeyspaceName;
use crate::transport::connection_pool::{NodeConnectionPool, PoolConfig};
use crate::transport::errors::QueryError;
use std::{
hash::{Hash, Hasher},
net::SocketAddr,
sync::{
atomic::{AtomicBool, Ordering},
Arc,
},
};
/// Node represents a cluster node along with it's data and connections
#[derive(Debug)]
pub struct Node {
pub address: SocketAddr,
pub datacenter: Option<String>,
pub rack: Option<String>,
// If the node is filtered out by the host filter, this will be None
pool: Option<NodeConnectionPool>,
down_marker: AtomicBool,
}
impl Node {
/// Creates new node which starts connecting in the background
/// # Arguments
///
/// `address` - address to connect to
/// `compression` - preferred compression to use
/// `datacenter` - optional datacenter name
/// `rack` - optional rack name
pub(crate) fn new(
address: SocketAddr,
pool_config: PoolConfig,
datacenter: Option<String>,
rack: Option<String>,
keyspace_name: Option<VerifiedKeyspaceName>,
enabled: bool,
) -> Self {
let pool = enabled.then(|| {
NodeConnectionPool::new(address.ip(), address.port(), pool_config, keyspace_name)
});
Node {
address,
datacenter,
rack,
pool,
down_marker: false.into(),
}
}
pub fn sharder(&self) -> Option<Sharder> {
self.pool.as_ref()?.sharder()
}
/// Get connection which should be used to connect using given token
/// If this connection is broken get any random connection to this Node
pub(crate) async fn connection_for_token(
&self,
token: Token,
) -> Result<Arc<Connection>, QueryError> {
self.get_pool()?.connection_for_token(token)
}
/// Get random connection
pub(crate) async fn random_connection(&self) -> Result<Arc<Connection>, QueryError> {
self.get_pool()?.random_connection()
}
pub fn is_down(&self) -> bool {
self.down_marker.load(Ordering::Relaxed)
}
/// Returns a boolean which indicates whether this node was is enabled.
/// Only enabled nodes will have connections open. For disabled nodes,
/// no connections will be opened.
pub fn is_enabled(&self) -> bool {
self.pool.is_some()
}
pub(crate) fn change_down_marker(&self, is_down: bool) {
self.down_marker.store(is_down, Ordering::Relaxed);
}
pub(crate) async fn use_keyspace(
&self,
keyspace_name: VerifiedKeyspaceName,
) -> Result<(), QueryError> {
if let Some(pool) = &self.pool {
pool.use_keyspace(keyspace_name).await?;
}
Ok(())
}
pub(crate) fn get_working_connections(&self) -> Result<Vec<Arc<Connection>>, QueryError> {
self.get_pool()?.get_working_connections()
}
pub(crate) async fn wait_until_pool_initialized(&self) {
if let Some(pool) = &self.pool {
pool.wait_until_initialized().await;
}
}
fn get_pool(&self) -> Result<&NodeConnectionPool, QueryError> {
self.pool.as_ref().ok_or_else(|| {
QueryError::IoError(Arc::new(std::io::Error::new(
std::io::ErrorKind::Other,
"No connections in the pool: the node has been disabled \
by the host filter",
)))
})
}
}
impl PartialEq for Node {
fn eq(&self, other: &Self) -> bool {
self.address == other.address
}
}
impl Eq for Node {}
impl Hash for Node {
fn hash<H: Hasher>(&self, state: &mut H) {
self.address.hash(state);
}
}