feat: reached mvp

Author: ynachi

.github/workflows/ci.yml

@@ -0,0 +1,22 @@
+name: Rust
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+
+env:
+  CARGO_TERM_COLOR: always
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v3
+      - name: Build
+        run: cargo build --verbose
+      - name: Run tests
+        run: cargo test --verbose

Cargo.lock

@@ -8,6 +8,7 @@ bytes = "1"
 tokio = {version = "1", features = ["full"]}
 tracing = "0.1"
 rustc-hash = "1"
+tracing-subscriber = "0.3.18"
 
 [dev-dependencies]
 criterion = "0.5.1"

Cargo.toml

@@ -4,7 +4,7 @@ use std::path::Path;
 use std::sync::Arc;
 use std::time::Duration;
 
-use criterion::{criterion_group, criterion_main, Criterion, black_box};
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
 use csv::ReaderBuilder;
 use rayon::prelude::*;
 

benches/bench_db.rs

@@ -0,0 +1,102 @@
+use std::time::Duration;
+
+use crate::frame::{Frame, FrameError};
+
+#[derive(Debug)]
+pub enum Command {
+    Ping(Option<String>),
+    Set(String, String, Duration),
+    Get(String),
+    Unknown(String),
+}
+
+impl Command {
+    /// Parse a command from a frame of array of bulk. This method should not be passed
+    /// another type or frame contents. Because RESP commands are represented as an array of bulk frames.
+    pub fn from_frame(frame: Frame) -> Result<Self, FrameError> {
+        // users are aware they should only give an array of bulk frame to this method but we have to
+        // check anyway. We opted to perform lazy validation as validating the whole array upfront
+        // can be expensive.
+        let frames = frame.get_array().ok_or(FrameError::Syntax(
+            "commands can only be frame array".to_string(),
+        ))?;
+        if let Some(name) = Self::get_name(frames) {
+            return match name.to_ascii_uppercase().as_str() {
+                "PING" => Command::parse_ping(&frames[1..]),
+                "GET" => Command::parse_get(&frames[1..]),
+                "SET" => Command::parse_set(&frames[1..]),
+                _ => Ok(Command::Unknown(name.to_string())),
+            };
+        }
+        Err(FrameError::Syntax(
+            "RESP command name should be of type bulk frame".to_string(),
+        ))
+    }
+
+    /// parse_get tries to retrieve get args from a slice of frames
+    fn parse_get(frames: &[Frame]) -> Result<Command, FrameError> {
+        if frames.len() != 1 {
+            return Err(FrameError::Syntax(
+                "PING command takes at most 1 argument".to_string(),
+            ));
+        }
+        Ok(Command::Get(Self::get_string(frames, 0)?))
+    }
+
+    fn get_name(frames: &[Frame]) -> Option<&String> {
+        if frames.is_empty() {
+            return None;
+        }
+        if let Some((_, name)) = frames[0].get_bulk() {
+            return Some(name);
+        }
+        None
+    }
+
+    /// parse_ping tries to retrieve ping args from a slice of frames
+    fn parse_ping(frames: &[Frame]) -> Result<Command, FrameError> {
+        match frames.len() {
+            0 => Ok(Command::Ping(None)),
+            1 => Ok(Command::Ping(Some(Self::get_string(frames, 0)?))),
+            _ => Err(FrameError::Syntax(
+                "PING command takes at most 1 argument".to_string(),
+            )),
+        }
+    }
+
+    fn get_string(frames: &[Frame], index: usize) -> Result<String, FrameError> {
+        match frames[index].get_bulk() {
+            Some(message) => Ok(message.1.to_owned()),
+            None => Err(FrameError::Syntax(
+                "RESP args should be of type bulk frame".to_string(),
+            )),
+        }
+    }
+
+    fn parse_set(frames: &[Frame]) -> Result<Command, FrameError> {
+        let len = frames.len();
+        if len != 2 && len != 4 {
+            return Err(FrameError::Syntax(
+                "SET should take 2 or 4 arguments".to_string(),
+            ));
+        }
+        let key = Self::get_string(frames, 0)?;
+        let value = Self::get_string(frames, 1)?;
+
+        let mut ttl = Duration::from_millis(0);
+
+        // check if we've got the right option to set the time in millis
+        if len > 2 && Self::get_string(frames, 2)?.to_uppercase() == "PX" {
+            let expiration = Self::get_u64(frames, 3)?;
+            ttl = Duration::from_millis(expiration);
+        }
+
+        Ok(Command::Set(key, value, ttl))
+    }
+
+    fn get_u64(frames: &[Frame], index: usize) -> Result<u64, FrameError> {
+        let content = Self::get_string(frames, index)?;
+        let content = content.parse::<u64>().map_err(|_| FrameError::UTF8ToInt)?;
+        Ok(content)
+    }
+}

src/command.rs

@@ -3,13 +3,13 @@
 //! We do not store the state for eviction. Time-based eviction is used and we perform lazy eviction.
 //! If an expired key is read, this key is deleted and no value is returned to the user.
 //! ...
-use std::collections::{BinaryHeap};
+use rustc_hash::FxHashMap;
+use std::collections::BinaryHeap;
 use std::fmt::Debug;
 use std::hash::{Hash, Hasher};
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::{Arc, Mutex};
 use std::time::{Duration, Instant};
-use rustc_hash::FxHashMap;
 
 struct Shard {
     storage: FxHashMap<String, String>,
@@ -42,7 +42,7 @@ impl Shard {
             0
         }
     }
-    
+
     fn latest_is_expired(&self) -> bool {
         if let Some((instant, _)) = self.eviction_state.peek() {
             if Instant::now() > *instant {
@@ -51,7 +51,7 @@ impl Shard {
         }
         false
     }
-    
+
     fn del_latest(&mut self) {
         if let Some((_, key)) = self.eviction_state.pop() {
             self.storage.remove(&key);
@@ -150,7 +150,6 @@ impl Storage {
     }
 }
 
-
 #[cfg(test)]
 mod tests {
     use super::*;

src/db.rs

@@ -8,9 +8,11 @@
 //! simple frames are valid at their creation. We do that because we want to pay the cost of checking
 //! this property only if needed as it is expensive.
 
-use crate::frame::FrameData::Nested;
-use tokio::io::ErrorKind;
-use tokio::io::{AsyncBufReadExt, AsyncReadExt, BufReader};
+use std::fmt::{self, Debug, Display, Formatter};
+use tokio::io::{
+    self, AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader, BufWriter, ErrorKind,
+};
+use tracing::error;
 
 const CR: u8 = b'\r';
 const LF: u8 = b'\n';
@@ -20,7 +22,7 @@ const LF: u8 = b'\n';
 ///  up front, but we have changed our mind.
 #[derive(Debug, Clone, Copy, Eq, PartialEq)]
 #[repr(u8)]
-pub(crate) enum FrameID {
+pub enum FrameID {
     Integer = 58, // ':'
     // @TODO: remove for now
     // Double = 44,       // ','
@@ -89,21 +91,156 @@ enum FrameData {
     Nested(Vec<Frame>),
 }
 
+impl FrameData {
+    fn get_integer(&self) -> Option<i64> {
+        match self {
+            FrameData::Integer(value) => Some(*value),
+            _ => None,
+        }
+    }
+    fn get_string(&self) -> Option<&String> {
+        match self {
+            FrameData::Simple(value) => Some(value),
+            _ => None,
+        }
+    }
+    fn get_bulk(&self) -> Option<(usize, &String)> {
+        match self {
+            FrameData::Bulk(size, data) => Some((*size, data)),
+            _ => None,
+        }
+    }
+    fn get_boolean(&self) -> Option<bool> {
+        match self {
+            FrameData::Boolean(value) => Some(*value),
+            _ => None,
+        }
+    }
+    pub fn get_nested(&self) -> Option<&Vec<Frame>> {
+        match self {
+            FrameData::Nested(value) => Some(value),
+            _ => None,
+        }
+    }
+}
+
 #[derive(Debug, Eq, PartialEq)]
-struct Frame {
+pub struct Frame {
     frame_type: FrameID,
     frame_data: FrameData,
 }
 
-fn validate_bool(data: String) -> Result<bool, FrameError> {
-    match data.as_str() {
+impl Frame {
+    pub fn get_id(&self) -> FrameID {
+        self.frame_type
+    }
+
+    pub fn get_array(&self) -> Option<&Vec<Frame>> {
+        if self.frame_type != FrameID::Array {
+            return None;
+        }
+        self.frame_data.get_nested()
+    }
+    pub fn get_bulk(&self) -> Option<(usize, &String)> {
+        match &self.frame_data {
+            FrameData::Bulk(size, data) => Some((*size, data)),
+            _ => None,
+        }
+    }
+
+    pub async fn write_flush_all<T>(&self, stream: &mut BufWriter<T>) -> io::Result<()>
+    where
+        T: AsyncWriteExt + Unpin,
+    {
+        stream.write_all(self.to_string().as_bytes()).await?;
+        stream.flush().await
+    }
+
+    pub fn new_bulk_error(inner: String) -> Frame {
+        Frame {
+            frame_type: FrameID::BulkError,
+            frame_data: FrameData::Bulk(inner.len(), inner),
+        }
+    }
+
+    pub fn new_simple_string(inner: String) -> Frame {
+        Frame {
+            frame_type: FrameID::SimpleString,
+            frame_data: FrameData::Simple(inner),
+        }
+    }
+
+    pub fn new_bulk_string(inner: String) -> Frame {
+        Frame {
+            frame_type: FrameID::BulkString,
+            frame_data: FrameData::Bulk(inner.len(), inner),
+        }
+    }
+
+    pub fn new_null() -> Frame {
+        Frame {
+            frame_type: FrameID::Null,
+            frame_data: FrameData::Null,
+        }
+    }
+}
+
+impl fmt::Display for Frame {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        match self.frame_type {
+            FrameID::Integer => {
+                let value = self.frame_data.get_integer().ok_or(fmt::Error)?;
+                write!(f, ":{}\r\n", value)
+            }
+            FrameID::SimpleString => {
+                let value = self.frame_data.get_string().ok_or(fmt::Error)?;
+                write!(f, "+{}\r\n", value)
+            }
+            FrameID::SimpleError => {
+                let value = self.frame_data.get_string().ok_or(fmt::Error)?;
+                write!(f, "-{}\r\n", value)
+            }
+            FrameID::BulkString => {
+                let (len, value) = self.get_bulk().ok_or(fmt::Error)?;
+                write!(f, "${}\r\n{}\r\n", len, value)
+            }
+            FrameID::BulkError => {
+                let (len, value) = self.frame_data.get_bulk().ok_or(fmt::Error)?;
+                write!(f, "!{}\r\n{}\r\n", len, value)
+            }
+            FrameID::Boolean => {
+                let value = self.frame_data.get_boolean().ok_or(fmt::Error)?;
+                let value = if value { "t" } else { "f" };
+                write!(f, "#{}\r\n", value)
+            }
+            FrameID::Null => {
+                write!(f, "_\r\n")
+            }
+            FrameID::BigNumber => {
+                let value = self.frame_data.get_string().ok_or(fmt::Error)?;
+                write!(f, "({}\r\n", value)
+            }
+            FrameID::Array => {
+                let frames = self.frame_data.get_nested().ok_or(fmt::Error)?;
+                write!(f, "*{}\r\n", frames.len())?;
+                for v in frames {
+                    write!(f, "{}", v)?;
+                }
+                Ok(())
+            }
+        }
+    }
+}
+
+fn validate_bool(data: &str) -> Result<bool, FrameError> {
+    match data {
         "t" => Ok(true),
         "f" => Ok(false),
         _ => Err(FrameError::Invalid),
     }
 }
 
-async fn decode<T>(stream: &mut BufReader<T>) -> Result<Frame, FrameError>
+pub async fn decode<T>(stream: &mut BufReader<T>) -> Result<Frame, FrameError>
 where
     T: AsyncReadExt + Unpin,
 {
@@ -122,7 +259,7 @@ where
             let frame_vec = process_aggregate_frames(id, stream).await?;
             Ok(Frame {
                 frame_type: FrameID::Array,
-                frame_data: Nested(frame_vec),
+                frame_data: FrameData::Nested(frame_vec),
             })
         }
     }
@@ -138,7 +275,7 @@ where
     let data = read_simple_string(stream).await?;
     match id {
         FrameID::Boolean => {
-            let bool = validate_bool(data)?;
+            let bool = validate_bool(&data)?;
             Ok(Frame {
                 frame_type: id,
                 frame_data: FrameData::Boolean(bool),
@@ -189,7 +326,9 @@ where
     T: AsyncReadExt + Unpin,
 {
     match id {
-        FrameID::Array => Err(FrameError::Syntax),
+        FrameID::Array => Err(FrameError::Syntax(
+            "received aggregate frame in non aggregate decoding".to_string(),
+        )),
         FrameID::BulkString | FrameID::BulkError => process_bulk_frames(id, stream).await,
         _ => process_simple_frames(id, stream).await,
     }
@@ -249,7 +388,7 @@ where
                         // to build the right aggregate.
                         frames.push(Frame {
                             frame_type: *id,
-                            frame_data: Nested(last_vec_of_frames),
+                            frame_data: FrameData::Nested(last_vec_of_frames),
                         });
                         *count -= 1;
                         if *count != 0 {
@@ -269,9 +408,9 @@ pub enum FrameError {
     // Frame is not correctly formatted
     Invalid,
     // Empty buffer should not be passed to get frame from, so this is an error.
-    EmptyBuffer,
+    // EmptyBuffer,
     // reached expected EOF
-    EOF,
+    Eof,
     // Connection unexpectedly reset
     ConnectionReset,
     // Unidentified IO error
@@ -281,7 +420,23 @@ pub enum FrameError {
     // Unknown frame type
     Unknown,
     // This is a programming error. It should not happen.
-    Syntax,
+    Syntax(String),
+}
+
+impl Display for FrameError {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        match self {
+            FrameError::Incomplete => write!(f, "not enough data to decode a full frame"),
+            FrameError::Invalid => write!(f, "frame is not correctly formatted"),
+            FrameError::Eof => write!(f, "seen EOF, this is generally a graceful disconnection"),
+            FrameError::ConnectionReset => write!(f, "unexpected connection reset"),
+            // this error should not happen in practice
+            FrameError::IOError => write!(f, "unexpected IO error"),
+            FrameError::UTF8ToInt => write!(f, "utf8 to int decoding error"),
+            FrameError::Unknown => write!(f, "unable to identify the frame type"),
+            FrameError::Syntax(message) => write!(f, "{}", message),
+        }
+    }
 }
 
 /// `read_simple_string` gets a simple string from the network. As a reminder, such string does
@@ -295,7 +450,7 @@ where
 {
     let mut buf = Vec::new();
     match stream.read_until(LF, &mut buf).await {
-        Ok(0) => Err(FrameError::EOF),
+        Ok(0) => Err(FrameError::Eof),
         Ok(size) => {
             if size < 2 {
                 return Err(FrameError::Incomplete);
@@ -308,8 +463,10 @@ where
             Ok(String::from_utf8_lossy(&buf[0..size - 2]).to_string())
         }
         Err(e) if e.kind() == ErrorKind::ConnectionReset => Err(FrameError::ConnectionReset),
-        // @TODO log e later
-        Err(e) => Err(FrameError::IOError),
+        Err(e) => {
+            error!("unexpected io error: {}", e);
+            Err(FrameError::IOError)
+        }
     }
 }
 
@@ -322,7 +479,7 @@ where
             Some(id) => Ok(id),
             None => Err(FrameError::Unknown),
         },
-        Err(e) if e.kind() == ErrorKind::UnexpectedEof => Err(FrameError::EOF),
+        Err(e) if e.kind() == ErrorKind::UnexpectedEof => Err(FrameError::Eof),
         // @TODO log e later
         Err(e) => Err(FrameError::IOError),
     }
@@ -360,7 +517,7 @@ where
             ))
         }
         // The caller will treat EOF differently, so it needs to be returned explicitly
-        Err(e) if e.kind() == ErrorKind::UnexpectedEof => Err(FrameError::EOF),
+        Err(e) if e.kind() == ErrorKind::UnexpectedEof => Err(FrameError::Eof),
         // @TODO log e later
         Err(e) => Err(FrameError::IOError),
     }
@@ -369,6 +526,7 @@ where
 #[cfg(test)]
 mod tests {
     use super::*;
+
     #[tokio::test]
     async fn decode_test() {
         //
@@ -394,7 +552,7 @@ mod tests {
         let frame = decode(&mut stream).await;
         assert_eq!(
             frame,
-            Err(FrameError::EOF),
+            Err(FrameError::Eof),
             "should return EOF error variant"
         );
 
@@ -441,7 +599,7 @@ mod tests {
         let mut stream = BufReader::new("*3\r\n:1\r\n+Two\r\n$5\r\nThree\r\n".as_bytes());
         let frame = decode(&mut stream).await.unwrap();
         assert_eq!(frame.frame_type, FrameID::Array);
-        let frame_data = Nested(vec![
+        let frame_data = FrameData::Nested(vec![
             Frame {
                 frame_type: FrameID::Integer,
                 frame_data: FrameData::Integer(1),
@@ -461,7 +619,7 @@ mod tests {
         let mut stream = BufReader::new("*2\r\n:1\r\n*1\r\n+Three\r\n".as_bytes());
         let frame = decode(&mut stream).await.unwrap();
         assert_eq!(frame.frame_type, FrameID::Array);
-        let frame_data = Nested(vec![
+        let frame_data = FrameData::Nested(vec![
             Frame {
                 frame_type: FrameID::Integer,
                 frame_data: FrameData::Integer(1),
@@ -480,7 +638,7 @@ mod tests {
         let mut stream = BufReader::new("*3\r\n:1\r\n*1\r\n+Three\r\n-Err\r\n".as_bytes());
         let frame = decode(&mut stream).await.unwrap();
         assert_eq!(frame.frame_type, FrameID::Array);
-        let frame_data = Nested(vec![
+        let frame_data = FrameData::Nested(vec![
             Frame {
                 frame_type: FrameID::Integer,
                 frame_data: FrameData::Integer(1),

src/frame.rs

@@ -1,2 +1,4 @@
+mod command;
 pub mod db;
-mod frame;
+pub mod frame;
+pub mod server;

src/lib.rs

@@ -1,5 +1,18 @@
+use mredis::server::{Config, Server};
+
 mod frame;
 
-fn main() {
-    println!("Hello, world!");
+#[tokio::main]
+pub async fn main() -> std::io::Result<()> {
+    // console_subscriber::init();
+    let cfg = Config {
+        ip_addr: "127.0.0.1".to_string(),
+        port: 6379,
+        capacity: 1_000_000,
+        shard_count: 32,
+    };
+    tracing_subscriber::fmt::try_init().expect("unable to initialize logging");
+    let server = Server::new(&cfg).await;
+    server.listen().await;
+    Ok(())
 }

src/main.rs

@@ -0,0 +1,149 @@
+use crate::command::Command;
+use crate::db::Storage;
+use crate::frame::{decode, Frame, FrameError};
+use std::sync::Arc;
+use tokio::io::{AsyncWriteExt, BufReader, BufWriter};
+use tokio::net::{TcpListener, TcpStream};
+use tracing::{debug, error};
+
+pub struct Config {
+    pub ip_addr: String,
+    pub port: u16,
+    pub capacity: usize,
+    pub shard_count: usize,
+}
+
+pub struct Server {
+    storage: Arc<Storage>,
+    tcp_listener: TcpListener,
+}
+
+impl Server {
+    pub async fn new(cfg: &Config) -> Self {
+        let tcp_listener = TcpListener::bind((cfg.ip_addr.to_owned(), cfg.port))
+            .await
+            .expect("failed to start TCP server");
+        let storage = Arc::new(Storage::new(cfg.capacity, cfg.shard_count));
+
+        Server {
+            storage,
+            tcp_listener,
+        }
+    }
+
+    pub async fn listen(&self) {
+        loop {
+            let conn_string = self.tcp_listener.accept().await;
+            match conn_string {
+                Ok((mut stream, addr)) => {
+                    debug!("new connection established: {}", addr);
+
+                    let state = self.storage.clone();
+
+                    tokio::spawn(async move { Self::process_stream(&mut stream, state).await });
+                }
+                Err(err) => {
+                    debug!("error accepting client connection: {:?}", err);
+                }
+            }
+        }
+    }
+
+    async fn process_stream(stream: &mut TcpStream, state: Arc<Storage>) {
+        let (reader_half, writer_half) = stream.split();
+
+        let mut reader = BufReader::with_capacity(8 * 1024, reader_half);
+        let mut writer = BufWriter::with_capacity(8 * 1024, writer_half);
+
+        loop {
+            let frame = decode(&mut reader).await;
+            match frame {
+                Ok(frame) => {
+                    debug!("command frame received!");
+                    process_frame(frame, &state, &mut writer).await;
+                }
+                Err(err) => {
+                    if seen_eof(&err, &mut writer).await {
+                        debug!("client gracefully closed connection");
+                        return;
+                    }
+                }
+            }
+        }
+    }
+}
+
+async fn process_frame<T>(frame: Frame, state: &Arc<Storage>, stream_writer: &mut BufWriter<T>)
+where
+    T: AsyncWriteExt + Unpin,
+{
+    match Command::from_frame(frame) {
+        Ok(cmd) => {
+            apply_command(&cmd, state, stream_writer).await;
+        }
+        Err(err) => send_error(&err, stream_writer).await,
+    }
+}
+
+async fn apply_command<T: AsyncWriteExt + Unpin>(command: &Command, state: &Arc<Storage>, stream_writer: &mut BufWriter<T>)
+{
+    let response_frame = match command {
+        Command::Ping(message) => {
+            if let Some(message) = message {
+                Frame::new_bulk_string(message.to_string())
+            } else {
+                Frame::new_simple_string("PONG".to_string())
+            }
+        }
+        Command::Get(key) => {
+            if let Some(ans) = state.get_v(key) {
+                Frame::new_simple_string(ans)
+            } else {
+                Frame::new_null()
+            }
+        }
+        Command::Set(key, value, ttl) => {
+            // if let Some(ans) = state.set_kv(key, value, *ttl) {
+            //     Frame::new_simple_string(ans)
+            // } else {
+            //     Frame::new_null()
+            // }
+            state.set_kv(key, value, *ttl);
+            Frame::new_simple_string("OK".to_string())
+        }
+        Command::Unknown(name) => Frame::new_bulk_error(format!("unknown command: {}", name)),
+    };
+    response_frame.write_flush_all(stream_writer).await.unwrap()
+}
+
+/// seen_eof filters FrameError because some errors need to be sent back to the client via
+/// the network, for instance, syntax errors. In this case, send the error to the client. If EOF
+/// return true to the caller. And, only log over error variants.
+async fn seen_eof<T: AsyncWriteExt + Unpin>(err: &FrameError, stream_writer: &mut BufWriter<T>) -> bool
+{
+    match err {
+        FrameError::Eof => true,
+        FrameError::Incomplete
+        | FrameError::Invalid
+        | FrameError::Unknown
+        | FrameError::UTF8ToInt
+        | FrameError::Syntax(_) => {
+            send_error(err, stream_writer).await;
+            false
+        }
+        _ => {
+            error!("error while decoding frame: {}", err);
+            false
+        }
+    }
+}
+
+/// send_error is a wrapper to send errors to the client over the network.
+/// These are mostly syntax errors.
+async fn send_error<T: AsyncWriteExt + Unpin>(err: &FrameError, stream: &mut BufWriter<T>)
+{
+    let err_frame = Frame::new_bulk_error(err.to_string());
+    if let Err(err) = err_frame.write_flush_all(stream).await {
+        error!("failed to write to network: {}", err);
+    }
+}