Adding more expiration related commands

src/cmd/key.rs

@@ -1,33 +1,97 @@
-use crate::{connection::Connection, error::Error, value::bytes_to_number, value::Value};
+use crate::{
+    check_arg, connection::Connection, error::Error, value::bytes_to_number, value::Value,
+};
 use bytes::Bytes;
+use std::time::{SystemTime, UNIX_EPOCH};
 use tokio::time::{Duration, Instant};
 
+pub fn now() -> Duration {
+    let start = SystemTime::now();
+    start
+        .duration_since(UNIX_EPOCH)
+        .expect("Time went backwards")
+}
+
 pub fn del(conn: &Connection, args: &[Bytes]) -> Result<Value, Error> {
     Ok(conn.db().del(&args[1..]))
 }
 
 pub fn expire(conn: &Connection, args: &[Bytes]) -> Result<Value, Error> {
-    let expires_at: i64 = bytes_to_number(&args[2])?;
+    let expires_in: i64 = bytes_to_number(&args[2])?;
 
-    if expires_at <= 0 {
+    if expires_in <= 0 {
+        // Delete key right away
         return Ok(conn.db().del(&args[1..2]));
     }
 
-    let expires_at = Duration::new(expires_at as u64, 0);
+    let expires_at = if check_arg!(args, 0, "EXPIRES") {
+        Duration::from_secs(expires_in as u64)
+    } else {
+        Duration::from_millis(expires_in as u64)
+    };
 
     Ok(conn.db().expire(&args[1], expires_at))
 }
 
-pub fn persist(conn: &Connection, args: &[Bytes]) -> Result<Value, Error> {
-    Ok(conn.db().persist(&args[1]))
+pub fn expire_at(conn: &Connection, args: &[Bytes]) -> Result<Value, Error> {
+    let secs = check_arg!(args, 0, "EXPIREAT");
+    let expires_at: i64 = bytes_to_number(&args[2])?;
+    let expires_in: i64 = if secs {
+        expires_at - now().as_secs() as i64
+    } else {
+        expires_at - now().as_millis() as i64
+    };
+
+    if expires_in <= 0 {
+        // Delete key right away
+        return Ok(conn.db().del(&args[1..2]));
+    }
+
+    let expires_at = if secs {
+        Duration::from_secs(expires_in as u64)
+    } else {
+        Duration::from_millis(expires_in as u64)
+    };
+
+    Ok(conn.db().expire(&args[1], expires_at))
 }
 
 pub fn ttl(conn: &Connection, args: &[Bytes]) -> Result<Value, Error> {
     let ttl = match conn.db().ttl(&args[1]) {
-        Some(Some(ttl)) => (ttl - Instant::now()).as_secs() as i64,
+        Some(Some(ttl)) => {
+            let ttl = ttl - Instant::now();
+            if check_arg!(args, 0, "TTL") {
+                ttl.as_secs() as i64
+            } else {
+                ttl.as_millis() as i64
+            }
+        }
         Some(None) => -1,
         None => -2,
     };
 
     Ok(ttl.into())
 }
+
+pub fn expire_time(conn: &Connection, args: &[Bytes]) -> Result<Value, Error> {
+    let ttl = match conn.db().ttl(&args[1]) {
+        Some(Some(ttl)) => {
+            // Is there a better way? There should be!
+            if check_arg!(args, 0, "EXPIRETIME") {
+                let secs: i64 = (ttl - Instant::now()).as_secs() as i64;
+                secs + (now().as_secs() as i64)
+            } else {
+                let secs: i64 = (ttl - Instant::now()).as_millis() as i64;
+                secs + (now().as_millis() as i64)
+            }
+        }
+        Some(None) => -1,
+        None => -2,
+    };
+
+    Ok(ttl.into())
+}
+
+pub fn persist(conn: &Connection, args: &[Bytes]) -> Result<Value, Error> {
+    Ok(conn.db().persist(&args[1]))
+}

src/dispatcher.rs

@@ -47,6 +47,16 @@ dispatcher! {
         ["read" "write" "fast"],
         3,
     },
+    expireat {
+        cmd::key::expire_at,
+        ["read" "write" "fast"],
+        3,
+    },
+    expiretime {
+        cmd::key::expire_time,
+        ["read" "write" "fast"],
+        2,
+    },
     del {
         cmd::key::del,
         ["random" "loading" "stale"],
@@ -77,6 +87,26 @@ dispatcher! {
         ["read" "read"],
         2,
     },
+    pexpire {
+        cmd::key::expire,
+        ["read" "write" "fast"],
+        3,
+    },
+    pexpireat {
+        cmd::key::expire_at,
+        ["read" "write" "fast"],
+        3,
+    },
+    pexpiretime {
+        cmd::key::expire_time,
+        ["read" "write" "fast"],
+        2,
+    },
+    pttl {
+        cmd::key::ttl,
+        ["read" "read"],
+        2,
+    },
     set {
         cmd::string::set,
         ["random" "loading" "stale"],

src/macros.rs

@@ -130,3 +130,16 @@ macro_rules! option {
         }
     }
 }
+
+#[macro_export]
+macro_rules! check_arg {
+    {$args: tt, $pos: tt, $command: tt} => {{
+        match $args.get($pos) {
+            Some(bytes) => {
+                let command = unsafe { std::str::from_utf8_unchecked(&bytes) };
+                command.to_uppercase() == $command
+            },
+            None => false,
+        }
+    }}
+}