diff --git a/Cargo.lock b/Cargo.lock index 0441fa82..cc2e7f63 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -762,6 +762,15 @@ dependencies = [ "parking_lot_core", ] +[[package]] +name = "deranged" +version = "0.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" +dependencies = [ + "powerfmt", +] + [[package]] name = "derive_more" version = "1.0.0" @@ -890,8 +899,12 @@ dependencies = [ "ember-protocol", "ember-server", "prost-reflect", + "rcgen", + "rustls", + "rustls-pki-types", "tempfile", "tokio", + "tokio-rustls", ] [[package]] @@ -1767,6 +1780,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "num-conv" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf97ec579c3c42f953ef76dbf8d55ac91fb219dde70e49aa4a6b7d74e9919050" + [[package]] name = "num-traits" version = "0.2.19" @@ -1880,6 +1899,16 @@ dependencies = [ "windows-link", ] +[[package]] +name = "pem" +version = "3.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d30c53c26bc5b31a98cd02d20f25a7c8567146caf63ed593a9d87b2775291be" +dependencies = [ + "base64", + "serde_core", +] + [[package]] name = "percent-encoding" version = "2.3.2" @@ -1987,6 +2016,12 @@ dependencies = [ "serde", ] +[[package]] +name = "powerfmt" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" + [[package]] name = "ppv-lite86" version = "0.2.21" @@ -2281,6 +2316,19 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "rcgen" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75e669e5202259b5314d1ea5397316ad400819437857b90861765f24c4cf80a2" +dependencies = [ + "pem", + "ring", + "rustls-pki-types", + "time", + "yasna", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -2889,6 +2937,25 @@ dependencies = [ "tikv-jemalloc-sys", ] +[[package]] +name = "time" +version = "0.3.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "743bd48c283afc0388f9b8827b976905fb217ad9e647fae3a379a9283c4def2c" +dependencies = [ + "deranged", + "num-conv", + "powerfmt", + "serde_core", + "time-core", +] + +[[package]] +name = "time-core" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7694e1cfe791f8d31026952abf09c69ca6f6fa4e1a1229e18988f06a04a12dca" + [[package]] name = "tinytemplate" version = "1.2.1" @@ -3792,6 +3859,15 @@ dependencies = [ "tap", ] +[[package]] +name = "yasna" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e17bb3549cc1321ae1296b9cdc2698e2b6cb1992adfa19a8c72e5b7a738f44cd" +dependencies = [ + "time", +] + [[package]] name = "zerocopy" version = "0.8.39" diff --git a/crates/ember-core/src/engine.rs b/crates/ember-core/src/engine.rs index 5473afed..5a9a915b 100644 --- a/crates/ember-core/src/engine.rs +++ b/crates/ember-core/src/engine.rs @@ -36,6 +36,8 @@ pub struct EngineConfig { /// [`ReplicationEvent`] so replication clients can stream it to /// replicas. pub replication_tx: Option>, + /// Optional channel to receive expired key names. Used for keyspace notifications. + pub expired_tx: Option>, /// Optional schema registry for protobuf value validation. /// When set, enables PROTO.* commands. #[cfg(feature = "protobuf")] @@ -53,6 +55,7 @@ pub struct EngineConfig { pub struct Engine { shards: Vec, replication_tx: Option>, + expired_tx: Option>, #[cfg(feature = "protobuf")] schema_registry: Option, } @@ -96,6 +99,7 @@ impl Engine { config.persistence.clone(), Some(drop_handle.clone()), config.replication_tx.clone(), + config.expired_tx.clone(), #[cfg(feature = "protobuf")] config.schema_registry.clone(), ) @@ -105,6 +109,7 @@ impl Engine { Self { shards, replication_tx: config.replication_tx, + expired_tx: config.expired_tx, #[cfg(feature = "protobuf")] schema_registry: config.schema_registry, } @@ -144,6 +149,7 @@ impl Engine { config.persistence.clone(), Some(drop_handle.clone()), config.replication_tx.clone(), + config.expired_tx.clone(), #[cfg(feature = "protobuf")] config.schema_registry.clone(), ); @@ -154,6 +160,7 @@ impl Engine { let engine = Self { shards: handles, replication_tx: config.replication_tx, + expired_tx: config.expired_tx, #[cfg(feature = "protobuf")] schema_registry: config.schema_registry, }; @@ -197,6 +204,14 @@ impl Engine { self.replication_tx.as_ref().map(|tx| tx.subscribe()) } + /// Creates a new broadcast receiver for expired key names. + /// + /// Returns `None` if no expired-key channel was configured. Used by the + /// server to subscribe a background task that fires keyspace notifications. + pub fn subscribe_expired(&self) -> Option> { + self.expired_tx.as_ref().map(|tx| tx.subscribe()) + } + /// Sends a request to a specific shard by index. /// /// Used by SCAN to iterate through shards sequentially. diff --git a/crates/ember-core/src/expiry.rs b/crates/ember-core/src/expiry.rs index 53d65446..2f919e7e 100644 --- a/crates/ember-core/src/expiry.rs +++ b/crates/ember-core/src/expiry.rs @@ -18,15 +18,16 @@ const MAX_ROUNDS: usize = 3; /// Runs one active expiration cycle on the keyspace. /// -/// Samples up to `SAMPLE_SIZE` random keys per round, removes expired -/// ones, and repeats if more than 25% of the sample was expired (up to -/// `MAX_ROUNDS` total). Returns the total number of keys removed. -pub fn run_expiration_cycle(ks: &mut Keyspace) -> usize { - let mut total_removed = 0; +/// Returns the keys removed this cycle. The caller can use this to fire +/// keyspace notifications. When nobody is listening, the caller drops +/// the returned Vec immediately (no allocation if empty). +pub fn run_expiration_cycle(ks: &mut Keyspace) -> Vec { + let mut expired_keys = Vec::new(); for _ in 0..MAX_ROUNDS { - let removed = ks.expire_sample(SAMPLE_SIZE); - total_removed += removed; + let prev_len = expired_keys.len(); + ks.expire_sample(SAMPLE_SIZE, &mut expired_keys); + let removed = expired_keys.len() - prev_len; // if we removed fewer than 25% of the sample, the keyspace // is reasonably clean — stop early @@ -35,7 +36,7 @@ pub fn run_expiration_cycle(ks: &mut Keyspace) -> usize { } } - total_removed + expired_keys } #[cfg(test)] @@ -52,7 +53,7 @@ mod tests { ks.set(format!("key:{i}"), Bytes::from("val"), None, false, false); } let removed = run_expiration_cycle(&mut ks); - assert_eq!(removed, 0); + assert!(removed.is_empty()); assert_eq!(ks.len(), 10); } @@ -77,7 +78,7 @@ mod tests { thread::sleep(Duration::from_millis(20)); let removed = run_expiration_cycle(&mut ks); - assert_eq!(removed, 10); + assert_eq!(removed.len(), 10); assert_eq!(ks.len(), 5); } @@ -94,7 +95,7 @@ mod tests { ); } let removed = run_expiration_cycle(&mut ks); - assert_eq!(removed, 0); + assert!(removed.is_empty()); assert_eq!(ks.len(), 10); } @@ -102,6 +103,6 @@ mod tests { fn empty_keyspace_is_fine() { let mut ks = Keyspace::new(); let removed = run_expiration_cycle(&mut ks); - assert_eq!(removed, 0); + assert!(removed.is_empty()); } } diff --git a/crates/ember-core/src/keyspace/mod.rs b/crates/ember-core/src/keyspace/mod.rs index fa40ff3b..93b16396 100644 --- a/crates/ember-core/src/keyspace/mod.rs +++ b/crates/ember-core/src/keyspace/mod.rs @@ -1262,9 +1262,11 @@ impl Keyspace { /// Randomly samples up to `count` keys and removes any that have expired. /// - /// Returns the number of keys actually removed. Used by the active - /// expiration cycle to clean up keys that no one is reading. - pub fn expire_sample(&mut self, count: usize) -> usize { + /// Samples up to `count` random keys and removes any that have expired. + /// + /// Expired key names are appended to `out` so the caller can emit + /// keyspace notifications. Returns the number of keys removed. + pub(crate) fn expire_sample(&mut self, count: usize, out: &mut Vec) -> usize { if self.entries.is_empty() { return 0; } @@ -1280,8 +1282,9 @@ impl Keyspace { .collect(); let mut removed = 0; - for key in &keys_to_check { - if self.remove_if_expired(key) { + for key in keys_to_check { + if self.remove_if_expired(&key) { + out.push(key); removed += 1; } } diff --git a/crates/ember-core/src/shard/mod.rs b/crates/ember-core/src/shard/mod.rs index fc4a139c..caf951f2 100644 --- a/crates/ember-core/src/shard/mod.rs +++ b/crates/ember-core/src/shard/mod.rs @@ -959,6 +959,8 @@ pub struct PreparedShard { persistence: Option, drop_handle: Option, replication_tx: Option>, + /// Optional channel to broadcast expired key names for keyspace notifications. + expired_tx: Option>, #[cfg(feature = "protobuf")] schema_registry: Option, } @@ -973,6 +975,7 @@ pub fn prepare_shard( persistence: Option, drop_handle: Option, replication_tx: Option>, + expired_tx: Option>, #[cfg(feature = "protobuf")] schema_registry: Option, ) -> (ShardHandle, PreparedShard) { let (tx, rx) = mpsc::channel(buffer); @@ -982,6 +985,7 @@ pub fn prepare_shard( persistence, drop_handle, replication_tx, + expired_tx, #[cfg(feature = "protobuf")] schema_registry, }; @@ -999,6 +1003,7 @@ pub async fn run_prepared(prepared: PreparedShard) { prepared.persistence, prepared.drop_handle, prepared.replication_tx, + prepared.expired_tx, #[cfg(feature = "protobuf")] prepared.schema_registry, ) @@ -1019,6 +1024,7 @@ pub fn spawn_shard( persistence: Option, drop_handle: Option, replication_tx: Option>, + expired_tx: Option>, #[cfg(feature = "protobuf")] schema_registry: Option, ) -> ShardHandle { let (handle, prepared) = prepare_shard( @@ -1027,6 +1033,7 @@ pub fn spawn_shard( persistence, drop_handle, replication_tx, + expired_tx, #[cfg(feature = "protobuf")] schema_registry, ); @@ -1042,6 +1049,7 @@ async fn run_shard( persistence: Option, drop_handle: Option, replication_tx: Option>, + expired_tx: Option>, #[cfg(feature = "protobuf")] schema_registry: Option, ) { let shard_id = config.shard_id; @@ -1226,7 +1234,14 @@ async fn run_shard( } } _ = expiry_tick.tick() => { - expiry::run_expiration_cycle(&mut keyspace); + let expired_keys = expiry::run_expiration_cycle(&mut keyspace); + if let Some(ref tx) = expired_tx { + if !expired_keys.is_empty() && tx.receiver_count() > 0 { + for key in &expired_keys { + let _ = tx.send(key.clone()); + } + } + } } _ = fsync_tick.tick(), if fsync_policy == FsyncPolicy::EverySec => { if let Some(ref mut writer) = aof_writer { diff --git a/crates/ember-server/src/config.rs b/crates/ember-server/src/config.rs index 50fa7721..c155a3c5 100644 --- a/crates/ember-server/src/config.rs +++ b/crates/ember-server/src/config.rs @@ -115,6 +115,7 @@ pub fn build_engine_config( }, persistence, replication_tx: None, + expired_tx: None, #[cfg(feature = "protobuf")] schema_registry: None, shard_channel_buffer, @@ -188,6 +189,10 @@ pub struct EmberConfig { #[serde(rename = "tls-auth-clients")] pub tls_auth_clients: String, + // -- notifications -- + #[serde(rename = "notify-keyspace-events")] + pub notify_keyspace_events: String, + // -- protocol limits -- #[serde(rename = "max-key-len")] pub max_key_len: String, @@ -243,6 +248,8 @@ impl Default for EmberConfig { tls_ca_cert_file: String::new(), tls_auth_clients: "no".into(), + notify_keyspace_events: String::new(), + max_key_len: "512kb".into(), max_value_len: "512mb".into(), max_subscriptions_per_connection: 10_000, @@ -451,6 +458,10 @@ impl EmberConfig { "save-interval-secs".into(), self.save_interval_secs.to_string(), ); + params.insert( + "notify-keyspace-events".into(), + self.notify_keyspace_events.clone(), + ); params.insert("max-key-len".into(), self.max_key_len.clone()); params.insert("max-value-len".into(), self.max_value_len.clone()); params.insert( @@ -593,6 +604,7 @@ const MUTABLE_PARAMS: &[&str] = &[ "max-pipeline-depth", "maxmemory", "maxmemory-policy", + "notify-keyspace-events", ]; impl ConfigRegistry { @@ -687,6 +699,7 @@ impl ConfigRegistry { )); } }, + "notify-keyspace-events" => {} // any string is valid; parsed at runtime _ => {} } diff --git a/crates/ember-server/src/connection/execute.rs b/crates/ember-server/src/connection/execute.rs index 6792f48e..1e8ef9ac 100644 --- a/crates/ember-server/src/connection/execute.rs +++ b/crates/ember-server/src/connection/execute.rs @@ -42,6 +42,27 @@ fn set_expire_to_duration(expire: SetExpire) -> Duration { } } +/// Emits keyspace/keyevent notifications for a successful write command. +/// +/// No-op when `notify-keyspace-events` is `""` / zero (the common case). +/// The `flags == 0` check is a single atomic load — true zero overhead +/// when notifications are disabled. +#[inline] +fn notify_write( + ctx: &Arc, + pubsub: &Arc, + event_flag: u32, + event: &str, + key: &str, +) { + let flags = ctx + .keyspace_event_flags + .load(std::sync::atomic::Ordering::Relaxed); + if flags != 0 { + crate::keyspace_notifications::notify_keyspace_event(flags, event_flag, event, key, pubsub); + } +} + /// Executes a parsed command and returns the response frame. /// /// Ping and Echo are handled inline (no shard routing needed). @@ -138,14 +159,23 @@ pub(super) async fn execute( let duration = expire.map(set_expire_to_duration); let idx = engine.shard_for_key(&key); let req = ShardRequest::Set { - key, + key: key.clone(), value, expire: duration, nx, xx, }; match engine.send_to_shard(idx, req).await { - Ok(ShardResponse::Ok) => Frame::Simple("OK".into()), + Ok(ShardResponse::Ok) => { + notify_write( + ctx, + pubsub, + crate::keyspace_notifications::FLAG_DOLLAR, + "set", + &key, + ); + Frame::Simple("OK".into()) + } Ok(ShardResponse::Value(None)) => Frame::Null, Ok(ShardResponse::OutOfMemory) => oom_error(), Ok(other) => Frame::Error(format!("ERR unexpected shard response: {other:?}")), @@ -155,9 +185,22 @@ pub(super) async fn execute( Command::Expire { key, seconds } => { let idx = engine.shard_for_key(&key); - let req = ShardRequest::Expire { key, seconds }; + let req = ShardRequest::Expire { + key: key.clone(), + seconds, + }; match engine.send_to_shard(idx, req).await { - Ok(ShardResponse::Bool(b)) => Frame::Integer(i64::from(b)), + Ok(ShardResponse::Bool(true)) => { + notify_write( + ctx, + pubsub, + crate::keyspace_notifications::FLAG_G, + "expire", + &key, + ); + Frame::Integer(1) + } + Ok(ShardResponse::Bool(false)) => Frame::Integer(0), Ok(other) => Frame::Error(format!("ERR unexpected shard response: {other:?}")), Err(e) => Frame::Error(format!("ERR {e}")), } @@ -509,6 +552,11 @@ pub(super) async fn execute( limit.unwrap_or(0) as u64, std::sync::atomic::Ordering::Relaxed, ); + } else if key == "notify-keyspace-events" { + let flags = + crate::keyspace_notifications::parse_keyspace_event_flags(&value); + ctx.keyspace_event_flags + .store(flags, std::sync::atomic::Ordering::Relaxed); } Frame::Simple("OK".into()) } @@ -782,9 +830,21 @@ pub(super) async fn execute( // -- list commands -- Command::LPush { key, values } => { let idx = engine.shard_for_key(&key); - let req = ShardRequest::LPush { key, values }; + let req = ShardRequest::LPush { + key: key.clone(), + values, + }; match engine.send_to_shard(idx, req).await { - Ok(ShardResponse::Len(n)) => Frame::Integer(n as i64), + Ok(ShardResponse::Len(n)) => { + notify_write( + ctx, + pubsub, + crate::keyspace_notifications::FLAG_L, + "lpush", + &key, + ); + Frame::Integer(n as i64) + } Ok(ShardResponse::WrongType) => wrongtype_error(), Ok(ShardResponse::OutOfMemory) => oom_error(), Ok(other) => Frame::Error(format!("ERR unexpected shard response: {other:?}")), @@ -794,9 +854,21 @@ pub(super) async fn execute( Command::RPush { key, values } => { let idx = engine.shard_for_key(&key); - let req = ShardRequest::RPush { key, values }; + let req = ShardRequest::RPush { + key: key.clone(), + values, + }; match engine.send_to_shard(idx, req).await { - Ok(ShardResponse::Len(n)) => Frame::Integer(n as i64), + Ok(ShardResponse::Len(n)) => { + notify_write( + ctx, + pubsub, + crate::keyspace_notifications::FLAG_L, + "rpush", + &key, + ); + Frame::Integer(n as i64) + } Ok(ShardResponse::WrongType) => wrongtype_error(), Ok(ShardResponse::OutOfMemory) => oom_error(), Ok(other) => Frame::Error(format!("ERR unexpected shard response: {other:?}")), @@ -977,7 +1049,7 @@ pub(super) async fn execute( } => { let idx = engine.shard_for_key(&key); let req = ShardRequest::ZAdd { - key, + key: key.clone(), members, nx: flags.nx, xx: flags.xx, @@ -986,7 +1058,18 @@ pub(super) async fn execute( ch: flags.ch, }; match engine.send_to_shard(idx, req).await { - Ok(ShardResponse::ZAddLen { count, .. }) => Frame::Integer(count as i64), + Ok(ShardResponse::ZAddLen { count, .. }) => { + if count > 0 { + notify_write( + ctx, + pubsub, + crate::keyspace_notifications::FLAG_Z, + "zadd", + &key, + ); + } + Frame::Integer(count as i64) + } Ok(ShardResponse::WrongType) => wrongtype_error(), Ok(ShardResponse::OutOfMemory) => oom_error(), Ok(other) => Frame::Error(format!("ERR unexpected shard response: {other:?}")), @@ -1250,9 +1333,21 @@ pub(super) async fn execute( // --- hash commands --- Command::HSet { key, fields } => { let idx = engine.shard_for_key(&key); - let req = ShardRequest::HSet { key, fields }; + let req = ShardRequest::HSet { + key: key.clone(), + fields, + }; match engine.send_to_shard(idx, req).await { - Ok(ShardResponse::Len(n)) => Frame::Integer(n as i64), + Ok(ShardResponse::Len(n)) => { + notify_write( + ctx, + pubsub, + crate::keyspace_notifications::FLAG_H, + "hset", + &key, + ); + Frame::Integer(n as i64) + } Ok(ShardResponse::WrongType) => wrongtype_error(), Ok(ShardResponse::OutOfMemory) => oom_error(), Ok(other) => Frame::Error(format!("ERR unexpected shard response: {other:?}")), @@ -1385,9 +1480,23 @@ pub(super) async fn execute( // --- set commands --- Command::SAdd { key, members } => { let idx = engine.shard_for_key(&key); - let req = ShardRequest::SAdd { key, members }; + let req = ShardRequest::SAdd { + key: key.clone(), + members, + }; match engine.send_to_shard(idx, req).await { - Ok(ShardResponse::Len(n)) => Frame::Integer(n as i64), + Ok(ShardResponse::Len(n)) => { + if n > 0 { + notify_write( + ctx, + pubsub, + crate::keyspace_notifications::FLAG_S, + "sadd", + &key, + ); + } + Frame::Integer(n as i64) + } Ok(ShardResponse::WrongType) => wrongtype_error(), Ok(ShardResponse::OutOfMemory) => oom_error(), Ok(other) => Frame::Error(format!("ERR unexpected shard response: {other:?}")), diff --git a/crates/ember-server/src/keyspace_notifications.rs b/crates/ember-server/src/keyspace_notifications.rs new file mode 100644 index 00000000..705c94b4 --- /dev/null +++ b/crates/ember-server/src/keyspace_notifications.rs @@ -0,0 +1,181 @@ +//! Keyspace notifications — publish events to pub/sub channels. +//! +//! When `notify-keyspace-events` is configured, Ember publishes a message +//! for each matching event: +//! +//! - `__keyspace@0__:` → message is the event name (e.g. "set") +//! - `__keyevent@0__:` → message is the key name +//! +//! Both channels are only published when the corresponding K (keyspace) or +//! E (keyevent) flag is set, and the event-type flag matches (e.g. `$` for +//! string commands, `g` for generic, `x` for expired). +//! +//! The guard `flags == 0` is a single AtomicU32 load — true zero overhead +//! when notifications are disabled. + +use bytes::Bytes; + +use crate::pubsub::PubSubManager; + +/// Emit keyspace events on `__keyspace@__:` channels. +pub const FLAG_K: u32 = 0x01; +/// Emit keyevent events on `__keyevent@__:` channels. +pub const FLAG_E: u32 = 0x02; +/// Generic commands: DEL, EXPIRE, RENAME, LPUSH (notify on key itself). +pub const FLAG_G: u32 = 0x04; +/// String commands: SET, GETSET, INCR, APPEND, etc. +pub const FLAG_DOLLAR: u32 = 0x08; +/// List commands: LPUSH, RPUSH, LPOP, RPOP, etc. +pub const FLAG_L: u32 = 0x10; +/// Sorted set commands: ZADD, ZINCRBY, ZREM, etc. +pub const FLAG_Z: u32 = 0x20; +/// Hash commands: HSET, HINCRBY, HDEL, etc. +pub const FLAG_H: u32 = 0x40; +/// Set commands: SADD, SREM, SPOP, etc. +pub const FLAG_S: u32 = 0x80; +/// Key expiration events. +pub const FLAG_X: u32 = 0x100; +/// Key eviction events (allkeys-lru policy). +pub const FLAG_D: u32 = 0x200; + +/// Parses a keyspace event flag string into a bitmask. +/// +/// Accepts any combination of: +/// - `K` — keyspace events +/// - `E` — keyevent events +/// - `g` — generic commands +/// - `$` — string commands +/// - `l` — list commands +/// - `z` — sorted set commands +/// - `h` — hash commands +/// - `s` — set commands +/// - `x` — expired events +/// - `d` — eviction events +/// - `A` — alias for `g$lzxhsd` (all event types) +/// - `""` or `"0"` — disable all notifications +/// +/// Unknown characters are silently ignored (Redis compat). +pub fn parse_keyspace_event_flags(s: &str) -> u32 { + if s.is_empty() || s == "0" { + return 0; + } + + let mut flags = 0u32; + for ch in s.chars() { + match ch { + 'K' => flags |= FLAG_K, + 'E' => flags |= FLAG_E, + 'g' => flags |= FLAG_G, + '$' => flags |= FLAG_DOLLAR, + 'l' => flags |= FLAG_L, + 'z' => flags |= FLAG_Z, + 'h' => flags |= FLAG_H, + 's' => flags |= FLAG_S, + 'x' => flags |= FLAG_X, + 'd' => flags |= FLAG_D, + 'A' => { + flags |= FLAG_G + | FLAG_DOLLAR + | FLAG_L + | FLAG_Z + | FLAG_X + | FLAG_H + | FLAG_S + | FLAG_D + } + _ => {} // unknown flags silently ignored + } + } + flags +} + +/// Publishes keyspace and keyevent notifications for one event. +/// +/// `flags` is the current `notify-keyspace-events` bitmask. +/// `event_flag` is the type flag for this specific event (e.g. `FLAG_X` for expired). +/// `event` is the event name (e.g. `"expired"`, `"set"`, `"del"`). +/// `key` is the key that was affected. +/// +/// Emits to: +/// - `__keyspace@0__:` with message `` (when FLAG_K is set) +/// - `__keyevent@0__:` with message `` (when FLAG_E is set) +/// +/// No-op when neither K nor E is set for this event type. +pub fn notify_keyspace_event( + flags: u32, + event_flag: u32, + event: &str, + key: &str, + pubsub: &PubSubManager, +) { + // both K and E are disabled — nothing to do + if flags & (FLAG_K | FLAG_E) == 0 { + return; + } + // event type not enabled + if flags & event_flag == 0 { + return; + } + + if flags & FLAG_K != 0 { + let channel = format!("__keyspace@0__:{key}"); + pubsub.publish(&channel, Bytes::copy_from_slice(event.as_bytes())); + } + + if flags & FLAG_E != 0 { + let channel = format!("__keyevent@0__:{event}"); + pubsub.publish(&channel, Bytes::copy_from_slice(key.as_bytes())); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn empty_string_returns_zero() { + assert_eq!(parse_keyspace_event_flags(""), 0); + assert_eq!(parse_keyspace_event_flags("0"), 0); + } + + #[test] + fn individual_flags_parsed() { + assert_eq!(parse_keyspace_event_flags("K"), FLAG_K); + assert_eq!(parse_keyspace_event_flags("E"), FLAG_E); + assert_eq!(parse_keyspace_event_flags("x"), FLAG_X); + assert_eq!(parse_keyspace_event_flags("$"), FLAG_DOLLAR); + } + + #[test] + fn uppercase_a_expands_all_events() { + let flags = parse_keyspace_event_flags("A"); + assert!(flags & FLAG_G != 0); + assert!(flags & FLAG_DOLLAR != 0); + assert!(flags & FLAG_L != 0); + assert!(flags & FLAG_Z != 0); + assert!(flags & FLAG_X != 0); + assert!(flags & FLAG_H != 0); + assert!(flags & FLAG_S != 0); + // K and E are NOT included in A + assert_eq!(flags & FLAG_K, 0); + assert_eq!(flags & FLAG_E, 0); + } + + #[test] + fn kex_is_common_config() { + // KEA = keyspace + keyevent + all events + let flags = parse_keyspace_event_flags("KEA"); + assert!(flags & FLAG_K != 0); + assert!(flags & FLAG_E != 0); + assert!(flags & FLAG_X != 0); + } + + #[test] + fn unknown_chars_ignored() { + // 'Q' is not a valid flag + assert_eq!( + parse_keyspace_event_flags("Kx"), + parse_keyspace_event_flags("KxQ") + ); + } +} diff --git a/crates/ember-server/src/main.rs b/crates/ember-server/src/main.rs index 13eae5d3..56f85c42 100644 --- a/crates/ember-server/src/main.rs +++ b/crates/ember-server/src/main.rs @@ -12,6 +12,7 @@ mod connection; mod connection_common; #[cfg(feature = "grpc")] mod grpc; +mod keyspace_notifications; mod metrics; mod pubsub; mod replication; diff --git a/crates/ember-server/src/metrics.rs b/crates/ember-server/src/metrics.rs index d6e6eb99..e062970e 100644 --- a/crates/ember-server/src/metrics.rs +++ b/crates/ember-server/src/metrics.rs @@ -217,7 +217,7 @@ async fn build_health_response(ctx: &ServerContext) -> Response> { } /// Spawns a background task that polls shard stats and publishes them -/// as prometheus gauges. +/// as prometheus gauges and counters. /// /// Keeps `ember-core` free of metrics dependencies — the poller /// pulls stats through the existing `ShardRequest::Stats` broadcast. @@ -226,6 +226,11 @@ pub fn spawn_stats_poller(engine: Engine, ctx: Arc, poll_interval let mut interval = tokio::time::interval(poll_interval); interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + // track cumulative totals from the previous poll so we can publish + // expired/evicted as monotonically-increasing counters rather than gauges. + let mut last_expired: u64 = 0; + let mut last_evicted: u64 = 0; + loop { interval.tick().await; @@ -256,12 +261,33 @@ pub fn spawn_stats_poller(engine: Engine, ctx: Arc, poll_interval gauge!("ember_keys_total").set(total.key_count as f64); gauge!("ember_memory_used_bytes").set(total.used_bytes as f64); - gauge!("ember_keys_expired_total").set(total.keys_expired as f64); - gauge!("ember_keys_evicted_total").set(total.keys_evicted as f64); gauge!("ember_oom_rejections_total").set(total.oom_rejections as f64); gauge!("ember_keyspace_hits_total").set(total.keyspace_hits as f64); gauge!("ember_keyspace_misses_total").set(total.keyspace_misses as f64); + // expired and evicted are cumulative totals from shards — + // publish the delta so prometheus sees a proper counter. + let expired = total.keys_expired as u64; + let evicted = total.keys_evicted as u64; + let delta_expired = expired.saturating_sub(last_expired); + let delta_evicted = evicted.saturating_sub(last_evicted); + if delta_expired > 0 { + counter!("ember_expired_keys_total").increment(delta_expired); + } + if delta_evicted > 0 { + counter!("ember_evicted_keys_total").increment(delta_evicted); + } + last_expired = expired; + last_evicted = evicted; + + // replication lag — how many records each replica is behind. + // published as max and count so alerting rules are straightforward. + let lags = ctx.replica_tracker.replica_lags(); + let replica_count = lags.len() as f64; + let max_lag = lags.iter().max().copied().unwrap_or(0) as f64; + gauge!("ember_replication_connected_replicas").set(replica_count); + gauge!("ember_replication_max_lag_records").set(max_lag); + // update atomic for /health endpoint ctx.memory_used_bytes.store(total.used_bytes as u64); } diff --git a/crates/ember-server/src/replication.rs b/crates/ember-server/src/replication.rs index 20bfeb45..35823cd2 100644 --- a/crates/ember-server/src/replication.rs +++ b/crates/ember-server/src/replication.rs @@ -122,6 +122,17 @@ impl ReplicaTracker { pub fn connected_count(&self) -> usize { self.offsets.lock().map(|map| map.len()).unwrap_or(0) } + + /// Returns the record lag for each connected replica. + /// + /// Lag is `write_offset - acked_offset`. A lag of 0 means fully caught up. + pub fn replica_lags(&self) -> Vec { + let write = self.write_offset.load(Ordering::Relaxed); + self.offsets + .lock() + .map(|map| map.values().map(|&ack| write.saturating_sub(ack)).collect()) + .unwrap_or_default() + } } // -- framed I/O primitives -- diff --git a/crates/ember-server/src/server.rs b/crates/ember-server/src/server.rs index ed525961..f8ac5b21 100644 --- a/crates/ember-server/src/server.rs +++ b/crates/ember-server/src/server.rs @@ -6,7 +6,7 @@ use std::collections::HashMap; use std::net::SocketAddr; use std::path::PathBuf; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::atomic::{AtomicU32, AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; @@ -97,6 +97,9 @@ pub struct ServerContext { /// Connected client metadata. Only locked on accept, disconnect, and /// CLIENT commands — never on the command hot path. pub clients: ClientRegistry, + /// Keyspace notification event flags. 0 = disabled (common case — zero overhead). + /// Updated atomically when CONFIG SET notify-keyspace-events is called. + pub keyspace_event_flags: AtomicU32, } /// Sets nodelay, acquires a semaphore permit, registers the client, and @@ -265,7 +268,7 @@ async fn reject_protected_mode(mut stream: TcpStream) { pub async fn run_concurrent( addr: SocketAddr, shard_count: usize, - config: EngineConfig, + mut config: EngineConfig, max_memory: Option, eviction_policy: EvictionPolicy, max_connections: Option, @@ -285,6 +288,10 @@ pub async fn run_concurrent( .map(|p| p.append_only) .unwrap_or(false); + // set up the expired-key channel for keyspace notifications + let (expired_kn_tx, _expired_kn_rx) = tokio::sync::broadcast::channel::(1024); + config.expired_tx = Some(expired_kn_tx.clone()); + // Create the concurrent keyspace let keyspace = Arc::new(ConcurrentKeyspace::new(max_memory, eviction_policy)); @@ -297,6 +304,15 @@ pub async fn run_concurrent( let tls_listener = setup_tls_listener(tls).await?; + // read the initial notify-keyspace-events flag from the registry + let initial_kn_flags = { + let pairs = config_registry.get_matching("notify-keyspace-events"); + pairs + .first() + .map(|(_, v)| crate::keyspace_notifications::parse_keyspace_event_flags(v)) + .unwrap_or(0) + }; + let metrics_enabled = metrics.is_some(); let stats_poll_interval = limits.stats_poll_interval; let clients: ClientRegistry = Arc::new(Mutex::new(HashMap::new())); @@ -326,6 +342,7 @@ pub async fn run_concurrent( last_save_timestamp: AtomicU64::new(0), next_client_id: AtomicU64::new(1), clients, + keyspace_event_flags: AtomicU32::new(initial_kn_flags), }); if let Some((metrics_addr, handle)) = metrics { @@ -336,6 +353,39 @@ pub async fn run_concurrent( let slow_log = Arc::new(SlowLog::new(slowlog_config)); let pubsub = Arc::new(PubSubManager::new()); + // spawn keyspace notification task for expired keys + { + let mut expired_rx = expired_kn_tx.subscribe(); + let pubsub_kn = Arc::clone(&pubsub); + let ctx_kn = Arc::clone(&ctx); + tokio::spawn(async move { + loop { + match expired_rx.recv().await { + Ok(key) => { + let flags = ctx_kn + .keyspace_event_flags + .load(std::sync::atomic::Ordering::Relaxed); + if flags != 0 { + crate::keyspace_notifications::notify_keyspace_event( + flags, + crate::keyspace_notifications::FLAG_X, + "expired", + &key, + &pubsub_kn, + ); + } + } + Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => { + tracing::warn!( + "keyspace notification subscriber lagged by {n} messages" + ); + } + Err(tokio::sync::broadcast::error::RecvError::Closed) => break, + } + } + }); + } + if save_interval_secs > 0 { let engine_snap = engine.clone(); let ctx_snap = Arc::clone(&ctx); @@ -535,7 +585,7 @@ fn pin_to_core(worker_id: usize) { pub async fn run_threaded( addr: SocketAddr, shard_count: usize, - config: EngineConfig, + mut config: EngineConfig, max_connections: Option, metrics: Option<(SocketAddr, metrics_exporter_prometheus::PrometheusHandle)>, slowlog_config: SlowLogConfig, @@ -563,6 +613,10 @@ pub async fn run_threaded( .max_memory .map(|per_shard| per_shard * shard_count); + // set up the expired-key channel for keyspace notifications + let (expired_kn_tx, _expired_kn_rx) = tokio::sync::broadcast::channel::(1024); + config.expired_tx = Some(expired_kn_tx.clone()); + let (engine, prepared_shards) = Engine::prepare(shard_count, config); let replica_tracker = Arc::new(crate::replication::ReplicaTracker::new()); @@ -579,6 +633,15 @@ pub async fn run_threaded( let max_conn = max_connections.unwrap_or(DEFAULT_MAX_CONNECTIONS); let semaphore = Arc::new(Semaphore::new(max_conn)); + // read the initial notify-keyspace-events flag from the registry + let initial_kn_flags = { + let pairs = config_registry.get_matching("notify-keyspace-events"); + pairs + .first() + .map(|(_, v)| crate::keyspace_notifications::parse_keyspace_event_flags(v)) + .unwrap_or(0) + }; + let metrics_enabled = metrics.is_some(); let stats_poll_interval = limits.stats_poll_interval; let clients: ClientRegistry = Arc::new(Mutex::new(HashMap::new())); @@ -607,6 +670,7 @@ pub async fn run_threaded( last_save_timestamp: AtomicU64::new(0), next_client_id: AtomicU64::new(1), clients, + keyspace_event_flags: AtomicU32::new(initial_kn_flags), }); // management tasks run on the caller's runtime (not the hot path) @@ -618,6 +682,39 @@ pub async fn run_threaded( let slow_log = Arc::new(SlowLog::new(slowlog_config)); let pubsub = Arc::new(PubSubManager::new()); + // spawn keyspace notification task for expired keys + { + let mut expired_rx = expired_kn_tx.subscribe(); + let pubsub_kn = Arc::clone(&pubsub); + let ctx_kn = Arc::clone(&ctx); + tokio::spawn(async move { + loop { + match expired_rx.recv().await { + Ok(key) => { + let flags = ctx_kn + .keyspace_event_flags + .load(std::sync::atomic::Ordering::Relaxed); + if flags != 0 { + crate::keyspace_notifications::notify_keyspace_event( + flags, + crate::keyspace_notifications::FLAG_X, + "expired", + &key, + &pubsub_kn, + ); + } + } + Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => { + tracing::warn!( + "keyspace notification subscriber lagged by {n} messages" + ); + } + Err(tokio::sync::broadcast::error::RecvError::Closed) => break, + } + } + }); + } + if save_interval_secs > 0 { let engine_snap = engine.clone(); let ctx_snap = Arc::clone(&ctx); diff --git a/tests/integration/Cargo.toml b/tests/integration/Cargo.toml index 62a6adf3..fc12b0a3 100644 --- a/tests/integration/Cargo.toml +++ b/tests/integration/Cargo.toml @@ -18,6 +18,10 @@ ember-protocol = { workspace = true } ember-server = { path = "../../crates/ember-server", default-features = false } prost-reflect = { workspace = true } tempfile = "3" +rcgen = "0.13" +rustls = { workspace = true } +rustls-pki-types = { workspace = true } +tokio-rustls = { workspace = true } [features] protobuf = ["ember-server/protobuf"] diff --git a/tests/integration/src/helpers.rs b/tests/integration/src/helpers.rs index d1570adb..c2032542 100644 --- a/tests/integration/src/helpers.rs +++ b/tests/integration/src/helpers.rs @@ -15,6 +15,8 @@ use tokio::net::TcpStream; pub struct TestServer { child: Child, pub port: u16, + /// TLS port, populated when the server was started with TLS options. + pub tls_port: Option, _data_dir: Option, /// Temp file holding server stderr for diagnostics on failure. stderr_path: Option, @@ -40,6 +42,10 @@ pub struct ServerOptions { pub shards: Option, /// Use concurrent (DashMap) mode instead of sharded channels. pub concurrent: bool, + /// PEM certificate file for the TLS listener (requires tls_key_file). + pub tls_cert_file: Option, + /// PEM private key file for the TLS listener (requires tls_cert_file). + pub tls_key_file: Option, } impl TestServer { @@ -125,6 +131,19 @@ impl TestServer { None }; + // TLS: allocate a separate port and pass cert/key files + let tls_port = if opts.tls_cert_file.is_some() && opts.tls_key_file.is_some() { + let tp = find_free_port(); + cmd.arg("--tls-port").arg(tp.to_string()); + cmd.arg("--tls-cert-file") + .arg(opts.tls_cert_file.as_ref().unwrap()); + cmd.arg("--tls-key-file") + .arg(opts.tls_key_file.as_ref().unwrap()); + Some(tp) + } else { + None + }; + // stderr is already redirected to a file for cluster servers if stderr_path.is_none() { cmd.stderr(std::process::Stdio::null()); @@ -151,6 +170,7 @@ impl TestServer { let server = Self { child, port, + tls_port, _data_dir: data_dir, stderr_path, }; diff --git a/tests/integration/src/main.rs b/tests/integration/src/main.rs index 70ef52cd..65b036b3 100644 --- a/tests/integration/src/main.rs +++ b/tests/integration/src/main.rs @@ -10,3 +10,4 @@ mod persistence; #[cfg(feature = "protobuf")] mod proto; mod pubsub; +mod tls; diff --git a/tests/integration/src/persistence.rs b/tests/integration/src/persistence.rs index 18830591..2f1f7b1b 100644 --- a/tests/integration/src/persistence.rs +++ b/tests/integration/src/persistence.rs @@ -6,6 +6,54 @@ use ember_protocol::Frame; use crate::helpers::{ServerOptions, TestServer}; +/// Simulates a hard crash: writes N keys with `appendfsync always`, then +/// SIGKILLs the server without any grace period. On restart the AOF tail +/// must be fully intact because every SET response guaranteed a prior fsync. +#[tokio::test] +async fn sigkill_crash_recovery() { + let data_dir = tempfile::tempdir().unwrap(); + let path = data_dir.path().to_path_buf(); + + const KEY_COUNT: usize = 50; + + { + let server = TestServer::start_with(ServerOptions { + appendonly: true, + data_dir_path: Some(path.clone()), + ..Default::default() + }); + let mut c = server.connect().await; + + for i in 0..KEY_COUNT { + // appendfsync=always means each OK guarantees a fsync — all + // of these must survive even a SIGKILL immediately after. + c.ok(&["SET", &format!("crash:{i}"), &format!("v{i}")]).await; + } + + // drop immediately — no sleep, no graceful shutdown. Child::kill() + // sends SIGKILL on unix, which is the worst-case crash scenario. + } + + // restart with same data directory + let server = TestServer::start_with(ServerOptions { + appendonly: true, + data_dir_path: Some(path), + ..Default::default() + }); + let mut c = server.connect().await; + + for i in 0..KEY_COUNT { + let val = c.get_bulk(&["GET", &format!("crash:{i}")]).await; + assert_eq!( + val, + Some(format!("v{i}")), + "key crash:{i} missing after crash recovery" + ); + } + + drop(data_dir); +} + #[tokio::test] async fn bgsave_and_snapshot_recovery() { let data_dir = tempfile::tempdir().unwrap(); diff --git a/tests/integration/src/tls.rs b/tests/integration/src/tls.rs new file mode 100644 index 00000000..95689206 --- /dev/null +++ b/tests/integration/src/tls.rs @@ -0,0 +1,171 @@ +//! TLS integration tests. +//! +//! Verifies that the server accepts TLS connections and processes commands +//! correctly over an encrypted transport. Uses a self-signed certificate +//! generated at test time via `rcgen`. + +use std::sync::Arc; + +use bytes::{Bytes, BytesMut}; +use ember_protocol::{parse_frame, Frame}; +use rustls::pki_types::{CertificateDer, ServerName}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio_rustls::TlsConnector; + +use crate::helpers::{ServerOptions, TestServer}; + +/// Generates a self-signed cert/key pair for `localhost` and writes PEM files +/// into the given directory. Returns the cert path, key path, and the raw DER +/// bytes needed to build a client-side trust store. +fn generate_test_cert( + dir: &std::path::Path, +) -> (std::path::PathBuf, std::path::PathBuf, Vec) { + let rcgen::CertifiedKey { cert, key_pair } = + rcgen::generate_simple_self_signed(vec!["localhost".to_string()]) + .expect("rcgen cert generation failed"); + + let cert_der = cert.der().to_vec(); + let cert_pem = cert.pem(); + let key_pem = key_pair.serialize_pem(); + + let cert_path = dir.join("server.crt"); + let key_path = dir.join("server.key"); + std::fs::write(&cert_path, cert_pem).expect("failed to write cert"); + std::fs::write(&key_path, key_pem).expect("failed to write key"); + + (cert_path, key_path, cert_der) +} + +/// Builds a rustls ClientConfig that trusts only the given DER certificate. +/// +/// This simulates a client that has been explicitly configured to trust the +/// server's self-signed cert — the same trust model as pinned certificates. +fn client_config(cert_der: Vec) -> rustls::ClientConfig { + let mut root_store = rustls::RootCertStore::empty(); + root_store + .add(CertificateDer::from(cert_der)) + .expect("failed to add cert to root store"); + + rustls::ClientConfig::builder() + .with_root_certificates(root_store) + .with_no_client_auth() +} + +/// Helper: send a command over a TLS stream and return the parsed response. +async fn tls_cmd( + stream: &mut tokio_rustls::client::TlsStream, + args: &[&str], +) -> Frame { + let parts: Vec = args + .iter() + .map(|a| Frame::Bulk(Bytes::copy_from_slice(a.as_bytes()))) + .collect(); + let mut out = BytesMut::new(); + Frame::Array(parts).serialize(&mut out); + stream.write_all(&out).await.expect("write failed"); + + let mut buf = BytesMut::with_capacity(1024); + loop { + let n = stream.read_buf(&mut buf).await.expect("read failed"); + assert!(n > 0 || !buf.is_empty(), "server closed connection unexpectedly"); + match parse_frame(&buf) { + Ok(Some((frame, consumed))) => { + let _ = buf.split_to(consumed); + return frame; + } + Ok(None) => continue, + Err(e) => panic!("protocol error: {e}"), + } + } +} + +/// Verifies that PING, SET, and GET work correctly over a TLS connection. +#[tokio::test] +async fn tls_basic_commands() { + let tmp = tempfile::tempdir().unwrap(); + let (cert_path, key_path, cert_der) = generate_test_cert(tmp.path()); + + let server = TestServer::start_with(ServerOptions { + tls_cert_file: Some(cert_path), + tls_key_file: Some(key_path), + ..Default::default() + }); + + let tls_port = server + .tls_port + .expect("server should have a TLS port after start_with with TLS options"); + + let config = Arc::new(client_config(cert_der)); + let connector = TlsConnector::from(config); + let server_name = ServerName::try_from("localhost").expect("invalid server name"); + + let tcp = tokio::net::TcpStream::connect(format!("127.0.0.1:{tls_port}")) + .await + .unwrap_or_else(|e| panic!("failed to connect to TLS port {tls_port}: {e}")); + + let mut stream = connector + .connect(server_name, tcp) + .await + .expect("TLS handshake failed"); + + // basic roundtrip: PING + let pong = tls_cmd(&mut stream, &["PING"]).await; + assert!( + matches!(pong, Frame::Simple(ref s) if s == "PONG"), + "expected PONG, got {pong:?}" + ); + + // write and read back a value + let ok = tls_cmd(&mut stream, &["SET", "tls:key", "hello"]).await; + assert!( + matches!(ok, Frame::Simple(ref s) if s == "OK"), + "expected OK from SET, got {ok:?}" + ); + + let val = tls_cmd(&mut stream, &["GET", "tls:key"]).await; + assert!( + matches!(val, Frame::Bulk(ref b) if b.as_ref() == b"hello"), + "expected 'hello' from GET, got {val:?}" + ); +} + +/// Verifies that a plain TCP connection to the TLS port is rejected cleanly +/// (the TLS handshake never completes, connection is closed). +#[tokio::test] +async fn plain_tcp_rejected_on_tls_port() { + let tmp = tempfile::tempdir().unwrap(); + let (cert_path, key_path, _cert_der) = generate_test_cert(tmp.path()); + + let server = TestServer::start_with(ServerOptions { + tls_cert_file: Some(cert_path), + tls_key_file: Some(key_path), + ..Default::default() + }); + + let tls_port = server.tls_port.unwrap(); + + // connect without TLS and attempt a plain PING — the server should + // close the connection when it sees non-TLS bytes, or return garbage. + let mut stream = tokio::net::TcpStream::connect(format!("127.0.0.1:{tls_port}")) + .await + .unwrap(); + + let mut out = BytesMut::new(); + Frame::Array(vec![Frame::Bulk(Bytes::from_static(b"PING"))]).serialize(&mut out); + let _ = stream.write_all(&out).await; + + // read whatever comes back — the server either closes the connection + // outright (0 bytes) or sends a TLS alert. either way it must NOT + // respond with a valid RESP3 PONG. + let mut buf = vec![0u8; 128]; + match stream.read(&mut buf).await { + Ok(0) | Err(_) => {} // closed — expected + Ok(n) => { + // got bytes back (likely a TLS alert record); it must not be PONG + assert!( + !String::from_utf8_lossy(&buf[..n]).contains("PONG"), + "server sent PONG on TLS port for a plain TCP connection" + ); + } + } +}