From 47c3d638bf202d713f213a1b9f8957ca8ffcab80 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Sun, 8 Feb 2026 08:23:46 -0500 Subject: [PATCH 1/8] feat: add pub/sub command parsing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit add SUBSCRIBE, UNSUBSCRIBE, PSUBSCRIBE, PUNSUBSCRIBE, and PUBLISH commands to the protocol layer. parsing only — no execution yet. --- crates/ember-protocol/src/command.rs | 74 ++++++++++++++++++++++++++++ 1 file changed, 74 insertions(+) diff --git a/crates/ember-protocol/src/command.rs b/crates/ember-protocol/src/command.rs index e565398b..d75b40ea 100644 --- a/crates/ember-protocol/src/command.rs +++ b/crates/ember-protocol/src/command.rs @@ -272,6 +272,23 @@ pub enum Command { /// SLOWLOG RESET. Clears the slow log. SlowLogReset, + // --- pub/sub commands --- + + /// SUBSCRIBE `channel` \[channel ...\]. Subscribe to one or more channels. + Subscribe { channels: Vec }, + + /// UNSUBSCRIBE \[channel ...\]. Unsubscribe from channels (all if none given). + Unsubscribe { channels: Vec }, + + /// PSUBSCRIBE `pattern` \[pattern ...\]. Subscribe to channels matching patterns. + PSubscribe { patterns: Vec }, + + /// PUNSUBSCRIBE \[pattern ...\]. Unsubscribe from patterns (all if none given). + PUnsubscribe { patterns: Vec }, + + /// PUBLISH `channel` `message`. Publish a message to a channel. + Publish { channel: String, message: Bytes }, + /// A command we don't recognize (yet). Unknown(String), } @@ -371,6 +388,11 @@ impl Command { Command::SlowLogGet { .. } => "slowlog", Command::SlowLogLen => "slowlog", Command::SlowLogReset => "slowlog", + Command::Subscribe { .. } => "subscribe", + Command::Unsubscribe { .. } => "unsubscribe", + Command::PSubscribe { .. } => "psubscribe", + Command::PUnsubscribe { .. } => "punsubscribe", + Command::Publish { .. } => "publish", Command::Unknown(_) => "unknown", } } @@ -452,6 +474,11 @@ impl Command { "ASKING" => parse_asking(&frames[1..]), "MIGRATE" => parse_migrate(&frames[1..]), "SLOWLOG" => parse_slowlog(&frames[1..]), + "SUBSCRIBE" => parse_subscribe(&frames[1..]), + "UNSUBSCRIBE" => parse_unsubscribe(&frames[1..]), + "PSUBSCRIBE" => parse_psubscribe(&frames[1..]), + "PUNSUBSCRIBE" => parse_punsubscribe(&frames[1..]), + "PUBLISH" => parse_publish(&frames[1..]), _ => Ok(Command::Unknown(name)), } } @@ -1450,6 +1477,53 @@ fn parse_migrate(args: &[Frame]) -> Result { }) } +fn parse_subscribe(args: &[Frame]) -> Result { + if args.is_empty() { + return Err(ProtocolError::WrongArity("SUBSCRIBE".into())); + } + let channels: Vec = args + .iter() + .map(extract_string) + .collect::>()?; + Ok(Command::Subscribe { channels }) +} + +fn parse_unsubscribe(args: &[Frame]) -> Result { + let channels: Vec = args + .iter() + .map(extract_string) + .collect::>()?; + Ok(Command::Unsubscribe { channels }) +} + +fn parse_psubscribe(args: &[Frame]) -> Result { + if args.is_empty() { + return Err(ProtocolError::WrongArity("PSUBSCRIBE".into())); + } + let patterns: Vec = args + .iter() + .map(extract_string) + .collect::>()?; + Ok(Command::PSubscribe { patterns }) +} + +fn parse_punsubscribe(args: &[Frame]) -> Result { + let patterns: Vec = args + .iter() + .map(extract_string) + .collect::>()?; + Ok(Command::PUnsubscribe { patterns }) +} + +fn parse_publish(args: &[Frame]) -> Result { + if args.len() != 2 { + return Err(ProtocolError::WrongArity("PUBLISH".into())); + } + let channel = extract_string(&args[0])?; + let message = extract_bytes(&args[1])?; + Ok(Command::Publish { channel, message }) +} + #[cfg(test)] mod tests { use super::*; From b8216e85891475811bee7e35ba1473e04fde96b8 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Sun, 8 Feb 2026 08:30:45 -0500 Subject: [PATCH 2/8] feat: add pubsub manager with broadcast channels introduces PubSubManager backed by DashMap + tokio broadcast for lock-free fan-out messaging. supports channel subscriptions, glob pattern matching (PSUBSCRIBE), and concurrent publish. includes comprehensive unit tests for glob matching and pub/sub semantics. --- Cargo.lock | 1 + crates/ember-server/Cargo.toml | 1 + crates/ember-server/src/main.rs | 1 + crates/ember-server/src/pubsub.rs | 407 ++++++++++++++++++++++++++++++ 4 files changed, 410 insertions(+) create mode 100644 crates/ember-server/src/pubsub.rs diff --git a/Cargo.lock b/Cargo.lock index c5de3c9f..f5d00eee 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -555,6 +555,7 @@ version = "0.3.0" dependencies = [ "bytes", "clap", + "dashmap", "ember-cluster", "ember-persistence", "ember-protocol", diff --git a/crates/ember-server/Cargo.toml b/crates/ember-server/Cargo.toml index ad96d653..7d1b81f4 100644 --- a/crates/ember-server/Cargo.toml +++ b/crates/ember-server/Cargo.toml @@ -26,6 +26,7 @@ clap = { workspace = true } metrics = { workspace = true } metrics-exporter-prometheus = { workspace = true } futures = "0.3" +dashmap = "6" # optional: better multi-threaded allocation performance tikv-jemallocator = { version = "0.6", optional = true } diff --git a/crates/ember-server/src/main.rs b/crates/ember-server/src/main.rs index f92c485f..7ec3c698 100644 --- a/crates/ember-server/src/main.rs +++ b/crates/ember-server/src/main.rs @@ -9,6 +9,7 @@ mod config; mod connection; mod connection_common; mod metrics; +mod pubsub; mod server; mod slowlog; diff --git a/crates/ember-server/src/pubsub.rs b/crates/ember-server/src/pubsub.rs new file mode 100644 index 00000000..6d1e5cd6 --- /dev/null +++ b/crates/ember-server/src/pubsub.rs @@ -0,0 +1,407 @@ +//! Pub/sub message broker for channel-based messaging. +//! +//! Manages subscriptions and broadcasts messages to all matching +//! subscribers. Supports both exact channel names and glob patterns. +//! Thread-safe via DashMap for lock-free concurrent access. + +use std::collections::HashSet; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; + +use bytes::Bytes; +use dashmap::DashMap; +use tokio::sync::broadcast; + +/// Maximum number of buffered messages per subscription before +/// slow consumers start missing messages. This is per-channel, +/// not global — a subscriber that falls behind on one busy channel +/// won't affect its other subscriptions. +const CHANNEL_CAPACITY: usize = 256; + +/// A message published to a channel. +#[derive(Debug, Clone)] +pub struct PubMessage { + /// The channel the message was published to. + pub channel: String, + /// The raw message data. + pub data: Bytes, + /// For pattern subscriptions, the pattern that matched. + /// None for exact channel subscriptions. + pub pattern: Option, +} + +/// Manages pub/sub state: channel subscriptions, pattern subscriptions, +/// and message broadcasting. +/// +/// Shared across all connection handlers via `Arc`. +/// Uses DashMap internally so all operations are lock-free. +pub struct PubSubManager { + /// Exact channel subscriptions: channel name → broadcast sender. + channels: DashMap>, + /// Pattern subscriptions: pattern string → broadcast sender. + patterns: DashMap>, + /// Total number of active subscriptions (channels + patterns). + subscription_count: AtomicUsize, +} + +impl PubSubManager { + pub fn new() -> Self { + Self { + channels: DashMap::new(), + patterns: DashMap::new(), + subscription_count: AtomicUsize::new(0), + } + } + + /// Subscribe to an exact channel. Returns a receiver for messages + /// on that channel. + pub fn subscribe(&self, channel: &str) -> broadcast::Receiver { + let entry = self + .channels + .entry(channel.to_string()) + .or_insert_with(|| { + let (tx, _) = broadcast::channel(CHANNEL_CAPACITY); + tx + }); + self.subscription_count.fetch_add(1, Ordering::Relaxed); + entry.subscribe() + } + + /// Unsubscribe from an exact channel. Returns true if the channel + /// had active subscriptions. + /// + /// Note: the actual receiver is dropped by the caller. This just + /// cleans up empty channels. + pub fn unsubscribe(&self, channel: &str) -> bool { + self.subscription_count.fetch_sub(1, Ordering::Relaxed); + if let Some(entry) = self.channels.get(channel) { + // if no receivers left, remove the channel entirely + if entry.receiver_count() <= 1 { + drop(entry); + self.channels.remove(channel); + } + true + } else { + false + } + } + + /// Subscribe to a glob pattern. Returns a receiver for messages + /// matching the pattern. + pub fn psubscribe(&self, pattern: &str) -> broadcast::Receiver { + let entry = self + .patterns + .entry(pattern.to_string()) + .or_insert_with(|| { + let (tx, _) = broadcast::channel(CHANNEL_CAPACITY); + tx + }); + self.subscription_count.fetch_add(1, Ordering::Relaxed); + entry.subscribe() + } + + /// Unsubscribe from a pattern. Returns true if the pattern had + /// active subscriptions. + pub fn punsubscribe(&self, pattern: &str) -> bool { + self.subscription_count.fetch_sub(1, Ordering::Relaxed); + if let Some(entry) = self.patterns.get(pattern) { + if entry.receiver_count() <= 1 { + drop(entry); + self.patterns.remove(pattern); + } + true + } else { + false + } + } + + /// Publish a message to a channel. Returns the total number of + /// subscribers that received the message (exact + pattern). + pub fn publish(&self, channel: &str, data: Bytes) -> usize { + let mut count = 0; + + // send to exact channel subscribers + if let Some(tx) = self.channels.get(channel) { + let msg = PubMessage { + channel: channel.to_string(), + data: data.clone(), + pattern: None, + }; + // send returns the number of receivers that got the message + count += tx.send(msg).unwrap_or(0); + } + + // send to pattern subscribers + for entry in self.patterns.iter() { + let pattern = entry.key(); + if glob_match(pattern, channel) { + let msg = PubMessage { + channel: channel.to_string(), + data: data.clone(), + pattern: Some(pattern.clone()), + }; + count += entry.value().send(msg).unwrap_or(0); + } + } + + count + } + + /// Returns the total number of active subscriptions. + pub fn total_subscriptions(&self) -> usize { + self.subscription_count.load(Ordering::Relaxed) + } + + /// Returns the number of active channels (with at least one subscriber). + pub fn active_channels(&self) -> usize { + self.channels.len() + } + + /// Returns the number of active patterns. + pub fn active_patterns(&self) -> usize { + self.patterns.len() + } +} + +/// Simple glob matching for pub/sub patterns. +/// +/// Supports: +/// - `*` matches any sequence of characters +/// - `?` matches any single character +/// - `[abc]` matches any character in the set +/// - `\x` escapes the next character +/// +/// This matches Redis behavior for PSUBSCRIBE patterns. +fn glob_match(pattern: &str, input: &str) -> bool { + let pat: Vec = pattern.chars().collect(); + let inp: Vec = input.chars().collect(); + glob_match_inner(&pat, &inp) +} + +fn glob_match_inner(pat: &[char], inp: &[char]) -> bool { + let (mut pi, mut ii) = (0, 0); + let (mut star_pi, mut star_ii) = (usize::MAX, usize::MAX); + + while ii < inp.len() { + if pi < pat.len() && pat[pi] == '\\' && pi + 1 < pat.len() { + // escaped character — must match literally + pi += 1; + if inp[ii] == pat[pi] { + pi += 1; + ii += 1; + continue; + } + } else if pi < pat.len() && pat[pi] == '?' { + pi += 1; + ii += 1; + continue; + } else if pi < pat.len() && pat[pi] == '*' { + star_pi = pi; + star_ii = ii; + pi += 1; + continue; + } else if pi < pat.len() && pat[pi] == '[' { + // character class + if let Some((matched, end)) = match_char_class(&pat[pi..], inp[ii]) { + if matched { + pi += end; + ii += 1; + continue; + } + } + } else if pi < pat.len() && pat[pi] == inp[ii] { + pi += 1; + ii += 1; + continue; + } + + // no match — backtrack to last star if possible + if star_pi != usize::MAX { + pi = star_pi + 1; + star_ii += 1; + ii = star_ii; + continue; + } + + return false; + } + + // consume trailing stars + while pi < pat.len() && pat[pi] == '*' { + pi += 1; + } + + pi == pat.len() +} + +/// Matches a `[...]` character class. Returns (matched, chars_consumed) +/// if the bracket expression is valid. +fn match_char_class(pat: &[char], ch: char) -> Option<(bool, usize)> { + if pat.is_empty() || pat[0] != '[' { + return None; + } + + let mut i = 1; + let negate = if i < pat.len() && pat[i] == '^' { + i += 1; + true + } else { + false + }; + + let mut matched = false; + while i < pat.len() && pat[i] != ']' { + if i + 2 < pat.len() && pat[i + 1] == '-' { + // range: [a-z] + if ch >= pat[i] && ch <= pat[i + 2] { + matched = true; + } + i += 3; + } else { + if ch == pat[i] { + matched = true; + } + i += 1; + } + } + + if i < pat.len() && pat[i] == ']' { + Some((matched ^ negate, i + 1)) + } else { + None // unterminated bracket + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn glob_exact_match() { + assert!(glob_match("hello", "hello")); + assert!(!glob_match("hello", "world")); + } + + #[test] + fn glob_star_match() { + assert!(glob_match("news.*", "news.sports")); + assert!(glob_match("news.*", "news.weather.today")); + assert!(!glob_match("news.*", "old.news")); + assert!(glob_match("*", "anything")); + assert!(glob_match("h*o", "hello")); + assert!(glob_match("h*o", "ho")); + } + + #[test] + fn glob_question_mark() { + assert!(glob_match("h?llo", "hello")); + assert!(glob_match("h?llo", "hallo")); + assert!(!glob_match("h?llo", "hllo")); + } + + #[test] + fn glob_char_class() { + assert!(glob_match("h[ae]llo", "hello")); + assert!(glob_match("h[ae]llo", "hallo")); + assert!(!glob_match("h[ae]llo", "hillo")); + } + + #[test] + fn glob_negated_class() { + assert!(glob_match("h[^ae]llo", "hillo")); + assert!(!glob_match("h[^ae]llo", "hello")); + } + + #[test] + fn glob_escaped_char() { + assert!(glob_match("hello\\*", "hello*")); + assert!(!glob_match("hello\\*", "helloX")); + } + + #[test] + fn subscribe_and_publish() { + let mgr = PubSubManager::new(); + let mut rx = mgr.subscribe("test"); + let count = mgr.publish("test", Bytes::from("hello")); + assert_eq!(count, 1); + + let msg = rx.try_recv().unwrap(); + assert_eq!(msg.channel, "test"); + assert_eq!(msg.data, Bytes::from("hello")); + assert!(msg.pattern.is_none()); + } + + #[test] + fn publish_to_empty_channel() { + let mgr = PubSubManager::new(); + let count = mgr.publish("nobody", Bytes::from("hello")); + assert_eq!(count, 0); + } + + #[test] + fn multiple_subscribers() { + let mgr = PubSubManager::new(); + let mut rx1 = mgr.subscribe("ch"); + let mut rx2 = mgr.subscribe("ch"); + + let count = mgr.publish("ch", Bytes::from("msg")); + assert_eq!(count, 2); + + assert_eq!(rx1.try_recv().unwrap().data, Bytes::from("msg")); + assert_eq!(rx2.try_recv().unwrap().data, Bytes::from("msg")); + } + + #[test] + fn pattern_subscribe_and_publish() { + let mgr = PubSubManager::new(); + let mut rx = mgr.psubscribe("news.*"); + + let count = mgr.publish("news.sports", Bytes::from("goal!")); + assert_eq!(count, 1); + + let msg = rx.try_recv().unwrap(); + assert_eq!(msg.channel, "news.sports"); + assert_eq!(msg.pattern, Some("news.*".to_string())); + + // shouldn't match + let count = mgr.publish("old.news", Bytes::from("nope")); + assert_eq!(count, 0); + } + + #[test] + fn exact_and_pattern_both_receive() { + let mgr = PubSubManager::new(); + let mut rx_exact = mgr.subscribe("news.sports"); + let mut rx_pattern = mgr.psubscribe("news.*"); + + let count = mgr.publish("news.sports", Bytes::from("goal!")); + assert_eq!(count, 2); + + assert!(rx_exact.try_recv().is_ok()); + assert!(rx_pattern.try_recv().is_ok()); + } + + #[test] + fn unsubscribe_stops_delivery() { + let mgr = PubSubManager::new(); + let rx = mgr.subscribe("ch"); + mgr.unsubscribe("ch"); + drop(rx); + + let count = mgr.publish("ch", Bytes::from("msg")); + assert_eq!(count, 0); + } + + #[test] + fn subscription_counts() { + let mgr = PubSubManager::new(); + assert_eq!(mgr.total_subscriptions(), 0); + + let _rx1 = mgr.subscribe("a"); + let _rx2 = mgr.subscribe("b"); + let _rx3 = mgr.psubscribe("c.*"); + assert_eq!(mgr.total_subscriptions(), 3); + assert_eq!(mgr.active_channels(), 2); + assert_eq!(mgr.active_patterns(), 1); + } +} From 1a23bb8ca8144b4fa2bdea57dbf82a9372be5d2f Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Sun, 8 Feb 2026 08:30:50 -0500 Subject: [PATCH 3/8] feat: wire pub/sub into connection handlers - server.rs: create and pass PubSubManager to all connections - connection.rs: subscriber mode state machine with select! loop, push message serialization, subscription cleanup on disconnect - concurrent_handler.rs: PUBLISH support, sub commands return error --- crates/ember-server/src/concurrent_handler.rs | 22 +- crates/ember-server/src/connection.rs | 366 +++++++++++++++++- crates/ember-server/src/server.rs | 9 +- 3 files changed, 387 insertions(+), 10 deletions(-) diff --git a/crates/ember-server/src/concurrent_handler.rs b/crates/ember-server/src/concurrent_handler.rs index 3e3d1f95..89774e4f 100644 --- a/crates/ember-server/src/concurrent_handler.rs +++ b/crates/ember-server/src/concurrent_handler.rs @@ -23,6 +23,7 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpStream; use crate::connection_common::{BUF_CAPACITY, IDLE_TIMEOUT, MAX_BUF_SIZE}; +use crate::pubsub::PubSubManager; use crate::server::ServerContext; use crate::slowlog::SlowLog; @@ -33,6 +34,7 @@ pub async fn handle( engine: Engine, // fallback for complex commands ctx: &Arc, slow_log: &Arc, + pubsub: &Arc, ) -> Result<(), Box> { stream.set_nodelay(true)?; @@ -60,7 +62,7 @@ pub async fn handle( match parse_frame(&buf) { Ok(Some((frame, consumed))) => { let _ = buf.split_to(consumed); - let response = process(frame, &keyspace, &engine, ctx, slow_log).await; + let response = process(frame, &keyspace, &engine, ctx, slow_log, pubsub).await; response.serialize(&mut out); } Ok(None) => break, @@ -85,6 +87,7 @@ async fn process( engine: &Engine, ctx: &Arc, slow_log: &Arc, + pubsub: &Arc, ) -> Frame { match Command::from_frame(frame) { Ok(cmd) => { @@ -96,7 +99,7 @@ async fn process( None }; - let response = execute_concurrent(cmd, keyspace, engine).await; + let response = execute_concurrent(cmd, keyspace, engine, pubsub).await; ctx.commands_processed.fetch_add(1, Ordering::Relaxed); if let Some(start) = start { @@ -119,6 +122,7 @@ async fn execute_concurrent( cmd: Command, keyspace: &Arc, _engine: &Engine, + pubsub: &Arc, ) -> Frame { match cmd { // Hot path: direct access without channels @@ -198,6 +202,20 @@ async fn execute_concurrent( Frame::Simple("OK".into()) } + // -- pub/sub -- + Command::Publish { channel, message } => { + let count = pubsub.publish(&channel, message); + Frame::Integer(count as i64) + } + + // subscribe commands are handled in the connection layer, not here + Command::Subscribe { .. } + | Command::Unsubscribe { .. } + | Command::PSubscribe { .. } + | Command::PUnsubscribe { .. } => { + Frame::Error("ERR pub/sub not supported in concurrent mode yet".into()) + } + // For unsupported commands, return an error Command::Unknown(name) => Frame::Error(format!("ERR unknown command '{name}'")), diff --git a/crates/ember-server/src/connection.rs b/crates/ember-server/src/connection.rs index c3b45644..890ec4d0 100644 --- a/crates/ember-server/src/connection.rs +++ b/crates/ember-server/src/connection.rs @@ -5,6 +5,7 @@ //! by dispatching multiple commands concurrently to shards using //! `join_all` for parallel execution. +use std::collections::HashMap; use std::sync::atomic::Ordering; use std::sync::Arc; use std::time::Instant; @@ -15,8 +16,10 @@ use ember_protocol::{parse_frame, Command, Frame, SetExpire}; use futures::future::join_all; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpStream; +use tokio::sync::broadcast; use crate::connection_common::{BUF_CAPACITY, IDLE_TIMEOUT, MAX_BUF_SIZE}; +use crate::pubsub::{PubMessage, PubSubManager}; use crate::server::ServerContext; use crate::slowlog::SlowLog; use std::time::Duration; @@ -31,6 +34,7 @@ pub async fn handle( engine: Engine, ctx: &Arc, slow_log: &Arc, + pubsub: &Arc, ) -> Result<(), Box> { // disable Nagle's algorithm — cache servers need low-latency writes, // and we already batch responses from pipelining into a single write @@ -79,14 +83,48 @@ pub async fn handle( } } - // dispatch all commands concurrently — this is the key optimization. - // instead of await-ing each command serially (16 round-trips for a - // pipeline of 16), we dispatch all at once and await them together - // (effectively 1 round-trip worth of latency for all 16). + // check if any frame is a subscribe command — if so, we need + // to enter subscriber mode which changes the connection loop + let mut enter_sub = false; + if !frames.is_empty() { + for frame in &frames { + if is_subscribe_frame(frame) { + enter_sub = true; + break; + } + } + } + + if enter_sub { + // process any non-subscribe commands that came before + let mut sub_frames = Vec::new(); + for frame in frames { + if is_subscribe_frame(&frame) { + sub_frames.push(frame); + } else { + let response = process(frame, &engine, ctx, slow_log, pubsub).await; + response.serialize(&mut out); + } + } + if !out.is_empty() { + stream.write_all(&out).await?; + out.clear(); + } + + // enter subscriber mode — this blocks until all subscriptions + // are removed or the client disconnects + handle_subscriber_mode( + &mut stream, &mut buf, &mut out, &engine, ctx, slow_log, pubsub, sub_frames, + ) + .await?; + return Ok(()); + } + + // normal command processing — dispatch concurrently if !frames.is_empty() { let futures: Vec<_> = frames .into_iter() - .map(|frame| process(frame, &engine, ctx, slow_log)) + .map(|frame| process(frame, &engine, ctx, slow_log, pubsub)) .collect(); let responses = join_all(futures).await; for response in responses { @@ -100,6 +138,305 @@ pub async fn handle( } } +/// Checks if a raw frame is a SUBSCRIBE or PSUBSCRIBE command. +fn is_subscribe_frame(frame: &Frame) -> bool { + if let Frame::Array(parts) = frame { + if let Some(Frame::Bulk(name)) = parts.first() { + let upper = String::from_utf8_lossy(name).to_ascii_uppercase(); + return matches!( + upper.as_str(), + "SUBSCRIBE" | "PSUBSCRIBE" | "UNSUBSCRIBE" | "PUNSUBSCRIBE" + ); + } + } + false +} + +/// Subscriber mode: listens for both broadcast messages and client commands. +/// +/// In this mode the connection can only process SUBSCRIBE, UNSUBSCRIBE, +/// PSUBSCRIBE, PUNSUBSCRIBE, and PING. All other commands return an error. +/// Returns to the caller when all subscriptions are removed or the client +/// disconnects. +async fn handle_subscriber_mode( + stream: &mut TcpStream, + buf: &mut BytesMut, + out: &mut BytesMut, + engine: &Engine, + ctx: &Arc, + slow_log: &Arc, + pubsub: &Arc, + initial_frames: Vec, +) -> Result<(), Box> { + // track subscriptions: channel/pattern -> receiver + let mut channel_rxs: HashMap> = HashMap::new(); + let mut pattern_rxs: HashMap> = HashMap::new(); + + // process the initial subscribe commands + for frame in initial_frames { + if let Ok(cmd) = Command::from_frame(frame) { + handle_sub_command( + cmd, pubsub, &mut channel_rxs, &mut pattern_rxs, out, + ); + } + } + + if !out.is_empty() { + stream.write_all(out).await?; + out.clear(); + } + + // main subscriber loop + loop { + let total_subs = channel_rxs.len() + pattern_rxs.len(); + if total_subs == 0 { + // no more subscriptions — exit subscriber mode + return Ok(()); + } + + tokio::select! { + // check for incoming messages from any subscription + msg = recv_any_message(&mut channel_rxs, &mut pattern_rxs) => { + if let Some(msg) = msg { + serialize_push_message(&msg, out); + stream.write_all(out).await?; + out.clear(); + } + } + + // check for new commands from the client + result = stream.read_buf(buf) => { + match result { + Ok(0) => { + // client disconnected — clean up subscriptions + cleanup_subscriptions(pubsub, &channel_rxs, &pattern_rxs); + return Ok(()); + } + Ok(_) => { + // parse and handle subscriber commands + loop { + match parse_frame(buf) { + Ok(Some((frame, consumed))) => { + let _ = buf.split_to(consumed); + match Command::from_frame(frame) { + Ok(cmd) => match &cmd { + Command::Subscribe { .. } + | Command::Unsubscribe { .. } + | Command::PSubscribe { .. } + | Command::PUnsubscribe { .. } => { + handle_sub_command( + cmd, pubsub, &mut channel_rxs, + &mut pattern_rxs, out, + ); + } + Command::Ping(msg) => { + let resp = match msg { + Some(m) => Frame::Bulk(m.clone()), + None => Frame::Simple("PONG".into()), + }; + resp.serialize(out); + } + _ => { + Frame::Error( + "ERR only (P)SUBSCRIBE / (P)UNSUBSCRIBE / PING are allowed in this context".into() + ).serialize(out); + } + } + Err(e) => { + Frame::Error(format!("ERR {e}")).serialize(out); + } + } + } + Ok(None) => break, + Err(e) => { + Frame::Error(format!("ERR protocol error: {e}")).serialize(out); + stream.write_all(out).await?; + cleanup_subscriptions(pubsub, &channel_rxs, &pattern_rxs); + return Ok(()); + } + } + } + + if !out.is_empty() { + stream.write_all(out).await?; + out.clear(); + } + } + Err(e) => { + cleanup_subscriptions(pubsub, &channel_rxs, &pattern_rxs); + return Err(e.into()); + } + } + } + } + } +} + +/// Processes a subscribe/unsubscribe command, updating the subscription maps +/// and writing RESP3 responses. +fn handle_sub_command( + cmd: Command, + pubsub: &PubSubManager, + channel_rxs: &mut HashMap>, + pattern_rxs: &mut HashMap>, + out: &mut BytesMut, +) { + match cmd { + Command::Subscribe { channels } => { + for ch in channels { + let rx = pubsub.subscribe(&ch); + channel_rxs.insert(ch.clone(), rx); + let count = channel_rxs.len() + pattern_rxs.len(); + serialize_sub_response("subscribe", &ch, count, out); + } + } + Command::Unsubscribe { channels } => { + if channels.is_empty() { + // unsubscribe from all channels + let names: Vec = channel_rxs.keys().cloned().collect(); + for ch in names { + channel_rxs.remove(&ch); + pubsub.unsubscribe(&ch); + let count = channel_rxs.len() + pattern_rxs.len(); + serialize_sub_response("unsubscribe", &ch, count, out); + } + if channel_rxs.is_empty() && pattern_rxs.is_empty() { + // send a final response with count 0 if we had nothing + serialize_sub_response("unsubscribe", "", 0, out); + } + } else { + for ch in channels { + channel_rxs.remove(&ch); + pubsub.unsubscribe(&ch); + let count = channel_rxs.len() + pattern_rxs.len(); + serialize_sub_response("unsubscribe", &ch, count, out); + } + } + } + Command::PSubscribe { patterns } => { + for pat in patterns { + let rx = pubsub.psubscribe(&pat); + pattern_rxs.insert(pat.clone(), rx); + let count = channel_rxs.len() + pattern_rxs.len(); + serialize_sub_response("psubscribe", &pat, count, out); + } + } + Command::PUnsubscribe { patterns } => { + if patterns.is_empty() { + let names: Vec = pattern_rxs.keys().cloned().collect(); + for pat in names { + pattern_rxs.remove(&pat); + pubsub.punsubscribe(&pat); + let count = channel_rxs.len() + pattern_rxs.len(); + serialize_sub_response("punsubscribe", &pat, count, out); + } + if channel_rxs.is_empty() && pattern_rxs.is_empty() { + serialize_sub_response("punsubscribe", "", 0, out); + } + } else { + for pat in patterns { + pattern_rxs.remove(&pat); + pubsub.punsubscribe(&pat); + let count = channel_rxs.len() + pattern_rxs.len(); + serialize_sub_response("punsubscribe", &pat, count, out); + } + } + } + _ => {} + } +} + +/// Receives a message from any active subscription (channels or patterns). +async fn recv_any_message( + channel_rxs: &mut HashMap>, + pattern_rxs: &mut HashMap>, +) -> Option { + // we need to poll all receivers — use a simple approach with select + // on the first available message from any receiver + if channel_rxs.is_empty() && pattern_rxs.is_empty() { + // no subscriptions — sleep forever (will be cancelled by select) + std::future::pending::>().await + } else { + // poll all receivers using tokio::select on a merged stream + // for simplicity, use a polling approach + loop { + for (_ch, rx) in channel_rxs.iter_mut() { + match rx.try_recv() { + Ok(msg) => return Some(msg), + Err(broadcast::error::TryRecvError::Lagged(n)) => { + tracing::warn!("subscriber lagged, missed {n} messages"); + // try again to get the next available + if let Ok(msg) = rx.try_recv() { + return Some(msg); + } + } + Err(_) => {} + } + } + for (_pat, rx) in pattern_rxs.iter_mut() { + match rx.try_recv() { + Ok(msg) => return Some(msg), + Err(broadcast::error::TryRecvError::Lagged(n)) => { + tracing::warn!("subscriber lagged, missed {n} messages"); + if let Ok(msg) = rx.try_recv() { + return Some(msg); + } + } + Err(_) => {} + } + } + // yield to avoid busy-spinning + tokio::task::yield_now().await; + } + } +} + +/// Serializes a subscribe/unsubscribe response: ["type", channel, count] +fn serialize_sub_response(kind: &str, channel: &str, count: usize, out: &mut BytesMut) { + let frame = Frame::Array(vec![ + Frame::Bulk(Bytes::from(kind.to_string())), + Frame::Bulk(Bytes::from(channel.to_string())), + Frame::Integer(count as i64), + ]); + frame.serialize(out); +} + +/// Serializes a pushed message for subscribers. +/// +/// For exact subscriptions: ["message", channel, data] +/// For pattern subscriptions: ["pmessage", pattern, channel, data] +fn serialize_push_message(msg: &PubMessage, out: &mut BytesMut) { + let frame = if let Some(ref pattern) = msg.pattern { + Frame::Array(vec![ + Frame::Bulk(Bytes::from("pmessage")), + Frame::Bulk(Bytes::from(pattern.clone())), + Frame::Bulk(Bytes::from(msg.channel.clone())), + Frame::Bulk(msg.data.clone()), + ]) + } else { + Frame::Array(vec![ + Frame::Bulk(Bytes::from("message")), + Frame::Bulk(Bytes::from(msg.channel.clone())), + Frame::Bulk(msg.data.clone()), + ]) + }; + frame.serialize(out); +} + +/// Cleans up all subscriptions when a subscriber disconnects. +fn cleanup_subscriptions( + pubsub: &PubSubManager, + channel_rxs: &HashMap>, + pattern_rxs: &HashMap>, +) { + for ch in channel_rxs.keys() { + pubsub.unsubscribe(ch); + } + for pat in pattern_rxs.keys() { + pubsub.punsubscribe(pat); + } +} + /// Converts a raw frame into a command and executes it. /// /// When metrics or slowlog are enabled, brackets the command with @@ -110,6 +447,7 @@ async fn process( engine: &Engine, ctx: &Arc, slow_log: &Arc, + pubsub: &Arc, ) -> Frame { match Command::from_frame(frame) { Ok(cmd) => { @@ -121,7 +459,7 @@ async fn process( None }; - let response = execute(cmd, engine, ctx, slow_log).await; + let response = execute(cmd, engine, ctx, slow_log, pubsub).await; ctx.commands_processed.fetch_add(1, Ordering::Relaxed); if let Some(start) = start { @@ -149,6 +487,7 @@ async fn execute( engine: &Engine, ctx: &Arc, slow_log: &Arc, + pubsub: &Arc, ) -> Frame { match cmd { // -- no shard needed -- @@ -943,6 +1282,21 @@ async fn execute( Frame::Simple("OK".into()) } + // -- pub/sub -- + Command::Publish { channel, message } => { + let count = pubsub.publish(&channel, message); + Frame::Integer(count as i64) + } + + // subscribe commands are handled in the connection loop, not here. + // if we reach this point, something went wrong. + Command::Subscribe { .. } + | Command::Unsubscribe { .. } + | Command::PSubscribe { .. } + | Command::PUnsubscribe { .. } => { + Frame::Error("ERR subscribe commands should not reach execute".into()) + } + Command::Unknown(name) => Frame::Error(format!("ERR unknown command '{name}'")), } } diff --git a/crates/ember-server/src/server.rs b/crates/ember-server/src/server.rs index 72c17b2d..7badf5af 100644 --- a/crates/ember-server/src/server.rs +++ b/crates/ember-server/src/server.rs @@ -14,6 +14,7 @@ use tokio::sync::Semaphore; use tracing::{error, info, warn}; use crate::connection; +use crate::pubsub::PubSubManager; use crate::slowlog::{SlowLog, SlowLogConfig}; /// Default maximum number of concurrent client connections. @@ -93,6 +94,7 @@ pub async fn run( }); let slow_log = Arc::new(SlowLog::new(slowlog_config)); + let pubsub = Arc::new(PubSubManager::new()); info!( "listening on {addr} with {} shards (max {max_conn} connections)", @@ -135,9 +137,10 @@ pub async fn run( let engine = engine.clone(); let ctx = Arc::clone(&ctx); let slow_log = Arc::clone(&slow_log); + let pubsub = Arc::clone(&pubsub); tokio::spawn(async move { - if let Err(e) = connection::handle(stream, engine, &ctx, &slow_log).await { + if let Err(e) = connection::handle(stream, engine, &ctx, &slow_log, &pubsub).await { error!("connection error from {peer}: {e}"); } ctx.connections_active.fetch_sub(1, Ordering::Relaxed); @@ -209,6 +212,7 @@ pub async fn run_concurrent( }); let slow_log = Arc::new(SlowLog::new(slowlog_config)); + let pubsub = Arc::new(PubSubManager::new()); info!("listening on {addr} with concurrent keyspace (max {max_conn} connections)"); @@ -249,10 +253,11 @@ pub async fn run_concurrent( let engine = engine.clone(); let ctx = Arc::clone(&ctx); let slow_log = Arc::clone(&slow_log); + let pubsub = Arc::clone(&pubsub); tokio::spawn(async move { if let Err(e) = crate::concurrent_handler::handle( - stream, keyspace, engine, &ctx, &slow_log + stream, keyspace, engine, &ctx, &slow_log, &pubsub ).await { error!("connection error from {peer}: {e}"); } From 8b9200dc509c2683a332ad03d9d445fea572aad9 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Sun, 8 Feb 2026 08:32:16 -0500 Subject: [PATCH 4/8] chore: fix clippy warnings and formatting - remove unused imports in pubsub.rs - remove unused parameters from handle_subscriber_mode - fix formatting in protocol parse functions --- crates/ember-protocol/src/command.rs | 21 ++++----------------- crates/ember-server/src/connection.rs | 12 ++---------- crates/ember-server/src/pubsub.rs | 27 +++++++++++---------------- 3 files changed, 17 insertions(+), 43 deletions(-) diff --git a/crates/ember-protocol/src/command.rs b/crates/ember-protocol/src/command.rs index d75b40ea..2424f90e 100644 --- a/crates/ember-protocol/src/command.rs +++ b/crates/ember-protocol/src/command.rs @@ -273,7 +273,6 @@ pub enum Command { SlowLogReset, // --- pub/sub commands --- - /// SUBSCRIBE `channel` \[channel ...\]. Subscribe to one or more channels. Subscribe { channels: Vec }, @@ -1481,18 +1480,12 @@ fn parse_subscribe(args: &[Frame]) -> Result { if args.is_empty() { return Err(ProtocolError::WrongArity("SUBSCRIBE".into())); } - let channels: Vec = args - .iter() - .map(extract_string) - .collect::>()?; + let channels: Vec = args.iter().map(extract_string).collect::>()?; Ok(Command::Subscribe { channels }) } fn parse_unsubscribe(args: &[Frame]) -> Result { - let channels: Vec = args - .iter() - .map(extract_string) - .collect::>()?; + let channels: Vec = args.iter().map(extract_string).collect::>()?; Ok(Command::Unsubscribe { channels }) } @@ -1500,18 +1493,12 @@ fn parse_psubscribe(args: &[Frame]) -> Result { if args.is_empty() { return Err(ProtocolError::WrongArity("PSUBSCRIBE".into())); } - let patterns: Vec = args - .iter() - .map(extract_string) - .collect::>()?; + let patterns: Vec = args.iter().map(extract_string).collect::>()?; Ok(Command::PSubscribe { patterns }) } fn parse_punsubscribe(args: &[Frame]) -> Result { - let patterns: Vec = args - .iter() - .map(extract_string) - .collect::>()?; + let patterns: Vec = args.iter().map(extract_string).collect::>()?; Ok(Command::PUnsubscribe { patterns }) } diff --git a/crates/ember-server/src/connection.rs b/crates/ember-server/src/connection.rs index 890ec4d0..cb983df4 100644 --- a/crates/ember-server/src/connection.rs +++ b/crates/ember-server/src/connection.rs @@ -113,10 +113,7 @@ pub async fn handle( // enter subscriber mode — this blocks until all subscriptions // are removed or the client disconnects - handle_subscriber_mode( - &mut stream, &mut buf, &mut out, &engine, ctx, slow_log, pubsub, sub_frames, - ) - .await?; + handle_subscriber_mode(&mut stream, &mut buf, &mut out, pubsub, sub_frames).await?; return Ok(()); } @@ -162,9 +159,6 @@ async fn handle_subscriber_mode( stream: &mut TcpStream, buf: &mut BytesMut, out: &mut BytesMut, - engine: &Engine, - ctx: &Arc, - slow_log: &Arc, pubsub: &Arc, initial_frames: Vec, ) -> Result<(), Box> { @@ -175,9 +169,7 @@ async fn handle_subscriber_mode( // process the initial subscribe commands for frame in initial_frames { if let Ok(cmd) = Command::from_frame(frame) { - handle_sub_command( - cmd, pubsub, &mut channel_rxs, &mut pattern_rxs, out, - ); + handle_sub_command(cmd, pubsub, &mut channel_rxs, &mut pattern_rxs, out); } } diff --git a/crates/ember-server/src/pubsub.rs b/crates/ember-server/src/pubsub.rs index 6d1e5cd6..7a942672 100644 --- a/crates/ember-server/src/pubsub.rs +++ b/crates/ember-server/src/pubsub.rs @@ -4,9 +4,7 @@ //! subscribers. Supports both exact channel names and glob patterns. //! Thread-safe via DashMap for lock-free concurrent access. -use std::collections::HashSet; use std::sync::atomic::{AtomicUsize, Ordering}; -use std::sync::Arc; use bytes::Bytes; use dashmap::DashMap; @@ -56,13 +54,10 @@ impl PubSubManager { /// Subscribe to an exact channel. Returns a receiver for messages /// on that channel. pub fn subscribe(&self, channel: &str) -> broadcast::Receiver { - let entry = self - .channels - .entry(channel.to_string()) - .or_insert_with(|| { - let (tx, _) = broadcast::channel(CHANNEL_CAPACITY); - tx - }); + let entry = self.channels.entry(channel.to_string()).or_insert_with(|| { + let (tx, _) = broadcast::channel(CHANNEL_CAPACITY); + tx + }); self.subscription_count.fetch_add(1, Ordering::Relaxed); entry.subscribe() } @@ -89,13 +84,10 @@ impl PubSubManager { /// Subscribe to a glob pattern. Returns a receiver for messages /// matching the pattern. pub fn psubscribe(&self, pattern: &str) -> broadcast::Receiver { - let entry = self - .patterns - .entry(pattern.to_string()) - .or_insert_with(|| { - let (tx, _) = broadcast::channel(CHANNEL_CAPACITY); - tx - }); + let entry = self.patterns.entry(pattern.to_string()).or_insert_with(|| { + let (tx, _) = broadcast::channel(CHANNEL_CAPACITY); + tx + }); self.subscription_count.fetch_add(1, Ordering::Relaxed); entry.subscribe() } @@ -148,16 +140,19 @@ impl PubSubManager { } /// Returns the total number of active subscriptions. + #[allow(dead_code)] // used in tests and future PUBSUB commands pub fn total_subscriptions(&self) -> usize { self.subscription_count.load(Ordering::Relaxed) } /// Returns the number of active channels (with at least one subscriber). + #[allow(dead_code)] pub fn active_channels(&self) -> usize { self.channels.len() } /// Returns the number of active patterns. + #[allow(dead_code)] pub fn active_patterns(&self) -> usize { self.patterns.len() } From 454cda7d549e0acde4bd4c1d8606b836be123b49 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Sun, 8 Feb 2026 08:33:14 -0500 Subject: [PATCH 5/8] test: add pub/sub command parsing tests covers subscribe, unsubscribe, psubscribe, punsubscribe, publish with edge cases: no args, multiple channels, wrong arity, case insensitivity. --- crates/ember-protocol/src/command.rs | 97 ++++++++++++++++++++++++++++ 1 file changed, 97 insertions(+) diff --git a/crates/ember-protocol/src/command.rs b/crates/ember-protocol/src/command.rs index 2424f90e..588ef345 100644 --- a/crates/ember-protocol/src/command.rs +++ b/crates/ember-protocol/src/command.rs @@ -3304,4 +3304,101 @@ mod tests { }, ); } + + // --- pub/sub --- + + #[test] + fn subscribe_single_channel() { + assert_eq!( + Command::from_frame(cmd(&["SUBSCRIBE", "news"])).unwrap(), + Command::Subscribe { + channels: vec!["news".into()] + }, + ); + } + + #[test] + fn subscribe_multiple_channels() { + assert_eq!( + Command::from_frame(cmd(&["SUBSCRIBE", "ch1", "ch2", "ch3"])).unwrap(), + Command::Subscribe { + channels: vec!["ch1".into(), "ch2".into(), "ch3".into()] + }, + ); + } + + #[test] + fn subscribe_no_args() { + let err = Command::from_frame(cmd(&["SUBSCRIBE"])).unwrap_err(); + assert!(matches!(err, ProtocolError::WrongArity(_))); + } + + #[test] + fn unsubscribe_all() { + assert_eq!( + Command::from_frame(cmd(&["UNSUBSCRIBE"])).unwrap(), + Command::Unsubscribe { channels: vec![] }, + ); + } + + #[test] + fn unsubscribe_specific() { + assert_eq!( + Command::from_frame(cmd(&["UNSUBSCRIBE", "news"])).unwrap(), + Command::Unsubscribe { + channels: vec!["news".into()] + }, + ); + } + + #[test] + fn psubscribe_pattern() { + assert_eq!( + Command::from_frame(cmd(&["PSUBSCRIBE", "news.*"])).unwrap(), + Command::PSubscribe { + patterns: vec!["news.*".into()] + }, + ); + } + + #[test] + fn psubscribe_no_args() { + let err = Command::from_frame(cmd(&["PSUBSCRIBE"])).unwrap_err(); + assert!(matches!(err, ProtocolError::WrongArity(_))); + } + + #[test] + fn punsubscribe_all() { + assert_eq!( + Command::from_frame(cmd(&["PUNSUBSCRIBE"])).unwrap(), + Command::PUnsubscribe { patterns: vec![] }, + ); + } + + #[test] + fn publish_basic() { + assert_eq!( + Command::from_frame(cmd(&["PUBLISH", "news", "hello world"])).unwrap(), + Command::Publish { + channel: "news".into(), + message: Bytes::from("hello world"), + }, + ); + } + + #[test] + fn publish_wrong_arity() { + let err = Command::from_frame(cmd(&["PUBLISH", "news"])).unwrap_err(); + assert!(matches!(err, ProtocolError::WrongArity(_))); + } + + #[test] + fn subscribe_case_insensitive() { + assert_eq!( + Command::from_frame(cmd(&["subscribe", "ch"])).unwrap(), + Command::Subscribe { + channels: vec!["ch".into()] + }, + ); + } } From 3bf381c78d983da2510e931260721a90fd2bf839 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Sun, 8 Feb 2026 08:43:32 -0500 Subject: [PATCH 6/8] feat: add PUBSUB CHANNELS/NUMSUB/NUMPAT introspection commands MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit completes the pub/sub command set with redis-compatible introspection: - PUBSUB CHANNELS [pattern] — list active channels with optional glob - PUBSUB NUMSUB [channel ...] — subscriber counts per channel - PUBSUB NUMPAT — count of active pattern subscriptions includes parsing tests and end-to-end verification with redis-cli. --- crates/ember-protocol/src/command.rs | 98 +++++++++++++++++++ crates/ember-server/src/concurrent_handler.rs | 17 ++++ crates/ember-server/src/connection.rs | 17 ++++ crates/ember-server/src/pubsub.rs | 39 ++++++-- 4 files changed, 163 insertions(+), 8 deletions(-) diff --git a/crates/ember-protocol/src/command.rs b/crates/ember-protocol/src/command.rs index 588ef345..7247e4ec 100644 --- a/crates/ember-protocol/src/command.rs +++ b/crates/ember-protocol/src/command.rs @@ -288,6 +288,15 @@ pub enum Command { /// PUBLISH `channel` `message`. Publish a message to a channel. Publish { channel: String, message: Bytes }, + /// PUBSUB CHANNELS \[pattern\]. List active channels, optionally matching a glob. + PubSubChannels { pattern: Option }, + + /// PUBSUB NUMSUB \[channel ...\]. Returns subscriber counts for given channels. + PubSubNumSub { channels: Vec }, + + /// PUBSUB NUMPAT. Returns the number of active pattern subscriptions. + PubSubNumPat, + /// A command we don't recognize (yet). Unknown(String), } @@ -392,6 +401,9 @@ impl Command { Command::PSubscribe { .. } => "psubscribe", Command::PUnsubscribe { .. } => "punsubscribe", Command::Publish { .. } => "publish", + Command::PubSubChannels { .. } => "pubsub", + Command::PubSubNumSub { .. } => "pubsub", + Command::PubSubNumPat => "pubsub", Command::Unknown(_) => "unknown", } } @@ -478,6 +490,7 @@ impl Command { "PSUBSCRIBE" => parse_psubscribe(&frames[1..]), "PUNSUBSCRIBE" => parse_punsubscribe(&frames[1..]), "PUBLISH" => parse_publish(&frames[1..]), + "PUBSUB" => parse_pubsub(&frames[1..]), _ => Ok(Command::Unknown(name)), } } @@ -1511,6 +1524,35 @@ fn parse_publish(args: &[Frame]) -> Result { Ok(Command::Publish { channel, message }) } +fn parse_pubsub(args: &[Frame]) -> Result { + if args.is_empty() { + return Err(ProtocolError::WrongArity("PUBSUB".into())); + } + + let subcmd = extract_string(&args[0])?.to_ascii_uppercase(); + match subcmd.as_str() { + "CHANNELS" => { + let pattern = if args.len() > 1 { + Some(extract_string(&args[1])?) + } else { + None + }; + Ok(Command::PubSubChannels { pattern }) + } + "NUMSUB" => { + let channels: Vec = args[1..] + .iter() + .map(extract_string) + .collect::>()?; + Ok(Command::PubSubNumSub { channels }) + } + "NUMPAT" => Ok(Command::PubSubNumPat), + other => Err(ProtocolError::InvalidCommandFrame(format!( + "unknown PUBSUB subcommand '{other}'" + ))), + } +} + #[cfg(test)] mod tests { use super::*; @@ -3401,4 +3443,60 @@ mod tests { }, ); } + + #[test] + fn pubsub_channels_no_pattern() { + assert_eq!( + Command::from_frame(cmd(&["PUBSUB", "CHANNELS"])).unwrap(), + Command::PubSubChannels { pattern: None }, + ); + } + + #[test] + fn pubsub_channels_with_pattern() { + assert_eq!( + Command::from_frame(cmd(&["PUBSUB", "CHANNELS", "news.*"])).unwrap(), + Command::PubSubChannels { + pattern: Some("news.*".into()) + }, + ); + } + + #[test] + fn pubsub_numsub_no_args() { + assert_eq!( + Command::from_frame(cmd(&["PUBSUB", "NUMSUB"])).unwrap(), + Command::PubSubNumSub { channels: vec![] }, + ); + } + + #[test] + fn pubsub_numsub_with_channels() { + assert_eq!( + Command::from_frame(cmd(&["PUBSUB", "NUMSUB", "ch1", "ch2"])).unwrap(), + Command::PubSubNumSub { + channels: vec!["ch1".into(), "ch2".into()] + }, + ); + } + + #[test] + fn pubsub_numpat() { + assert_eq!( + Command::from_frame(cmd(&["PUBSUB", "NUMPAT"])).unwrap(), + Command::PubSubNumPat, + ); + } + + #[test] + fn pubsub_no_subcommand() { + let err = Command::from_frame(cmd(&["PUBSUB"])).unwrap_err(); + assert!(matches!(err, ProtocolError::WrongArity(_))); + } + + #[test] + fn pubsub_unknown_subcommand() { + let err = Command::from_frame(cmd(&["PUBSUB", "BOGUS"])).unwrap_err(); + assert!(matches!(err, ProtocolError::InvalidCommandFrame(_))); + } } diff --git a/crates/ember-server/src/concurrent_handler.rs b/crates/ember-server/src/concurrent_handler.rs index 89774e4f..05a231d4 100644 --- a/crates/ember-server/src/concurrent_handler.rs +++ b/crates/ember-server/src/concurrent_handler.rs @@ -208,6 +208,23 @@ async fn execute_concurrent( Frame::Integer(count as i64) } + Command::PubSubChannels { pattern } => { + let names = pubsub.channel_names(pattern.as_deref()); + Frame::Array(names.into_iter().map(|n| Frame::Bulk(n.into())).collect()) + } + + Command::PubSubNumSub { channels } => { + let pairs = pubsub.numsub(&channels); + let mut frames = Vec::with_capacity(pairs.len() * 2); + for (ch, count) in pairs { + frames.push(Frame::Bulk(ch.into())); + frames.push(Frame::Integer(count as i64)); + } + Frame::Array(frames) + } + + Command::PubSubNumPat => Frame::Integer(pubsub.active_patterns() as i64), + // subscribe commands are handled in the connection layer, not here Command::Subscribe { .. } | Command::Unsubscribe { .. } diff --git a/crates/ember-server/src/connection.rs b/crates/ember-server/src/connection.rs index cb983df4..b5594230 100644 --- a/crates/ember-server/src/connection.rs +++ b/crates/ember-server/src/connection.rs @@ -1280,6 +1280,23 @@ async fn execute( Frame::Integer(count as i64) } + Command::PubSubChannels { pattern } => { + let names = pubsub.channel_names(pattern.as_deref()); + Frame::Array(names.into_iter().map(|n| Frame::Bulk(n.into())).collect()) + } + + Command::PubSubNumSub { channels } => { + let pairs = pubsub.numsub(&channels); + let mut frames = Vec::with_capacity(pairs.len() * 2); + for (ch, count) in pairs { + frames.push(Frame::Bulk(ch.into())); + frames.push(Frame::Integer(count as i64)); + } + Frame::Array(frames) + } + + Command::PubSubNumPat => Frame::Integer(pubsub.active_patterns() as i64), + // subscribe commands are handled in the connection loop, not here. // if we reach this point, something went wrong. Command::Subscribe { .. } diff --git a/crates/ember-server/src/pubsub.rs b/crates/ember-server/src/pubsub.rs index 7a942672..a09f3d7a 100644 --- a/crates/ember-server/src/pubsub.rs +++ b/crates/ember-server/src/pubsub.rs @@ -140,19 +140,42 @@ impl PubSubManager { } /// Returns the total number of active subscriptions. - #[allow(dead_code)] // used in tests and future PUBSUB commands + #[allow(dead_code)] // used in tests pub fn total_subscriptions(&self) -> usize { self.subscription_count.load(Ordering::Relaxed) } - /// Returns the number of active channels (with at least one subscriber). - #[allow(dead_code)] - pub fn active_channels(&self) -> usize { - self.channels.len() + /// Returns active channel names, optionally filtered by a glob pattern. + /// Used by PUBSUB CHANNELS [pattern]. + pub fn channel_names(&self, pattern: Option<&str>) -> Vec { + self.channels + .iter() + .map(|entry| entry.key().clone()) + .filter(|name| match pattern { + Some(pat) => glob_match(pat, name), + None => true, + }) + .collect() } - /// Returns the number of active patterns. - #[allow(dead_code)] + /// Returns (channel, subscriber_count) pairs for the given channels. + /// Used by PUBSUB NUMSUB [channel ...]. + pub fn numsub(&self, channels: &[String]) -> Vec<(String, usize)> { + channels + .iter() + .map(|ch| { + let count = self + .channels + .get(ch) + .map(|tx| tx.receiver_count()) + .unwrap_or(0); + (ch.clone(), count) + }) + .collect() + } + + /// Returns the number of active pattern subscriptions. + /// Used by PUBSUB NUMPAT. pub fn active_patterns(&self) -> usize { self.patterns.len() } @@ -396,7 +419,7 @@ mod tests { let _rx2 = mgr.subscribe("b"); let _rx3 = mgr.psubscribe("c.*"); assert_eq!(mgr.total_subscriptions(), 3); - assert_eq!(mgr.active_channels(), 2); + assert_eq!(mgr.channel_names(None).len(), 2); assert_eq!(mgr.active_patterns(), 1); } } From 1fe6fde6a94641bbedd79b1c1f0975d32367ce5a Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Sun, 8 Feb 2026 08:54:11 -0500 Subject: [PATCH 7/8] refactor: audit fixes for pub/sub implementation - fix race condition: unsubscribe/punsubscribe now only decrement subscription_count when the channel actually exists in the registry - replace busy-wait polling in recv_any_message with FuturesUnordered for proper async multiplexing across all broadcast receivers - remove unnecessary String allocations in serialize functions (use Bytes::copy_from_slice and Bytes::from_static instead) - simplify is_subscribe_frame using eq_ignore_ascii_case on byte slices instead of allocating a String for case comparison - simplify enter_sub detection with iter().any() - organize imports (group std::time together) --- crates/ember-server/src/connection.rs | 119 +++++++++++++------------- crates/ember-server/src/pubsub.rs | 12 +-- 2 files changed, 66 insertions(+), 65 deletions(-) diff --git a/crates/ember-server/src/connection.rs b/crates/ember-server/src/connection.rs index b5594230..fd1e1e17 100644 --- a/crates/ember-server/src/connection.rs +++ b/crates/ember-server/src/connection.rs @@ -8,7 +8,7 @@ use std::collections::HashMap; use std::sync::atomic::Ordering; use std::sync::Arc; -use std::time::Instant; +use std::time::{Duration, Instant}; use bytes::{Bytes, BytesMut}; use ember_core::{Engine, KeyspaceStats, ShardRequest, ShardResponse, TtlResult, Value}; @@ -22,7 +22,6 @@ use crate::connection_common::{BUF_CAPACITY, IDLE_TIMEOUT, MAX_BUF_SIZE}; use crate::pubsub::{PubMessage, PubSubManager}; use crate::server::ServerContext; use crate::slowlog::SlowLog; -use std::time::Duration; /// Drives a single client connection to completion. /// @@ -85,15 +84,7 @@ pub async fn handle( // check if any frame is a subscribe command — if so, we need // to enter subscriber mode which changes the connection loop - let mut enter_sub = false; - if !frames.is_empty() { - for frame in &frames { - if is_subscribe_frame(frame) { - enter_sub = true; - break; - } - } - } + let enter_sub = frames.iter().any(is_subscribe_frame); if enter_sub { // process any non-subscribe commands that came before @@ -135,15 +126,14 @@ pub async fn handle( } } -/// Checks if a raw frame is a SUBSCRIBE or PSUBSCRIBE command. +/// Checks if a raw frame is a SUBSCRIBE/PSUBSCRIBE/UNSUBSCRIBE/PUNSUBSCRIBE command. fn is_subscribe_frame(frame: &Frame) -> bool { if let Frame::Array(parts) = frame { if let Some(Frame::Bulk(name)) = parts.first() { - let upper = String::from_utf8_lossy(name).to_ascii_uppercase(); - return matches!( - upper.as_str(), - "SUBSCRIBE" | "PSUBSCRIBE" | "UNSUBSCRIBE" | "PUNSUBSCRIBE" - ); + return name.eq_ignore_ascii_case(b"SUBSCRIBE") + || name.eq_ignore_ascii_case(b"PSUBSCRIBE") + || name.eq_ignore_ascii_case(b"UNSUBSCRIBE") + || name.eq_ignore_ascii_case(b"PUNSUBSCRIBE"); } } false @@ -339,58 +329,69 @@ fn handle_sub_command( } /// Receives a message from any active subscription (channels or patterns). +/// +/// Uses `FuturesUnordered` to efficiently await all broadcast receivers +/// concurrently, avoiding busy-wait polling. Returns `None` only when +/// there are no active subscriptions. async fn recv_any_message( channel_rxs: &mut HashMap>, pattern_rxs: &mut HashMap>, ) -> Option { - // we need to poll all receivers — use a simple approach with select - // on the first available message from any receiver + use std::pin::Pin; + + use futures::stream::{FuturesUnordered, StreamExt}; + if channel_rxs.is_empty() && pattern_rxs.is_empty() { // no subscriptions — sleep forever (will be cancelled by select) - std::future::pending::>().await - } else { - // poll all receivers using tokio::select on a merged stream - // for simplicity, use a polling approach - loop { - for (_ch, rx) in channel_rxs.iter_mut() { - match rx.try_recv() { - Ok(msg) => return Some(msg), - Err(broadcast::error::TryRecvError::Lagged(n)) => { - tracing::warn!("subscriber lagged, missed {n} messages"); - // try again to get the next available - if let Ok(msg) = rx.try_recv() { - return Some(msg); - } - } - Err(_) => {} - } + return std::future::pending::>().await; + } + + type RecvFuture<'a> = Pin< + Box< + dyn std::future::Future> + + Send + + 'a, + >, + >; + + // collect all receivers into a FuturesUnordered so we await them + // all concurrently without spinning. each future resolves when its + // broadcast channel has a message (or reports lag). + let mut pending: FuturesUnordered> = FuturesUnordered::new(); + + for rx in channel_rxs.values_mut() { + pending.push(Box::pin(rx.recv())); + } + for rx in pattern_rxs.values_mut() { + pending.push(Box::pin(rx.recv())); + } + + while let Some(result) = pending.next().await { + match result { + Ok(msg) => return Some(msg), + Err(broadcast::error::RecvError::Lagged(n)) => { + tracing::warn!("subscriber lagged, missed {n} messages"); + // the receiver auto-advances past the gap, so the next + // call to recv() will return the oldest available message. + // we drop through to re-poll on the next loop iteration. } - for (_pat, rx) in pattern_rxs.iter_mut() { - match rx.try_recv() { - Ok(msg) => return Some(msg), - Err(broadcast::error::TryRecvError::Lagged(n)) => { - tracing::warn!("subscriber lagged, missed {n} messages"); - if let Ok(msg) = rx.try_recv() { - return Some(msg); - } - } - Err(_) => {} - } + Err(broadcast::error::RecvError::Closed) => { + // sender was dropped — channel was removed. skip it. } - // yield to avoid busy-spinning - tokio::task::yield_now().await; } } + + None } /// Serializes a subscribe/unsubscribe response: ["type", channel, count] fn serialize_sub_response(kind: &str, channel: &str, count: usize, out: &mut BytesMut) { - let frame = Frame::Array(vec![ - Frame::Bulk(Bytes::from(kind.to_string())), - Frame::Bulk(Bytes::from(channel.to_string())), + Frame::Array(vec![ + Frame::Bulk(Bytes::copy_from_slice(kind.as_bytes())), + Frame::Bulk(Bytes::copy_from_slice(channel.as_bytes())), Frame::Integer(count as i64), - ]); - frame.serialize(out); + ]) + .serialize(out); } /// Serializes a pushed message for subscribers. @@ -400,15 +401,15 @@ fn serialize_sub_response(kind: &str, channel: &str, count: usize, out: &mut Byt fn serialize_push_message(msg: &PubMessage, out: &mut BytesMut) { let frame = if let Some(ref pattern) = msg.pattern { Frame::Array(vec![ - Frame::Bulk(Bytes::from("pmessage")), - Frame::Bulk(Bytes::from(pattern.clone())), - Frame::Bulk(Bytes::from(msg.channel.clone())), + Frame::Bulk(Bytes::from_static(b"pmessage")), + Frame::Bulk(Bytes::copy_from_slice(pattern.as_bytes())), + Frame::Bulk(Bytes::copy_from_slice(msg.channel.as_bytes())), Frame::Bulk(msg.data.clone()), ]) } else { Frame::Array(vec![ - Frame::Bulk(Bytes::from("message")), - Frame::Bulk(Bytes::from(msg.channel.clone())), + Frame::Bulk(Bytes::from_static(b"message")), + Frame::Bulk(Bytes::copy_from_slice(msg.channel.as_bytes())), Frame::Bulk(msg.data.clone()), ]) }; diff --git a/crates/ember-server/src/pubsub.rs b/crates/ember-server/src/pubsub.rs index a09f3d7a..362b9e3b 100644 --- a/crates/ember-server/src/pubsub.rs +++ b/crates/ember-server/src/pubsub.rs @@ -63,13 +63,13 @@ impl PubSubManager { } /// Unsubscribe from an exact channel. Returns true if the channel - /// had active subscriptions. + /// existed in the registry. /// /// Note: the actual receiver is dropped by the caller. This just - /// cleans up empty channels. + /// cleans up empty channels and adjusts the subscription count. pub fn unsubscribe(&self, channel: &str) -> bool { - self.subscription_count.fetch_sub(1, Ordering::Relaxed); if let Some(entry) = self.channels.get(channel) { + self.subscription_count.fetch_sub(1, Ordering::Relaxed); // if no receivers left, remove the channel entirely if entry.receiver_count() <= 1 { drop(entry); @@ -92,11 +92,11 @@ impl PubSubManager { entry.subscribe() } - /// Unsubscribe from a pattern. Returns true if the pattern had - /// active subscriptions. + /// Unsubscribe from a pattern. Returns true if the pattern existed + /// in the registry. pub fn punsubscribe(&self, pattern: &str) -> bool { - self.subscription_count.fetch_sub(1, Ordering::Relaxed); if let Some(entry) = self.patterns.get(pattern) { + self.subscription_count.fetch_sub(1, Ordering::Relaxed); if entry.receiver_count() <= 1 { drop(entry); self.patterns.remove(pattern); From 5e16e10f126badbe5be72c8f3fa460de4f68a41a Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Sun, 8 Feb 2026 08:54:15 -0500 Subject: [PATCH 8/8] docs: update command count to 76, add pub/sub to feature list MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - README.md: 65+ → 76 commands, 609 → 639 tests, ~14k → ~21k LOC - README.md: add pub/sub to feature list - bench/README.md: update command count, remove pub/sub from dragonfly advantages (ember now has it) --- README.md | 3 ++- bench/README.md | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 61ae1f92..7a277658 100644 --- a/README.md +++ b/README.md @@ -16,6 +16,7 @@ a low-latency, memory-efficient, distributed cache written in Rust. designed to - **sets** — SADD, SREM, SMEMBERS, SISMEMBER, SCARD - **key commands** — DEL, EXISTS, EXPIRE, TTL, PEXPIRE, PTTL, PERSIST, TYPE, SCAN - **server commands** — PING, ECHO, INFO, DBSIZE, FLUSHDB, BGSAVE, BGREWRITEAOF +- **pub/sub** — SUBSCRIBE, UNSUBSCRIBE, PSUBSCRIBE, PUNSUBSCRIBE, PUBLISH, plus PUBSUB introspection - **observability** — prometheus metrics (`--metrics-port`), enriched INFO with 6 sections, SLOWLOG command - **sharded engine** — shared-nothing, thread-per-core design with no cross-shard locking - **concurrent mode** — experimental DashMap-backed keyspace for lock-free GET/SET (2x faster than Redis) @@ -183,7 +184,7 @@ contributions welcome — see [CONTRIBUTING.md](CONTRIBUTING.md). | 4 | clustering (raft, gossip, slots, migration) | ✅ complete | | 5 | developer experience (observability, CLI, clients) | 🚧 in progress | -**current**: 65+ commands, 609 tests, ~14k lines of code +**current**: 76 commands, 639 tests, ~21k lines of code ## security diff --git a/bench/README.md b/bench/README.md index 0b77a75d..bf9916b6 100644 --- a/bench/README.md +++ b/bench/README.md @@ -31,12 +31,12 @@ tested on GCP c2-standard-8 (8 vCPU Intel Xeon @ 3.10GHz), Ubuntu 22.04. **important caveat**: this comparison is not apples-to-apples. dragonfly is a production-ready Redis replacement with features ember doesn't have: -- full Redis API compatibility (100+ commands vs ember's 65) +- full Redis API compatibility (100+ commands vs ember's 76) - sophisticated memory management (dashtable for ~25% of Redis memory usage) - transactional semantics (MULTI/EXEC, Lua scripting) - fork-free snapshotting - replication and clustering -- streams, pub/sub, and more +- streams and more ember's concurrent mode wins on raw GET/SET throughput because it's architecturally simpler — essentially a concurrent hashmap with RESP3 parsing. this simplicity comes at the cost of features. for production Redis replacement, dragonfly is likely the better choice. ember is best suited for simple caching workloads where raw throughput matters more than feature completeness.