Pub/sub messaging
Step 7 of 10. Step 6 ended at the tag mini-ch06 with the store on disk: SAVE dumps it, the
next boot loads it back, and deadlines survive the restart. This step starts from that tag and
ends at mini-ch07. A connection can now SUBSCRIBE to channels, PUBLISH reaches every
subscriber of a channel at once, and the server pushes messages to a client instead of only
answering it — the first time the clone sends a frame nobody asked for. The store is untouched by
any of it: a published message lives nowhere, which is exactly what the real server does too.
What the real server puts on the wire
Section titled “What the real server puts on the wire”Captured on redis-server 7.2.4, raw bytes, before any of the Rust below was written:
| Event | Reply |
|---|---|
SUBSCRIBE news |
*3\r\n$9\r\nsubscribe\r\n$4\r\nnews\r\n:1\r\n |
second SUBSCRIBE sports, same connection |
same shape, :2 — the running count of channels |
PUBLISH news hello with 2 subscribers |
:2\r\n — clients reached, not subscribers |
| each subscriber receives, unsolicited | *3\r\n$7\r\nmessage\r\n$4\r\nnews\r\n$5\r\nhello\r\n |
PUBLISH news hello after one hangs up |
:1 — the departed client stops counting |
PUBLISH empty x |
:0 — a channel nobody is on answers zero, not an error |
GET x while subscribed |
`-ERR Can’t execute ‘get’: only (P |
PING while subscribed |
*2\r\n$4\r\npong\r\n$0\r\n\r\n — a two-part pong, not +PONG |
UNSUBSCRIBE with no arguments |
one confirmation per channel, in subscription order |
Three details are worth pinning down before designing anything. The subscribe confirmation
carries a count of this connection’s channels, which is why it climbs 1, 2, 3. The publish
reply carries a count of deliveries — the real server counts the clients the message actually
reached, and a connection that has hung up is not one of them. And a subscribed connection is a
different kind of client: store commands are refused with a message that names the mode, and even
PING changes shape. The clone reproduces all of this byte for byte — every row above was
checked against the finished server at the tag, not just the first one.
One mailbox per connection
Section titled “One mailbox per connection”Everything until now has been a request and a reply on one socket. Pub/sub breaks that symmetry: a message arrives whenever someone publishes, possibly while the connection sits quietly between commands. The connection task has two things to wait for, so it needs two queues:
- commands arrive on the socket, as always;
- messages arrive in a mailbox — one
tokio::sync::mpscchannel per connection.
PUBLISH never touches a socket. It drops a message into each subscriber’s mailbox and returns
the count. The subscriber’s task notices the mailbox has work and writes the frame itself. The
queue is unbounded on purpose: a publish must not block on a slow reader, the way the real
server’s output buffer would grow rather than stall the publisher.
The standalone version of that design is small enough to run on its own:
[dependencies]tokio = { version = "1", features = ["macros", "rt-multi-thread", "net", "io-util", "time", "sync"] }use std::collections::HashMap;use tokio::sync::mpsc;
type Mailbox = mpsc::UnboundedSender<(Vec<u8>, Vec<u8>)>;
/// channel name -> every connection's mailbox, tagged with its id so an/// unsubscribe can find just that connection's entry.#[derive(Default)]struct Hub { subs: HashMap<Vec<u8>, Vec<(u64, Mailbox)>>,}
impl Hub { fn subscribe(&mut self, id: u64, channel: &[u8], mail: Mailbox) { self.subs.entry(channel.to_vec()).or_default().push((id, mail)); }
fn unsubscribe(&mut self, id: u64, channel: &[u8]) { if let Some(list) = self.subs.get_mut(channel) { list.retain(|(other, _)| *other != id); } }
/// Deliver `message` to every live subscriber of `channel` and count /// the deliveries. A send fails only when the receiving end is gone — /// the connection closed without unsubscribing — and those entries /// are pruned on the spot. fn publish(&mut self, channel: &[u8], message: &[u8]) -> i64 { match self.subs.get_mut(channel) { None => 0, Some(list) => { list.retain(|(_, mail)| { mail.send((channel.to_vec(), message.to_vec())).is_ok() }); let delivered = list.len() as i64; if list.is_empty() { self.subs.remove(channel); } delivered } } }}
#[tokio::main]async fn main() { let mut hub = Hub::default(); let (alice_mail, alice) = mpsc::unbounded_channel(); hub.subscribe(1, b"news", alice_mail); let (bob_mail, mut bob) = mpsc::unbounded_channel(); hub.subscribe(2, b"news", bob_mail.clone()); hub.subscribe(2, b"sports", bob_mail);
let got = hub.publish(b"news", b"hello"); println!("{got} subscriber(s) got it");
// Bob leaves the news channel; sports still delivers to him. hub.unsubscribe(2, b"news"); let got = hub.publish(b"news", b"again"); println!("{got} subscriber(s) got it");
// Alice hangs up without unsubscribing. Her mailbox's receiving end // is gone; the next publish prunes her entry and counts one fewer. drop(alice); let got = hub.publish(b"news", b"third"); println!("{got} subscriber(s) got it");
let got = hub.publish(b"empty", b"nobody"); println!("{got} subscriber(s) got it");
let got = hub.publish(b"sports", b"goal"); println!("{got} subscriber(s) got it");
// What Bob has queued, in order. while let Ok((channel, message)) = bob.try_recv() { println!( "bob <- {} / {}", String::from_utf8_lossy(&channel), String::from_utf8_lossy(&message) ); }}2 subscriber(s) got it1 subscriber(s) got it0 subscriber(s) got it0 subscriber(s) got it1 subscriber(s) got itbob <- news / hellobob <- sports / goalThe third line is the quiet one: Alice never unsubscribed, and the next publish still answered
0 subscribers on news before it. A send through UnboundedSender fails exactly when the
receiving half is dropped, so a hang-up costs one failed send and the hub prunes itself — no
background reaper, no reference counting. The channel rides along inside the message because a
connection can listen to several channels and has to be told which one a message came from.
The commands
Section titled “The commands”Three new variants, at the end of the enum:
Ttl(Vec<u8>), Save, Subscribe(Vec<Vec<u8>>), Unsubscribe(Vec<Vec<u8>>), Publish(Vec<u8>, Vec<u8>),}Arity follows the captures: SUBSCRIBE needs at least one channel, PUBLISH exactly a channel
and a payload, and UNSUBSCRIBE with no arguments means leave everything:
// At least one channel. Captured on redis-server 7.2.4 — // `SUBSCRIBE` with none answers the wrong-arity error. b"SUBSCRIBE" => match args { [] => Err(wrong_args(name)), channels => Ok(Command::Subscribe( channels.iter().map(|channel| channel.to_vec()).collect(), )), }, // No arguments means every channel this connection is // subscribed to — the real server confirms each one, in // subscription order. Captured on 7.2.4. b"UNSUBSCRIBE" => Ok(Command::Unsubscribe( args.iter().map(|channel| channel.to_vec()).collect(), )), // Exactly a channel and a payload. Captured on 7.2.4 — // `PUBLISH` with one argument answers the wrong-arity error. b"PUBLISH" => match args { [channel, message] => { Ok(Command::Publish(channel.to_vec(), message.to_vec())) } _ => Err(wrong_args(name)), },The hub
Section titled “The hub”Same shape as the store: a map behind one shared lock, one clone of the Arc per owner. The
entries are tagged with a connection id so an unsubscribe can find exactly its own mailbox and
nobody else’s:
/// One message on its way to a subscriber: the channel it was published/// on, and the payload. The channel rides along because a connection can/// subscribe to several and must be told which one a message came from.type Push = (Vec<u8>, Vec<u8>);
/// A connection's mailbox: the sending half of the queue its messages/// land on. Unbounded on purpose — a publish must never block on a slow/// reader — and dropped by the connection task when the connection ends.type Mailbox = tokio::sync::mpsc::UnboundedSender<Push>;
/// Every channel that has subscribers, and each subscriber's mailbox,/// tagged with the connection's id so an unsubscribe can find exactly its/// own entry. Same shape as `Db`: the map lives behind one shared lock,/// one clone of the `Arc` per owner.type Hub = Arc<Mutex<HashMap<Vec<u8>, Vec<(u64, Mailbox)>>>>;
fn new_hub() -> Hub { Arc::new(Mutex::new(HashMap::new()))}
/// One connection's pub/sub state: its id in the hub, its mailbox, and/// the channels it is subscribed to, in subscription order — the order an/// argument-less `UNSUBSCRIBE` confirms them in. A `Vec`, not a `HashSet`:/// a connection has a handful of channels at most, and order is part of/// what the real server's replies show.struct Sub { id: u64, mail: Mailbox, channels: Vec<Vec<u8>>,}
/// The reply a subscribe or unsubscribe is confirmed with: the kind, the/// channel, and how many channels the connection is on afterwards. Three/// bulks and an integer, on the wire. Captured on redis-server 7.2.4.fn confirmation(kind: &str, channel: &[u8], count: usize) -> Frame { Frame::Array(vec![ Frame::Bulk(kind.as_bytes().to_vec()), Frame::Bulk(channel.to_vec()), Frame::Integer(count as i64), ])}
/// The frame a subscriber is pushed: `message`, the channel, the payload./// Captured on redis-server 7.2.4.fn message_frame(channel: &[u8], payload: &[u8]) -> Frame { Frame::Array(vec![ Frame::Bulk(b"message".to_vec()), Frame::Bulk(channel.to_vec()), Frame::Bulk(payload.to_vec()), ])}Then the commands themselves. execute keeps its signature and its sixty-two tests; pub/sub
lives in its own function beside it, which answers with a Vec<Frame> because the real server
confirms one frame per channel as separate frames on the wire — a two-channel UNSUBSCRIBE is
two replies, not one nested array:
/// Run a pub/sub command, or say there was none.////// An empty `Vec` means "this frame is not mine" and the caller falls/// through to `execute`. Non-empty means the replies to write — a `Vec`/// because the real server confirms one frame per channel, as separate/// frames on the wire, and a two-channel `UNSUBSCRIBE` answers with two.fn execute_pubsub(hub: &Hub, me: &mut Sub, frame: &Frame) -> Vec<Frame> { let command = match Command::from_frame(frame) { Ok(command) => command, Err(reply) => return vec![reply], }; match command { // While subscribed, PING answers a two-part array — the pong and // an empty payload. Captured on redis-server 7.2.4. Command::Ping if !me.channels.is_empty() => { vec![Frame::Array(vec![ Frame::Bulk(b"pong".to_vec()), Frame::Bulk(Vec::new()), ])] } Command::Subscribe(channels) => { let mut replies = Vec::with_capacity(channels.len()); let mut subs = hub.lock().unwrap(); for channel in channels { if !me.channels.contains(&channel) { subs.entry(channel.clone()) .or_default() .push((me.id, me.mail.clone())); me.channels.push(channel.clone()); } // Subscribing twice to one channel is not an error: the // real server re-confirms it without doubling the // delivery. replies.push(confirmation( "subscribe", &channel, me.channels.len(), )); } replies } Command::Unsubscribe(named) => { // No names: leave everything, confirmed in subscription order. let leaving = if named.is_empty() { me.channels.clone() } else { named }; let mut replies = Vec::with_capacity(leaving.len()); let mut subs = hub.lock().unwrap(); for channel in leaving { me.channels.retain(|mine| *mine != channel); if let Some(list) = subs.get_mut(&channel) { list.retain(|(other, _)| *other != me.id); if list.is_empty() { subs.remove(&channel); } } // Naming a channel the connection was not on is confirmed // too, with the count it has left. replies.push(confirmation( "unsubscribe", &channel, me.channels.len(), )); } replies } Command::Publish(channel, message) => { let delivered = { let mut subs = hub.lock().unwrap(); match subs.get_mut(&channel) { None => 0, Some(list) => { // Fan-out: one send per subscriber, none of them // waited for. A send fails only when the receiving // end is gone — the connection closed without // unsubscribing — and those entries are pruned on // the spot, so a departed client costs one failed // send and nothing more. list.retain(|(_, mail)| { mail.send((channel.clone(), message.clone())).is_ok() }); let delivered = list.len() as i64; // The last live subscriber's failure takes the // channel out with it: an empty list would sit in // the map forever. if list.is_empty() { subs.remove(&channel); } delivered } } }; // Delivered, not subscribed: the real server counts the // clients the message reached. Captured on 7.2.4. vec![Frame::Integer(delivered)] } // Everything else, from a subscribed connection, is refused with // the real server's message, naming the command as sent and // lowercased. Captured on redis-server 7.2.4. _ if !me.channels.is_empty() => vec![cannot_execute(frame)], _ => Vec::new(), }}The guard is the last arm, and it is what makes the mode real: once a connection has channels,
GET, SET and friends never reach the store — the real server’s message comes back instead,
with the command lowercased the way 7.2.4 lowercases it. The two helpers:
/// The reply to a store command sent from a subscribed connection.fn cannot_execute(frame: &Frame) -> Frame { let name = match frame { Frame::Array(parts) => match parts.first() { Some(Frame::Bulk(name)) => String::from_utf8_lossy(name).to_lowercase(), _ => String::new(), }, _ => String::new(), }; Frame::Error(format!( "ERR Can't execute '{name}': only (P|S)SUBSCRIBE / (P|S)UNSUBSCRIBE / \ PING / QUIT / RESET are allowed in this context" ))}
/// Every hub entry that belongs to connection `id`, out. Runs when the/// connection ends; an entry left behind would be paid for as a failed/// send on every publish.fn detach(hub: &Hub, id: u64) { hub.lock().unwrap().retain(|_, list| { list.retain(|(other, _)| *other != id); !list.is_empty() });}Two clean-up paths, deliberately different. A connection that unsubscribes removes its entry
at once. A connection that hangs up without unsubscribing is caught by the prune inside
Publish — one failed send, then gone — and detach runs when the task ends anyway, so the
only entries that can outlive a connection are ones no publish has touched since.
serve learns to wait on two things
Section titled “serve learns to wait on two things”The loop from step 4 was while let Some(frame) = conn.read_frame(). With a mailbox, waiting on
the socket alone would leave pushed messages queued until the next command — a subscriber that
goes quiet would stop receiving. tokio::select! takes whichever source stirs first:
/// The connection counter. The hub stores mailboxes tagged with an id, so/// a connection leaving a channel removes its own entry and nobody else's.static NEXT_CONNECTION: AtomicU64 = AtomicU64::new(1);
async fn serve(stream: TcpStream, db: Db, hub: Hub) -> std::io::Result<()> { let mut conn = Connection::new(stream); // This connection's mailbox. `inbox` is polled by the select below; // publishes land there from whichever task ran the command. let (mail, mut inbox) = tokio::sync::mpsc::unbounded_channel(); let mut me = Sub { id: NEXT_CONNECTION.fetch_add(1, Ordering::Relaxed), mail, channels: Vec::new(), }; loop { // A subscribed connection has two things to wait for: its next // command, and the next message fanned out to it. `select!` takes // whichever comes first, so a push never waits behind a client // that has gone quiet. tokio::select! { frame = conn.read_frame() => { match frame? { Some(frame) => { let replies = execute_pubsub(&hub, &mut me, &frame); let mut out = Vec::new(); if replies.is_empty() { execute(&db, &frame).encode(&mut out); } else { for reply in replies { reply.encode(&mut out); } } conn.stream.write_all(&out).await?; } // Ok(None): the client closed the connection cleanly. None => break, } } Some((channel, payload)) = inbox.recv() => { let mut out = Vec::new(); message_frame(&channel, &payload).encode(&mut out); conn.stream.write_all(&out).await?; } } } detach(&hub, me.id); Ok(())}main creates the hub next to the store and hands every connection a clone, exactly as it does
with db. select! cancellation is safe here by construction: read_frame only ever consumes
buffer bytes after a full frame decodes, and a half-read stays in Connection::buffer for the
next poll — the same survival rule step 3 built the buffer for.
Seeing it run
Section titled “Seeing it run”The binary hardcodes 127.0.0.1:6379, which the real redis-server owns, so the same binary was rebuilt on a scratch port and run from an empty directory. One terminal subscribes to two channels:
$ redis-cli -p 7783 subscribe news sports1) "subscribe"2) "news"3) (integer) 11) "subscribe"2) "sports"3) (integer) 21) "message"2) "news"3) "hello"The confirmations climb 1, 2 — and then, without anyone typing GET, the third block arrives:
the pushed message. A second terminal does the publishing:
$ redis-cli -p 7783 publish news hello(integer) 1
$ redis-cli -p 7783 publish empty x(integer) 0empty answers 0, not an error — the channel was never created, and Redis pub/sub has no
memory of channels that nobody is on. The server’s own stdout is one line, as it has been since
step 4:
listening on 127.0.0.1:7783Seventy-four, up from sixty-two. The twelve new ones pin the wire shapes and the hub’s
bookkeeping: the confirmation frame, fan-out to two subscribers over real sockets, the count
dropping to zero after an unsubscribe and for an empty channel, the argument-less
UNSUBSCRIBE leaving every channel and returning the connection to normal commands, the
subscribe-mode guard, the two-part pong, the double subscribe that delivers once, a departed
connection pruned on the next publish, and a publish that lands in the mailbox with its channel
and payload.
$ cargo test Compiling mini-redis v0.1.0 (/Users/rango/www/hackingwithrust/code/mini-redis) Finished `test` profile [unoptimized + debuginfo] target(s) in 0.65s Running unittests src/main.rs (/Volumes/RustBuild/tracks/debug/deps/mini_redis-d1b449e02a3de456)
running 74 teststest tests::a_non_positive_ttl_on_a_missing_key_answers_zero ... oktest tests::a_deadline_that_has_arrived_expires_the_key_on_read ... oktest tests::a_frame_in_two_pieces_is_incomplete_not_invalid ... oktest tests::a_deadline_survives_the_round_trip ... oktest tests::a_key_expired_while_down_is_dropped_at_load ... oktest tests::a_non_positive_ttl_deletes_the_key ... oktest tests::a_key_with_no_deadline_answers_minus_one ... oktest tests::a_departed_connection_is_pruned_on_the_next_publish ... oktest tests::a_snapshot_round_trips_every_value_kind ... oktest tests::a_buffered_second_frame_needs_no_read ... oktest tests::a_store_command_is_untouched_while_not_subscribed ... oktest tests::a_snapshot_with_the_wrong_magic_is_an_error ... oktest tests::a_truncated_snapshot_is_an_error ... oktest tests::a_subscribed_connection_cannot_run_store_commands ... oktest tests::a_subscription_is_confirmed_with_the_channel_and_the_count ... oktest tests::an_empty_array_is_not_a_command ... oktest tests::an_expire_time_past_the_cutoff_is_refused ... oktest tests::an_expired_key_is_not_written ... oktest tests::an_unknown_command_caps_its_argument_list_at_128_bytes ... oktest tests::an_unknown_command_lists_its_arguments ... oktest tests::a_protocol_error_ends_the_connection_without_a_reply ... oktest tests::an_unsupported_version_is_an_error ... oktest tests::bytes_that_are_not_resp_are_invalid ... oktest tests::del_counts_only_the_keys_that_existed ... oktest tests::downtime_shrinks_the_ttl_the_way_it_does_on_the_real_server ... oktest tests::echo_round_trips_through_the_codec ... oktest tests::a_half_sent_command_does_not_block_another_connection ... oktest tests::echo_with_the_wrong_number_of_arguments_is_refused ... oktest tests::error_messages_quote_the_command_name_as_it_was_sent ... oktest tests::every_argument_must_be_a_bulk_string ... oktest tests::every_frame_type_round_trips ... oktest tests::expire_on_a_missing_key_answers_zero ... oktest tests::expire_reads_its_argument_the_way_the_real_server_does ... oktest tests::a_message_fans_out_to_every_subscriber ... oktest tests::expire_replaces_an_earlier_deadline ... oktest tests::expire_then_ttl_reports_the_seconds_left ... oktest tests::expire_with_the_wrong_number_of_arguments_is_refused ... oktest tests::an_argument_less_unsubscribe_leaves_every_channel ... oktest tests::expire_works_on_a_list_key ... oktest tests::get_on_a_list_is_a_wrongtype_error ... oktest tests::get_on_a_missing_key_answers_null ... oktest tests::lpush_needs_at_least_one_value ... oktest tests::load_snapshot_answers_none_when_the_file_is_missing ... oktest tests::lpush_on_a_string_key_is_a_wrongtype_error ... oktest tests::lpush_on_an_expired_key_starts_a_new_list ... oktest tests::lpush_prepends_each_value_in_order ... oktest tests::lrange_on_a_missing_key_answers_an_empty_array ... oktest tests::lrange_on_a_string_key_is_a_wrongtype_error ... oktest tests::lrange_reads_a_slice_the_way_the_real_server_clamps_it ... oktest tests::one_decode_consumes_exactly_one_frame ... oktest tests::lrange_with_a_non_integer_index_is_refused ... oktest tests::ping_while_subscribed_answers_the_two_part_pong ... oktest tests::ping_with_an_argument_echoes_it ... oktest tests::publish_reaches_the_mailbox_with_channel_and_payload ... oktest tests::save_with_an_argument_is_refused ... oktest tests::set_clears_the_deadline ... oktest tests::set_replaces_a_list_without_an_error ... oktest tests::set_then_get_round_trips ... oktest tests::ping_over_a_real_socket ... oktest tests::pipelined_commands_are_all_answered ... oktest tests::subscribe_with_no_arguments_is_refused ... oktest tests::subscribing_twice_does_not_double_the_delivery ... oktest tests::the_exact_bytes_the_real_redis_cli_sends_for_ping ... oktest tests::the_sweeper_removes_only_the_keys_whose_deadline_has_passed ... oktest tests::ttl_on_a_missing_key_answers_minus_two ... oktest tests::ttl_rounds_to_the_nearest_second ... oktest tests::ttl_with_the_wrong_number_of_arguments_is_refused ... oktest tests::unsubscribe_removes_only_the_channel_named ... oktest tests::save_answers_ok_and_writes_the_file ... oktest tests::publish_counts_what_it_delivered ... oktest tests::save_then_load_round_trips_through_a_file ... oktest tests::two_connections_share_one_store ... oktest tests::a_frame_split_across_two_reads_is_reassembled ... oktest tests::the_sweeper_task_reclaims_an_expired_key_on_its_own ... ok
test result: ok. 74 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.11sThe order the tests print in varies between runs; the set does not. The Running line names
this machine’s shared build cache; elsewhere it reads target/debug/deps/mini_redis-….
What is deliberately not here
Section titled “What is deliberately not here”- Patterns.
PSUBSCRIBEmatches channels withnews.*-style globs; the hub would need to try every pattern against every published channel. The real server’s counts in the confirmations get fiddly there too — one connection can receive a message through several patterns at once. - Key-space notifications. The real server can publish events like
expiredonto channels you opt into withCONFIG SET notify-keyspace-events; the clone has no configuration layer. - Sharded pub/sub. Redis 7 added
SSUBSCRIBE/SUNSUSCRIBE, which route channels by slot in a cluster. There is no cluster here, and plain pub/sub already has no routing to do. QUITandRESETin subscribe mode. The guard’s message names both; the clone accepts neither yet — a subscribed connection ends only by hanging up. Both are small, and both belong with the connection-lifecycle work.- Message delivery guarantees. None, by design: a publish goes to the subscribers that are
there now, and a slow reader’s unbounded queue is the only buffering. The real server has
the same shape — pub/sub is fire-and-forget, and keys with
TTLare the durable half of this server, not channels.
Step 8, Transactions with MULTI and EXEC, queues several commands behind one connection and runs them as a group. It is not written yet, so it gets no link — the track overview carries the full list as articles land.
Related
Section titled “Related”- Handling concurrent connections —
step 4, the one-task-per-connection shape that
select!plugs into. - Persistence with a snapshot — step 6, the tag this step starts from.
- Async —
select!and unbounded channels are the two Tokio tools this step leans on.