Skip to main content

netget/client/redis/
cluster.rs

1//! Cluster mode for the Redis client: one connection per master, each command sent to the
2//! node that owns its hash slot, and redirections followed the way `redis-cli -c` and every
3//! cluster library follow them.
4//!
5//! * The slot map comes from `CLUSTER SLOTS` on the seed node, read before the model is told
6//!   the client is connected.
7//! * A command goes to the owner of its keys' slot (`utils::redis_keys::slot_keys`, the same
8//!   table NetGet's server routes by); a keyless one to the seed.
9//! * `-MOVED <slot> <addr>`: the slot now lives at `addr`. The map is corrected for that slot
10//!   at once, refreshed from `addr` in full, and the command is sent again there.
11//! * `-ASK <slot> <addr>`: this one command, and only it, goes to `addr` preceded by `ASKING`;
12//!   the map is left alone, because the slot is still mid-migration.
13//! * At most [`MAX_REDIRECTS`] redirections per command. A command still being redirected after
14//!   that is answered with the last redirection error and `redirects_exhausted`, rather than
15//!   chased forever between two nodes that disagree.
16//!
17//! Commands run one at a time, each to completion, so pub/sub is not available here: a
18//! subscribed connection never completes.
19use 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
30/// Redirections followed for one command before giving up.
31pub const MAX_REDIRECTS: usize = 5;
32/// Distinct nodes the router will open connections to. A slot map naming more is refused: it
33/// is not a cluster this client was asked to talk to, it is a fan-out.
34pub const MAX_NODES: usize = 64;
35/// How long one node may take to answer one command.
36pub const NODE_REPLY_TIMEOUT: Duration = Duration::from_secs(30);
37
38struct NodeConn {
39    reader: BufReader<OwnedReadHalf>,
40    writer: OwnedWriteHalf,
41}
42
43/// A redirection a node answered with.
44#[derive(Debug, Clone, PartialEq)]
45pub enum Redirect {
46    Moved(u16, String),
47    Ask(u16, String),
48}
49
50/// Read a redirection out of an error reply.
51pub 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
67/// One command's outcome.
68pub struct Routed {
69    pub reply: Reply,
70    /// The node that gave the final answer.
71    pub node: String,
72    pub slot: Option<u16>,
73    /// Each redirection followed, as the node wrote it (`MOVED 3999 127.0.0.1:7001`).
74    pub redirects: Vec<String>,
75    /// The command was still being redirected after [`MAX_REDIRECTS`].
76    pub exhausted: bool,
77}
78
79/// The slot map and the node connections.
80pub struct Router {
81    seed: String,
82    slots: Vec<Option<String>>,
83    conns: HashMap<String, NodeConn>,
84}
85
86impl Router {
87    /// Connect to the seed and read its slot map. Fails if the seed is not a cluster node.
88    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    /// The seed connection's local address.
102    pub fn local_addr(&self) -> Option<std::net::SocketAddr> {
103        self.conns.get(&self.seed)?.writer.local_addr().ok()
104    }
105
106    /// Distinct nodes in the slot map.
107    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    /// Owner of a slot, as the map has it.
115    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    /// One command to one node, one reply back. A node that fails mid-exchange is forgotten so
144    /// the next command reconnects rather than reading a desynchronised stream.
145    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    /// Replace the slot map with `CLUSTER SLOTS` as `via` reports it.
163    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            // An empty host means "the address you reached me on".
196            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        // Connections to nodes that left the map are closed.
228        self.conns
229            .retain(|addr, _| nodes.contains(addr) || *addr == self.seed);
230        Ok(())
231    }
232
233    /// Run one command against the cluster.
234    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}