Set key with expiry

Author: fangpin

Cargo.lock

@@ -2,7 +2,7 @@
 name = "redis-rs"
 version = "0.0.1"
 authors = ["Pin Fang <[email protected]>"]
-edition = "2024"
+edition = "2021"
 
 [dependencies]
 anyhow = "1.0.59"

Cargo.toml

@@ -1,6 +1,6 @@
 # Build Your Own Redis
 
-To build a toy Redis-Server clone that's capable of handling
+This project is to build a toy Redis-Server clone that's capable of handling
 basic commands like `PING`, `SET` and `GET`. Also implement the event loops, the Redis protocol and more.
 
 ## Prerequisites
@@ -23,6 +23,10 @@ redis-cli ECHO hey
 redis-cli SET foo bar
 ```
 
+```sh
+redis-cli SET foo bar px 100
+```
+
 ```sh
 redis-cli GET foo
 ```

README.md

@@ -8,6 +8,8 @@ pub enum Cmd {
     Echo(String),
     Get(String),
     Set(String, String),
+    SetPx(String, String, u128),
+    SetEx(String, String, u128),
 }
 
 impl Cmd {
@@ -23,7 +25,15 @@ impl Cmd {
                     "echo" => Cmd::Echo(cmd[1].clone()),
                     "ping" => Cmd::Ping,
                     "get" => Cmd::Get(cmd[1].clone()),
-                    "set" => Cmd::Set(cmd[1].clone(), cmd[2].clone()),
+                    "set" => {
+                        if cmd.len() == 5 && cmd[3] == "px" {
+                            Cmd::SetPx(cmd[1].clone(), cmd[2].clone(), cmd[4].parse().unwrap())
+                        } else if cmd.len() == 5 && cmd[3] == "ex" {
+                            Cmd::SetEx(cmd[1].clone(), cmd[2].clone(), cmd[4].parse().unwrap())
+                        } else {
+                            Cmd::Set(cmd[1].clone(), cmd[2].clone())
+                        }
+                    }
                     _ => return Err(anyhow::anyhow!("unknown cmd {:?}", cmd[0])),
                 })
             }
@@ -42,14 +52,28 @@ impl Cmd {
                 } else {
                     Protocol::Null
                 })
-            },
+            }
             Cmd::Set(k, v) => {
                 {
                     let mut s = storage.lock().unwrap();
                     s.set(k.clone(), v.clone());
                 }
                 Ok(Protocol::ok())
             }
+            Cmd::SetPx(k, v, x) => {
+                {
+                    let mut s = storage.lock().unwrap();
+                    s.setx(k.clone(), v.clone(), *x);
+                }
+                Ok(Protocol::ok())
+            }
+            Cmd::SetEx(k, v, x) => {
+                {
+                    let mut s = storage.lock().unwrap();
+                    s.setx(k.clone(), v.clone(), *x * 1000);
+                }
+                Ok(Protocol::ok())
+            }
         }
     }
 }

src/cmd.rs

@@ -1,6 +1,6 @@
 #![allow(unused_imports)]
 
-use redis_starter_rust::server;
+use redis_rs::server;
 
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::net::TcpListener;

src/main.rs

@@ -1,7 +1,9 @@
-use std::collections::HashMap;
+use core::time;
+use std::{collections::HashMap, time::Instant};
 
 pub struct Storage {
-    set: HashMap<String, String>,
+    // key -> (value, (insert/update time, expire milli seconds))
+    set: HashMap<String, (String, Option<(Instant, u128)>)>,
 }
 
 impl Storage {
@@ -12,10 +14,26 @@ impl Storage {
     }
 
     pub fn get(self: &Self, k: &str) -> Option<&String> {
-        self.set.get(k)
+        match self.set.get(k) {
+            Some((ss, time_info)) => match time_info {
+                Some((instant, ms)) => {
+                    if instant.elapsed().as_millis() > *ms {
+                        None
+                    } else {
+                        Some(ss)
+                    }
+                }
+                _ => Some(ss),
+            },
+            _ => None,
+        }
     }
 
     pub fn set(self: &mut Self, k: String, v: String) {
-        self.set.insert(k, v);
+        self.set.insert(k, (v, None));
+    }
+
+    pub fn setx(self: &mut Self, k: String, v: String, expire_ms: u128) {
+        self.set.insert(k, (v, Some((Instant::now(), expire_ms))));
     }
 }