feat: EXPIREAT, PEXPIREAT, GETSET, MSETNX, LPOP/RPOP count

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]

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<Option<Bytes>, 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(&[]));
+    }
 }

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.