netget/client/redis/
pubsub.rs

1//! Pub/sub on the client: what this connection is subscribed to, telling a delivered message
2//! from an ordinary reply, and publishing while subscribed.
3//!
4//! # Telling a message from a reply
5//!
6//! On RESP3 a message is a push (`>`), unambiguous. On RESP2 it is an ordinary array whose first
7//! element is `message`, `pmessage` or `smessage` — which an `LRANGE` reply could also be. So an
8//! array is read as pub/sub traffic only while the connection is *engaged*: from the moment a
9//! `(P|S)SUBSCRIBE` is written until a confirmation reports that nothing is left. That is the
10//! same rule every Redis client library applies, for the same reason.
11//!
12//! # Publishing while subscribed
13//!
14//! A RESP2 connection holding a subscription may send nothing but the pub/sub family, `PING`,
15//! `QUIT` and `RESET`; Redis refuses a `PUBLISH` there. So `redis_publish` on such a connection
16//! opens a second, short-lived connection to the same server, publishes, reads the receiver
17//! count and closes — what redis-rs, redis-py and Jedis all require of their users. On RESP3, or
18//! when not subscribed, it publishes on the connection itself and the `:N` reply is an ordinary
19//! `redis_response_received`.
20use super::resp::{self, Reply};
21use anyhow::{bail, Context, Result};
22use serde_json::{json, Value};
23use std::collections::BTreeSet;
24use std::sync::atomic::{AtomicBool, Ordering};
25use std::sync::Mutex;
26use std::time::Duration;
27use tokio::io::{AsyncWriteExt, BufReader};
28use tokio::net::TcpStream;
29
30/// How long the side connection a subscribed RESP2 client publishes over may take, end to end.
31pub const SIDE_PUBLISH_TIMEOUT: Duration = Duration::from_secs(10);
32
33/// The pub/sub family: commands whose replies are confirmations rather than ordinary replies.
34const SUBSCRIBE_FAMILY: &[&str] = &[
35    "SUBSCRIBE",
36    "PSUBSCRIBE",
37    "SSUBSCRIBE",
38    "UNSUBSCRIBE",
39    "PUNSUBSCRIBE",
40    "SUNSUBSCRIBE",
41];
42
43#[derive(Default)]
44struct Subscriptions {
45    channels: BTreeSet<String>,
46    patterns: BTreeSet<String>,
47    shard: BTreeSet<String>,
48}
49
50/// Connection-wide state the read loop and the write path share.
51pub struct Session {
52    /// The server, for the side connection.
53    pub remote_addr: String,
54    resp3: AtomicBool,
55    engaged: AtomicBool,
56    subs: Mutex<Subscriptions>,
57}
58
59/// What one reply turned out to be.
60#[derive(Debug, PartialEq)]
61pub enum Classified {
62    /// A delivered message: `kind` is `message`, `pmessage` or `smessage`.
63    Message {
64        kind: String,
65        pattern: Option<String>,
66        channel: String,
67        message: String,
68    },
69    /// A (un)subscribe confirmation.
70    Subscription {
71        kind: String,
72        channel: Option<String>,
73        count: i64,
74    },
75    /// Anything else.
76    Reply,
77}
78
79impl Session {
80    pub fn new(remote_addr: String) -> Self {
81        Session {
82            remote_addr,
83            resp3: AtomicBool::new(false),
84            engaged: AtomicBool::new(false),
85            subs: Mutex::new(Subscriptions::default()),
86        }
87    }
88
89    fn subs(&self) -> std::sync::MutexGuard<'_, Subscriptions> {
90        self.subs.lock().unwrap_or_else(|p| p.into_inner())
91    }
92
93    pub fn is_resp3(&self) -> bool {
94        self.resp3.load(Ordering::Relaxed)
95    }
96
97    pub fn set_resp3(&self, on: bool) {
98        self.resp3.store(on, Ordering::Relaxed);
99    }
100
101    /// A RESP2 connection in subscribed mode, which cannot publish on itself.
102    pub fn subscribed_resp2(&self) -> bool {
103        !self.is_resp3() && self.engaged.load(Ordering::Relaxed)
104    }
105
106    /// Note a command about to be written.
107    pub fn note_sent(&self, args: &[Vec<u8>]) {
108        let name = args
109            .first()
110            .map(|n| String::from_utf8_lossy(n).to_ascii_uppercase())
111            .unwrap_or_default();
112        if SUBSCRIBE_FAMILY.contains(&name.as_str()) {
113            self.engaged.store(true, Ordering::Relaxed);
114        }
115    }
116
117    /// The `UNSUBSCRIBE`-family commands that drop every subscription held.
118    pub fn unsubscribe_all_commands(&self) -> Vec<Vec<Vec<u8>>> {
119        let subs = self.subs();
120        let mut out = Vec::new();
121        if !subs.channels.is_empty() {
122            out.push(vec![b"UNSUBSCRIBE".to_vec()]);
123        }
124        if !subs.patterns.is_empty() {
125            out.push(vec![b"PUNSUBSCRIBE".to_vec()]);
126        }
127        if !subs.shard.is_empty() {
128            out.push(vec![b"SUNSUBSCRIBE".to_vec()]);
129        }
130        out
131    }
132
133    /// What this connection holds, for the model.
134    pub fn snapshot(&self) -> Value {
135        let subs = self.subs();
136        json!({
137            "channels": subs.channels.iter().collect::<Vec<_>>(),
138            "patterns": subs.patterns.iter().collect::<Vec<_>>(),
139            "shard_channels": subs.shard.iter().collect::<Vec<_>>(),
140        })
141    }
142
143    /// Read a reply as pub/sub traffic where it is one, updating the subscription state from
144    /// confirmations. Also notices a `HELLO` reply and records the version it agreed.
145    pub fn classify(&self, reply: &Reply) -> Classified {
146        if let Some(proto) = hello_proto(reply) {
147            self.set_resp3(proto == 3);
148            return Classified::Reply;
149        }
150        let pubsub_shaped = match reply.reply_type {
151            "push" => true,
152            "array" => self.engaged.load(Ordering::Relaxed),
153            _ => false,
154        };
155        if !pubsub_shaped {
156            return Classified::Reply;
157        }
158        let Some(items) = reply.value.as_array() else {
159            return Classified::Reply;
160        };
161        let text = |v: &Value| v.as_str().map(str::to_string);
162        let kind = items.first().and_then(text).unwrap_or_default();
163        match (kind.as_str(), items.len()) {
164            ("message" | "smessage", 3) => Classified::Message {
165                kind,
166                pattern: None,
167                channel: text(&items[1]).unwrap_or_default(),
168                message: text(&items[2]).unwrap_or_default(),
169            },
170            ("pmessage", 4) => Classified::Message {
171                kind,
172                pattern: text(&items[1]),
173                channel: text(&items[2]).unwrap_or_default(),
174                message: text(&items[3]).unwrap_or_default(),
175            },
176            (
177                "subscribe" | "psubscribe" | "ssubscribe" | "unsubscribe" | "punsubscribe"
178                | "sunsubscribe",
179                3,
180            ) => {
181                let channel = text(&items[1]);
182                let count = items[2].as_i64().unwrap_or(0);
183                let mut subs = self.subs();
184                if let Some(name) = &channel {
185                    let set = match kind.as_str() {
186                        "subscribe" | "unsubscribe" => &mut subs.channels,
187                        "psubscribe" | "punsubscribe" => &mut subs.patterns,
188                        _ => &mut subs.shard,
189                    };
190                    if kind.ends_with("unsubscribe") {
191                        set.remove(name);
192                    } else {
193                        set.insert(name.clone());
194                    }
195                }
196                if kind.ends_with("unsubscribe")
197                    && count == 0
198                    && subs.channels.is_empty()
199                    && subs.patterns.is_empty()
200                    && subs.shard.is_empty()
201                {
202                    self.engaged.store(false, Ordering::Relaxed);
203                }
204                Classified::Subscription {
205                    kind,
206                    channel,
207                    count,
208                }
209            }
210            _ => Classified::Reply,
211        }
212    }
213}
214
215/// The protocol version a `HELLO` reply agreed: a RESP3 map, or the RESP2 flat array, carrying
216/// both `server` and `proto`.
217fn hello_proto(reply: &Reply) -> Option<i64> {
218    match &reply.value {
219        Value::Object(map) if map.contains_key("server") => map.get("proto")?.as_i64(),
220        Value::Array(items) if reply.reply_type == "array" => {
221            let at = |key: &str| items.iter().position(|v| v.as_str() == Some(key));
222            at("server")?;
223            items.get(at("proto")? + 1)?.as_i64()
224        }
225        _ => None,
226    }
227}
228
229/// `PUBLISH` (or `SPUBLISH`) over a fresh connection, for a RESP2 connection in subscribed mode.
230/// Returns the receiver count the server answered.
231pub async fn publish_on_side_connection(addr: &str, args: &[Vec<u8>]) -> Result<i64> {
232    let work = async {
233        let stream = TcpStream::connect(addr)
234            .await
235            .with_context(|| format!("side connection to {addr}"))?;
236        let (read_half, mut write_half) = tokio::io::split(stream);
237        write_half.write_all(&resp::encode_command(args)).await?;
238        write_half.flush().await?;
239        let mut reader = BufReader::new(read_half);
240        let reply = resp::read_reply(&mut reader)
241            .await?
242            .context("the server closed the side connection without answering")?;
243        let _ = write_half.shutdown().await;
244        match (reply.reply_type, reply.value.as_i64()) {
245            ("integer", Some(n)) => Ok(n),
246            _ => bail!("the server answered the publish with {}", reply.response),
247        }
248    };
249    tokio::time::timeout(SIDE_PUBLISH_TIMEOUT, work)
250        .await
251        .context("the side connection's publish timed out")?
252}
netget/client/redis/pubsub.rs

netget/client/redis/
pubsub.rs
