diff --git a/crates/ember-core/src/keyspace/mod.rs b/crates/ember-core/src/keyspace/mod.rs index 18122cf8..18036204 100644 --- a/crates/ember-core/src/keyspace/mod.rs +++ b/crates/ember-core/src/keyspace/mod.rs @@ -967,6 +967,48 @@ impl Keyspace { } } + /// Sets an expiration at an absolute Unix timestamp (seconds). + /// + /// Returns `true` if the key exists and the expiry was set, + /// `false` if the key doesn't exist. + pub fn expireat(&mut self, key: &str, unix_secs: u64) -> bool { + if self.remove_if_expired(key) { + return false; + } + match self.entries.get_mut(key) { + Some(entry) => { + if entry.expires_at_ms == 0 { + self.expiry_count += 1; + } + entry.expires_at_ms = time::unix_ms_to_monotonic_ms(unix_secs.saturating_mul(1000)); + self.bump_version(key); + true + } + None => false, + } + } + + /// Sets an expiration at an absolute Unix timestamp (milliseconds). + /// + /// Returns `true` if the key exists and the expiry was set, + /// `false` if the key doesn't exist. + pub fn pexpireat(&mut self, key: &str, unix_ms: u64) -> bool { + if self.remove_if_expired(key) { + return false; + } + match self.entries.get_mut(key) { + Some(entry) => { + if entry.expires_at_ms == 0 { + self.expiry_count += 1; + } + entry.expires_at_ms = time::unix_ms_to_monotonic_ms(unix_ms); + self.bump_version(key); + true + } + None => false, + } + } + /// Returns the absolute Unix timestamp (seconds) when the key expires. /// /// Returns `-2` if the key doesn't exist, `-1` if it has no expiry. @@ -2191,6 +2233,70 @@ mod tests { assert_eq!(ks.stats().keys_with_expiry, 1); } + // --- expireat / pexpireat --- + + #[test] + fn expireat_sets_expiry_on_existing_key() { + use std::time::{SystemTime, UNIX_EPOCH}; + let mut ks = Keyspace::new(); + ks.set("k".into(), Bytes::from("v"), None, false, false); + let future_secs = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs() + + 60; + assert!(ks.expireat("k", future_secs)); + assert!(matches!(ks.ttl("k"), TtlResult::Seconds(_))); + assert_eq!(ks.stats().keys_with_expiry, 1); + } + + #[test] + fn expireat_missing_key_returns_false() { + let mut ks = Keyspace::new(); + assert!(!ks.expireat("missing", 9_999_999_999)); + } + + #[test] + fn expireat_does_not_double_count_expiry() { + use std::time::{SystemTime, UNIX_EPOCH}; + let mut ks = Keyspace::new(); + let base = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs(); + ks.set( + "k".into(), + Bytes::from("v"), + Some(Duration::from_secs(30)), + false, + false, + ); + assert_eq!(ks.stats().keys_with_expiry, 1); + assert!(ks.expireat("k", base + 120)); + assert_eq!(ks.stats().keys_with_expiry, 1); + } + + #[test] + fn pexpireat_sets_expiry_in_ms() { + use std::time::{SystemTime, UNIX_EPOCH}; + let mut ks = Keyspace::new(); + ks.set("k".into(), Bytes::from("v"), None, false, false); + let future_ms = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_millis() as u64 + + 60_000; + assert!(ks.pexpireat("k", future_ms)); + assert!(matches!(ks.pttl("k"), TtlResult::Milliseconds(_))); + assert_eq!(ks.stats().keys_with_expiry, 1); + } + + #[test] + fn pexpireat_missing_key_returns_false() { + let mut ks = Keyspace::new(); + assert!(!ks.pexpireat("missing", 9_999_999_999_000)); + } + // --- keys tests --- #[test] diff --git a/crates/ember-core/src/keyspace/string.rs b/crates/ember-core/src/keyspace/string.rs index bc65af98..52ecb4c9 100644 --- a/crates/ember-core/src/keyspace/string.rs +++ b/crates/ember-core/src/keyspace/string.rs @@ -262,6 +262,59 @@ impl Keyspace { Ok(Some(bytes)) } + /// Atomically sets a key to a new value and returns the old value. + /// + /// Equivalent to `GET` followed by `SET` in a single operation. The new + /// value is always stored as a plain string with no expiry (any existing + /// TTL is cleared, matching Redis behaviour for GETSET). + /// + /// Returns `Ok(None)` if the key did not exist or had expired. + /// Returns `Err(WrongType)` if the key holds a non-string value. + pub fn getset(&mut self, key: &str, value: Bytes) -> Result, WrongType> { + if self.remove_if_expired(key) { + // key was expired — treat as missing, fall through to set below + } else { + match self.entries.get(key) { + None => {} + Some(e) if !matches!(e.value, Value::String(_)) => return Err(WrongType), + _ => {} + } + } + + let old = match self.entries.get(key) { + Some(e) => match &e.value { + Value::String(b) => Some(b.clone()), + _ => return Err(WrongType), + }, + None => None, + }; + + // Always set with no TTL — GETSET clears any existing expiry. + self.set(key.to_owned(), value, None, false, false); + Ok(old) + } + + /// Sets multiple keys only if none of them already exist. + /// + /// Checks all keys atomically before writing any. Returns `true` and + /// writes all pairs if no key exists, `false` and writes nothing if any + /// key is already present (including keys with a TTL that hasn't expired + /// yet). + pub fn msetnx(&mut self, pairs: &[(String, Bytes)]) -> bool { + // first pass: check that no key exists + for (key, _) in pairs { + self.remove_if_expired(key); + if self.entries.contains_key(key.as_str()) { + return false; + } + } + // second pass: write all pairs + for (key, value) in pairs { + self.set(key.clone(), value.clone(), None, false, false); + } + true + } + /// Returns the value of a key and optionally updates its expiry. /// /// `expire` controls what happens to the TTL: @@ -1047,4 +1100,86 @@ mod tests { let mut ks = Keyspace::new(); assert_eq!(ks.getex("nope", None).unwrap(), None); } + + // --- getset --- + + #[test] + fn getset_returns_old_value_and_sets_new() { + let mut ks = Keyspace::new(); + ks.set("k".into(), Bytes::from("old"), None, false, false); + let old = ks.getset("k", Bytes::from("new")).unwrap(); + assert_eq!(old, Some(Bytes::from("old"))); + assert_eq!( + ks.get("k").unwrap(), + Some(Value::String(Bytes::from("new"))) + ); + } + + #[test] + fn getset_missing_key_returns_none_and_sets_value() { + let mut ks = Keyspace::new(); + let old = ks.getset("k", Bytes::from("v")).unwrap(); + assert_eq!(old, None); + assert_eq!(ks.get("k").unwrap(), Some(Value::String(Bytes::from("v")))); + } + + #[test] + fn getset_clears_existing_ttl() { + let mut ks = Keyspace::new(); + ks.set( + "k".into(), + Bytes::from("old"), + Some(Duration::from_secs(60)), + false, + false, + ); + assert!(matches!(ks.ttl("k"), TtlResult::Seconds(_))); + let _ = ks.getset("k", Bytes::from("new")).unwrap(); + assert!(matches!(ks.ttl("k"), TtlResult::NoExpiry)); + } + + #[test] + fn getset_wrong_type_returns_error() { + let mut ks = Keyspace::new(); + ks.zadd("z", &[(1.0, "a".into())], &ZAddFlags::default()) + .unwrap(); + assert!(ks.getset("z", Bytes::from("v")).is_err()); + } + + // --- msetnx --- + + #[test] + fn msetnx_all_new_keys_returns_true_and_sets() { + let mut ks = Keyspace::new(); + let pairs = vec![ + ("a".to_owned(), Bytes::from("1")), + ("b".to_owned(), Bytes::from("2")), + ]; + assert!(ks.msetnx(&pairs)); + assert_eq!(ks.get("a").unwrap(), Some(Value::String(Bytes::from("1")))); + assert_eq!(ks.get("b").unwrap(), Some(Value::String(Bytes::from("2")))); + } + + #[test] + fn msetnx_any_existing_returns_false_and_no_changes() { + let mut ks = Keyspace::new(); + ks.set("a".into(), Bytes::from("existing"), None, false, false); + let pairs = vec![ + ("a".to_owned(), Bytes::from("new")), + ("b".to_owned(), Bytes::from("2")), + ]; + assert!(!ks.msetnx(&pairs)); + // "a" unchanged, "b" not created + assert_eq!( + ks.get("a").unwrap(), + Some(Value::String(Bytes::from("existing"))) + ); + assert_eq!(ks.get("b").unwrap(), None); + } + + #[test] + fn msetnx_empty_pairs_returns_true() { + let mut ks = Keyspace::new(); + assert!(ks.msetnx(&[])); + } } diff --git a/crates/ember-core/src/shard/aof.rs b/crates/ember-core/src/shard/aof.rs index 00d3d656..937219ab 100644 --- a/crates/ember-core/src/shard/aof.rs +++ b/crates/ember-core/src/shard/aof.rs @@ -44,6 +44,19 @@ pub(super) fn to_aof_records( (ShardRequest::Expire { key, seconds }, ShardResponse::Bool(true)) => { smallvec![AofRecord::Expire { key, seconds }] } + // EXPIREAT: convert the absolute unix timestamp to a pexpire record (ms). + // On replay, pexpire will recompute the monotonic deadline from the stored ms. + // We store as Pexpireat so replay sets the same absolute deadline regardless + // of when recovery runs. + (ShardRequest::Expireat { key, timestamp }, ShardResponse::Bool(true)) => { + smallvec![AofRecord::Pexpireat { + key, + timestamp_ms: timestamp.saturating_mul(1000), + }] + } + (ShardRequest::Pexpireat { key, timestamp_ms }, ShardResponse::Bool(true)) => { + smallvec![AofRecord::Pexpireat { key, timestamp_ms }] + } (ShardRequest::LPush { key, values }, ShardResponse::Len(_)) => { smallvec![AofRecord::LPush { key, values }] } @@ -56,6 +69,23 @@ pub(super) fn to_aof_records( (ShardRequest::RPop { key }, ShardResponse::Value(Some(_))) => { smallvec![AofRecord::RPop { key }] } + // LPopCount/RPopCount: emit one pop record per element removed. + (ShardRequest::LPopCount { key, .. }, ShardResponse::Array(items)) if !items.is_empty() => { + let n = items.len(); + let mut records = SmallVec::with_capacity(n); + for _ in 0..n { + records.push(AofRecord::LPop { key: key.clone() }); + } + records + } + (ShardRequest::RPopCount { key, .. }, ShardResponse::Array(items)) if !items.is_empty() => { + let n = items.len(); + let mut records = SmallVec::with_capacity(n); + for _ in 0..n { + records.push(AofRecord::RPop { key: key.clone() }); + } + records + } (ShardRequest::LSet { key, index, value }, ShardResponse::Ok) => { smallvec![AofRecord::LSet { key, index, value }] } @@ -408,6 +438,23 @@ pub(super) fn to_aof_records( (ShardRequest::GetDel { key }, ShardResponse::Value(Some(_))) => { smallvec![AofRecord::Del { key }] } + // GETSET: persist as SET with the new value and no expiry. + (ShardRequest::GetSet { key, value }, ShardResponse::Value(_)) => { + smallvec![AofRecord::Set { + key, + value, + expire_ms: -1, + }] + } + // MSETNX: when all keys were new (Bool(true)), persist as individual SET records. + (ShardRequest::MSetNx { pairs }, ShardResponse::Bool(true)) => pairs + .into_iter() + .map(|(key, value)| AofRecord::Set { + key, + value, + expire_ms: -1, + }) + .collect(), // 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. diff --git a/crates/ember-core/src/shard/mod.rs b/crates/ember-core/src/shard/mod.rs index 7c76786b..7dd1715c 100644 --- a/crates/ember-core/src/shard/mod.rs +++ b/crates/ember-core/src/shard/mod.rs @@ -244,6 +244,16 @@ pub enum ShardRequest { key: String, milliseconds: u64, }, + /// EXPIREAT: set expiry at an absolute Unix timestamp (seconds). + Expireat { + key: String, + timestamp: u64, + }, + /// PEXPIREAT: set expiry at an absolute Unix timestamp (milliseconds). + Pexpireat { + key: String, + timestamp_ms: u64, + }, LPush { key: String, values: Vec, @@ -258,6 +268,16 @@ pub enum ShardRequest { RPop { key: String, }, + /// LPOP key count — pop up to `count` elements from the list head, returning an array. + LPopCount { + key: String, + count: usize, + }, + /// RPOP key count — pop up to `count` elements from the list tail, returning an array. + RPopCount { + key: String, + count: usize, + }, /// Blocking left-pop. If the list has elements, pops immediately and sends /// the result on `waiter`. If empty, the shard registers the waiter to be /// woken when an element is pushed. Uses an mpsc sender so multiple shards @@ -516,6 +536,15 @@ pub enum ShardRequest { GetDel { key: String, }, + /// GETSET: atomically sets key to a new value and returns the old value. + GetSet { + key: String, + value: Bytes, + }, + /// MSETNX: sets multiple keys only if none already exist (atomic all-or-nothing). + MSetNx { + pairs: Vec<(String, Bytes)>, + }, /// 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. @@ -735,12 +764,16 @@ impl ShardRequest { | ShardRequest::Rename { .. } | ShardRequest::Copy { .. } | ShardRequest::Expire { .. } + | ShardRequest::Expireat { .. } | ShardRequest::Persist { .. } | ShardRequest::Pexpire { .. } + | ShardRequest::Pexpireat { .. } | ShardRequest::LPush { .. } | ShardRequest::RPush { .. } | ShardRequest::LPop { .. } | ShardRequest::RPop { .. } + | ShardRequest::LPopCount { .. } + | ShardRequest::RPopCount { .. } | ShardRequest::LSet { .. } | ShardRequest::LTrim { .. } | ShardRequest::LInsert { .. } @@ -767,6 +800,8 @@ impl ShardRequest { | ShardRequest::LMove { .. } | ShardRequest::GetDel { .. } | ShardRequest::GetEx { .. } + | ShardRequest::GetSet { .. } + | ShardRequest::MSetNx { .. } | ShardRequest::FlushDb | ShardRequest::FlushDbAsync | ShardRequest::RestoreKey { .. } => true, @@ -1718,12 +1753,18 @@ fn dispatch( Err(_) => ShardResponse::WrongType, }, ShardRequest::Expire { key, seconds } => ShardResponse::Bool(ks.expire(key, *seconds)), + ShardRequest::Expireat { key, timestamp } => { + ShardResponse::Bool(ks.expireat(key, *timestamp)) + } ShardRequest::Ttl { key } => ShardResponse::Ttl(ks.ttl(key)), ShardRequest::Persist { key } => ShardResponse::Bool(ks.persist(key)), ShardRequest::Pttl { key } => ShardResponse::Ttl(ks.pttl(key)), ShardRequest::Pexpire { key, milliseconds } => { ShardResponse::Bool(ks.pexpire(key, *milliseconds)) } + ShardRequest::Pexpireat { key, timestamp_ms } => { + ShardResponse::Bool(ks.pexpireat(key, *timestamp_ms)) + } ShardRequest::LPush { key, values } => write_result_len(ks.lpush(key, values)), ShardRequest::RPush { key, values } => write_result_len(ks.rpush(key, values)), ShardRequest::LPop { key } => match ks.lpop(key) { @@ -1734,6 +1775,16 @@ fn dispatch( Ok(val) => ShardResponse::Value(val.map(Value::String)), Err(_) => ShardResponse::WrongType, }, + ShardRequest::LPopCount { key, count } => match ks.lpop_count(key, *count) { + Ok(Some(items)) => ShardResponse::Array(items), + Ok(None) => ShardResponse::Value(None), + Err(_) => ShardResponse::WrongType, + }, + ShardRequest::RPopCount { key, count } => match ks.rpop_count(key, *count) { + Ok(Some(items)) => ShardResponse::Array(items), + Ok(None) => ShardResponse::Value(None), + Err(_) => ShardResponse::WrongType, + }, ShardRequest::LRange { key, start, stop } => match ks.lrange(key, *start, *stop) { Ok(items) => ShardResponse::Array(items), Err(_) => ShardResponse::WrongType, @@ -2051,6 +2102,14 @@ fn dispatch( Err(_) => ShardResponse::WrongType, } } + ShardRequest::GetSet { key, value } => match ks.getset(key, value.clone()) { + Ok(old) => ShardResponse::Value(old.map(Value::String)), + Err(_) => ShardResponse::WrongType, + }, + ShardRequest::MSetNx { pairs } => { + let result = ks.msetnx(pairs); + ShardResponse::Bool(result) + } ShardRequest::ZDiff { keys } => match ks.zdiff(keys) { Ok(pairs) => ShardResponse::ScoredArray(pairs), Err(_) => ShardResponse::WrongType, diff --git a/crates/ember-core/src/time.rs b/crates/ember-core/src/time.rs index c2428e24..e1180610 100644 --- a/crates/ember-core/src/time.rs +++ b/crates/ember-core/src/time.rs @@ -30,6 +30,33 @@ pub fn now_secs() -> u32 { start.elapsed().as_secs() as u32 } +/// Anchors the monotonic clock to wall time once and caches it for all +/// subsequent conversions. Both `monotonic_to_unix_ms` and +/// `unix_ms_to_monotonic_ms` share this anchor so the two operations are +/// perfectly invertible. +struct ClockAnchor { + /// Unix epoch ms at the moment we captured the anchor. + unix_ms_at_capture: u64, + /// Monotonic ms at the moment we captured the anchor. + mono_ms_at_capture: u64, +} + +fn clock_anchor() -> &'static ClockAnchor { + static ANCHOR: OnceLock = OnceLock::new(); + ANCHOR.get_or_init(|| { + let unix_ms_at_capture = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis() + .min(u64::MAX as u128) as u64; + let mono_ms_at_capture = now_ms(); + ClockAnchor { + unix_ms_at_capture, + mono_ms_at_capture, + } + }) +} + /// Converts a monotonic expiry timestamp (ms since process start) to a Unix /// epoch timestamp in milliseconds. /// @@ -37,41 +64,37 @@ pub fn now_secs() -> u32 { /// a single `SystemTime::now()` sample. Subsequent calls use only the fast /// monotonic clock and arithmetic — no system call. /// -/// Returns `None` if the system clock predates the Unix epoch (shouldn't -/// happen on any real machine) or if `expires_at_ms` is `NO_EXPIRY`. +/// Returns `None` if `expires_at_ms` is `NO_EXPIRY`. #[inline] pub fn monotonic_to_unix_ms(expires_at_ms: u64) -> Option { if expires_at_ms == NO_EXPIRY { return None; } - - // Capture the relationship between monotonic and wall-clock time once. - struct Anchor { - /// Unix epoch ms at the moment we captured the anchor. - unix_ms_at_capture: u64, - /// Monotonic ms at the moment we captured the anchor. - mono_ms_at_capture: u64, - } - - static ANCHOR: OnceLock = OnceLock::new(); - let anchor = ANCHOR.get_or_init(|| { - let unix_ms_at_capture = SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_millis() - .min(u64::MAX as u128) as u64; - let mono_ms_at_capture = now_ms(); - Anchor { - unix_ms_at_capture, - mono_ms_at_capture, - } - }); - + let anchor = clock_anchor(); // unix_ms = unix_at_capture + (mono_expiry - mono_at_capture) let offset = expires_at_ms.saturating_sub(anchor.mono_ms_at_capture); Some(anchor.unix_ms_at_capture.saturating_add(offset)) } +/// Converts a Unix epoch timestamp in milliseconds to a monotonic expiry +/// value suitable for storage in `Entry::expires_at_ms`. +/// +/// The inverse of `monotonic_to_unix_ms`. Used by EXPIREAT and PEXPIREAT +/// to convert an absolute wall-clock timestamp into the internal monotonic +/// representation. Both functions share the same clock anchor so the +/// conversion is coherent. +/// +/// A timestamp in the past results in a value less than or equal to +/// `now_ms()`, which means the key will be treated as already expired on +/// the next access. +#[inline] +pub fn unix_ms_to_monotonic_ms(unix_ms: u64) -> u64 { + let anchor = clock_anchor(); + // mono_expiry = mono_at_capture + (unix_ms - unix_at_capture) + let offset = unix_ms.saturating_sub(anchor.unix_ms_at_capture); + anchor.mono_ms_at_capture.saturating_add(offset) +} + /// Sentinel value meaning "no expiry". pub const NO_EXPIRY: u64 = 0; diff --git a/crates/ember-persistence/src/aof.rs b/crates/ember-persistence/src/aof.rs index 6872d8d9..5949fd9a 100644 --- a/crates/ember-persistence/src/aof.rs +++ b/crates/ember-persistence/src/aof.rs @@ -97,6 +97,7 @@ const TAG_DEL: u8 = 2; const TAG_EXPIRE: u8 = 3; const TAG_PERSIST: u8 = 10; const TAG_PEXPIRE: u8 = 11; +const TAG_PEXPIREAT: u8 = 35; const TAG_RENAME: u8 = 22; const TAG_COPY: u8 = 27; const TAG_LSET: u8 = 28; @@ -174,6 +175,12 @@ pub enum AofRecord { Persist { key: String }, /// PEXPIRE key milliseconds. Pexpire { key: String, milliseconds: u64 }, + /// PEXPIREAT key timestamp-ms — set expiry at an absolute Unix timestamp (ms). + /// + /// Used to persist EXPIREAT and PEXPIREAT commands so that after recovery + /// the expiry deadline is the same absolute point in time rather than + /// being re-anchored to the moment of replay. + Pexpireat { key: String, timestamp_ms: u64 }, /// INCR key. Incr { key: String }, /// DECR key. @@ -282,6 +289,7 @@ impl AofRecord { AofRecord::ZRem { .. } => TAG_ZREM, AofRecord::Persist { .. } => TAG_PERSIST, AofRecord::Pexpire { .. } => TAG_PEXPIRE, + AofRecord::Pexpireat { .. } => TAG_PEXPIREAT, AofRecord::Incr { .. } => TAG_INCR, AofRecord::Decr { .. } => TAG_DECR, AofRecord::HSet { .. } => TAG_HSET, @@ -331,10 +339,10 @@ impl AofRecord { | AofRecord::Persist { key } | AofRecord::Incr { key } | AofRecord::Decr { key } => 1 + LEN_PREFIX + key.len(), - // 1 tag + 4 key-len + key + 8 seconds/millis - AofRecord::Expire { key, .. } | AofRecord::Pexpire { key, .. } => { - 1 + LEN_PREFIX + key.len() + 8 - } + // 1 tag + 4 key-len + key + 8 seconds/millis/timestamp + AofRecord::Expire { key, .. } + | AofRecord::Pexpire { key, .. } + | AofRecord::Pexpireat { key, .. } => 1 + LEN_PREFIX + key.len() + 8, // 1 tag + 4 key-len + key + 4 count + (4 value-len + value) * n AofRecord::LPush { key, values } | AofRecord::RPush { key, values } => { let values_size: usize = values.iter().map(|v| LEN_PREFIX + v.len()).sum(); @@ -477,6 +485,10 @@ impl AofRecord { format::write_bytes(&mut buf, key.as_bytes())?; format::write_i64(&mut buf, (*milliseconds).min(i64::MAX as u64) as i64)?; } + AofRecord::Pexpireat { key, timestamp_ms } => { + format::write_bytes(&mut buf, key.as_bytes())?; + format::write_i64(&mut buf, (*timestamp_ms).min(i64::MAX as u64) as i64)?; + } AofRecord::IncrBy { key, delta } | AofRecord::DecrBy { key, delta } => { format::write_bytes(&mut buf, key.as_bytes())?; format::write_i64(&mut buf, *delta)?; @@ -777,6 +789,16 @@ impl AofRecord { })?; Ok(AofRecord::Pexpire { key, milliseconds }) } + TAG_PEXPIREAT => { + let key = read_string(&mut cursor, "key")?; + let raw = format::read_i64(&mut cursor)?; + let timestamp_ms = u64::try_from(raw).map_err(|_| { + FormatError::InvalidData(format!( + "PEXPIREAT timestamp_ms is negative ({raw}) in AOF record" + )) + })?; + Ok(AofRecord::Pexpireat { key, timestamp_ms }) + } TAG_INCR => { let key = read_string(&mut cursor, "key")?; Ok(AofRecord::Incr { key }) diff --git a/crates/ember-persistence/src/recovery.rs b/crates/ember-persistence/src/recovery.rs index d3f48cc2..0cd8fcfd 100644 --- a/crates/ember-persistence/src/recovery.rs +++ b/crates/ember-persistence/src/recovery.rs @@ -510,6 +510,24 @@ fn replay_aof( entry.1 = milliseconds.min(i64::MAX as u64) as i64; } } + AofRecord::Pexpireat { key, timestamp_ms } => { + if let Some(entry) = map.get_mut(&key) { + // Convert the absolute unix timestamp to a remaining TTL + // relative to now. This preserves the exact wall-clock + // deadline across restarts. + let now_ms = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64; + if timestamp_ms <= now_ms { + // Already expired — mark as 0 so the filter removes it. + entry.1 = 0; + } else { + let remaining = timestamp_ms.saturating_sub(now_ms); + entry.1 = remaining.min(i64::MAX as u64) as i64; + } + } + } AofRecord::Incr { key } => { apply_incr(map, key, 1); } diff --git a/crates/ember-protocol/src/command/attributes.rs b/crates/ember-protocol/src/command/attributes.rs index fc8b1e22..20f723ca 100644 --- a/crates/ember-protocol/src/command/attributes.rs +++ b/crates/ember-protocol/src/command/attributes.rs @@ -64,10 +64,12 @@ impl Command { Command::ZScan { .. } => "zscan", Command::Type { .. } => "type", Command::Expire { .. } => "expire", + Command::Expireat { .. } => "expireat", Command::Ttl { .. } => "ttl", Command::Persist { .. } => "persist", Command::Pttl { .. } => "pttl", Command::Pexpire { .. } => "pexpire", + Command::Pexpireat { .. } => "pexpireat", Command::Expiretime { .. } => "expiretime", Command::Pexpiretime { .. } => "pexpiretime", @@ -108,6 +110,8 @@ impl Command { Command::LMove { .. } => "lmove", Command::GetDel { .. } => "getdel", Command::GetEx { .. } => "getex", + Command::GetSet { .. } => "getset", + Command::MSetNx { .. } => "msetnx", // sorted set Command::ZAdd { .. } => "zadd", @@ -247,6 +251,8 @@ impl Command { // string mutations Command::Set { .. } | Command::MSet { .. } + | Command::MSetNx { .. } + | Command::GetSet { .. } | Command::Append { .. } | Command::Incr { .. } | Command::Decr { .. } @@ -261,7 +267,9 @@ impl Command { | Command::Rename { .. } | Command::Copy { .. } | Command::Expire { .. } + | Command::Expireat { .. } | Command::Pexpire { .. } + | Command::Pexpireat { .. } | Command::Persist { .. } // list | Command::LPush { .. } @@ -373,6 +381,8 @@ impl Command { // string — writes Command::Set { .. } | Command::MSet { .. } + | Command::MSetNx { .. } + | Command::GetSet { .. } | Command::Append { .. } | Command::SetRange { .. } | Command::SetBit { .. } @@ -404,9 +414,11 @@ impl Command { // keyspace — writes Command::Del { .. } | Command::Unlink { .. } => WRITE | KEYSPACE | FAST, Command::Rename { .. } | Command::Copy { .. } => WRITE | KEYSPACE | SLOW, - Command::Expire { .. } | Command::Pexpire { .. } | Command::Persist { .. } => { - WRITE | KEYSPACE | FAST - } + Command::Expire { .. } + | Command::Expireat { .. } + | Command::Pexpire { .. } + | Command::Pexpireat { .. } + | Command::Persist { .. } => WRITE | KEYSPACE | FAST, // list — reads Command::LRange { .. } | Command::LLen { .. } => READ | LIST | SLOW, @@ -590,7 +602,9 @@ impl Command { | Command::SetRange { key, .. } | Command::Persist { key } | Command::Expire { key, .. } + | Command::Expireat { key, .. } | Command::Pexpire { key, .. } + | Command::Pexpireat { key, .. } | Command::Ttl { key } | Command::Pttl { key } | Command::Expiretime { key } @@ -599,10 +613,11 @@ impl Command { | Command::Rename { key, .. } | Command::ObjectEncoding { key } | Command::ObjectRefcount { key } + | Command::GetSet { key, .. } | Command::LPush { key, .. } | Command::RPush { key, .. } - | Command::LPop { key } - | Command::RPop { key } + | Command::LPop { key, .. } + | Command::RPop { key, .. } | Command::LRange { key, .. } | Command::LLen { key } | Command::LIndex { key, .. } @@ -678,7 +693,9 @@ impl Command { Command::SUnionStore { dest, .. } | Command::SInterStore { dest, .. } | Command::SDiffStore { dest, .. } => Some(dest), - Command::MSet { pairs } => pairs.first().map(|(k, _)| k.as_str()), + Command::MSet { pairs } | Command::MSetNx { pairs } => { + pairs.first().map(|(k, _)| k.as_str()) + } _ => None, } } diff --git a/crates/ember-protocol/src/command/mod.rs b/crates/ember-protocol/src/command/mod.rs index 09e644d4..1dfec04e 100644 --- a/crates/ember-protocol/src/command/mod.rs +++ b/crates/ember-protocol/src/command/mod.rs @@ -129,9 +129,20 @@ pub enum Command { /// MSET `key` `value` \[key value ...\]. Sets multiple key-value pairs. MSet { pairs: Vec<(String, Bytes)> }, + /// MSETNX `key` `value` \[key value ...\]. Sets multiple keys only if none already exist. + /// Returns 1 if all keys were set, 0 if any key already existed (atomic: all-or-nothing). + MSetNx { pairs: Vec<(String, Bytes)> }, + + /// GETSET `key` `value`. Atomically sets `key` to `value` and returns the old value. + /// Deprecated in Redis 6.2 but widely used; kept for compatibility. + GetSet { key: String, value: Bytes }, + /// EXPIRE `key` `seconds`. Sets a TTL on an existing key. Expire { key: String, seconds: u64 }, + /// EXPIREAT `key` `timestamp`. Sets expiry at an absolute Unix timestamp (seconds). + Expireat { key: String, timestamp: u64 }, + /// TTL `key`. Returns remaining time-to-live in seconds. Ttl { key: String }, @@ -144,6 +155,9 @@ pub enum Command { /// PEXPIRE `key` `milliseconds`. Sets a TTL in milliseconds on an existing key. Pexpire { key: String, milliseconds: u64 }, + /// PEXPIREAT `key` `timestamp-ms`. Sets expiry at an absolute Unix timestamp (milliseconds). + Pexpireat { key: String, timestamp_ms: u64 }, + /// DBSIZE. Returns the number of keys in the database. DbSize, @@ -214,11 +228,17 @@ pub enum Command { /// RPUSH `key` `value` \[value ...\]. Pushes values to the tail of a list. RPush { key: String, values: Vec }, - /// LPOP `key`. Pops a value from the head of a list. - LPop { key: String }, + /// LPOP `key` \[count\]. Pops one or more values from the head of a list. + /// + /// Without `count`: returns a bulk string (or nil). With `count`: returns + /// an array of up to `count` elements (Redis 6.2+ semantics). + LPop { key: String, count: Option }, - /// RPOP `key`. Pops a value from the tail of a list. - RPop { key: String }, + /// RPOP `key` \[count\]. Pops one or more values from the tail of a list. + /// + /// Without `count`: returns a bulk string (or nil). With `count`: returns + /// an array of up to `count` elements (Redis 6.2+ semantics). + RPop { key: String, count: Option }, /// LRANGE `key` `start` `stop`. Returns a range of elements by index. LRange { key: String, start: i64, stop: i64 }, diff --git a/crates/ember-protocol/src/command/parse.rs b/crates/ember-protocol/src/command/parse.rs index bd7ad62a..73fb91a4 100644 --- a/crates/ember-protocol/src/command/parse.rs +++ b/crates/ember-protocol/src/command/parse.rs @@ -101,11 +101,15 @@ impl Command { "EXISTS" => parse_exists(&frames[1..]), "MGET" => parse_mget(&frames[1..]), "MSET" => parse_mset(&frames[1..]), + "MSETNX" => parse_msetnx(&frames[1..]), + "GETSET" => parse_getset(&frames[1..]), "EXPIRE" => parse_expire(&frames[1..]), + "EXPIREAT" => parse_expireat(&frames[1..]), "TTL" => parse_ttl(&frames[1..]), "PERSIST" => parse_persist(&frames[1..]), "PTTL" => parse_pttl(&frames[1..]), "PEXPIRE" => parse_pexpire(&frames[1..]), + "PEXPIREAT" => parse_pexpireat_cmd(&frames[1..]), "DBSIZE" => parse_dbsize(&frames[1..]), "INFO" => parse_info(&frames[1..]), "BGSAVE" => parse_bgsave(&frames[1..]), @@ -328,6 +332,14 @@ fn parse_u64(frame: &Frame, cmd: &str) -> Result { }) } +/// Parses a frame's bytes as a `usize`. Used for optional count arguments. +fn parse_usize(frame: &Frame, cmd: &str) -> Result { + let n = parse_u64(frame, cmd)?; + usize::try_from(n).map_err(|_| { + ProtocolError::InvalidCommandFrame(format!("count is out of range for '{cmd}'")) + }) +} + /// Parses an unsigned integer directly from a byte slice. fn parse_u64_bytes(buf: &[u8]) -> Option { if buf.is_empty() { @@ -803,6 +815,28 @@ fn parse_mset(args: &[Frame]) -> Result { Ok(Command::MSet { pairs }) } +fn parse_msetnx(args: &[Frame]) -> Result { + if args.is_empty() || !args.len().is_multiple_of(2) { + return Err(wrong_arity("MSETNX")); + } + let mut pairs = Vec::with_capacity(args.len() / 2); + for chunk in args.chunks(2) { + let key = extract_string(&chunk[0])?; + let value = extract_bytes(&chunk[1])?; + pairs.push((key, value)); + } + Ok(Command::MSetNx { pairs }) +} + +fn parse_getset(args: &[Frame]) -> Result { + if args.len() != 2 { + return Err(wrong_arity("GETSET")); + } + let key = extract_string(&args[0])?; + let value = extract_bytes(&args[1])?; + Ok(Command::GetSet { key, value }) +} + fn parse_expire(args: &[Frame]) -> Result { if args.len() != 2 { return Err(wrong_arity("EXPIRE")); @@ -859,6 +893,24 @@ fn parse_pexpire(args: &[Frame]) -> Result { Ok(Command::Pexpire { key, milliseconds }) } +fn parse_expireat(args: &[Frame]) -> Result { + if args.len() != 2 { + return Err(wrong_arity("EXPIREAT")); + } + let key = extract_string(&args[0])?; + let timestamp = parse_u64(&args[1], "EXPIREAT")?; + Ok(Command::Expireat { key, timestamp }) +} + +fn parse_pexpireat_cmd(args: &[Frame]) -> Result { + if args.len() != 2 { + return Err(wrong_arity("PEXPIREAT")); + } + let key = extract_string(&args[0])?; + let timestamp_ms = parse_u64(&args[1], "PEXPIREAT")?; + Ok(Command::Pexpireat { key, timestamp_ms }) +} + fn parse_expiretime(args: &[Frame]) -> Result { if args.len() != 1 { return Err(wrong_arity("EXPIRETIME")); @@ -1161,19 +1213,29 @@ fn parse_rpush(args: &[Frame]) -> Result { } fn parse_lpop(args: &[Frame]) -> Result { - if args.len() != 1 { + if args.is_empty() || args.len() > 2 { return Err(wrong_arity("LPOP")); } let key = extract_string(&args[0])?; - Ok(Command::LPop { key }) + let count = if args.len() == 2 { + Some(parse_usize(&args[1], "LPOP")?) + } else { + None + }; + Ok(Command::LPop { key, count }) } fn parse_rpop(args: &[Frame]) -> Result { - if args.len() != 1 { + if args.is_empty() || args.len() > 2 { return Err(wrong_arity("RPOP")); } let key = extract_string(&args[0])?; - Ok(Command::RPop { key }) + let count = if args.len() == 2 { + Some(parse_usize(&args[1], "RPOP")?) + } else { + None + }; + Ok(Command::RPop { key, count }) } fn parse_lrange(args: &[Frame]) -> Result { diff --git a/crates/ember-protocol/src/command/tests.rs b/crates/ember-protocol/src/command/tests.rs index 348bfdc4..979c0dac 100644 --- a/crates/ember-protocol/src/command/tests.rs +++ b/crates/ember-protocol/src/command/tests.rs @@ -661,7 +661,10 @@ fn rpush_no_value() { fn lpop_basic() { assert_eq!( Command::from_frame(cmd(&["LPOP", "list"])).unwrap(), - Command::LPop { key: "list".into() }, + Command::LPop { + key: "list".into(), + count: None, + }, ); } @@ -677,7 +680,10 @@ fn lpop_wrong_arity() { fn rpop_basic() { assert_eq!( Command::from_frame(cmd(&["RPOP", "list"])).unwrap(), - Command::RPop { key: "list".into() }, + Command::RPop { + key: "list".into(), + count: None, + }, ); } diff --git a/crates/ember-server/src/acl.rs b/crates/ember-server/src/acl.rs index b0ddedfe..cb8857ea 100644 --- a/crates/ember-server/src/acl.rs +++ b/crates/ember-server/src/acl.rs @@ -401,8 +401,8 @@ fn extract_keys(cmd: &ember_protocol::Command) -> Vec<&str> { | Command::ObjectRefcount { key } | Command::LPush { key, .. } | Command::RPush { key, .. } - | Command::LPop { key } - | Command::RPop { key } + | Command::LPop { key, .. } + | Command::RPop { key, .. } | Command::LRange { key, .. } | Command::LLen { key } | Command::ZAdd { key, .. } @@ -894,8 +894,14 @@ fn commands_in_category(flag: u64) -> Vec<&'static str> { key: String::new(), values: vec![], }, - Command::LPop { key: String::new() }, - Command::RPop { key: String::new() }, + Command::LPop { + key: String::new(), + count: None, + }, + Command::RPop { + key: String::new(), + count: None, + }, Command::LRange { key: String::new(), start: 0, diff --git a/crates/ember-server/src/connection/dispatch.rs b/crates/ember-server/src/connection/dispatch.rs index a523441d..f4d77e53 100644 --- a/crates/ember-server/src/connection/dispatch.rs +++ b/crates/ember-server/src/connection/dispatch.rs @@ -420,12 +420,32 @@ pub(super) async fn prepare_command( ResponseTag::LenResultOom ) } - Command::LPop { key } => { + Command::LPop { key, count: None } => { route!(key, ShardRequest::LPop { key }, ResponseTag::PopResult) } - Command::RPop { key } => { + Command::LPop { + key, + count: Some(count), + } => { + route!( + key, + ShardRequest::LPopCount { key, count }, + ResponseTag::ArrayResult + ) + } + Command::RPop { key, count: None } => { route!(key, ShardRequest::RPop { key }, ResponseTag::PopResult) } + Command::RPop { + key, + count: Some(count), + } => { + route!( + key, + ShardRequest::RPopCount { key, count }, + ResponseTag::ArrayResult + ) + } Command::LRange { key, start, stop } => { route!( key, @@ -1200,11 +1220,14 @@ pub(super) async fn cluster_slot_check( | Command::Persist { ref key } | Command::Pttl { ref key } | Command::Pexpire { ref key, .. } + | Command::Expireat { ref key, .. } + | Command::Pexpireat { ref key, .. } | Command::Type { ref key } + | Command::GetSet { ref key, .. } | Command::LPush { ref key, .. } | Command::RPush { ref key, .. } - | Command::LPop { ref key } - | Command::RPop { ref key } + | Command::LPop { ref key, .. } + | Command::RPop { ref key, .. } | Command::LRange { ref key, .. } | Command::LLen { ref key } | Command::LIndex { ref key, .. } @@ -1336,8 +1359,8 @@ pub(super) async fn cluster_slot_check( .await } - // mset: extract keys from pairs for crossslot check - Command::MSet { ref pairs } => { + // mset / msetnx: extract keys from pairs for crossslot check + Command::MSet { ref pairs } | Command::MSetNx { ref pairs } => { let keys: Vec<&str> = pairs.iter().map(|(k, _)| k.as_str()).collect(); if let Err(err) = cluster.check_crossslot(&keys) { return Some(err); diff --git a/crates/ember-server/src/connection/execute.rs b/crates/ember-server/src/connection/execute.rs index b732e7c6..3578fae2 100644 --- a/crates/ember-server/src/connection/execute.rs +++ b/crates/ember-server/src/connection/execute.rs @@ -206,6 +206,52 @@ pub(super) async fn execute( } } + Command::Expireat { key, timestamp } => { + let idx = engine.shard_for_key(&key); + let req = ShardRequest::Expireat { + key: key.clone(), + timestamp, + }; + match engine.send_to_shard(idx, req).await { + Ok(ShardResponse::Bool(true)) => { + notify_write( + ctx, + pubsub, + crate::keyspace_notifications::FLAG_G, + "expireat", + &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}")), + } + } + + Command::Pexpireat { key, timestamp_ms } => { + let idx = engine.shard_for_key(&key); + let req = ShardRequest::Pexpireat { + key: key.clone(), + timestamp_ms, + }; + match engine.send_to_shard(idx, req).await { + Ok(ShardResponse::Bool(true)) => { + notify_write( + ctx, + pubsub, + crate::keyspace_notifications::FLAG_G, + "pexpireat", + &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}")), + } + } + Command::Ttl { key } => { let idx = engine.shard_for_key(&key); let req = ShardRequest::Ttl { key }; @@ -513,6 +559,77 @@ pub(super) async fn execute( } } + Command::MSetNx { pairs } => { + // MSETNX is all-or-nothing: set all keys only if none exist. + // + // We implement this with two fan-out passes: + // 1. Check existence of every key across all shards. + // 2. If all are absent, write all pairs. + // + // This is not atomic across shards (no distributed transaction), + // but matches Redis cluster semantics where MSETNX pairs must + // share a hash slot. For single-node mode it is correct. + if pairs.is_empty() { + return Frame::Error("ERR wrong number of arguments for 'MSETNX'".into()); + } + + let keys: Vec = pairs.iter().map(|(k, _)| k.clone()).collect(); + + // phase 1: check existence + let exists_responses = match engine + .route_multi(&keys, |k| ShardRequest::Exists { key: k }) + .await + { + Ok(r) => r, + Err(e) => return Frame::Error(format!("ERR {e}")), + }; + + let any_exists = exists_responses + .iter() + .any(|r| matches!(r, ShardResponse::Bool(true))); + if any_exists { + return Frame::Integer(0); + } + + // phase 2: write all pairs + let values: std::collections::HashMap = pairs.into_iter().collect(); + match engine + .route_multi(&keys, |k| { + let value = values.get(&k).cloned().unwrap_or_default(); + ShardRequest::Set { + key: k, + value, + expire: None, + nx: false, + xx: false, + } + }) + .await + { + Ok(responses) => { + for r in &responses { + if matches!(r, ShardResponse::OutOfMemory) { + return oom_error(); + } + } + Frame::Integer(1) + } + Err(e) => Frame::Error(format!("ERR {e}")), + } + } + + Command::GetSet { key, value } => { + let idx = engine.shard_for_key(&key); + let req = ShardRequest::GetSet { key, value }; + 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::RandomKey => match engine.broadcast(|| ShardRequest::RandomKey).await { Ok(responses) => { let keys: Vec = responses @@ -952,7 +1069,7 @@ pub(super) async fn execute( } } - Command::LPop { key } => { + Command::LPop { key, count: None } => { let idx = engine.shard_for_key(&key); let req = ShardRequest::LPop { key }; match engine.send_to_shard(idx, req).await { @@ -964,7 +1081,25 @@ pub(super) async fn execute( } } - Command::RPop { key } => { + Command::LPop { + key, + count: Some(count), + } => { + let idx = engine.shard_for_key(&key); + let req = ShardRequest::LPopCount { key, count }; + match engine.send_to_shard(idx, req).await { + Ok(ShardResponse::Array(items)) => { + let frames = items.into_iter().map(Frame::Bulk).collect(); + Frame::Array(frames) + } + 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::RPop { key, count: None } => { let idx = engine.shard_for_key(&key); let req = ShardRequest::RPop { key }; match engine.send_to_shard(idx, req).await { @@ -976,6 +1111,24 @@ pub(super) async fn execute( } } + Command::RPop { + key, + count: Some(count), + } => { + let idx = engine.shard_for_key(&key); + let req = ShardRequest::RPopCount { key, count }; + match engine.send_to_shard(idx, req).await { + Ok(ShardResponse::Array(items)) => { + let frames = items.into_iter().map(Frame::Bulk).collect(); + Frame::Array(frames) + } + 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::LRange { key, start, stop } => { let idx = engine.shard_for_key(&key); let req = ShardRequest::LRange { key, start, stop }; diff --git a/crates/ember-server/src/replication.rs b/crates/ember-server/src/replication.rs index a201a9c9..61df6008 100644 --- a/crates/ember-server/src/replication.rs +++ b/crates/ember-server/src/replication.rs @@ -782,6 +782,10 @@ pub fn aof_record_to_shard_request(record: &AofRecord) -> Option { key: key.clone(), milliseconds: *milliseconds, }), + AofRecord::Pexpireat { key, timestamp_ms } => Some(ShardRequest::Pexpireat { + key: key.clone(), + timestamp_ms: *timestamp_ms, + }), AofRecord::Incr { key } => Some(ShardRequest::Incr { key: key.clone() }), AofRecord::Decr { key } => Some(ShardRequest::Decr { key: key.clone() }), AofRecord::HSet { key, fields } => Some(ShardRequest::HSet { diff --git a/tests/integration/src/basic_operations.rs b/tests/integration/src/basic_operations.rs index ce56d83e..1c914366 100644 --- a/tests/integration/src/basic_operations.rs +++ b/tests/integration/src/basic_operations.rs @@ -324,3 +324,114 @@ async fn unknown_command() { let msg = c.err(&["NOTACOMMAND"]).await; assert!(msg.contains("unknown command")); } + +// --- EXPIREAT / PEXPIREAT --- + +#[tokio::test] +async fn expireat_sets_absolute_expiry() { + let server = TestServer::start(); + let mut c = server.connect().await; + + c.ok(&["SET", "k", "v"]).await; + + // timestamp 100 seconds in the future + let future = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs() + + 100; + let future_str = future.to_string(); + + assert_eq!(c.get_int(&["EXPIREAT", "k", &future_str]).await, 1); + let ttl = c.get_int(&["TTL", "k"]).await; + assert!(ttl > 0 && ttl <= 100, "expected TTL in (0,100], got {ttl}"); +} + +#[tokio::test] +async fn expireat_missing_key_returns_zero() { + let server = TestServer::start(); + let mut c = server.connect().await; + + assert_eq!(c.get_int(&["EXPIREAT", "nope", "9999999999"]).await, 0); +} + +#[tokio::test] +async fn pexpireat_sets_absolute_expiry_ms() { + let server = TestServer::start(); + let mut c = server.connect().await; + + c.ok(&["SET", "k", "v"]).await; + + let future_ms = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_millis() as u64 + + 60_000; + let future_str = future_ms.to_string(); + + assert_eq!(c.get_int(&["PEXPIREAT", "k", &future_str]).await, 1); + let pttl = c.get_int(&["PTTL", "k"]).await; + assert!( + pttl > 0 && pttl <= 60_000, + "expected PTTL in (0,60000], got {pttl}" + ); +} + +// --- GETSET --- + +#[tokio::test] +async fn getset_returns_old_value() { + let server = TestServer::start(); + let mut c = server.connect().await; + + c.ok(&["SET", "k", "old"]).await; + + let old = c.get_bulk(&["GETSET", "k", "new"]).await; + assert_eq!(old, Some("old".into())); + + let current = c.get_bulk(&["GET", "k"]).await; + assert_eq!(current, Some("new".into())); +} + +#[tokio::test] +async fn getset_missing_key_returns_nil() { + let server = TestServer::start(); + let mut c = server.connect().await; + + let resp = c.cmd(&["GETSET", "nope", "v"]).await; + assert!(matches!(resp, Frame::Null)); + + // key should now exist + let current = c.get_bulk(&["GET", "nope"]).await; + assert_eq!(current, Some("v".into())); +} + +// --- MSETNX --- + +#[tokio::test] +async fn msetnx_all_new_returns_one() { + let server = TestServer::start(); + let mut c = server.connect().await; + + assert_eq!(c.get_int(&["MSETNX", "a", "1", "b", "2"]).await, 1); + assert_eq!(c.get_bulk(&["GET", "a"]).await, Some("1".into())); + assert_eq!(c.get_bulk(&["GET", "b"]).await, Some("2".into())); +} + +#[tokio::test] +async fn msetnx_any_existing_returns_zero_and_no_changes() { + let server = TestServer::start(); + let mut c = server.connect().await; + + c.ok(&["SET", "a", "existing"]).await; + + // should fail atomically — neither "a" nor "b" should change + assert_eq!(c.get_int(&["MSETNX", "a", "new", "b", "2"]).await, 0); + + // "a" keeps its original value + assert_eq!(c.get_bulk(&["GET", "a"]).await, Some("existing".into())); + + // "b" was not created + let resp = c.cmd(&["GET", "b"]).await; + assert!(matches!(resp, Frame::Null)); +} diff --git a/tests/integration/src/data_types.rs b/tests/integration/src/data_types.rs index 665ed898..239be965 100644 --- a/tests/integration/src/data_types.rs +++ b/tests/integration/src/data_types.rs @@ -535,3 +535,67 @@ async fn zmpop_all_empty_returns_nil() { let resp = c.cmd(&["ZMPOP", "2", "nope1", "nope2", "MIN"]).await; assert!(matches!(resp, Frame::Null)); } + +// --- lpop / rpop with count --- + +#[tokio::test] +async fn lpop_no_count_returns_bulk_string() { + let server = TestServer::start(); + let mut c = server.connect().await; + + c.cmd(&["RPUSH", "list", "a", "b", "c"]).await; + + // no count — returns a bulk string, not an array + let val = c.get_bulk(&["LPOP", "list"]).await; + assert_eq!(val, Some("a".into())); +} + +#[tokio::test] +async fn lpop_count_returns_array() { + let server = TestServer::start(); + let mut c = server.connect().await; + + c.cmd(&["RPUSH", "list", "a", "b", "c", "d"]).await; + + let resp = c.cmd(&["LPOP", "list", "2"]).await; + match resp { + Frame::Array(frames) => { + assert_eq!(frames.len(), 2); + assert!(matches!(&frames[0], Frame::Bulk(b) if b == &b"a"[..])); + assert!(matches!(&frames[1], Frame::Bulk(b) if b == &b"b"[..])); + } + other => panic!("expected Array, got {other:?}"), + } +} + +#[tokio::test] +async fn rpop_count_returns_array() { + let server = TestServer::start(); + let mut c = server.connect().await; + + c.cmd(&["RPUSH", "list", "a", "b", "c", "d"]).await; + + let resp = c.cmd(&["RPOP", "list", "2"]).await; + match resp { + Frame::Array(frames) => { + assert_eq!(frames.len(), 2); + assert!(matches!(&frames[0], Frame::Bulk(b) if b == &b"d"[..])); + assert!(matches!(&frames[1], Frame::Bulk(b) if b == &b"c"[..])); + } + other => panic!("expected Array, got {other:?}"), + } +} + +#[tokio::test] +async fn lpop_count_exceeding_list_returns_all() { + let server = TestServer::start(); + let mut c = server.connect().await; + + c.cmd(&["RPUSH", "list", "x", "y"]).await; + + let resp = c.cmd(&["LPOP", "list", "10"]).await; + match resp { + Frame::Array(frames) => assert_eq!(frames.len(), 2), + other => panic!("expected Array, got {other:?}"), + } +}