netget/client/redis/
cluster.rs1use super::resp::{self, Reply};
20use crate::utils::redis_keys::{key_hash_slot, slot_keys, CLUSTER_SLOTS};
21use anyhow::{bail, ensure, Context, Result};
22use serde_json::Value;
23use std::collections::HashMap;
24use std::time::Duration;
25use tokio::io::{AsyncWriteExt, BufReader};
26use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
27use tokio::net::TcpStream;
28use tracing::{debug, info, warn};
29
30pub const MAX_REDIRECTS: usize = 5;
32pub const MAX_NODES: usize = 64;
35pub const NODE_REPLY_TIMEOUT: Duration = Duration::from_secs(30);
37
38struct NodeConn {
39 reader: BufReader<OwnedReadHalf>,
40 writer: OwnedWriteHalf,
41}
42
43#[derive(Debug, Clone, PartialEq)]
45pub enum Redirect {
46 Moved(u16, String),
47 Ask(u16, String),
48}
49
50pub fn parse_redirect(reply: &Reply) -> Option<Redirect> {
52 let text = reply.value.get("error")?.as_str()?;
53 let mut parts = text.split_whitespace();
54 let word = parts.next()?;
55 let slot: u16 = parts.next()?.parse().ok()?;
56 let addr = parts.next()?.to_string();
57 if slot >= CLUSTER_SLOTS || !addr.contains(':') {
58 return None;
59 }
60 match word {
61 "MOVED" => Some(Redirect::Moved(slot, addr)),
62 "ASK" => Some(Redirect::Ask(slot, addr)),
63 _ => None,
64 }
65}
66
67pub struct Routed {
69 pub reply: Reply,
70 pub node: String,
72 pub slot: Option<u16>,
73 pub redirects: Vec<String>,
75 pub exhausted: bool,
77}
78
79pub struct Router {
81 seed: String,
82 slots: Vec<Option<String>>,
83 conns: HashMap<String, NodeConn>,
84}
85
86impl Router {
87 pub async fn connect(seed: &str) -> Result<Router> {
89 let mut router = Router {
90 seed: seed.to_string(),
91 slots: vec![None; usize::from(CLUSTER_SLOTS)],
92 conns: HashMap::new(),
93 };
94 router
95 .refresh(seed)
96 .await
97 .with_context(|| format!("reading the cluster's slot map from {seed}"))?;
98 Ok(router)
99 }
100
101 pub fn local_addr(&self) -> Option<std::net::SocketAddr> {
103 self.conns.get(&self.seed)?.writer.local_addr().ok()
104 }
105
106 pub fn nodes(&self) -> Vec<String> {
108 let mut nodes: Vec<String> = self.slots.iter().flatten().cloned().collect();
109 nodes.sort();
110 nodes.dedup();
111 nodes
112 }
113
114 pub fn owner(&self, slot: u16) -> Option<&str> {
116 self.slots[usize::from(slot)].as_deref()
117 }
118
119 async fn conn(&mut self, addr: &str) -> Result<&mut NodeConn> {
120 if !self.conns.contains_key(addr) {
121 ensure!(
122 self.conns.len() < MAX_NODES,
123 "refusing to open a connection to a {}th cluster node",
124 MAX_NODES + 1
125 );
126 let stream = tokio::time::timeout(NODE_REPLY_TIMEOUT, TcpStream::connect(addr))
127 .await
128 .with_context(|| format!("connecting to cluster node {addr} timed out"))?
129 .with_context(|| format!("connecting to cluster node {addr}"))?;
130 let (r, w) = stream.into_split();
131 info!("Redis cluster client connected to node {}", addr);
132 self.conns.insert(
133 addr.to_string(),
134 NodeConn {
135 reader: BufReader::new(r),
136 writer: w,
137 },
138 );
139 }
140 Ok(self.conns.get_mut(addr).expect("inserted above"))
141 }
142
143 async fn roundtrip(&mut self, addr: &str, args: &[Vec<u8>]) -> Result<Reply> {
146 let result = async {
147 let conn = self.conn(addr).await?;
148 conn.writer.write_all(&resp::encode_command(args)).await?;
149 conn.writer.flush().await?;
150 tokio::time::timeout(NODE_REPLY_TIMEOUT, resp::read_reply(&mut conn.reader))
151 .await
152 .with_context(|| format!("cluster node {addr} did not answer"))??
153 .with_context(|| format!("cluster node {addr} closed the connection"))
154 }
155 .await;
156 if result.is_err() {
157 self.conns.remove(addr);
158 }
159 result
160 }
161
162 pub async fn refresh(&mut self, via: &str) -> Result<()> {
164 let reply = self
165 .roundtrip(via, &[b"CLUSTER".to_vec(), b"SLOTS".to_vec()])
166 .await?;
167 let ranges = reply
168 .value
169 .as_array()
170 .with_context(|| format!("CLUSTER SLOTS answered {}", reply.response))?;
171 let via_host = via.rsplit_once(':').map(|(h, _)| h).unwrap_or(via);
172 let mut slots: Vec<Option<String>> = vec![None; usize::from(CLUSTER_SLOTS)];
173 let mut nodes = std::collections::BTreeSet::new();
174 for range in ranges {
175 let parts = range
176 .as_array()
177 .context("a CLUSTER SLOTS entry is an array")?;
178 let start = parts
179 .first()
180 .and_then(Value::as_u64)
181 .context("slot range start")?;
182 let end = parts
183 .get(1)
184 .and_then(Value::as_u64)
185 .context("slot range end")?;
186 ensure!(
187 start <= end && end < u64::from(CLUSTER_SLOTS),
188 "CLUSTER SLOTS names slots {start}-{end}"
189 );
190 let master = parts
191 .get(2)
192 .and_then(Value::as_array)
193 .context("a slot range names its master")?;
194 let host = master.first().and_then(Value::as_str).unwrap_or("");
195 let host = if host.is_empty() || host == "?" {
197 via_host
198 } else {
199 host
200 };
201 let port = master
202 .get(1)
203 .and_then(Value::as_u64)
204 .filter(|p| (1..=65535).contains(p))
205 .context("a master's port")?;
206 let addr = format!("{host}:{port}");
207 nodes.insert(addr.clone());
208 ensure!(
209 nodes.len() <= MAX_NODES,
210 "the slot map names more than {MAX_NODES} nodes"
211 );
212 for slot in start..=end {
213 slots[slot as usize] = Some(addr.clone());
214 }
215 }
216 ensure!(
217 !nodes.is_empty(),
218 "CLUSTER SLOTS assigned no slot to any node"
219 );
220 debug!(
221 "Redis cluster slot map from {}: {} node(s), {} slot(s) assigned",
222 via,
223 nodes.len(),
224 slots.iter().filter(|s| s.is_some()).count()
225 );
226 self.slots = slots;
227 self.conns
229 .retain(|addr, _| nodes.contains(addr) || *addr == self.seed);
230 Ok(())
231 }
232
233 pub async fn execute(&mut self, args: &[Vec<u8>]) -> Result<Routed> {
235 if args.is_empty() {
236 bail!("an empty command");
237 }
238 let slot = slot_keys(args).first().map(|k| key_hash_slot(k));
239 let mut target = slot
240 .and_then(|s| self.slots[usize::from(s)].clone())
241 .unwrap_or_else(|| self.seed.clone());
242 let mut redirects = Vec::new();
243 let mut asking = false;
244 let mut redirects_left = MAX_REDIRECTS;
245 loop {
246 if asking {
247 let ok = self.roundtrip(&target, &[b"ASKING".to_vec()]).await?;
248 if ok.reply_type == "error" {
249 warn!(
250 "Redis cluster node {} refused ASKING: {}",
251 target, ok.response
252 );
253 }
254 }
255 let reply = self.roundtrip(&target, args).await?;
256 let Some(redirect) = parse_redirect(&reply) else {
257 return Ok(Routed {
258 reply,
259 node: target,
260 slot,
261 redirects,
262 exhausted: false,
263 });
264 };
265 if redirects_left == 0 {
266 warn!(
267 "Redis cluster: still redirected after {} redirections; giving up on this \
268 command",
269 MAX_REDIRECTS
270 );
271 return Ok(Routed {
272 reply,
273 node: target,
274 slot,
275 redirects,
276 exhausted: true,
277 });
278 }
279 redirects_left -= 1;
280 redirects.push(
281 reply
282 .value
283 .get("error")
284 .and_then(Value::as_str)
285 .unwrap_or_default()
286 .to_string(),
287 );
288 match redirect {
289 Redirect::Moved(moved_slot, addr) => {
290 info!(
291 "Redis cluster: slot {} moved to {}; refreshing the slot map",
292 moved_slot, addr
293 );
294 self.slots[usize::from(moved_slot)] = Some(addr.clone());
295 if let Err(e) = self.refresh(&addr).await {
296 warn!(
297 "Redis cluster: refreshing the slot map from {} failed: {:#}",
298 addr, e
299 );
300 self.slots[usize::from(moved_slot)] = Some(addr.clone());
301 }
302 target = addr;
303 asking = false;
304 }
305 Redirect::Ask(_, addr) => {
306 target = addr;
307 asking = true;
308 }
309 }
310 }
311 }
312}