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/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);