diff --git a/crates/ember-core/src/concurrent.rs b/crates/ember-core/src/concurrent.rs index 40bf2a8f..d0e3a6d4 100644 --- a/crates/ember-core/src/concurrent.rs +++ b/crates/ember-core/src/concurrent.rs @@ -9,10 +9,50 @@ use std::time::Duration; use bytes::Bytes; use dashmap::DashMap; -use crate::keyspace::{EvictionPolicy, TtlResult}; +use crate::keyspace::{format_float, EvictionPolicy, TtlResult}; use crate::memory; use crate::time; +/// Errors from integer/float operations on the concurrent keyspace. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ConcurrentOpError { + /// Value cannot be parsed as a number. + NotAnInteger, + /// Increment or decrement would overflow i64. + Overflow, +} + +impl std::fmt::Display for ConcurrentOpError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::NotAnInteger => write!(f, "ERR value is not an integer or out of range"), + Self::Overflow => write!(f, "ERR increment or decrement would overflow"), + } + } +} + +impl std::error::Error for ConcurrentOpError {} + +/// Errors from float operations on the concurrent keyspace. +#[derive(Debug, Clone, PartialEq)] +pub enum ConcurrentFloatError { + /// Value cannot be parsed as a float. + NotAFloat, + /// Result would be NaN or Infinity. + NanOrInfinity, +} + +impl std::fmt::Display for ConcurrentFloatError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::NotAFloat => write!(f, "ERR value is not a valid float"), + Self::NanOrInfinity => write!(f, "ERR increment would produce NaN or Infinity"), + } + } +} + +impl std::error::Error for ConcurrentFloatError {} + /// An entry in the concurrent keyspace. /// Optimized for memory: 40 bytes (down from 56). #[derive(Debug, Clone)] @@ -176,6 +216,299 @@ impl ConcurrentKeyspace { } } + /// Increments the integer value of a key by 1. + /// If the key doesn't exist, it's initialized to 0 before incrementing. + pub fn incr(&self, key: &str) -> Result { + self.incr_by(key, 1) + } + + /// Decrements the integer value of a key by 1. + /// If the key doesn't exist, it's initialized to 0 before decrementing. + pub fn decr(&self, key: &str) -> Result { + self.incr_by(key, -1) + } + + /// Adds `delta` to the integer value of a key, creating it if missing. + /// Preserves existing TTL when updating. + pub fn incr_by(&self, key: &str, delta: i64) -> Result { + self.ops_count.fetch_add(1, Ordering::Relaxed); + + // try to update in-place via get_mut + if let Some(mut entry) = self.data.get_mut(key) { + if entry.is_expired() { + let key_len = entry.key().len(); + let old_size = entry.size(key_len); + drop(entry); + if self.data.remove(key).is_some() { + self.memory_used.fetch_sub(old_size, Ordering::Relaxed); + } + // treat as missing — fall through to insert below + } else { + let s = std::str::from_utf8(&entry.value) + .map_err(|_| ConcurrentOpError::NotAnInteger)?; + let current: i64 = s.parse().map_err(|_| ConcurrentOpError::NotAnInteger)?; + let new_val = current + .checked_add(delta) + .ok_or(ConcurrentOpError::Overflow)?; + let new_bytes = Bytes::from(new_val.to_string()); + + let key_len = entry.key().len(); + let old_size = entry.size(key_len); + entry.value = new_bytes; + let new_size = entry.size(key_len); + let diff = new_size as isize - old_size as isize; + if diff > 0 { + self.memory_used.fetch_add(diff as usize, Ordering::Relaxed); + } else if diff < 0 { + self.memory_used + .fetch_sub((-diff) as usize, Ordering::Relaxed); + } + return Ok(new_val); + } + } + + // key doesn't exist — treat as 0 + let new_val = (0i64) + .checked_add(delta) + .ok_or(ConcurrentOpError::Overflow)?; + self.set(key.to_owned(), Bytes::from(new_val.to_string()), None); + Ok(new_val) + } + + /// Adds `delta` to the float value of a key, creating it if missing. + /// Preserves existing TTL when updating. + pub fn incr_by_float(&self, key: &str, delta: f64) -> Result { + self.ops_count.fetch_add(1, Ordering::Relaxed); + + if let Some(mut entry) = self.data.get_mut(key) { + if entry.is_expired() { + let key_len = entry.key().len(); + let old_size = entry.size(key_len); + drop(entry); + if self.data.remove(key).is_some() { + self.memory_used.fetch_sub(old_size, Ordering::Relaxed); + } + } else { + let s = std::str::from_utf8(&entry.value) + .map_err(|_| ConcurrentFloatError::NotAFloat)?; + let current: f64 = s.parse().map_err(|_| ConcurrentFloatError::NotAFloat)?; + let new_val = current + delta; + if new_val.is_nan() || new_val.is_infinite() { + return Err(ConcurrentFloatError::NanOrInfinity); + } + let new_bytes = Bytes::from(format_float(new_val)); + + let key_len = entry.key().len(); + let old_size = entry.size(key_len); + entry.value = new_bytes; + let new_size = entry.size(key_len); + let diff = new_size as isize - old_size as isize; + if diff > 0 { + self.memory_used.fetch_add(diff as usize, Ordering::Relaxed); + } else if diff < 0 { + self.memory_used + .fetch_sub((-diff) as usize, Ordering::Relaxed); + } + return Ok(new_val); + } + } + + // key doesn't exist — treat as 0.0 + let new_val = delta; + if new_val.is_nan() || new_val.is_infinite() { + return Err(ConcurrentFloatError::NanOrInfinity); + } + self.set(key.to_owned(), Bytes::from(format_float(new_val)), None); + Ok(new_val) + } + + /// Appends a value to an existing string key, or creates a new key. + /// Returns the new string length. + pub fn append(&self, key: &str, suffix: &[u8]) -> usize { + self.ops_count.fetch_add(1, Ordering::Relaxed); + + if let Some(mut entry) = self.data.get_mut(key) { + if !entry.is_expired() { + let mut new_data = Vec::with_capacity(entry.value.len() + suffix.len()); + new_data.extend_from_slice(&entry.value); + new_data.extend_from_slice(suffix); + let new_len = new_data.len(); + + let key_len = entry.key().len(); + let old_size = entry.size(key_len); + entry.value = Bytes::from(new_data); + let new_size = entry.size(key_len); + let diff = new_size as isize - old_size as isize; + if diff > 0 { + self.memory_used.fetch_add(diff as usize, Ordering::Relaxed); + } else if diff < 0 { + self.memory_used + .fetch_sub((-diff) as usize, Ordering::Relaxed); + } + return new_len; + } + // expired — remove and fall through to create + let key_len = entry.key().len(); + let old_size = entry.size(key_len); + drop(entry); + if self.data.remove(key).is_some() { + self.memory_used.fetch_sub(old_size, Ordering::Relaxed); + } + } + + // key doesn't exist — create with just the suffix + let new_len = suffix.len(); + self.set(key.to_owned(), Bytes::copy_from_slice(suffix), None); + new_len + } + + /// Returns the length of the string value stored at key. + /// Returns 0 if the key doesn't exist. + pub fn strlen(&self, key: &str) -> usize { + match self.get(key) { + Some(data) => data.len(), + None => 0, + } + } + + /// Removes the expiration from a key. + /// Returns true if the timeout was successfully removed. + pub fn persist(&self, key: &str) -> bool { + self.ops_count.fetch_add(1, Ordering::Relaxed); + + if let Some(mut entry) = self.data.get_mut(key) { + if entry.is_expired() { + return false; + } + if entry.expires_at_ms != 0 { + entry.expires_at_ms = 0; + true + } else { + false + } + } else { + false + } + } + + /// Sets expiration in milliseconds on an existing key. + pub fn pexpire(&self, key: &str, millis: u64) -> bool { + self.ops_count.fetch_add(1, Ordering::Relaxed); + + if let Some(mut entry) = self.data.get_mut(key) { + if entry.is_expired() { + return false; + } + entry.expires_at_ms = time::now_ms().saturating_add(millis); + true + } else { + false + } + } + + /// Returns the remaining TTL in milliseconds. + pub fn pttl(&self, key: &str) -> TtlResult { + match self.data.get(key) { + None => TtlResult::NotFound, + Some(entry) => { + if entry.is_expired() { + TtlResult::NotFound + } else { + match time::remaining_ms(entry.expires_at_ms) { + None => TtlResult::NoExpiry, + Some(ms) => TtlResult::Milliseconds(ms), + } + } + } + } + } + + /// Returns all keys matching a glob pattern. + pub fn keys(&self, pattern: &str) -> Vec { + let len = self.data.len(); + if len > 10_000 { + tracing::warn!( + key_count = len, + "KEYS on large keyspace, consider SCAN instead" + ); + } + self.data + .iter() + .filter(|entry| !entry.value().is_expired()) + .filter(|entry| crate::keyspace::glob_match(pattern, entry.key())) + .map(|entry| entry.key().to_string()) + .collect() + } + + /// Iterates keys using a cursor. Returns (next_cursor, keys). + /// A next_cursor of 0 means the iteration is complete. + pub fn scan_keys( + &self, + cursor: u64, + count: usize, + pattern: Option<&str>, + ) -> (u64, Vec) { + let target_count = if count == 0 { 10 } else { count }; + let mut keys = Vec::with_capacity(target_count); + let mut position = 0u64; + + for entry in self.data.iter() { + if entry.value().is_expired() { + continue; + } + if position < cursor { + position += 1; + continue; + } + if let Some(pat) = pattern { + if !crate::keyspace::glob_match(pat, entry.key()) { + position += 1; + continue; + } + } + keys.push(entry.key().to_string()); + position += 1; + if keys.len() >= target_count { + // there may be more keys — return position as next cursor + return (position, keys); + } + } + + // iteration complete + (0, keys) + } + + /// Renames a key. Returns true if the source key existed. + pub fn rename(&self, key: &str, newkey: &str) -> Result<(), &'static str> { + self.ops_count.fetch_add(1, Ordering::Relaxed); + + let (_, entry) = self.data.remove(key).ok_or("ERR no such key")?; + + if entry.is_expired() { + let size = entry.size(key.len()); + self.memory_used.fetch_sub(size, Ordering::Relaxed); + return Err("ERR no such key"); + } + + // remove destination if it exists + if let Some((k, old_dest)) = self.data.remove(newkey) { + self.memory_used + .fetch_sub(old_dest.size(k.len()), Ordering::Relaxed); + } + + // adjust memory: old key removed, new key added + let old_key_len = key.len(); + let new_key_len = newkey.len(); + let old_mem = old_key_len + entry.value.len() + 48; + let new_mem = new_key_len + entry.value.len() + 48; + + self.memory_used.fetch_sub(old_mem, Ordering::Relaxed); + self.memory_used.fetch_add(new_mem, Ordering::Relaxed); + + self.data.insert(newkey.into(), entry); + Ok(()) + } + /// Returns the number of keys. pub fn len(&self) -> usize { self.data.len() @@ -285,6 +618,237 @@ mod tests { assert_eq!(ks.get("key"), None); } + #[test] + fn incr_new_key() { + let ks = ConcurrentKeyspace::default(); + assert_eq!(ks.incr("counter").unwrap(), 1); + assert_eq!(ks.get("counter"), Some(Bytes::from("1"))); + } + + #[test] + fn incr_existing_key() { + let ks = ConcurrentKeyspace::default(); + ks.set("counter".into(), Bytes::from("10"), None); + assert_eq!(ks.incr("counter").unwrap(), 11); + } + + #[test] + fn decr_below_zero() { + let ks = ConcurrentKeyspace::default(); + assert_eq!(ks.decr("counter").unwrap(), -1); + assert_eq!(ks.decr("counter").unwrap(), -2); + } + + #[test] + fn incr_by_delta() { + let ks = ConcurrentKeyspace::default(); + assert_eq!(ks.incr_by("counter", 5).unwrap(), 5); + assert_eq!(ks.incr_by("counter", -3).unwrap(), 2); + } + + #[test] + fn incr_non_integer_value() { + let ks = ConcurrentKeyspace::default(); + ks.set("key".into(), Bytes::from("not_a_number"), None); + assert_eq!(ks.incr("key"), Err(ConcurrentOpError::NotAnInteger)); + } + + #[test] + fn incr_overflow() { + let ks = ConcurrentKeyspace::default(); + ks.set("key".into(), Bytes::from(i64::MAX.to_string()), None); + assert_eq!(ks.incr("key"), Err(ConcurrentOpError::Overflow)); + } + + #[test] + fn incr_by_float_new_key() { + let ks = ConcurrentKeyspace::default(); + let val = ks.incr_by_float("key", 2.5).unwrap(); + assert!((val - 2.5).abs() < f64::EPSILON); + } + + #[test] + fn incr_by_float_existing() { + let ks = ConcurrentKeyspace::default(); + ks.set("key".into(), Bytes::from("10.5"), None); + let val = ks.incr_by_float("key", 1.5).unwrap(); + assert!((val - 12.0).abs() < f64::EPSILON); + } + + #[test] + fn incr_by_float_not_a_float() { + let ks = ConcurrentKeyspace::default(); + ks.set("key".into(), Bytes::from("hello"), None); + assert_eq!( + ks.incr_by_float("key", 1.0), + Err(ConcurrentFloatError::NotAFloat) + ); + } + + #[test] + fn incr_by_float_infinity() { + let ks = ConcurrentKeyspace::default(); + ks.set("key".into(), Bytes::from(f64::MAX.to_string()), None); + assert_eq!( + ks.incr_by_float("key", f64::MAX), + Err(ConcurrentFloatError::NanOrInfinity) + ); + } + + #[test] + fn append_new_key() { + let ks = ConcurrentKeyspace::default(); + assert_eq!(ks.append("key", b"hello"), 5); + assert_eq!(ks.get("key"), Some(Bytes::from("hello"))); + } + + #[test] + fn append_existing_key() { + let ks = ConcurrentKeyspace::default(); + ks.set("key".into(), Bytes::from("hello"), None); + assert_eq!(ks.append("key", b" world"), 11); + assert_eq!(ks.get("key"), Some(Bytes::from("hello world"))); + } + + #[test] + fn strlen_existing() { + let ks = ConcurrentKeyspace::default(); + ks.set("key".into(), Bytes::from("hello"), None); + assert_eq!(ks.strlen("key"), 5); + } + + #[test] + fn strlen_missing() { + let ks = ConcurrentKeyspace::default(); + assert_eq!(ks.strlen("missing"), 0); + } + + #[test] + fn persist_removes_ttl() { + let ks = ConcurrentKeyspace::default(); + ks.set("key".into(), Bytes::from("val"), Some(Duration::from_secs(60))); + assert!(ks.persist("key")); + assert!(matches!(ks.ttl("key"), TtlResult::NoExpiry)); + } + + #[test] + fn persist_no_ttl() { + let ks = ConcurrentKeyspace::default(); + ks.set("key".into(), Bytes::from("val"), None); + assert!(!ks.persist("key")); // no TTL to remove + } + + #[test] + fn persist_missing_key() { + let ks = ConcurrentKeyspace::default(); + assert!(!ks.persist("missing")); + } + + #[test] + fn pexpire_and_pttl() { + let ks = ConcurrentKeyspace::default(); + ks.set("key".into(), Bytes::from("val"), None); + assert!(ks.pexpire("key", 5000)); + match ks.pttl("key") { + TtlResult::Milliseconds(ms) => assert!(ms > 0 && ms <= 5000), + other => panic!("expected Milliseconds, got {other:?}"), + } + } + + #[test] + fn pttl_no_expiry() { + let ks = ConcurrentKeyspace::default(); + ks.set("key".into(), Bytes::from("val"), None); + assert!(matches!(ks.pttl("key"), TtlResult::NoExpiry)); + } + + #[test] + fn pttl_missing() { + let ks = ConcurrentKeyspace::default(); + assert!(matches!(ks.pttl("missing"), TtlResult::NotFound)); + } + + #[test] + fn keys_match_pattern() { + let ks = ConcurrentKeyspace::default(); + ks.set("user:1".into(), Bytes::from("a"), None); + ks.set("user:2".into(), Bytes::from("b"), None); + ks.set("item:1".into(), Bytes::from("c"), None); + let mut result = ks.keys("user:*"); + result.sort(); + assert_eq!(result, vec!["user:1", "user:2"]); + } + + #[test] + fn keys_match_all() { + let ks = ConcurrentKeyspace::default(); + ks.set("a".into(), Bytes::from("1"), None); + ks.set("b".into(), Bytes::from("2"), None); + let result = ks.keys("*"); + assert_eq!(result.len(), 2); + } + + #[test] + fn scan_basic() { + let ks = ConcurrentKeyspace::default(); + ks.set("a".into(), Bytes::from("1"), None); + ks.set("b".into(), Bytes::from("2"), None); + ks.set("c".into(), Bytes::from("3"), None); + let (cursor, keys) = ks.scan_keys(0, 10, None); + assert_eq!(cursor, 0); // complete in one pass + assert_eq!(keys.len(), 3); + } + + #[test] + fn scan_with_pattern() { + let ks = ConcurrentKeyspace::default(); + ks.set("user:1".into(), Bytes::from("a"), None); + ks.set("user:2".into(), Bytes::from("b"), None); + ks.set("item:1".into(), Bytes::from("c"), None); + let (_, keys) = ks.scan_keys(0, 10, Some("user:*")); + assert_eq!(keys.len(), 2); + } + + #[test] + fn scan_with_count() { + let ks = ConcurrentKeyspace::default(); + for i in 0..10 { + ks.set(format!("k{i}"), Bytes::from("v"), None); + } + let (cursor, keys) = ks.scan_keys(0, 3, None); + assert!(keys.len() <= 3); + // if cursor > 0, there are more keys + if cursor > 0 { + let (_, keys2) = ks.scan_keys(cursor, 3, None); + assert!(!keys2.is_empty()); + } + } + + #[test] + fn rename_basic() { + let ks = ConcurrentKeyspace::default(); + ks.set("old".into(), Bytes::from("value"), None); + ks.rename("old", "new").unwrap(); + assert_eq!(ks.get("old"), None); + assert_eq!(ks.get("new"), Some(Bytes::from("value"))); + } + + #[test] + fn rename_missing_key() { + let ks = ConcurrentKeyspace::default(); + assert!(ks.rename("missing", "new").is_err()); + } + + #[test] + fn rename_overwrites_destination() { + let ks = ConcurrentKeyspace::default(); + ks.set("src".into(), Bytes::from("new_val"), None); + ks.set("dst".into(), Bytes::from("old_val"), None); + ks.rename("src", "dst").unwrap(); + assert_eq!(ks.get("src"), None); + assert_eq!(ks.get("dst"), Some(Bytes::from("new_val"))); + } + #[test] fn concurrent_access() { use std::sync::Arc; diff --git a/crates/ember-core/src/keyspace.rs b/crates/ember-core/src/keyspace.rs index 4dc40a92..1771f74c 100644 --- a/crates/ember-core/src/keyspace.rs +++ b/crates/ember-core/src/keyspace.rs @@ -2184,7 +2184,7 @@ impl Default for Keyspace { /// /// Uses up to 17 significant digits and strips unnecessary trailing zeros, /// but always keeps at least one decimal place for non-integer results. -fn format_float(val: f64) -> String { +pub(crate) fn format_float(val: f64) -> String { if val == 0.0 { return "0".into(); } @@ -2212,7 +2212,7 @@ fn format_float(val: f64) -> String { /// /// Uses an iterative two-pointer algorithm with backtracking for O(n*m) /// worst-case performance, where n is pattern length and m is text length. -fn glob_match(pattern: &str, text: &str) -> bool { +pub(crate) fn glob_match(pattern: &str, text: &str) -> bool { let pat: Vec = pattern.chars().collect(); let txt: Vec = text.chars().collect(); diff --git a/crates/ember-core/src/lib.rs b/crates/ember-core/src/lib.rs index ce16145f..254c234e 100644 --- a/crates/ember-core/src/lib.rs +++ b/crates/ember-core/src/lib.rs @@ -18,7 +18,7 @@ pub mod types; #[cfg(feature = "protobuf")] pub mod schema; -pub use concurrent::ConcurrentKeyspace; +pub use concurrent::{ConcurrentFloatError, ConcurrentKeyspace, ConcurrentOpError}; pub use engine::{Engine, EngineConfig}; pub use error::ShardError; pub use keyspace::{ diff --git a/crates/ember-server/src/concurrent_handler.rs b/crates/ember-server/src/concurrent_handler.rs index 79f247bb..81122eb9 100644 --- a/crates/ember-server/src/concurrent_handler.rs +++ b/crates/ember-server/src/concurrent_handler.rs @@ -16,7 +16,6 @@ use std::sync::atomic::Ordering; use std::sync::Arc; use std::time::{Duration, Instant}; -#[cfg(feature = "protobuf")] use bytes::Bytes; use bytes::BytesMut; use ember_core::{ConcurrentKeyspace, Engine, TtlResult}; @@ -233,6 +232,56 @@ async fn execute_concurrent( TtlResult::Milliseconds(ms) => Frame::Integer((ms / 1000) as i64), }, + Command::Incr { key } => match keyspace.incr(&key) { + Ok(val) => Frame::Integer(val), + Err(e) => Frame::Error(e.to_string()), + }, + + Command::Decr { key } => match keyspace.decr(&key) { + Ok(val) => Frame::Integer(val), + Err(e) => Frame::Error(e.to_string()), + }, + + Command::IncrBy { key, delta } => match keyspace.incr_by(&key, delta) { + Ok(val) => Frame::Integer(val), + Err(e) => Frame::Error(e.to_string()), + }, + + Command::DecrBy { key, delta } => match keyspace.incr_by(&key, -delta) { + Ok(val) => Frame::Integer(val), + Err(e) => Frame::Error(e.to_string()), + }, + + Command::IncrByFloat { key, delta } => match keyspace.incr_by_float(&key, delta) { + Ok(val) => Frame::Bulk(Bytes::from(val.to_string())), + Err(e) => Frame::Error(e.to_string()), + }, + + Command::Append { key, value } => { + let new_len = keyspace.append(&key, &value); + Frame::Integer(new_len as i64) + } + + Command::Strlen { key } => Frame::Integer(keyspace.strlen(&key) as i64), + + Command::Persist { key } => Frame::Integer(if keyspace.persist(&key) { 1 } else { 0 }), + + Command::Pexpire { + key, + milliseconds, + } => Frame::Integer(if keyspace.pexpire(&key, milliseconds) { + 1 + } else { + 0 + }), + + Command::Pttl { key } => match keyspace.pttl(&key) { + TtlResult::Milliseconds(ms) => Frame::Integer(ms as i64), + TtlResult::NoExpiry => Frame::Integer(-1), + TtlResult::NotFound => Frame::Integer(-2), + TtlResult::Seconds(s) => Frame::Integer(s as i64 * 1000), + }, + Command::Ping(None) => Frame::Simple("PONG".into()), Command::Ping(Some(msg)) => Frame::Bulk(msg), Command::Echo(msg) => Frame::Bulk(msg), @@ -244,6 +293,68 @@ async fn execute_concurrent( Frame::Simple("OK".into()) } + Command::MGet { keys } => { + let mut frames = Vec::with_capacity(keys.len()); + for key in keys { + match keyspace.get(&key) { + Some(data) => frames.push(Frame::Bulk(data)), + None => frames.push(Frame::Null), + } + } + Frame::Array(frames) + } + + Command::MSet { pairs } => { + for (key, value) in pairs { + keyspace.set(key, value, None); + } + Frame::Simple("OK".into()) + } + + Command::Type { key } => { + if keyspace.exists(&key) { + Frame::Simple("string".into()) + } else { + Frame::Simple("none".into()) + } + } + + Command::Keys { pattern } => { + let matched = keyspace.keys(&pattern); + Frame::Array( + matched + .into_iter() + .map(|k| Frame::Bulk(Bytes::from(k))) + .collect(), + ) + } + + Command::Scan { + cursor, + pattern, + count, + } => { + let (next_cursor, keys) = + keyspace.scan_keys(cursor, count.unwrap_or(10), pattern.as_deref()); + Frame::Array(vec![ + Frame::Bulk(Bytes::from(next_cursor.to_string())), + Frame::Array( + keys.into_iter() + .map(|k| Frame::Bulk(Bytes::from(k))) + .collect(), + ), + ]) + } + + Command::Rename { key, newkey } => match keyspace.rename(&key, &newkey) { + Ok(()) => Frame::Simple("OK".into()), + Err(msg) => Frame::Error(msg.into()), + }, + + Command::Info { section } => { + render_concurrent_info(keyspace, ctx, section.as_deref()) + } + // -- pub/sub -- Command::Publish { channel, message } => { let count = pubsub.publish(&channel, message); @@ -566,3 +677,75 @@ async fn execute_concurrent( _ => Frame::Error("ERR command not supported in concurrent mode".into()), } } + +/// Renders INFO output for the concurrent keyspace. +fn render_concurrent_info( + keyspace: &ConcurrentKeyspace, + ctx: &ServerContext, + section: Option<&str>, +) -> Frame { + let section_upper = section.map(|s| s.to_ascii_uppercase()); + let want_all = section_upper.is_none(); + let want = |name: &str| want_all || section_upper.as_deref() == Some(name); + + let mut out = String::with_capacity(512); + + if want("SERVER") { + let uptime = ctx.start_time.elapsed().as_secs(); + out.push_str("# Server\r\n"); + out.push_str(&format!("ember_version:{}\r\n", ctx.version)); + out.push_str(&format!("process_id:{}\r\n", std::process::id())); + out.push_str(&format!("uptime_in_seconds:{uptime}\r\n")); + out.push_str("mode:concurrent\r\n"); + out.push_str("\r\n"); + } + + if want("CLIENTS") { + let connected = ctx.connections_active.load(Ordering::Relaxed); + out.push_str("# Clients\r\n"); + out.push_str(&format!("connected_clients:{connected}\r\n")); + out.push_str(&format!("max_clients:{}\r\n", ctx.max_connections)); + out.push_str("\r\n"); + } + + if want("MEMORY") { + let used = keyspace.memory_used(); + out.push_str("# Memory\r\n"); + out.push_str(&format!("used_memory:{used}\r\n")); + if let Some(max) = ctx.max_memory { + out.push_str(&format!("max_memory:{max}\r\n")); + } else { + out.push_str("max_memory:0\r\n"); + } + out.push_str("\r\n"); + } + + if want("STATS") { + let total_cmds = ctx.commands_processed.load(Ordering::Relaxed); + let total_conns = ctx.connections_accepted.load(Ordering::Relaxed); + out.push_str("# Stats\r\n"); + out.push_str(&format!("total_connections_received:{total_conns}\r\n")); + out.push_str(&format!("total_commands_processed:{total_cmds}\r\n")); + out.push_str("\r\n"); + } + + if want("KEYSPACE") { + let key_count = keyspace.len(); + out.push_str("# Keyspace\r\n"); + if key_count > 0 { + out.push_str(&format!( + "db0:keys={},used_bytes={}\r\n", + key_count, + keyspace.memory_used() + )); + } + out.push_str("\r\n"); + } + + // trim trailing blank line + if out.ends_with("\r\n\r\n") { + out.truncate(out.len() - 2); + } + + Frame::Bulk(Bytes::from(out)) +}