diff --git a/crates/ember-core/src/keyspace/list.rs b/crates/ember-core/src/keyspace/list.rs index d33456dd..1aa8da06 100644 --- a/crates/ember-core/src/keyspace/list.rs +++ b/crates/ember-core/src/keyspace/list.rs @@ -502,6 +502,45 @@ impl Keyspace { Ok(len) } + /// Atomically pops a value from `source` and pushes it to `destination`. + /// + /// `src_left` controls whether to pop from the head (true) or tail (false) + /// of the source. `dst_left` controls whether to push to the head (true) + /// or tail (false) of the destination. Source and destination may be the + /// same key (rotate). Returns `Ok(None)` if the source list is missing. + pub fn lmove( + &mut self, + source: &str, + destination: &str, + src_left: bool, + dst_left: bool, + ) -> Result, WriteError> { + if self.remove_if_expired(source) { + return Ok(None); + } + match self.entries.get(source) { + None => return Ok(None), + Some(e) if !matches!(e.value, Value::List(_)) => return Err(WriteError::WrongType), + _ => {} + } + + if source != destination { + self.remove_if_expired(destination); + if let Some(e) = self.entries.get(destination) { + if !matches!(e.value, Value::List(_)) { + return Err(WriteError::WrongType); + } + } + } + + let popped = self.list_pop(source, src_left).map_err(WriteError::from)?; + let Some(value) = popped else { + return Ok(None); + }; + self.list_push(destination, &[value.clone()], dst_left)?; + Ok(Some(value)) + } + /// Internal pop implementation shared by lpop/rpop. pub(super) fn list_pop(&mut self, key: &str, left: bool) -> Result, WrongType> { if self.remove_if_expired(key) { @@ -1238,4 +1277,44 @@ mod tests { ks.set("s".into(), Bytes::from("val"), None, false, false); assert!(ks.lpos("s", b"a", 1, 1, 0).is_err()); } + + #[test] + fn lmove_left_to_right() { + let mut ks = Keyspace::new(); + ks.rpush("src", &[Bytes::from("a"), Bytes::from("b"), Bytes::from("c")]) + .unwrap(); + let moved = ks.lmove("src", "dst", true, false).unwrap(); + assert_eq!(moved, Some(Bytes::from("a"))); + // src should now be [b, c] + assert_eq!(ks.lrange("src", 0, -1).unwrap(), vec![Bytes::from("b"), Bytes::from("c")]); + // dst should be [a] + assert_eq!(ks.lrange("dst", 0, -1).unwrap(), vec![Bytes::from("a")]); + } + + #[test] + fn lmove_rotate_same_key() { + let mut ks = Keyspace::new(); + ks.rpush("q", &[Bytes::from("1"), Bytes::from("2"), Bytes::from("3")]) + .unwrap(); + // rotate: pop from left, push to right + let moved = ks.lmove("q", "q", true, false).unwrap(); + assert_eq!(moved, Some(Bytes::from("1"))); + let items = ks.lrange("q", 0, -1).unwrap(); + assert_eq!(items, vec![Bytes::from("2"), Bytes::from("3"), Bytes::from("1")]); + } + + #[test] + fn lmove_missing_source() { + let mut ks = Keyspace::new(); + let moved = ks.lmove("missing", "dst", true, true).unwrap(); + assert_eq!(moved, None); + assert!(!ks.exists("dst")); + } + + #[test] + fn lmove_wrong_type_returns_error() { + let mut ks = Keyspace::new(); + ks.set("s".into(), Bytes::from("hello"), None, false, false); + assert!(ks.lmove("s", "dst", true, true).is_err()); + } } diff --git a/crates/ember-core/src/keyspace/string.rs b/crates/ember-core/src/keyspace/string.rs index 7e53e1d2..cfa6368a 100644 --- a/crates/ember-core/src/keyspace/string.rs +++ b/crates/ember-core/src/keyspace/string.rs @@ -231,6 +231,81 @@ impl Keyspace { } } + /// Returns the value of a key and deletes it atomically. + /// + /// Returns `Ok(None)` if the key does not exist or has expired. + /// Returns `Err(WrongType)` if the key holds a non-string value. + pub fn getdel(&mut self, key: &str) -> Result, WrongType> { + if self.remove_if_expired(key) { + return Ok(None); + } + match self.entries.get(key) { + None => return Ok(None), + Some(e) if !matches!(e.value, Value::String(_)) => return Err(WrongType), + _ => {} + } + let entry = self.entries.remove(key).expect("verified above"); + let bytes = match entry.value { + Value::String(ref b) => b.clone(), + _ => unreachable!("type checked above"), + }; + self.memory + .remove_with_size(entry.cached_value_size as usize + key.len() + memory::ENTRY_OVERHEAD); + self.decrement_expiry_if_set(&entry); + self.remove_version(key); + Ok(Some(bytes)) + } + + /// Returns the value of a key and optionally updates its expiry. + /// + /// `expire` controls what happens to the TTL: + /// - `None` — leave the TTL unchanged (plain GET semantics) + /// - `Some(None)` — remove the TTL (PERSIST semantics) + /// - `Some(Some(duration))` — set a new TTL from now + /// + /// Returns `Ok(None)` if the key does not exist or has expired. + /// Returns `Err(WrongType)` if the key holds a non-string value. + pub fn getex( + &mut self, + key: &str, + expire: Option>, + ) -> Result, WrongType> { + if self.remove_if_expired(key) { + return Ok(None); + } + + let (bytes, had_expiry) = match self.entries.get(key) { + None => return Ok(None), + Some(e) => match &e.value { + Value::String(b) => (b.clone(), e.expires_at_ms != 0), + _ => return Err(WrongType), + }, + }; + + if let Some(new_expire) = expire { + let entry = self.entries.get_mut(key).expect("verified above"); + match new_expire { + Some(duration) => { + entry.expires_at_ms = + time::now_ms().saturating_add(duration.as_millis() as u64); + if !had_expiry { + self.expiry_count += 1; + } + } + None => { + entry.expires_at_ms = 0; + if had_expiry { + self.expiry_count = self.expiry_count.saturating_sub(1); + } + } + } + entry.touch(self.track_access); + self.bump_version(key); + } + + Ok(Some(bytes)) + } + /// Appends a value to an existing string key, or creates a new key if /// it doesn't exist. Returns the new string length. pub fn append(&mut self, key: &str, value: &[u8]) -> Result { @@ -893,4 +968,65 @@ mod tests { assert_eq!(super::format_float(2.72), "2.72"); assert_eq!(super::format_float(10.5), "10.5"); } + + #[test] + fn getdel_returns_value_and_removes_key() { + let mut ks = Keyspace::new(); + ks.set("k".into(), Bytes::from("hello"), None, false, false); + let val = ks.getdel("k").unwrap(); + assert_eq!(val, Some(Bytes::from("hello"))); + assert!(!ks.exists("k")); + assert_eq!(ks.stats().used_bytes, 0); + } + + #[test] + fn getdel_missing_key_returns_none() { + let mut ks = Keyspace::new(); + assert_eq!(ks.getdel("nope").unwrap(), None); + } + + #[test] + fn getdel_wrong_type_returns_error() { + let mut ks = Keyspace::new(); + ks.zadd("z", &[(1.0, "a".into())], &ZAddFlags::default()) + .unwrap(); + assert!(ks.getdel("z").is_err()); + } + + #[test] + fn getex_no_option_leaves_ttl_unchanged() { + let mut ks = Keyspace::new(); + ks.set("k".into(), Bytes::from("v"), Some(Duration::from_secs(60)), false, false); + let ttl_before = ks.ttl("k"); + let val = ks.getex("k", None).unwrap(); + assert_eq!(val, Some(Bytes::from("v"))); + let ttl_after = ks.ttl("k"); + // ttl should still be set (both should be positive) + assert!(matches!(ttl_before, TtlResult::Seconds(_))); + assert!(matches!(ttl_after, TtlResult::Seconds(_))); + } + + #[test] + fn getex_persist_clears_ttl() { + let mut ks = Keyspace::new(); + ks.set("k".into(), Bytes::from("v"), Some(Duration::from_secs(60)), false, false); + let val = ks.getex("k", Some(None)).unwrap(); + assert_eq!(val, Some(Bytes::from("v"))); + assert!(matches!(ks.ttl("k"), TtlResult::NoExpiry)); + } + + #[test] + fn getex_set_new_ttl() { + let mut ks = Keyspace::new(); + ks.set("k".into(), Bytes::from("v"), None, false, false); + let val = ks.getex("k", Some(Some(Duration::from_secs(30)))).unwrap(); + assert_eq!(val, Some(Bytes::from("v"))); + assert!(matches!(ks.ttl("k"), TtlResult::Seconds(_))); + } + + #[test] + fn getex_missing_key_returns_none() { + let mut ks = Keyspace::new(); + assert_eq!(ks.getex("nope", None).unwrap(), None); + } } diff --git a/crates/ember-core/src/keyspace/zset.rs b/crates/ember-core/src/keyspace/zset.rs index 1d4fb6f4..9ff74bfe 100644 --- a/crates/ember-core/src/keyspace/zset.rs +++ b/crates/ember-core/src/keyspace/zset.rs @@ -437,6 +437,120 @@ impl Keyspace { }, } } + + /// Returns members in the first sorted set that are not in any of the others. + /// + /// Missing keys are treated as empty sets. Results are ordered by score + /// then by member name. + pub fn zdiff(&mut self, keys: &[String]) -> Result, WrongType> { + if keys.is_empty() { + return Ok(vec![]); + } + for key in keys { + self.remove_if_expired(key); + } + let first: Vec<(String, f64)> = match self.entries.get(keys[0].as_str()) { + None => return Ok(vec![]), + Some(e) => match &e.value { + Value::SortedSet(ss) => ss.iter().map(|(m, s)| (m.to_owned(), s)).collect(), + _ => return Err(WrongType), + }, + }; + let mut excluded: AHashMap = AHashMap::new(); + for key in &keys[1..] { + match self.entries.get(key.as_str()) { + None => {} + Some(e) => match &e.value { + Value::SortedSet(ss) => { + for (member, _) in ss.iter() { + excluded.insert(member.to_owned(), ()); + } + } + _ => return Err(WrongType), + }, + } + } + Ok(first + .into_iter() + .filter(|(m, _)| !excluded.contains_key(m)) + .collect()) + } + + /// Returns members present in all of the given sorted sets, with scores summed. + /// + /// If any key is missing the result is empty. Results are ordered by + /// score then member name. + pub fn zinter(&mut self, keys: &[String]) -> Result, WrongType> { + if keys.is_empty() { + return Ok(vec![]); + } + for key in keys { + self.remove_if_expired(key); + } + let mut candidates: Vec<(String, f64)> = match self.entries.get(keys[0].as_str()) { + None => return Ok(vec![]), + Some(e) => match &e.value { + Value::SortedSet(ss) => ss.iter().map(|(m, s)| (m.to_owned(), s)).collect(), + _ => return Err(WrongType), + }, + }; + for key in &keys[1..] { + match self.entries.get(key.as_str()) { + None => return Ok(vec![]), + Some(e) => match &e.value { + Value::SortedSet(ss) => { + let lookup: AHashMap = + ss.iter().map(|(m, s)| (m.to_owned(), s)).collect(); + candidates = candidates + .into_iter() + .filter_map(|(m, score)| lookup.get(&m).map(|&s| (m, score + s))) + .collect(); + } + _ => return Err(WrongType), + }, + } + } + candidates.sort_by(|(am, as_), (bm, bs)| { + as_.partial_cmp(bs) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| am.cmp(bm)) + }); + Ok(candidates) + } + + /// Returns the union of all given sorted sets, with scores summed across keys. + /// + /// Missing keys contribute no members. Results are ordered by score then + /// member name. + pub fn zunion(&mut self, keys: &[String]) -> Result, WrongType> { + if keys.is_empty() { + return Ok(vec![]); + } + for key in keys { + self.remove_if_expired(key); + } + let mut totals: AHashMap = AHashMap::new(); + for key in keys { + match self.entries.get(key.as_str()) { + None => {} + Some(e) => match &e.value { + Value::SortedSet(ss) => { + for (member, score) in ss.iter() { + *totals.entry(member.to_owned()).or_insert(0.0) += score; + } + } + _ => return Err(WrongType), + }, + } + } + let mut result: Vec<(String, f64)> = totals.into_iter().collect(); + result.sort_by(|(am, as_), (bm, bs)| { + as_.partial_cmp(bs) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| am.cmp(bm)) + }); + Ok(result) + } } #[cfg(test)] @@ -841,4 +955,71 @@ mod tests { // original key should be untouched assert!(ks.exists("a")); } + + #[test] + fn zdiff_returns_members_unique_to_first() { + let mut ks = Keyspace::new(); + ks.zadd("a", &[(1.0, "x".into()), (2.0, "y".into()), (3.0, "z".into())], &ZAddFlags::default()).unwrap(); + ks.zadd("b", &[(1.0, "y".into()), (1.0, "w".into())], &ZAddFlags::default()).unwrap(); + + let keys = vec!["a".to_owned(), "b".to_owned()]; + let diff = ks.zdiff(&keys).unwrap(); + let members: Vec<&str> = diff.iter().map(|(m, _)| m.as_str()).collect(); + assert!(members.contains(&"x")); + assert!(members.contains(&"z")); + assert!(!members.contains(&"y")); + } + + #[test] + fn zdiff_with_missing_second_key_returns_all() { + let mut ks = Keyspace::new(); + ks.zadd("a", &[(1.0, "x".into())], &ZAddFlags::default()).unwrap(); + let keys = vec!["a".to_owned(), "missing".to_owned()]; + let diff = ks.zdiff(&keys).unwrap(); + assert_eq!(diff.len(), 1); + assert_eq!(diff[0].0, "x"); + } + + #[test] + fn zinter_returns_common_members_with_summed_scores() { + let mut ks = Keyspace::new(); + ks.zadd("a", &[(1.0, "x".into()), (2.0, "y".into())], &ZAddFlags::default()).unwrap(); + ks.zadd("b", &[(3.0, "x".into()), (4.0, "z".into())], &ZAddFlags::default()).unwrap(); + + let keys = vec!["a".to_owned(), "b".to_owned()]; + let inter = ks.zinter(&keys).unwrap(); + assert_eq!(inter.len(), 1); + assert_eq!(inter[0].0, "x"); + assert!((inter[0].1 - 4.0).abs() < f64::EPSILON); // 1.0 + 3.0 + } + + #[test] + fn zinter_empty_when_no_common_members() { + let mut ks = Keyspace::new(); + ks.zadd("a", &[(1.0, "x".into())], &ZAddFlags::default()).unwrap(); + ks.zadd("b", &[(1.0, "y".into())], &ZAddFlags::default()).unwrap(); + let keys = vec!["a".to_owned(), "b".to_owned()]; + assert!(ks.zinter(&keys).unwrap().is_empty()); + } + + #[test] + fn zunion_combines_all_members_with_summed_scores() { + let mut ks = Keyspace::new(); + ks.zadd("a", &[(1.0, "x".into()), (2.0, "y".into())], &ZAddFlags::default()).unwrap(); + ks.zadd("b", &[(3.0, "x".into()), (4.0, "z".into())], &ZAddFlags::default()).unwrap(); + + let keys = vec!["a".to_owned(), "b".to_owned()]; + let union = ks.zunion(&keys).unwrap(); + assert_eq!(union.len(), 3); // x, y, z + let x = union.iter().find(|(m, _)| m == "x").unwrap(); + assert!((x.1 - 4.0).abs() < f64::EPSILON); // 1.0 + 3.0 + } + + #[test] + fn zdiff_wrong_type_returns_error() { + let mut ks = Keyspace::new(); + ks.set("s".into(), Bytes::from("v"), None, false, false); + let keys = vec!["s".to_owned()]; + assert!(ks.zdiff(&keys).is_err()); + } } diff --git a/crates/ember-core/src/shard/aof.rs b/crates/ember-core/src/shard/aof.rs index 4d159254..656dd1ec 100644 --- a/crates/ember-core/src/shard/aof.rs +++ b/crates/ember-core/src/shard/aof.rs @@ -311,6 +311,61 @@ pub(super) fn to_aof_records( (ShardRequest::VRem { key, element }, ShardResponse::Bool(true)) => { smallvec![AofRecord::VRem { key, element }] } + // LMOVE: persist as lpop from source + lpush to destination. + // Only written when a value was actually moved (Value is Some). + ( + ShardRequest::LMove { + source, + destination, + src_left, + dst_left, + }, + ShardResponse::Value(Some(v)), + ) => { + let val = match v { + Value::String(b) => b.clone(), + _ => return SmallVec::new(), + }; + let pop_record = if src_left { + AofRecord::LPop { key: source } + } else { + AofRecord::RPop { key: source } + }; + let push_record = if dst_left { + AofRecord::LPush { + key: destination, + values: vec![val], + } + } else { + AofRecord::RPush { + key: destination, + values: vec![val], + } + }; + smallvec![pop_record, push_record] + } + // GETDEL: persist as DEL when the key actually existed. + (ShardRequest::GetDel { key }, ShardResponse::Value(Some(_))) => { + smallvec![AofRecord::Del { key }] + } + // GETEX with a new TTL: persist as Expire (seconds) or Pexpire (ms). + // PERSIST (expire = Some(None)) is represented as Pexpire with 0. + // No TTL change (expire = None): nothing to persist. + ( + ShardRequest::GetEx { + key, + expire: Some(new_expire), + }, + ShardResponse::Value(Some(_)), + ) => match new_expire { + Some(ms) if ms > 0 => { + smallvec![AofRecord::Pexpire { key, milliseconds: ms }] + } + // PERSIST — expire set to 0 in the keyspace; record as pexpire 0 + // so replay calls persist. We use a negative sentinel to signal + // PERSIST on replay: store as Expire with seconds = 0. + _ => smallvec![AofRecord::Persist { key }], + }, _ => SmallVec::new(), } } diff --git a/crates/ember-core/src/shard/mod.rs b/crates/ember-core/src/shard/mod.rs index b3a2ee45..d5877f7b 100644 --- a/crates/ember-core/src/shard/mod.rs +++ b/crates/ember-core/src/shard/mod.rs @@ -445,6 +445,36 @@ pub enum ShardRequest { key: String, members: Vec, }, + /// LMOVE: atomically pops from source and pushes to destination. + LMove { + source: String, + destination: String, + src_left: bool, + dst_left: bool, + }, + /// GETDEL: returns the value at key and deletes it. + GetDel { + key: String, + }, + /// GETEX: returns the value at key and optionally updates its TTL. + /// + /// `expire`: `None` = no change, `Some(None)` = persist, `Some(Some(ms))` = new TTL in ms. + GetEx { + key: String, + expire: Option>, + }, + /// ZDIFF: returns members in the first sorted set not in the others. + ZDiff { + keys: Vec, + }, + /// ZINTER: returns members present in all sorted sets, scores summed. + ZInter { + keys: Vec, + }, + /// ZUNION: returns the union of all sorted sets, scores summed. + ZUnion { + keys: Vec, + }, /// Returns the key count for this shard. DbSize, /// Returns keyspace stats for this shard. @@ -661,6 +691,9 @@ impl ShardRequest { | ShardRequest::SUnionStore { .. } | ShardRequest::SInterStore { .. } | ShardRequest::SDiffStore { .. } + | ShardRequest::LMove { .. } + | ShardRequest::GetDel { .. } + | ShardRequest::GetEx { .. } | ShardRequest::FlushDb | ShardRequest::FlushDbAsync | ShardRequest::RestoreKey { .. } => true, @@ -1823,6 +1856,44 @@ fn dispatch( Ok(results) => ShardResponse::BoolArray(results), Err(_) => ShardResponse::WrongType, }, + ShardRequest::LMove { + source, + destination, + src_left, + dst_left, + } => match ks.lmove(source, destination, *src_left, *dst_left) { + Ok(Some(v)) => ShardResponse::Value(Some(Value::String(v))), + Ok(None) => ShardResponse::Value(None), + Err(e) => match e { + WriteError::WrongType => ShardResponse::WrongType, + WriteError::OutOfMemory => ShardResponse::OutOfMemory, + }, + }, + ShardRequest::GetDel { key } => match ks.getdel(key) { + Ok(Some(v)) => ShardResponse::Value(Some(Value::String(v))), + Ok(None) => ShardResponse::Value(None), + Err(_) => ShardResponse::WrongType, + }, + ShardRequest::GetEx { key, expire } => { + let dur = expire.map(|opt| opt.map(Duration::from_millis)); + match ks.getex(key, dur) { + Ok(Some(v)) => ShardResponse::Value(Some(Value::String(v))), + Ok(None) => ShardResponse::Value(None), + Err(_) => ShardResponse::WrongType, + } + } + ShardRequest::ZDiff { keys } => match ks.zdiff(keys) { + Ok(pairs) => ShardResponse::ScoredArray(pairs), + Err(_) => ShardResponse::WrongType, + }, + ShardRequest::ZInter { keys } => match ks.zinter(keys) { + Ok(pairs) => ShardResponse::ScoredArray(pairs), + Err(_) => ShardResponse::WrongType, + }, + ShardRequest::ZUnion { keys } => match ks.zunion(keys) { + Ok(pairs) => ShardResponse::ScoredArray(pairs), + Err(_) => ShardResponse::WrongType, + }, ShardRequest::SScan { key, cursor, diff --git a/crates/ember-protocol/src/command/attributes.rs b/crates/ember-protocol/src/command/attributes.rs index ff936c58..0658e3eb 100644 --- a/crates/ember-protocol/src/command/attributes.rs +++ b/crates/ember-protocol/src/command/attributes.rs @@ -95,6 +95,9 @@ impl Command { Command::LInsert { .. } => "linsert", Command::LRem { .. } => "lrem", Command::LPos { .. } => "lpos", + Command::LMove { .. } => "lmove", + Command::GetDel { .. } => "getdel", + Command::GetEx { .. } => "getex", // sorted set Command::ZAdd { .. } => "zadd", @@ -111,6 +114,9 @@ impl Command { Command::ZRevRangeByScore { .. } => "zrevrangebyscore", Command::ZPopMin { .. } => "zpopmin", Command::ZPopMax { .. } => "zpopmax", + Command::ZDiff { .. } => "zdiff", + Command::ZInter { .. } => "zinter", + Command::ZUnion { .. } => "zunion", // hash Command::HSet { .. } => "hset", @@ -252,6 +258,10 @@ impl Command { | Command::LTrim { .. } | Command::LInsert { .. } | Command::LRem { .. } + | Command::LMove { .. } + // string extras + | Command::GetDel { .. } + | Command::GetEx { .. } // sorted set | Command::ZAdd { .. } | Command::ZRem { .. } @@ -406,6 +416,17 @@ impl Command { | Command::ZPopMin { .. } | Command::ZPopMax { .. } => WRITE | SORTEDSET | SLOW, + // sorted set — reads (Redis 6.2+) + Command::ZDiff { .. } | Command::ZInter { .. } | Command::ZUnion { .. } => { + READ | SORTEDSET | SLOW + } + + // string extras (Redis 6.2+) + Command::GetDel { .. } | Command::GetEx { .. } => WRITE | STRING | FAST, + + // list extras (Redis 6.2+) + Command::LMove { .. } => WRITE | LIST | FAST, + // hash — reads Command::HGet { .. } | Command::HExists { .. } | Command::HLen { .. } => { READ | HASH | FAST @@ -596,7 +617,10 @@ impl Command { | Command::ProtoSetField { key, .. } | Command::ProtoDelField { key, .. } | Command::Restore { key, .. } - | Command::Sort { key, .. } => Some(key), + | Command::Sort { key, .. } + | Command::GetDel { key } + | Command::GetEx { key, .. } => Some(key), + Command::LMove { source, .. } => Some(source), Command::Copy { source, .. } => Some(source), Command::Del { keys } | Command::Unlink { keys } @@ -607,7 +631,10 @@ impl Command { | Command::BRPop { keys, .. } | Command::SUnion { keys } | Command::SInter { keys } - | Command::SDiff { keys } => keys.first().map(String::as_str), + | Command::SDiff { keys } + | Command::ZDiff { keys, .. } + | Command::ZInter { keys, .. } + | Command::ZUnion { keys, .. } => keys.first().map(String::as_str), Command::SUnionStore { dest, .. } | Command::SInterStore { dest, .. } | Command::SDiffStore { dest, .. } => Some(dest), diff --git a/crates/ember-protocol/src/command/mod.rs b/crates/ember-protocol/src/command/mod.rs index dab6a215..57105b95 100644 --- a/crates/ember-protocol/src/command/mod.rs +++ b/crates/ember-protocol/src/command/mod.rs @@ -6,13 +6,17 @@ use bytes::Bytes; -/// Expiration option for the SET command. +/// Expiration option for the SET and GETEX commands. #[derive(Debug, Clone, PartialEq, Eq)] pub enum SetExpire { /// EX seconds — expire after N seconds. Ex(u64), /// PX milliseconds — expire after N milliseconds. Px(u64), + /// EXAT unix-seconds — expire at an absolute unix timestamp (seconds). + ExAt(u64), + /// PXAT unix-milliseconds — expire at an absolute unix timestamp (milliseconds). + PxAt(u64), } /// A parsed client command, ready for execution. @@ -238,6 +242,40 @@ pub enum Command { maxlen: usize, }, + /// LMOVE `source` `destination` LEFT|RIGHT LEFT|RIGHT. + /// Atomically pops from the source list and pushes to the destination list. + LMove { + source: String, + destination: String, + /// Pop from the left (head) if true, right (tail) if false. + src_left: bool, + /// Push to the left (head) if true, right (tail) if false. + dst_left: bool, + }, + + /// GETDEL `key`. Returns the value of a key and deletes it atomically. + GetDel { key: String }, + + /// GETEX `key` \[EX seconds | PX ms | EXAT timestamp | PXAT timestamp-ms | PERSIST\]. + /// Returns the value of a key and optionally updates its expiry. + GetEx { + key: String, + /// `None` — no change; `Some(None)` — remove TTL (PERSIST); `Some(Some(_))` — set TTL. + expire: Option>, + }, + + /// ZDIFF `numkeys` `key` \[key ...\] \[WITHSCORES\]. + /// Returns members in the first sorted set not present in the others. + ZDiff { keys: Vec, with_scores: bool }, + + /// ZINTER `numkeys` `key` \[key ...\] \[WITHSCORES\]. + /// Returns members present in all of the given sorted sets. + ZInter { keys: Vec, with_scores: bool }, + + /// ZUNION `numkeys` `key` \[key ...\] \[WITHSCORES\]. + /// Returns the union of all given sorted sets. + ZUnion { keys: Vec, with_scores: bool }, + /// TYPE `key`. Returns the type of the value stored at key. Type { key: String }, diff --git a/crates/ember-protocol/src/command/parse.rs b/crates/ember-protocol/src/command/parse.rs index 9b5ff889..f4d2d7ac 100644 --- a/crates/ember-protocol/src/command/parse.rs +++ b/crates/ember-protocol/src/command/parse.rs @@ -124,6 +124,9 @@ impl Command { "LINSERT" => parse_linsert(&frames[1..]), "LREM" => parse_lrem(&frames[1..]), "LPOS" => parse_lpos(&frames[1..]), + "LMOVE" => parse_lmove(&frames[1..]), + "GETDEL" => parse_getdel(&frames[1..]), + "GETEX" => parse_getex(&frames[1..]), "TYPE" => parse_type(&frames[1..]), "ZADD" => parse_zadd(&frames[1..]), "ZREM" => parse_zrem(&frames[1..]), @@ -145,6 +148,9 @@ impl Command { let (key, count) = parse_zpop_args(&frames[1..], "ZPOPMAX")?; Ok(Command::ZPopMax { key, count }) } + "ZDIFF" => parse_zset_multi("ZDIFF", &frames[1..]), + "ZINTER" => parse_zset_multi("ZINTER", &frames[1..]), + "ZUNION" => parse_zset_multi("ZUNION", &frames[1..]), "HSET" => parse_hset(&frames[1..]), "HGET" => parse_hget(&frames[1..]), "HGETALL" => parse_hgetall(&frames[1..]), @@ -3046,3 +3052,144 @@ fn parse_sort(args: &[Frame]) -> Result { store, }) } + +// --- Redis 6.2+ commands --- + +fn parse_lmove(args: &[Frame]) -> Result { + if args.len() != 4 { + return Err(wrong_arity("LMOVE")); + } + let source = extract_string(&args[0])?; + let destination = extract_string(&args[1])?; + + let mut kw = [0u8; MAX_KEYWORD_LEN]; + let src_left = match uppercase_arg(&args[2], &mut kw)? { + "LEFT" => true, + "RIGHT" => false, + other => { + return Err(ProtocolError::InvalidCommandFrame(format!( + "LMOVE: invalid wherefrom '{other}', expected LEFT or RIGHT" + ))); + } + }; + let mut kw = [0u8; MAX_KEYWORD_LEN]; + let dst_left = match uppercase_arg(&args[3], &mut kw)? { + "LEFT" => true, + "RIGHT" => false, + other => { + return Err(ProtocolError::InvalidCommandFrame(format!( + "LMOVE: invalid whereto '{other}', expected LEFT or RIGHT" + ))); + } + }; + + Ok(Command::LMove { + source, + destination, + src_left, + dst_left, + }) +} + +fn parse_getdel(args: &[Frame]) -> Result { + if args.len() != 1 { + return Err(wrong_arity("GETDEL")); + } + let key = extract_string(&args[0])?; + Ok(Command::GetDel { key }) +} + +fn parse_getex(args: &[Frame]) -> Result { + if args.is_empty() { + return Err(wrong_arity("GETEX")); + } + let key = extract_string(&args[0])?; + let rest = &args[1..]; + + let expire = if rest.is_empty() { + // no options — TTL unchanged + None + } else { + let mut kw = [0u8; MAX_KEYWORD_LEN]; + match uppercase_arg(&rest[0], &mut kw)? { + "PERSIST" => Some(None), + "EX" => { + if rest.len() < 2 { + return Err(wrong_arity("GETEX")); + } + let n = parse_u64(&rest[1], "GETEX")?; + if n == 0 { + return Err(ProtocolError::InvalidCommandFrame( + "invalid expire time in 'GETEX' command".into(), + )); + } + Some(Some(SetExpire::Ex(n))) + } + "PX" => { + if rest.len() < 2 { + return Err(wrong_arity("GETEX")); + } + let n = parse_u64(&rest[1], "GETEX")?; + if n == 0 { + return Err(ProtocolError::InvalidCommandFrame( + "invalid expire time in 'GETEX' command".into(), + )); + } + Some(Some(SetExpire::Px(n))) + } + "EXAT" => { + if rest.len() < 2 { + return Err(wrong_arity("GETEX")); + } + let n = parse_u64(&rest[1], "GETEX")?; + Some(Some(SetExpire::ExAt(n))) + } + "PXAT" => { + if rest.len() < 2 { + return Err(wrong_arity("GETEX")); + } + let n = parse_u64(&rest[1], "GETEX")?; + Some(Some(SetExpire::PxAt(n))) + } + other => { + return Err(ProtocolError::InvalidCommandFrame(format!( + "GETEX: unsupported option '{other}'" + ))); + } + } + }; + + Ok(Command::GetEx { key, expire }) +} + +/// Parses ZDIFF/ZINTER/ZUNION: `numkeys key [key ...] [WITHSCORES]`. +fn parse_zset_multi(cmd: &'static str, args: &[Frame]) -> Result { + if args.is_empty() { + return Err(wrong_arity(cmd)); + } + let numkeys = parse_u64(&args[0], cmd)? as usize; + if numkeys == 0 { + return Err(ProtocolError::InvalidCommandFrame(format!( + "{cmd}: numkeys must be positive" + ))); + } + if args.len() < 1 + numkeys { + return Err(wrong_arity(cmd)); + } + let keys = extract_strings(&args[1..1 + numkeys])?; + + let mut with_scores = false; + for frame in &args[1 + numkeys..] { + let mut kw = [0u8; MAX_KEYWORD_LEN]; + if let Ok("WITHSCORES") = uppercase_arg(frame, &mut kw) { + with_scores = true; + } + } + + match cmd { + "ZDIFF" => Ok(Command::ZDiff { keys, with_scores }), + "ZINTER" => Ok(Command::ZInter { keys, with_scores }), + "ZUNION" => Ok(Command::ZUnion { keys, with_scores }), + _ => Err(wrong_arity(cmd)), + } +} diff --git a/crates/ember-server/src/concurrent_handler.rs b/crates/ember-server/src/concurrent_handler.rs index fe018c19..09d40cb9 100644 --- a/crates/ember-server/src/concurrent_handler.rs +++ b/crates/ember-server/src/concurrent_handler.rs @@ -468,9 +468,26 @@ async fn execute_concurrent( return Frame::Null; } - let ttl = expire.map(|e| match e { - SetExpire::Ex(secs) => Duration::from_secs(secs), - SetExpire::Px(millis) => Duration::from_millis(millis), + let ttl = expire.map(|e| { + use std::time::{SystemTime, UNIX_EPOCH}; + match e { + SetExpire::Ex(secs) => Duration::from_secs(secs), + SetExpire::Px(millis) => Duration::from_millis(millis), + SetExpire::ExAt(ts) => { + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + Duration::from_secs(ts.saturating_sub(now)) + } + SetExpire::PxAt(ts_ms) => { + let now_ms = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64; + Duration::from_millis(ts_ms.saturating_sub(now_ms)) + } + } }); if keyspace.set(key, value, ttl) { diff --git a/crates/ember-server/src/connection/dispatch.rs b/crates/ember-server/src/connection/dispatch.rs index facb18a1..639f0cc7 100644 --- a/crates/ember-server/src/connection/dispatch.rs +++ b/crates/ember-server/src/connection/dispatch.rs @@ -266,9 +266,26 @@ pub(super) async fn prepare_command( nx, xx, } => { - let duration = expire.map(|e| match e { - SetExpire::Ex(secs) => Duration::from_secs(secs), - SetExpire::Px(millis) => Duration::from_millis(millis), + let duration = expire.map(|e| { + use std::time::{SystemTime, UNIX_EPOCH}; + match e { + SetExpire::Ex(secs) => Duration::from_secs(secs), + SetExpire::Px(millis) => Duration::from_millis(millis), + SetExpire::ExAt(ts) => { + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + Duration::from_secs(ts.saturating_sub(now)) + } + SetExpire::PxAt(ts_ms) => { + let now_ms = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64; + Duration::from_millis(ts_ms.saturating_sub(now_ms)) + } + } }); route!( key, diff --git a/crates/ember-server/src/connection/execute.rs b/crates/ember-server/src/connection/execute.rs index bccb5da1..d2e7b38d 100644 --- a/crates/ember-server/src/connection/execute.rs +++ b/crates/ember-server/src/connection/execute.rs @@ -15,6 +15,33 @@ use ember_protocol::{parse_frame, Command, Frame, SetExpire}; use subtle::ConstantTimeEq; use tokio::io::{AsyncReadExt, AsyncWriteExt}; +/// Converts a [`SetExpire`] option to a [`Duration`] relative to now. +/// +/// EX/PX are relative; EXAT/PXAT are unix timestamps that are converted to +/// a duration by subtracting the current wall time. A past timestamp results +/// in a zero duration (the key expires immediately). +fn set_expire_to_duration(expire: SetExpire) -> Duration { + use std::time::{SystemTime, UNIX_EPOCH}; + match expire { + SetExpire::Ex(secs) => Duration::from_secs(secs), + SetExpire::Px(ms) => Duration::from_millis(ms), + SetExpire::ExAt(ts) => { + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + Duration::from_secs(ts.saturating_sub(now)) + } + SetExpire::PxAt(ts_ms) => { + let now_ms = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64; + Duration::from_millis(ts_ms.saturating_sub(now_ms)) + } + } +} + /// Executes a parsed command and returns the response frame. /// /// Ping and Echo are handled inline (no shard routing needed). @@ -108,10 +135,7 @@ pub(super) async fn execute( nx, xx, } => { - let duration = expire.map(|e| match e { - SetExpire::Ex(secs) => Duration::from_secs(secs), - SetExpire::Px(millis) => Duration::from_millis(millis), - }); + let duration = expire.map(|e| set_expire_to_duration(e)); let idx = engine.shard_for_key(&key); let req = ShardRequest::Set { key, @@ -1535,6 +1559,140 @@ pub(super) async fn execute( } } + Command::LMove { + source, + destination, + src_left, + dst_left, + } => { + // route to the source key's shard + let idx = engine.shard_for_key(&source); + let req = ShardRequest::LMove { + source, + destination, + src_left, + dst_left, + }; + match engine.send_to_shard(idx, req).await { + Ok(ShardResponse::Value(Some(Value::String(data)))) => Frame::Bulk(data), + Ok(ShardResponse::Value(None)) => Frame::Null, + Ok(ShardResponse::WrongType) => wrongtype_error(), + Ok(ShardResponse::OutOfMemory) => oom_error(), + Ok(other) => Frame::Error(format!("ERR unexpected shard response: {other:?}")), + Err(e) => Frame::Error(format!("ERR {e}")), + } + } + + Command::GetDel { key } => { + let idx = engine.shard_for_key(&key); + let req = ShardRequest::GetDel { key }; + match engine.send_to_shard(idx, req).await { + Ok(ShardResponse::Value(Some(Value::String(data)))) => Frame::Bulk(data), + Ok(ShardResponse::Value(None)) => Frame::Null, + Ok(ShardResponse::WrongType) => wrongtype_error(), + Ok(other) => Frame::Error(format!("ERR unexpected shard response: {other:?}")), + Err(e) => Frame::Error(format!("ERR {e}")), + } + } + + Command::GetEx { key, expire } => { + let idx = engine.shard_for_key(&key); + // convert SetExpire into an Option> (milliseconds from now) + let expire_ms: Option> = expire.map(|opt| { + opt.map(|se| match se { + SetExpire::Ex(s) => Duration::from_secs(s).as_millis() as u64, + SetExpire::Px(ms) => ms, + SetExpire::ExAt(ts) => { + use std::time::{SystemTime, UNIX_EPOCH}; + let now_s = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + Duration::from_secs(ts.saturating_sub(now_s)).as_millis() as u64 + } + SetExpire::PxAt(ts_ms) => { + use std::time::{SystemTime, UNIX_EPOCH}; + let now_ms = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64; + ts_ms.saturating_sub(now_ms) + } + }) + }); + let req = ShardRequest::GetEx { key, expire: expire_ms }; + match engine.send_to_shard(idx, req).await { + Ok(ShardResponse::Value(Some(Value::String(data)))) => Frame::Bulk(data), + Ok(ShardResponse::Value(None)) => Frame::Null, + Ok(ShardResponse::WrongType) => wrongtype_error(), + Ok(other) => Frame::Error(format!("ERR unexpected shard response: {other:?}")), + Err(e) => Frame::Error(format!("ERR {e}")), + } + } + + Command::ZDiff { keys, with_scores } => { + let key = keys.first().cloned().unwrap_or_default(); + let idx = engine.shard_for_key(&key); + let req = ShardRequest::ZDiff { keys }; + match engine.send_to_shard(idx, req).await { + Ok(ShardResponse::ScoredArray(items)) => { + let mut frames = Vec::new(); + for (member, score) in items { + frames.push(Frame::Bulk(Bytes::from(member))); + if with_scores { + frames.push(Frame::Bulk(Bytes::from(format!("{score}")))); + } + } + Frame::Array(frames) + } + Ok(ShardResponse::WrongType) => wrongtype_error(), + Ok(other) => Frame::Error(format!("ERR unexpected shard response: {other:?}")), + Err(e) => Frame::Error(format!("ERR {e}")), + } + } + + Command::ZInter { keys, with_scores } => { + let key = keys.first().cloned().unwrap_or_default(); + let idx = engine.shard_for_key(&key); + let req = ShardRequest::ZInter { keys }; + match engine.send_to_shard(idx, req).await { + Ok(ShardResponse::ScoredArray(items)) => { + let mut frames = Vec::new(); + for (member, score) in items { + frames.push(Frame::Bulk(Bytes::from(member))); + if with_scores { + frames.push(Frame::Bulk(Bytes::from(format!("{score}")))); + } + } + Frame::Array(frames) + } + Ok(ShardResponse::WrongType) => wrongtype_error(), + Ok(other) => Frame::Error(format!("ERR unexpected shard response: {other:?}")), + Err(e) => Frame::Error(format!("ERR {e}")), + } + } + + Command::ZUnion { keys, with_scores } => { + let key = keys.first().cloned().unwrap_or_default(); + let idx = engine.shard_for_key(&key); + let req = ShardRequest::ZUnion { keys }; + match engine.send_to_shard(idx, req).await { + Ok(ShardResponse::ScoredArray(items)) => { + let mut frames = Vec::new(); + for (member, score) in items { + frames.push(Frame::Bulk(Bytes::from(member))); + if with_scores { + frames.push(Frame::Bulk(Bytes::from(format!("{score}")))); + } + } + Frame::Array(frames) + } + Ok(ShardResponse::WrongType) => wrongtype_error(), + Ok(other) => Frame::Error(format!("ERR unexpected shard response: {other:?}")), + Err(e) => Frame::Error(format!("ERR {e}")), + } + } + Command::SScan { key, cursor,