Implement the SET & GET commands

Author: fangpin

src/cmd.rs

@@ -1,10 +1,13 @@
-use crate::protocol::Protocol;
+use std::sync::{Arc, Mutex};
+
+use crate::{protocol::Protocol, storage::Storage};
 use anyhow::Result;
 
 pub enum Cmd {
     Ping,
     Echo(String),
-    Unknown(),
+    Get(String),
+    Set(String, String),
 }
 
 impl Cmd {
@@ -19,18 +22,34 @@ impl Cmd {
                 Ok(match cmd[0].as_str() {
                     "echo" => Cmd::Echo(cmd[1].clone()),
                     "ping" => Cmd::Ping,
-                    _ => Cmd::Unknown(),
+                    "get" => Cmd::Get(cmd[1].clone()),
+                    "set" => Cmd::Set(cmd[1].clone(), cmd[2].clone()),
+                    _ => return Err(anyhow::anyhow!("unknown cmd {:?}", cmd[0])),
                 })
             }
             _ => Err(anyhow::anyhow!("fail to parse as cmd for {:?}", protocol.0)),
         }
     }
 
-    pub fn run(self: &Self) -> Result<Protocol> {
+    pub fn run(self: &Self, storage: &mut Arc<Mutex<Storage>>) -> Result<Protocol> {
         match self {
-            Cmd::Echo(s) => Ok(Protocol::SimpleString(s.clone())),
             Cmd::Ping => Ok(Protocol::SimpleString("PONG".to_string())),
-            _ => Err(anyhow::anyhow!("unknown cmd")),
+            Cmd::Echo(s) => Ok(Protocol::SimpleString(s.clone())),
+            Cmd::Get(k) => {
+                let s = storage.lock().unwrap();
+                Ok(if let Some(v) = s.get(k) {
+                    Protocol::SimpleString(v.clone())
+                } else {
+                    Protocol::Null
+                })
+            },
+            Cmd::Set(k, v) => {
+                {
+                    let mut s = storage.lock().unwrap();
+                    s.set(k.clone(), v.clone());
+                }
+                Ok(Protocol::ok())
+            }
         }
     }
 }

src/handler.rs

@@ -1,3 +1,4 @@
 mod cmd;
-pub mod handler;
 mod protocol;
+pub mod server;
+mod storage;

src/lib.rs

@@ -1,6 +1,6 @@
 #![allow(unused_imports)]
 
-use redis_starter_rust::handler;
+use redis_starter_rust::server;
 
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::net::TcpListener;
@@ -10,15 +10,17 @@ async fn main() {
     println!("Logs from your program will appear here!");
 
     let listener = TcpListener::bind("127.0.0.1:6379").await.unwrap();
+    let server = server::Server::new();
 
     loop {
         let stream = listener.accept().await;
         match stream {
             Ok((stream, _)) => {
                 println!("accepted new connection");
 
+                let mut sc = server.clone();
                 tokio::spawn(async move {
-                    handler::handle(stream).await;
+                    sc.handle(stream).await;
                 });
             }
             Err(e) => {

src/main.rs

@@ -4,6 +4,7 @@ use anyhow::Result;
 pub enum Protocol {
     SimpleString(String),
     BulkString(String),
+    Null,
     Array(Vec<Protocol>),
 }
 
@@ -24,10 +25,16 @@ impl Protocol {
         }
     }
 
+    #[inline]
+    pub fn ok() -> Self {
+        Protocol::SimpleString("ok".to_string())
+    }
+
     pub fn decode(self: &Self) -> String {
         match self {
             Protocol::SimpleString(s) => s.to_string(),
             Protocol::BulkString(s) => s.to_string(),
+            Protocol::Null => "".to_string(),
             Protocol::Array(s) => s
                 .into_iter()
                 .map(|x| x.decode())
@@ -48,6 +55,7 @@ impl Protocol {
                         .join("")
                         .as_str()
             }
+            Protocol::Null => "$-1\r\n".to_string(),
         }
     }
 

src/protocol.rs

@@ -0,0 +1,42 @@
+use core::str;
+use std::sync::Arc;
+use std::sync::Mutex;
+use tokio::io::AsyncReadExt;
+use tokio::io::AsyncWriteExt;
+
+use crate::cmd::Cmd;
+use crate::storage::Storage;
+
+#[derive(Clone)]
+pub struct Server {
+    storage: Arc<Mutex<Storage>>,
+}
+
+impl Server {
+    pub fn new() -> Self {
+        Server {
+            storage: Arc::new(Mutex::new(Storage::new())),
+        }
+    }
+
+    pub async fn handle(self: &mut Self, mut stream: tokio::net::TcpStream) {
+        let mut buf = [0; 512];
+        loop {
+            if let Ok(len) = stream.read(&mut buf).await {
+                if len == 0 {
+                    println!("[handle] connection closed");
+                    return;
+                }
+                let s = str::from_utf8(&buf[..len]).unwrap();
+                let cmd = Cmd::from(s).unwrap();
+                let res = cmd.run(&mut self.storage).unwrap();
+                println!("going to send response {}", res.encode());
+                stream.write_all(res.encode().as_bytes()).await.unwrap();
+                println!("finish processing");
+            } else {
+                println!("[handle] going to break");
+                break;
+            }
+        }
+    }
+}

src/server.rs

@@ -0,0 +1,21 @@
+use std::collections::HashMap;
+
+pub struct Storage {
+    set: HashMap<String, String>,
+}
+
+impl Storage {
+    pub fn new() -> Self {
+        Storage {
+            set: HashMap::new(),
+        }
+    }
+
+    pub fn get(self: &Self, k: &str) -> Option<&String> {
+        self.set.get(k)
+    }
+
+    pub fn set(self: &mut Self, k: String, v: String) {
+        self.set.insert(k, v);
+    }
+}