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/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. diff --git a/crates/ember-protocol/src/command.rs b/crates/ember-protocol/src/command.rs index e565398b..7247e4ec 100644 --- a/crates/ember-protocol/src/command.rs +++ b/crates/ember-protocol/src/command.rs @@ -272,6 +272,31 @@ 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 }, + + /// 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), } @@ -371,6 +396,14 @@ 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::PubSubChannels { .. } => "pubsub", + Command::PubSubNumSub { .. } => "pubsub", + Command::PubSubNumPat => "pubsub", Command::Unknown(_) => "unknown", } } @@ -452,6 +485,12 @@ 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..]), + "PUBSUB" => parse_pubsub(&frames[1..]), _ => Ok(Command::Unknown(name)), } } @@ -1450,6 +1489,70 @@ 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 }) +} + +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::*; @@ -3243,4 +3346,157 @@ 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()] + }, + ); + } + + #[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/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/concurrent_handler.rs b/crates/ember-server/src/concurrent_handler.rs index 3e3d1f95..05a231d4 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,37 @@ 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) + } + + 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 { .. } + | 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..fd1e1e17 100644 --- a/crates/ember-server/src/connection.rs +++ b/crates/ember-server/src/connection.rs @@ -5,9 +5,10 @@ //! 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; +use std::time::{Duration, Instant}; use bytes::{Bytes, BytesMut}; use ember_core::{Engine, KeyspaceStats, ShardRequest, ShardResponse, TtlResult, Value}; @@ -15,11 +16,12 @@ 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; /// Drives a single client connection to completion. /// @@ -31,6 +33,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 +82,37 @@ 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 enter_sub = frames.iter().any(is_subscribe_frame); + + 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, 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 +126,310 @@ pub async fn handle( } } +/// 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() { + 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 +} + +/// 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, + 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). +/// +/// 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 { + 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) + 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. + } + Err(broadcast::error::RecvError::Closed) => { + // sender was dropped — channel was removed. skip it. + } + } + } + + None +} + +/// Serializes a subscribe/unsubscribe response: ["type", channel, count] +fn serialize_sub_response(kind: &str, channel: &str, count: usize, out: &mut BytesMut) { + 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), + ]) + .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_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_static(b"message")), + Frame::Bulk(Bytes::copy_from_slice(msg.channel.as_bytes())), + 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 +440,7 @@ async fn process( engine: &Engine, ctx: &Arc, slow_log: &Arc, + pubsub: &Arc, ) -> Frame { match Command::from_frame(frame) { Ok(cmd) => { @@ -121,7 +452,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 +480,7 @@ async fn execute( engine: &Engine, ctx: &Arc, slow_log: &Arc, + pubsub: &Arc, ) -> Frame { match cmd { // -- no shard needed -- @@ -943,6 +1275,38 @@ async fn execute( Frame::Simple("OK".into()) } + // -- pub/sub -- + Command::Publish { channel, message } => { + let count = pubsub.publish(&channel, message); + 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 { .. } + | 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/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..362b9e3b --- /dev/null +++ b/crates/ember-server/src/pubsub.rs @@ -0,0 +1,425 @@ +//! 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::sync::atomic::{AtomicUsize, Ordering}; + +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 + /// existed in the registry. + /// + /// Note: the actual receiver is dropped by the caller. This just + /// cleans up empty channels and adjusts the subscription count. + pub fn unsubscribe(&self, channel: &str) -> bool { + 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); + 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 existed + /// in the registry. + pub fn punsubscribe(&self, pattern: &str) -> bool { + 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); + } + 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. + #[allow(dead_code)] // used in tests + pub fn total_subscriptions(&self) -> usize { + self.subscription_count.load(Ordering::Relaxed) + } + + /// 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 (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() + } +} + +/// 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.channel_names(None).len(), 2); + assert_eq!(mgr.active_patterns(), 1); + } +} 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}"); }