Simple Redis implementation in Rust

Question

This is a very tiny implementation of Rust that uses redis serialization protocol (RESP). I implemented it as one of code crafters challenges. I have less than a few month experience in Rust, so it may not seems to be a very advanced Rust program. But I appreciate any feedback to improve my implementation and also to make my Rust code more advanced.

use std::io::{Read, Write};
use std::net::{TcpListener, TcpStream};
use std::{
    collections,
    sync::{Arc, Mutex},
    thread, time,
};

fn main() {
    let listener = TcpListener::bind("127.0.0.1:6379").unwrap();

    let data_storage = Arc::new(Mutex::new(collections::HashMap::<
        String,
        (Option<time::Instant>, Vec<u8>),
    >::new()));

    for stream in listener.incoming() {
        match stream {
            Ok(s) => {
                let cloned_storage = data_storage.clone();
                thread::spawn(move || handle(s, cloned_storage));
            }
            Err(e) => {
                println!("error: {}", e);
            }
        }
    }
}

fn handle(
    mut stream: TcpStream,
    storage: Arc<Mutex<collections::HashMap<String, (Option<time::Instant>, Vec<u8>)>>>,
) {
    let mut buf = [0u8; 64];
    loop {
        let read_count = stream.read(&mut buf).expect("Could not read from client");
        if read_count == 0 {
            return;
        }
        let mut new_buf = Vec::new();
        for i in 0..read_count {
            new_buf.push(buf[i]);
        }
        match Command::from_buffer(&new_buf.as_slice()) {
            Ok(Command::Ping) => {
                stream.write_all(b"+PONG\r\n");
            }
            Ok(Command::Echo(s)) => {
                let out = serialize_to_bulk_string(s.as_bytes());
                stream.write_all(out.as_slice());
            }
            Ok(Command::Set(key, value, expiry)) => {
                let mut storage = storage.lock().unwrap();
                let expiry =
                    expiry.map(|t| time::Instant::now() + time::Duration::from_millis(t as u64));
                storage.insert(key, (expiry, value));
                let out = serialize_to_simple_string("OK".as_bytes());
                stream.write_all(out.as_slice());
            }
            Ok(Command::Get(key)) => {
                let mut storage = storage.lock().unwrap();
                match storage.get(&key) {
                    Some((expiry, v)) => {
                        if let Some(expiry) = expiry {
                            if time::Instant::now() >= *expiry {
                                stream.write_all(b"$-1\r\n");
                            } else {
                                let out = serialize_to_bulk_string(v);
                                stream.write_all(out.as_slice());
                            }
                        } else {
                            let out = serialize_to_bulk_string(v);
                            stream.write_all(out.as_slice());
                        }
                    }
                    None => {
                        stream.write_all(b"$-1\r\n");
                    }
                }
            }
            Err(_) => {
                stream.write_all(b"-Error\r\n");
            }
        }
    }
}

fn serialize_to_simple_string(s: &[u8]) -> Vec<u8> {
    [b"+", s, b"\r\n"].concat()
}

fn serialize_to_bulk_string(s: &[u8]) -> Vec<u8> {
    [b"$", format!("{}", s.len()).as_bytes(), b"\r\n", s, b"\r\n"].concat()
}

#[derive(Debug)]
enum Command {
    Ping,
    Echo(String),
    Set(String, Vec<u8>, Option<u64>),
    Get(String),
}

enum DataType {
    SimpleString,
    SimpleErr,
    Integer,
    BulkString,
    Array,
}

impl DataType {
    fn from_byte(b: u8) -> Self {
        match b {
            b'+' => Self::SimpleString,
            b'-' => Self::SimpleErr,
            b':' => Self::Integer,
            b'$' => Self::BulkString,
            b'*' => Self::Array,
            _ => unimplemented!(),
        }
    }
}

#[derive(Debug)]
enum RedisObject {
    SimpleString(String),
    SimpleErr(String),
    Integer(i32),
    BulkString(usize, String),
    Array(Vec<RedisObject>),
}

struct Parser<'a> {
    stream: &'a [u8],
}

impl<'a> Parser<'a> {
    fn new(stream: &'a [u8]) -> Self {
        Self { stream }
    }

    fn parse(&mut self) -> Result<RedisObject, ()> {
        match Self::parse_object(self.stream) {
            Ok((Some(object), _)) => Ok(object),
            Ok(_) => Err(()),
            Err(_) => Err(()),
        }
    }

    fn parse_object(stream: &[u8]) -> Result<(Option<RedisObject>, usize), ()> {
        if stream[0..2] == *b"\r\n" {
            return Ok((None, 2));
        }
        match DataType::from_byte(stream[0]) {
            DataType::Array => Parser::parse_array(&stream[1..]),
            DataType::SimpleString => {
                let parts = split_by_line(&stream[1..]);
                Ok((
                    Some(RedisObject::SimpleString(
                        String::from_utf8(parts[0].clone()).unwrap(),
                    )),
                    parts[0].len(),
                ))
            }
            DataType::Integer => {
                let parts = split_by_line(&stream[1..]);
                Ok((
                    Some(RedisObject::Integer(
                        String::from_utf8(parts[0].clone())
                            .unwrap()
                            .parse::<i32>()
                            .unwrap(),
                    )),
                    parts[0].len(),
                ))
            }
            DataType::BulkString => {
                let parts = split_by_line(&stream[1..]);
                let Ok(size) = String::from_utf8(parts[0].clone())
                    .unwrap()
                    .parse::<usize>()
                else {
                    panic!("invalid string");
                };
                let string = String::from_utf8(parts[1].clone()).unwrap();
                assert!(string.len() == size as usize);
                Ok((
                    Some(RedisObject::BulkString(size, string)),
                    parts[0].len() + parts[1].len() + 3,
                ))
            }
            _ => unimplemented!("type not implemented"),
        }
    }

    fn parse_array(stream: &[u8]) -> Result<(Option<RedisObject>, usize), ()> {
        let parts = split_by_line(stream);
        let _size = String::from_utf8(parts[0].clone())
            .expect("invalid string")
            .parse::<usize>()
            .expect("invalid string");
        let mut objects = vec![];
        let mut pos: usize = parts[0].len() + 2;
        loop {
            match Parser::parse_object(&stream[pos..]) {
                Ok((Some(object), consumed)) => {
                    objects.push(object);
                    pos += consumed;
                    if pos > stream.len() {
                        break;
                    }
                }
                Ok((None, consumed)) => {
                    pos += consumed;
                    if pos >= stream.len() {
                        break;
                    }
                }
                Err(_) => {
                    return Err(());
                }
            }
        }
        println!("{objects:?}");
        Ok((Some(RedisObject::Array(objects)), pos))
    }
}

fn split_by_line(stream: &[u8]) -> Vec<Vec<u8>> {
    let line_positions = stream
        .windows(2)
        .enumerate()
        .filter(|(_, w)| w == b"\r\n")
        .map(|(i, _)| i)
        .collect::<Vec<_>>();
    let mut lines = vec![stream[..line_positions[0]].to_vec()];
    lines.extend(
        line_positions
            .windows(2)
            .map(|i| stream[i[0] + 2..i[1]].to_vec())
            .collect::<Vec<_>>(),
    );
    lines.push(stream[*line_positions.last().unwrap() + 2..].to_vec());
    lines
        .into_iter()
        .filter(|l| !l.is_empty())
        .collect::<Vec<_>>()
}

impl Command {
    fn from_buffer(buf: &[u8]) -> Result<Self, ()> {
        let mut p = Parser::new(buf);
        match p.parse() {
            Ok(object) => match object {
                RedisObject::Array(arr) => match arr.as_slice() {
                    [RedisObject::BulkString(4, s)] => {
                        if s.to_uppercase() == "PING".to_string() {
                            Ok(Command::Ping)
                        } else {
                            Err(())
                        }
                    }
                    [RedisObject::BulkString(4, s), RedisObject::BulkString(_, o)] => {
                        if s.to_uppercase() == "ECHO".to_string() {
                            Ok(Command::Echo(o.to_string()))
                        } else {
                            Err(())
                        }
                    }
                    [RedisObject::BulkString(3, s), RedisObject::BulkString(_, key)] => {
                        if s.to_uppercase() == "GET".to_string() {
                            Ok(Command::Get(key.to_string()))
                        } else {
                            Err(())
                        }
                    }
                    [RedisObject::BulkString(3, s), RedisObject::BulkString(_, key), RedisObject::BulkString(_, value), RedisObject::BulkString(2, ex), RedisObject::BulkString(_, duration)] => {
                        if s.to_uppercase() == "SET".to_string()
                            && ex.to_uppercase() == "PX"
                            && duration.parse::<u64>().is_ok()
                        {
                            Ok(Command::Set(
                                key.to_string(),
                                value.as_bytes().to_vec(),
                                Some(duration.parse::<u64>().unwrap()),
                            ))
                        } else {
                            Err(())
                        }
                    }
                    [RedisObject::BulkString(3, s), RedisObject::BulkString(_, key), RedisObject::BulkString(_, value)] => {
                        if s.to_uppercase() == "SET".to_string() {
                            Ok(Command::Set(
                                key.to_string(),
                                value.as_bytes().to_vec(),
                                None,
                            ))
                        } else {
                            Err(())
                        }
                    }
                    _ => Err(()),
                },
                _ => Err(()),
            },
            Err(_) => Err(()),
        }
    }
}

Answer 1

Don’t Panic!

(On Runtime Errors)

I admit I’ve sometimes put an .expect() or .unwrap() in my code examples, for simplicity. But I’m mostly on the side that believes end users should never see a runtime panic. Those should only be used for logic errors, or debugging. But a network error, especially, is something that often happens in the real world, so a program really ought to handle it more gracefully than panicking.

In general, you might want to have each function on a stream that could fail, return a value wrapped in a std::io::Result. Computations that shouldn’t have to worry or know about TCP and errors should be called on the unwrapped values, with syntax sugar like the ? operator. Or you can use railway-style with calls like and_then.

Let’s see what happens when we change main to return a std::io::result, allowing it to short-circuit and return the first Err:

fn main() -> std::io::Result<()> {
    let listener = TcpListener::bind("127.0.0.1:6379")?;

    let data_storage = Arc::new(Mutex::new(collections::HashMap::<
        String,
        (Option<time::Instant>, Vec<u8>),
    >::new()));

    for stream in listener.incoming() {
       let s = stream?;
       let cloned_storage = data_storage.clone();
       thread::spawn(move || handle(s, cloned_storage));
    }
    Ok(())
}

This error-handling is still bare-bones: any I/O error that gets returned from main makes the program print an error message and return with an exit status of 1. In this implementation, the error message is something like:

Error: Os { code: 98, kind: AddrInUse, message: "Address already in use" }

Unlike the default panic message from Result::expect, this doesn’t have any explanation of where the error happened or what you were trying to do (or the ability to get a backtrace). You might rename main to run and write a new main that inspects the return value of run, and handles it, for example by logging the error and trying again to listen on the port if possible. For an alternative that lets you add some context to the error, explaining where it happened and what you were doing, you might look at the error_chain crate.

Refactor to Remove False Dependencies

And to Reduce Cyclomatic Complexity

You have a large match block in handle, currently, which is a good candidate for refactoring. It really takes a line of input and a shared database (by mutable reference), and returns some output. It doesn’t really need to know about sockets. This means it doesn’t need to know about, or handle, errors: we can check for them in one place, when we write the output string we returned. If we move this block into a helper function, we get (Warning: untested!):

fn execute(
    s: &str,
    storage: &mut Arc<Mutex<collections::HashMap<String, (Option<time::Instant>, String)>>>,
) -> String {
    match Command::from_buffer(s.as_bytes()) {
        Ok(Command::Ping) => "+PONG\r\n".to_string(),
        Ok(Command::Echo(cmd)) => "$".to_string() + &cmd + "\r\n",
        Ok(Command::Set(key, value, expiry)) => {
            let mut storage = storage.lock().unwrap();
            let expiry =
                expiry.map(|t| time::Instant::now() + time::Duration::from_millis(t as u64));
            storage.insert(key, (expiry, value));
            "+OK\r\n".to_string()
        }
        Ok(Command::Get(key)) => {
            let storage = storage.lock().unwrap();
            match storage.get(&key) {
                Some((expiry, v)) => {
                    if let Some(expiry) = expiry {
                        if time::Instant::now() >= *expiry {
                            "$-1\r\n".to_string()
                        } else {
                            "$".to_string() + &v + "\r\n"
                        }
                    } else {
                        "$".to_string() + &v + "\r\n"
                    }
                }
                None => "$-1\r\n".to_string(),
            } // end match
        }
        Err(_) => "-Error\r\n".to_string(),
    } // end match
}

This fixes one minor bug (the locked storage in the Get branch was mut when it didn’t need to be) and almost certainly introduces several more.

This also simplifies the next piece:

Use a BufReader to Split Lines

You currently have some complex code to SplitByLine, SerializeToSimpleString, SerializeToBulkString, etc. While I haven’t checked the requirements or tested to see if what I’m about to recommend works robustly for your use case, it looks like you could be reinventing the wheel.

There’s already a class in the standard library that turns anything with the Read trait into a buffered reader. The BufRead::lines method turns an input stream into an iterator over its lines. (More precisely, an iterator over std::io::Result objects which hold either a line of input, or a std::io::Error.) This is how I’d split into lines whenever possible. There’s also a more-general split method to delimit buffered input at an arbitrary byte.

A wrinkle here is that a BufReader can’t be written to, nor can our stream be mutably borrowed twice, so we actually need to wrap a clone of the TcpStream for buffered reading, and use the original for writing.

fn handle(
    mut stream: TcpStream,
    mut storage: Arc<Mutex<collections::HashMap<String, (Option<time::Instant>, String)>>>,
) {
    let buffered_stream = BufReader::new(stream.try_clone().expect("Cloning socket"));
    for line in buffered_stream.lines() {
        match line {
            Ok(s) => {
                if let Err(e) = stream.write_all(execute(&s, &mut storage).as_bytes()) {
                    let _ = std::io::stdout().flush(); // Don't cross the streams!
                    eprintln!("Network write error: {e}\n");
                    let _ = std::io::stderr().flush(); // Could be writing from multiple threads.
                    return;
                }
            }
            Err(e) => {
                let _ = std::io::stdout().flush(); // Don't cross the streams!
                eprintln!("Network read error: {e}\n");
                let _ = std::io::stderr().flush(); // Could be writing from multiple threads.
                return;
            }
        }
    }
}

You’ll notice a bit of hypocrisy in how I handled errors. Cloning the network socket, to get separate unbuffered writer and buffered reader streams, should “always” work, so I fell back to expect for error handling. (Don’t quote me on that: some OS might run out of file descriptors, or something.) Other errors are handled by printing a message to standard error and terminating the thread (which closes the network connection).

I flush the standard output and error streams, and ignore the error conditions. (If printing an error message failed, there’s no point trying to print an error message about it!) This reduces the odds that a bunch of threads will garble their error messages together when they all crash at once, but if you want a robust solution, factor this code out into a locked critical section.

Putting it All Together

Because a lot of little changes were needed to get this tweaked version to compile, here's a full listing. There’s plenty of work left to do. You’ll notice for example that it still calls split_by_line even though the lines have already been split.

use std::io::{BufRead, BufReader, Write};
use std::net::{TcpListener, TcpStream};
use std::{
    collections,
    sync::{Arc, Mutex},
    thread, time,
};

fn main() -> std::io::Result<()> {
    let listener = TcpListener::bind("127.0.0.1:6379")?;

    let data_storage = Arc::new(Mutex::new(collections::HashMap::<
        String,
        (Option<time::Instant>, String),
    >::new()));

    for stream in listener.incoming() {
        let s = stream?;
        let cloned_storage = data_storage.clone();
        thread::spawn(move || handle(s, cloned_storage));
    }
    Ok(())
}

fn handle(
    mut stream: TcpStream,
    mut storage: Arc<Mutex<collections::HashMap<String, (Option<time::Instant>, String)>>>,
) {
    let buffered_stream = BufReader::new(stream.try_clone().expect("Cloning socket"));
    for line in buffered_stream.lines() {
        match line {
            Ok(s) => {
                if let Err(e) = stream.write_all(execute(&s, &mut storage).as_bytes()) {
                    let _ = std::io::stdout().flush(); // Don't cross the streams!
                    eprintln!("Network write error: {e}\n");
                    let _ = std::io::stderr().flush(); // Could be writing from multiple threads.
                    return;
                }
            }
            Err(e) => {
                let _ = std::io::stdout().flush(); // Don't cross the streams!
                eprintln!("Network read error: {e}\n");
                let _ = std::io::stderr().flush(); // Could be writing from multiple threads.
                return;
            }
        }
    }
}

fn execute(
    s: &str,
    storage: &mut Arc<Mutex<collections::HashMap<String, (Option<time::Instant>, String)>>>,
) -> String {
    match Command::from_buffer(s.as_bytes()) {
        Ok(Command::Ping) => "+PONG\r\n".to_string(),
        Ok(Command::Echo(cmd)) => "$".to_string() + &cmd + "\r\n",
        Ok(Command::Set(key, value, expiry)) => {
            let mut storage = storage.lock().unwrap();
            let expiry =
                expiry.map(|t| time::Instant::now() + time::Duration::from_millis(t as u64));
            storage.insert(key, (expiry, value));
            "+OK\r\n".to_string()
        }
        Ok(Command::Get(key)) => {
            let storage = storage.lock().unwrap();
            match storage.get(&key) {
                Some((expiry, v)) => {
                    if let Some(expiry) = expiry {
                        if time::Instant::now() >= *expiry {
                            "$-1\r\n".to_string()
                        } else {
                            "$".to_string() + &v + "\r\n"
                        }
                    } else {
                        "$".to_string() + &v + "\r\n"
                    }
                }
                None => "$-1\r\n".to_string(),
            } // end match
        }
        Err(_) => "-Error\r\n".to_string(),
    } // end match
}

/*
fn serialize_to_simple_string(s: &[u8]) -> Vec<u8> {
    [b"+", s, b"\r\n"].concat()
}

fn serialize_to_bulk_string(s: &[u8]) -> Vec<u8> {
    [b"$", format!("{}", s.len()).as_bytes(), b"\r\n", s, b"\r\n"].concat()
}
 */

#[derive(Debug)]
enum Command {
    Ping,
    Echo(String),
    Set(String, String, Option<u64>),
    Get(String),
}

enum DataType {
    SimpleString,
    SimpleErr,
    Integer,
    BulkString,
    Array,
}

impl DataType {
    fn from_byte(b: u8) -> Self {
        match b {
            b'+' => Self::SimpleString,
            b'-' => Self::SimpleErr,
            b':' => Self::Integer,
            b'$' => Self::BulkString,
            b'*' => Self::Array,
            _ => unimplemented!(),
        }
    }
}

#[derive(Debug)]
enum RedisObject {
    SimpleString(String),
    SimpleErr(String),
    Integer(i32),
    BulkString(usize, String),
    Array(Vec<RedisObject>),
}

struct Parser<'a> {
    stream: &'a [u8],
}

impl<'a> Parser<'a> {
    fn new(stream: &'a [u8]) -> Self {
        Self { stream }
    }

    fn parse(&mut self) -> Result<RedisObject, ()> {
        match Self::parse_object(self.stream) {
            Ok((Some(object), _)) => Ok(object),
            Ok(_) => Err(()),
            Err(_) => Err(()),
        }
    }

    fn parse_object(stream: &[u8]) -> Result<(Option<RedisObject>, usize), ()> {
        if stream[0..2] == *b"\r\n" {
            return Ok((None, 2));
        }
        match DataType::from_byte(stream[0]) {
            DataType::Array => Parser::parse_array(&stream[1..]),
            DataType::SimpleString => {
                let parts = split_by_line(&stream[1..]);
                Ok((
                    Some(RedisObject::SimpleString(
                        String::from_utf8(parts[0].clone()).unwrap(),
                    )),
                    parts[0].len(),
                ))
            }
            DataType::Integer => {
                let parts = split_by_line(&stream[1..]);
                Ok((
                    Some(RedisObject::Integer(
                        String::from_utf8(parts[0].clone())
                            .unwrap()
                            .parse::<i32>()
                            .unwrap(),
                    )),
                    parts[0].len(),
                ))
            }
            DataType::BulkString => {
                let parts = split_by_line(&stream[1..]);
                let Ok(size) = String::from_utf8(parts[0].clone())
                    .unwrap()
                    .parse::<usize>()
                else {
                    panic!("invalid string");
                };
                let string = String::from_utf8(parts[1].clone()).unwrap();
                assert!(string.len() == size as usize);
                Ok((
                    Some(RedisObject::BulkString(size, string)),
                    parts[0].len() + parts[1].len() + 3,
                ))
            }
            _ => unimplemented!("type not implemented"),
        }
    }

    fn parse_array(stream: &[u8]) -> Result<(Option<RedisObject>, usize), ()> {
        let parts = split_by_line(stream);
        let _size = String::from_utf8(parts[0].clone())
            .expect("invalid string")
            .parse::<usize>()
            .expect("invalid string");
        let mut objects = vec![];
        let mut pos: usize = parts[0].len() + 2;
        loop {
            match Parser::parse_object(&stream[pos..]) {
                Ok((Some(object), consumed)) => {
                    objects.push(object);
                    pos += consumed;
                    if pos > stream.len() {
                        break;
                    }
                }
                Ok((None, consumed)) => {
                    pos += consumed;
                    if pos >= stream.len() {
                        break;
                    }
                }
                Err(_) => {
                    return Err(());
                }
            }
        }
        println!("{objects:?}");
        Ok((Some(RedisObject::Array(objects)), pos))
    }
}

fn split_by_line(stream: &[u8]) -> Vec<Vec<u8>> {
    let line_positions = stream
        .windows(2)
        .enumerate()
        .filter(|(_, w)| w == b"\r\n")
        .map(|(i, _)| i)
        .collect::<Vec<_>>();
    let mut lines = vec![stream[..line_positions[0]].to_vec()];
    lines.extend(
        line_positions
            .windows(2)
            .map(|i| stream[i[0] + 2..i[1]].to_vec())
            .collect::<Vec<_>>(),
    );
    lines.push(stream[*line_positions.last().unwrap() + 2..].to_vec());
    lines
        .into_iter()
        .filter(|l| !l.is_empty())
        .collect::<Vec<_>>()
}

impl Command {
    fn from_buffer(buf: &[u8]) -> Result<Self, ()> {
        let mut p = Parser::new(buf);
        match p.parse() {
            Ok(object) => match object {
                RedisObject::Array(arr) => match arr.as_slice() {
                    [RedisObject::BulkString(4, s)] => {
                        if s.to_uppercase() == "PING".to_string() {
                            Ok(Command::Ping)
                        } else {
                            Err(())
                        }
                    }
                    [RedisObject::BulkString(4, s), RedisObject::BulkString(_, o)] => {
                        if s.to_uppercase() == "ECHO".to_string() {
                            Ok(Command::Echo(o.to_string()))
                        } else {
                            Err(())
                        }
                    }
                    [RedisObject::BulkString(3, s), RedisObject::BulkString(_, key)] => {
                        if s.to_uppercase() == "GET".to_string() {
                            Ok(Command::Get(key.to_string()))
                        } else {
                            Err(())
                        }
                    }
                    [RedisObject::BulkString(3, s), RedisObject::BulkString(_, key), RedisObject::BulkString(_, value), RedisObject::BulkString(2, ex), RedisObject::BulkString(_, duration)] => {
                        if s.to_uppercase() == "SET".to_string()
                            && ex.to_uppercase() == "PX"
                            && duration.parse::<u64>().is_ok()
                        {
                            Ok(Command::Set(
                                key.to_string(),
                                value.to_string(),
                                Some(duration.parse::<u64>().unwrap()),
                            ))
                        } else {
                            Err(())
                        }
                    }
                    [RedisObject::BulkString(3, s), RedisObject::BulkString(_, key), RedisObject::BulkString(_, value)] => {
                        if s.to_uppercase() == "SET".to_string() {
                            Ok(Command::Set(key.to_string(), value.to_string(), None))
                        } else {
                            Err(())
                        }
                    }
                    _ => Err(()),
                },
                _ => Err(()),
            },
            Err(_) => Err(()),
        }
    }
}