diff --git a/crates/ember-core/src/keyspace/list.rs b/crates/ember-core/src/keyspace/list.rs index 70e6e77a..34e2c0dd 100644 --- a/crates/ember-core/src/keyspace/list.rs +++ b/crates/ember-core/src/keyspace/list.rs @@ -53,6 +53,46 @@ impl Keyspace { self.list_pop(key, false) } + /// Pops up to `count` values from the head of a list. + /// + /// Returns `Ok(None)` if the key doesn't exist or is empty. Returns + /// `Ok(Some(items))` with 1–count elements otherwise. Removes the key + /// when the list becomes empty. Returns `Err(WrongType)` on type mismatch. + pub fn lpop_count(&mut self, key: &str, count: usize) -> Result>, WrongType> { + let mut items = Vec::with_capacity(count); + for _ in 0..count { + match self.lpop(key)? { + Some(v) => items.push(v), + None => break, + } + } + if items.is_empty() { + Ok(None) + } else { + Ok(Some(items)) + } + } + + /// Pops up to `count` values from the tail of a list. + /// + /// Returns `Ok(None)` if the key doesn't exist or is empty. Returns + /// `Ok(Some(items))` with 1–count elements otherwise. Removes the key + /// when the list becomes empty. Returns `Err(WrongType)` on type mismatch. + pub fn rpop_count(&mut self, key: &str, count: usize) -> Result>, WrongType> { + let mut items = Vec::with_capacity(count); + for _ in 0..count { + match self.rpop(key)? { + Some(v) => items.push(v), + None => break, + } + } + if items.is_empty() { + Ok(None) + } else { + Ok(Some(items)) + } + } + /// Returns a range of elements from a list by index. /// /// Supports negative indices (e.g. -1 = last element). Out-of-bounds @@ -1326,4 +1366,61 @@ mod tests { ks.set("s".into(), Bytes::from("hello"), None, false, false); assert!(ks.lmove("s", "dst", true, true).is_err()); } + + #[test] + fn lpop_count_pops_multiple_from_head() { + let mut ks = Keyspace::new(); + ks.rpush( + "l", + &[ + Bytes::from("a"), + Bytes::from("b"), + Bytes::from("c"), + Bytes::from("d"), + ], + ) + .unwrap(); + let result = ks.lpop_count("l", 3).unwrap(); + assert_eq!( + result, + Some(vec![Bytes::from("a"), Bytes::from("b"), Bytes::from("c")]) + ); + // one item left + assert_eq!(ks.llen("l").unwrap(), 1); + } + + #[test] + fn rpop_count_pops_multiple_from_tail() { + let mut ks = Keyspace::new(); + ks.rpush("l", &[Bytes::from("a"), Bytes::from("b"), Bytes::from("c")]) + .unwrap(); + let result = ks.rpop_count("l", 2).unwrap(); + assert_eq!(result, Some(vec![Bytes::from("c"), Bytes::from("b")])); + assert_eq!(ks.llen("l").unwrap(), 1); + } + + #[test] + fn lpop_count_missing_key_returns_none() { + let mut ks = Keyspace::new(); + assert_eq!(ks.lpop_count("missing", 5).unwrap(), None); + } + + #[test] + fn lpop_count_capped_at_list_size() { + let mut ks = Keyspace::new(); + ks.rpush("l", &[Bytes::from("x"), Bytes::from("y")]) + .unwrap(); + // request more than available + let result = ks.lpop_count("l", 10).unwrap(); + assert_eq!(result, Some(vec![Bytes::from("x"), Bytes::from("y")])); + // key deleted after emptied + assert!(!ks.exists("l")); + } + + #[test] + fn lpop_count_wrong_type_returns_error() { + let mut ks = Keyspace::new(); + ks.set("s".into(), Bytes::from("hello"), None, false, false); + assert!(ks.lpop_count("s", 1).is_err()); + } } diff --git a/crates/ember-core/src/shard/aof.rs b/crates/ember-core/src/shard/aof.rs index b529dfba..00d3d656 100644 --- a/crates/ember-core/src/shard/aof.rs +++ b/crates/ember-core/src/shard/aof.rs @@ -114,6 +114,30 @@ pub(super) fn to_aof_records( members: items.iter().map(|(m, _)| m.clone()).collect(), }] } + // LMPOP single-key: persist as individual lpop/rpop records + (ShardRequest::LmpopSingle { key, left, .. }, ShardResponse::Array(items)) + if !items.is_empty() => + { + let n = items.len(); + let mut records = SmallVec::with_capacity(n); + for _ in 0..n { + if left { + records.push(AofRecord::LPop { key: key.clone() }); + } else { + records.push(AofRecord::RPop { key: key.clone() }); + } + } + records + } + // ZMPOP single-key: persist as ZREM of the popped members + (ShardRequest::ZmpopSingle { key, .. }, ShardResponse::ZPopResult(items)) + if !items.is_empty() => + { + smallvec![AofRecord::ZRem { + key, + members: items.iter().map(|(m, _)| m.clone()).collect(), + }] + } (ShardRequest::Incr { key }, ShardResponse::Integer(_)) => { smallvec![AofRecord::Incr { key }] } diff --git a/crates/ember-core/src/shard/mod.rs b/crates/ember-core/src/shard/mod.rs index 764f7611..7c76786b 100644 --- a/crates/ember-core/src/shard/mod.rs +++ b/crates/ember-core/src/shard/mod.rs @@ -386,6 +386,18 @@ pub enum ShardRequest { key: String, count: usize, }, + /// LMPOP single-key sub-request: pop up to `count` items from one list. + LmpopSingle { + key: String, + left: bool, + count: usize, + }, + /// ZMPOP single-key sub-request: pop up to `count` items from one sorted set. + ZmpopSingle { + key: String, + min: bool, + count: usize, + }, HSet { key: String, fields: Vec<(String, Bytes)>, @@ -740,6 +752,8 @@ impl ShardRequest { | ShardRequest::ZIncrBy { .. } | ShardRequest::ZPopMin { .. } | ShardRequest::ZPopMax { .. } + | ShardRequest::LmpopSingle { .. } + | ShardRequest::ZmpopSingle { .. } | ShardRequest::HSet { .. } | ShardRequest::HDel { .. } | ShardRequest::HIncrBy { .. } @@ -1873,6 +1887,30 @@ fn dispatch( Ok(items) => ShardResponse::ZPopResult(items), Err(_) => ShardResponse::WrongType, }, + ShardRequest::LmpopSingle { key, left, count } => { + let result = if *left { + ks.lpop_count(key, *count) + } else { + ks.rpop_count(key, *count) + }; + match result { + Ok(Some(items)) => ShardResponse::Array(items), + Ok(None) => ShardResponse::Value(None), + Err(_) => ShardResponse::WrongType, + } + } + ShardRequest::ZmpopSingle { key, min, count } => { + let result = if *min { + ks.zpopmin(key, *count) + } else { + ks.zpopmax(key, *count) + }; + match result { + Ok(items) if !items.is_empty() => ShardResponse::ZPopResult(items), + Ok(_) => ShardResponse::Value(None), + Err(_) => ShardResponse::WrongType, + } + } ShardRequest::DbSize => ShardResponse::KeyCount(ks.len()), ShardRequest::Stats => ShardResponse::Stats(ks.stats()), ShardRequest::KeyVersion { ref key } => ShardResponse::Version(ks.key_version(key)), diff --git a/crates/ember-protocol/src/command/attributes.rs b/crates/ember-protocol/src/command/attributes.rs index af072621..fc8b1e22 100644 --- a/crates/ember-protocol/src/command/attributes.rs +++ b/crates/ember-protocol/src/command/attributes.rs @@ -124,6 +124,8 @@ impl Command { Command::ZRevRangeByScore { .. } => "zrevrangebyscore", Command::ZPopMin { .. } => "zpopmin", Command::ZPopMax { .. } => "zpopmax", + Command::Lmpop { .. } => "lmpop", + Command::Zmpop { .. } => "zmpop", Command::ZDiff { .. } => "zdiff", Command::ZInter { .. } => "zinter", Command::ZUnion { .. } => "zunion", @@ -282,6 +284,8 @@ impl Command { | Command::ZIncrBy { .. } | Command::ZPopMin { .. } | Command::ZPopMax { .. } + | Command::Lmpop { .. } + | Command::Zmpop { .. } // hash | Command::HSet { .. } | Command::HDel { .. } @@ -436,7 +440,11 @@ impl Command { | Command::ZRem { .. } | Command::ZIncrBy { .. } | Command::ZPopMin { .. } - | Command::ZPopMax { .. } => WRITE | SORTEDSET | SLOW, + | Command::ZPopMax { .. } + | Command::Zmpop { .. } => WRITE | SORTEDSET | SLOW, + + // list — multi-key pop (Redis 7.0+) + Command::Lmpop { .. } => WRITE | LIST | SLOW, // sorted set — reads (Redis 6.2+) Command::ZDiff { .. } | Command::ZInter { .. } | Command::ZUnion { .. } => { @@ -664,7 +672,9 @@ impl Command { | Command::SInterCard { keys, .. } | Command::ZDiff { keys, .. } | Command::ZInter { keys, .. } - | Command::ZUnion { keys, .. } => keys.first().map(String::as_str), + | Command::ZUnion { keys, .. } + | Command::Lmpop { keys, .. } + | Command::Zmpop { 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 317679d8..09e644d4 100644 --- a/crates/ember-protocol/src/command/mod.rs +++ b/crates/ember-protocol/src/command/mod.rs @@ -401,6 +401,24 @@ pub enum Command { /// ZPOPMAX `key` \[count\]. Removes and returns the highest scored members. ZPopMax { key: String, count: usize }, + /// LMPOP `numkeys` `key` \[key ...\] LEFT|RIGHT \[COUNT n\]. + /// Tries keys left-to-right, popping up to `count` elements from the first + /// non-empty list. Returns `[key_name, [elem, ...]]` or nil if all empty. + Lmpop { + keys: Vec, + left: bool, + count: usize, + }, + + /// ZMPOP `numkeys` `key` \[key ...\] MIN|MAX \[COUNT n\]. + /// Tries keys left-to-right, popping up to `count` elements from the first + /// non-empty sorted set. Returns `[key_name, [[member, score], ...]]` or nil. + Zmpop { + keys: Vec, + min: bool, + count: usize, + }, + /// HSET `key` `field` `value` \[field value ...\]. Sets field-value pairs in a hash. HSet { key: String, diff --git a/crates/ember-protocol/src/command/parse.rs b/crates/ember-protocol/src/command/parse.rs index 4bc6eb08..bd7ad62a 100644 --- a/crates/ember-protocol/src/command/parse.rs +++ b/crates/ember-protocol/src/command/parse.rs @@ -153,6 +153,8 @@ impl Command { let (key, count) = parse_zpop_args(&frames[1..], "ZPOPMAX")?; Ok(Command::ZPopMax { key, count }) } + "LMPOP" => parse_lmpop(&frames[1..]), + "ZMPOP" => parse_zmpop(&frames[1..]), "ZDIFF" => parse_zset_multi("ZDIFF", &frames[1..]), "ZINTER" => parse_zset_multi("ZINTER", &frames[1..]), "ZUNION" => parse_zset_multi("ZUNION", &frames[1..]), @@ -1747,6 +1749,100 @@ fn parse_zpop_args(args: &[Frame], cmd: &'static str) -> Result<(String, usize), Ok((key, count)) } +// LMPOP numkeys key [key ...] LEFT|RIGHT [COUNT n] +fn parse_lmpop(args: &[Frame]) -> Result { + if args.len() < 3 { + return Err(wrong_arity("LMPOP")); + } + let numkeys = parse_u64(&args[0], "LMPOP")? as usize; + if numkeys == 0 || args.len() < 1 + numkeys + 1 { + return Err(ProtocolError::InvalidCommandFrame( + "LMPOP numkeys must match key count".into(), + )); + } + let keys: Vec = args[1..=numkeys] + .iter() + .map(extract_string) + .collect::>()?; + let dir = extract_string(&args[numkeys + 1])?.to_ascii_uppercase(); + let left = match dir.as_str() { + "LEFT" => true, + "RIGHT" => false, + _ => { + return Err(ProtocolError::InvalidCommandFrame( + "LMPOP: direction must be LEFT or RIGHT".into(), + )) + } + }; + let count = if args.len() == numkeys + 4 { + let tag = extract_string(&args[numkeys + 2])?.to_ascii_uppercase(); + if tag != "COUNT" { + return Err(ProtocolError::InvalidCommandFrame( + "LMPOP: expected COUNT".into(), + )); + } + let n = parse_u64(&args[numkeys + 3], "LMPOP")? as usize; + if n == 0 { + return Err(ProtocolError::InvalidCommandFrame( + "LMPOP: COUNT must be positive".into(), + )); + } + n + } else if args.len() == numkeys + 2 { + 1 + } else { + return Err(wrong_arity("LMPOP")); + }; + Ok(Command::Lmpop { keys, left, count }) +} + +// ZMPOP numkeys key [key ...] MIN|MAX [COUNT n] +fn parse_zmpop(args: &[Frame]) -> Result { + if args.len() < 3 { + return Err(wrong_arity("ZMPOP")); + } + let numkeys = parse_u64(&args[0], "ZMPOP")? as usize; + if numkeys == 0 || args.len() < 1 + numkeys + 1 { + return Err(ProtocolError::InvalidCommandFrame( + "ZMPOP numkeys must match key count".into(), + )); + } + let keys: Vec = args[1..=numkeys] + .iter() + .map(extract_string) + .collect::>()?; + let order = extract_string(&args[numkeys + 1])?.to_ascii_uppercase(); + let min = match order.as_str() { + "MIN" => true, + "MAX" => false, + _ => { + return Err(ProtocolError::InvalidCommandFrame( + "ZMPOP: order must be MIN or MAX".into(), + )) + } + }; + let count = if args.len() == numkeys + 4 { + let tag = extract_string(&args[numkeys + 2])?.to_ascii_uppercase(); + if tag != "COUNT" { + return Err(ProtocolError::InvalidCommandFrame( + "ZMPOP: expected COUNT".into(), + )); + } + let n = parse_u64(&args[numkeys + 3], "ZMPOP")? as usize; + if n == 0 { + return Err(ProtocolError::InvalidCommandFrame( + "ZMPOP: COUNT must be positive".into(), + )); + } + n + } else if args.len() == numkeys + 2 { + 1 + } else { + return Err(wrong_arity("ZMPOP")); + }; + Ok(Command::Zmpop { keys, min, count }) +} + // --- hash commands --- fn parse_hset(args: &[Frame]) -> Result { diff --git a/crates/ember-server/src/connection/execute.rs b/crates/ember-server/src/connection/execute.rs index c498986b..b732e7c6 100644 --- a/crates/ember-server/src/connection/execute.rs +++ b/crates/ember-server/src/connection/execute.rs @@ -1406,6 +1406,65 @@ pub(super) async fn execute( } } + Command::Lmpop { keys, left, count } => { + for key in &keys { + let idx = engine.shard_for_key(key); + let req = ShardRequest::LmpopSingle { + key: key.clone(), + left, + count, + }; + match engine.send_to_shard(idx, req).await { + Ok(ShardResponse::Array(items)) if !items.is_empty() => { + let elems = Frame::Array(items.into_iter().map(Frame::Bulk).collect()); + return Frame::Array(vec![Frame::Bulk(Bytes::from(key.clone())), elems]); + } + Ok(ShardResponse::Array(_)) | Ok(ShardResponse::Value(None)) => continue, + Ok(ShardResponse::WrongType) => return wrongtype_error(), + Ok(other) => { + return Frame::Error(format!("ERR unexpected shard response: {other:?}")) + } + Err(e) => return Frame::Error(format!("ERR {e}")), + } + } + Frame::Null + } + + Command::Zmpop { keys, min, count } => { + for key in &keys { + let idx = engine.shard_for_key(key); + let req = ShardRequest::ZmpopSingle { + key: key.clone(), + min, + count, + }; + match engine.send_to_shard(idx, req).await { + Ok(ShardResponse::ZPopResult(members)) if !members.is_empty() => { + let pairs: Vec = members + .into_iter() + .flat_map(|(m, s)| { + vec![ + Frame::Bulk(Bytes::from(m)), + Frame::Bulk(Bytes::from(format!("{s}"))), + ] + }) + .collect(); + return Frame::Array(vec![ + Frame::Bulk(Bytes::from(key.clone())), + Frame::Array(pairs), + ]); + } + Ok(ShardResponse::ZPopResult(_)) | Ok(ShardResponse::Value(None)) => continue, + Ok(ShardResponse::WrongType) => return wrongtype_error(), + Ok(other) => { + return Frame::Error(format!("ERR unexpected shard response: {other:?}")) + } + Err(e) => return Frame::Error(format!("ERR {e}")), + } + } + Frame::Null + } + // --- hash commands --- Command::HSet { key, fields } => { let idx = engine.shard_for_key(&key); diff --git a/tests/integration/src/data_types.rs b/tests/integration/src/data_types.rs index e115bb4c..665ed898 100644 --- a/tests/integration/src/data_types.rs +++ b/tests/integration/src/data_types.rs @@ -408,3 +408,130 @@ async fn sintercard_missing_key_returns_zero() { assert_eq!(c.get_int(&["SINTERCARD", "2", "s", "missing"]).await, 0); } + +// --- LMPOP --- + +#[tokio::test] +async fn lmpop_returns_first_nonempty_list() { + let server = TestServer::start(); + let mut c = server.connect().await; + + // first key is empty, second has data + c.cmd(&["RPUSH", "b", "x", "y", "z"]).await; + + let resp = c.cmd(&["LMPOP", "2", "a", "b", "LEFT"]).await; + match resp { + Frame::Array(outer) => { + assert_eq!(outer.len(), 2); + // first element is the key name + assert!(matches!(&outer[0], Frame::Bulk(k) if k == &b"b"[..])); + // second is an array of popped elements (default count=1) + match &outer[1] { + Frame::Array(items) => { + assert_eq!(items.len(), 1); + assert!(matches!(&items[0], Frame::Bulk(v) if v == &b"x"[..])); + } + other => panic!("expected inner Array, got {other:?}"), + } + } + other => panic!("expected outer Array, got {other:?}"), + } +} + +#[tokio::test] +async fn lmpop_count_pops_multiple() { + let server = TestServer::start(); + let mut c = server.connect().await; + + c.cmd(&["RPUSH", "lst", "1", "2", "3", "4"]).await; + + let resp = c.cmd(&["LMPOP", "1", "lst", "RIGHT", "COUNT", "3"]).await; + match resp { + Frame::Array(outer) => { + assert_eq!(outer.len(), 2); + match &outer[1] { + Frame::Array(items) => { + // pops 3 from the right: 4, 3, 2 + assert_eq!(items.len(), 3); + assert!(matches!(&items[0], Frame::Bulk(v) if v == &b"4"[..])); + assert!(matches!(&items[1], Frame::Bulk(v) if v == &b"3"[..])); + assert!(matches!(&items[2], Frame::Bulk(v) if v == &b"2"[..])); + } + other => panic!("expected inner Array, got {other:?}"), + } + } + other => panic!("expected outer Array, got {other:?}"), + } +} + +#[tokio::test] +async fn lmpop_all_empty_returns_nil() { + let server = TestServer::start(); + let mut c = server.connect().await; + + let resp = c.cmd(&["LMPOP", "2", "empty1", "empty2", "LEFT"]).await; + assert!(matches!(resp, Frame::Null)); +} + +// --- ZMPOP --- + +#[tokio::test] +async fn zmpop_min_pops_lowest_score() { + let server = TestServer::start(); + let mut c = server.connect().await; + + // first key missing, second has data + c.cmd(&["ZADD", "scores", "1", "alice", "2", "bob", "3", "carol"]) + .await; + + let resp = c.cmd(&["ZMPOP", "2", "missing", "scores", "MIN"]).await; + match resp { + Frame::Array(outer) => { + assert_eq!(outer.len(), 2); + assert!(matches!(&outer[0], Frame::Bulk(k) if k == &b"scores"[..])); + match &outer[1] { + Frame::Array(pairs) => { + // one (member, score) pair flattened + assert_eq!(pairs.len(), 2); + assert!(matches!(&pairs[0], Frame::Bulk(m) if m == &b"alice"[..])); + assert!(matches!(&pairs[1], Frame::Bulk(s) if s == &b"1"[..])); + } + other => panic!("expected pairs Array, got {other:?}"), + } + } + other => panic!("expected outer Array, got {other:?}"), + } +} + +#[tokio::test] +async fn zmpop_count_pops_multiple() { + let server = TestServer::start(); + let mut c = server.connect().await; + + c.cmd(&["ZADD", "z", "10", "a", "20", "b", "30", "c"]).await; + + let resp = c.cmd(&["ZMPOP", "1", "z", "MAX", "COUNT", "2"]).await; + match resp { + Frame::Array(outer) => { + match &outer[1] { + Frame::Array(pairs) => { + // 2 members popped from MAX: c(30) then b(20), flattened + assert_eq!(pairs.len(), 4); + assert!(matches!(&pairs[0], Frame::Bulk(m) if m == &b"c"[..])); + assert!(matches!(&pairs[2], Frame::Bulk(m) if m == &b"b"[..])); + } + other => panic!("expected pairs Array, got {other:?}"), + } + } + other => panic!("expected outer Array, got {other:?}"), + } +} + +#[tokio::test] +async fn zmpop_all_empty_returns_nil() { + let server = TestServer::start(); + let mut c = server.connect().await; + + let resp = c.cmd(&["ZMPOP", "2", "nope1", "nope2", "MIN"]).await; + assert!(matches!(resp, Frame::Null)); +}