Day30:a redis-memcache server with tokio - 100DayOfRust
Tags: rust, redis
Published: 2020-01-05 20:27:13

I will write a memcache server with tokio by using redis protocol. This memcache server only support 2 command from redis: 
* `get <key>`
* `set <key> <value>`

The goal is, after we run the server, we can use the standard `redis-cli` to connect to it and use `get` & `set` command.

### Basic code structure

We utilize `tokio` to write this high performance server software:

In `Cargo.toml`:

```
[dependencies]
tokio = { version = "0.2", features = ["full"] }
```

In `main.rs`:

```rust
use std::collections::HashMap;
use std::error::Error;

use tokio::net::{TcpListener, TcpStream};
use tokio::prelude::*;
use tokio::sync::Mutex;
use std::sync::Arc;


#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let addr = "127.0.0.1:7777";
    let mut listener = TcpListener::bind(addr).await?;

    println!("Listen on {}", addr);

    let dict = Arc::new(Mutex::new(HashMap::new()));

    loop {
        let (mut sock, _) = listener.accept().await?;
        let dict = dict.clone();
        tokio::spawn(async move {
            // todo
        });
    }
}
```

The `dict` variable is for saving the key-value pair later. This dict we wrapped it with `Mutex` and `Arc`, because it will be used in multi-threads in tokio.

### Redis Protocol

Redis use `\r\n` as separator in its [protocol](https://redis.io/topics/protocol). To know the details of its protocol needs another post, so here we just introduce the basics:

> For Simple Strings the first byte of the reply is "+"
> For Errors the first byte of the reply is "-"
> For Integers the first byte of the reply is ":"
> For Bulk Strings the first byte of the reply is "$"
> For Arrays the first byte of the reply is "*"

For example, for the `get <key>` command, redis-cli will send the content like this (suppose our key is "hello" in this case):

```
*2
$3
get
$5
hello
``` 

The leading `*2` indicates this is an array which contains 2 elements, the first one's length is 3, and the second one's length is 5.

Let's write 2 util functions to get the length integer and from the length integer get the followed string.


```rust
async fn read_till_crlf(stream: &mut TcpStream, skip: u8) -> Vec<u8> {
    let mut ret: Vec<u8> = vec![];
    let mut skip_num = skip;
    loop {
        let mut buf = [0; 1];
        stream.read_exact(&mut buf).await.unwrap();
        // LF's ascii number is 10
        if skip_num == 0 && buf[0] == 10 {
            break;
        }
        if skip_num > 0 {
            skip_num -= 1;
        } else {
            ret.push(buf[0]);
        }
    }
    // pop the last CR
    ret.pop();
    ret
}


async fn read_nbytes(stream: &mut TcpStream, nbytes: usize) -> Vec<u8> {
    let mut ret: Vec<u8> = vec![0; nbytes];
    stream.read_exact(&mut ret).await.unwrap();
    ret
}


async fn get_next_len(stream: &mut TcpStream) -> usize {
    let vlen = read_till_crlf(stream, 1).await;
    let slen = String::from_utf8(vlen).unwrap();
    let len:usize = slen.parse().unwrap();
    len
}


async fn get_next_string(stream: &mut TcpStream) -> String {
    let len = get_next_len(stream).await;
    let vs = read_nbytes(stream, len).await;
    // consume the followed \r\n
    let _ = read_nbytes(stream, 2).await;
    // build string and return
    let s = String::from_utf8(vs).unwrap();
    s
}
```

The `get_next_len` will return the number from the stream like "*2\r\n" and "$3\r\n", etc. The `get_next_string` will return the string, for example, for "$5\r\nhello\r\n", it will return "hello".

### Handle unknown command and syntax error

Since we only support `get <key>` and `set <key> <value>` command, let's write other 2 util functions:

```rust
async fn handle_unknown(stream: &mut TcpStream) {
    stream.write_all(b"-Unknown command\r\n").await.unwrap();
}

async fn handle_syntax_err(stream: &mut TcpStream) {
    stream.write_all(b"-ERR syntax error\r\n").await.unwrap();
}
```

### Handle GET command

```rust
async fn handle_get(stream: &mut TcpStream,
                    dict: &Arc<Mutex<HashMap<String, String>>>) {
    let key = get_next_string(stream).await;
    let map = dict.lock().await;
    let s = match map.get(key.as_str()) {
        Some(v) => {
            format!("${}\r\n{}\r\n", v.len(), v)
        },
        None => {
            "$-1\r\n".to_owned()
        },
    };
    stream.write_all(s.as_bytes()).await.unwrap();
}
```

First we get the key, if key is in our dictionary, then return the value. Remember the returned value also need to follow the redis protocol.


### Handle SET command

```rust
async fn handle_set(stream: &mut TcpStream,
                    dict: &Arc<Mutex<HashMap<String, String>>>) {
    let key = get_next_string(stream).await;
    let val = get_next_string(stream).await;
    let mut map = dict.lock().await;
    map.insert(key, val);
    stream.write_all(b"+OK\r\n").await.unwrap();
}
```

For the `set` command, we get the key then the value, then set it into our dictionary. We don't care if the key already exists: if it does, we will just overwrite it.


### Put it together

Full code:

```rust
use std::collections::HashMap;
use std::error::Error;

use tokio::net::{TcpListener, TcpStream};
use tokio::prelude::*;
use tokio::sync::Mutex;
use std::sync::Arc;


#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let addr = "127.0.0.1:7777";
    let mut listener = TcpListener::bind(addr).await?;

    println!("Listen on {}", addr);

    let dict = Arc::new(Mutex::new(HashMap::new()));

    loop {
        let (mut sock, _) = listener.accept().await?;
        let dict = dict.clone();
        tokio::spawn(async move {
            // get arg array length like *2, *3
            let arg_len = get_next_len(&mut sock).await;
            let cmd = get_next_string(&mut sock).await;
            let cmd = cmd.to_lowercase();
            if cmd == "get" {
                if arg_len != 2 {
                    handle_syntax_err(&mut sock).await;
                } else {
                    handle_get(&mut sock, &dict).await;
                }
            } else if cmd == "set" {
                if arg_len != 3 {
                    handle_syntax_err(&mut sock).await;
                } else {
                    handle_set(&mut sock, &dict).await;
                }
            } else {
                handle_unknown(&mut sock).await;
            }
        });
    }
}


async fn read_till_crlf(stream: &mut TcpStream, skip: u8) -> Vec<u8> {
    let mut ret: Vec<u8> = vec![];
    let mut skip_num = skip;
    loop {
        let mut buf = [0; 1];
        stream.read_exact(&mut buf).await.unwrap();
        // LF's ascii number is 10
        if skip_num == 0 && buf[0] == 10 {
            break;
        }
        if skip_num > 0 {
            skip_num -= 1;
        } else {
            ret.push(buf[0]);
        }
    }
    // pop the last CR
    ret.pop();
    ret
}


async fn read_nbytes(stream: &mut TcpStream, nbytes: usize) -> Vec<u8> {
    let mut ret: Vec<u8> = vec![0; nbytes];
    stream.read_exact(&mut ret).await.unwrap();
    ret
}


async fn get_next_len(stream: &mut TcpStream) -> usize {
    let vlen = read_till_crlf(stream, 1).await;
    let slen = String::from_utf8(vlen).unwrap();
    let len:usize = slen.parse().unwrap();
    len
}


async fn get_next_string(stream: &mut TcpStream) -> String {
    let len = get_next_len(stream).await;
    let vs = read_nbytes(stream, len).await;
    // consume the followed \r\n
    let _ = read_nbytes(stream, 2).await;
    // build string and return
    let s = String::from_utf8(vs).unwrap();
    s
}


async fn handle_get(stream: &mut TcpStream,
                    dict: &Arc<Mutex<HashMap<String, String>>>) {
    let key = get_next_string(stream).await;
    let map = dict.lock().await;
    let s = match map.get(key.as_str()) {
        Some(v) => {
            format!("${}\r\n{}\r\n", v.len(), v)
        },
        None => {
            "$-1\r\n".to_owned()
        },
    };
    stream.write_all(s.as_bytes()).await.unwrap();
}


async fn handle_set(stream: &mut TcpStream,
                    dict: &Arc<Mutex<HashMap<String, String>>>) {
    let key = get_next_string(stream).await;
    let val = get_next_string(stream).await;
    let mut map = dict.lock().await;
    map.insert(key, val);
    stream.write_all(b"+OK\r\n").await.unwrap();
}


async fn handle_unknown(stream: &mut TcpStream) {
    stream.write_all(b"-Unknown command\r\n").await.unwrap();
}

async fn handle_syntax_err(stream: &mut TcpStream) {
    stream.write_all(b"-ERR syntax error\r\n").await.unwrap();
}
```

Let's run it with `cargo run`:

```bash
$ cargo run
Listen on 127.0.0.1:7777

```

In another terminal window, use `redis-cli` to connect to it:

```bash
$ redis-cli -p 7777
127.0.0.1:7777> get hello
(nil)
127.0.0.1:7777> set hello world
OK
127.0.0.1:7777> get hello
"world"
127.0.0.1:7777> set hello world1 world2 world3
(error) ERR syntax error
127.0.0.1:7777> get hello
"world"
127.0.0.1:7777> set hello "world overwrite"
OK
127.0.0.1:7777> get hello
"world overwrite"
127.0.0.1:7777> command
(error) Unknown command
127.0.0.1:7777>
```

It works! Thanks to `tokio`, with only around 100 lines code, we now have a multi-threads async memcache server. Even more, because we are using redis protocol, we can use the exisiting client library. For example, with Python's redis package (`pip install redis`), we can get the value associate with "hello" we just set:


```python
>>> import redis
>>> r = redis.Redis(port=7777)
>>> r.get("hello")
b'world overwrite'
```


### Reference
* https://redis.io/topics/protocol
* https://redisgreen.net/blog/beginners-guide-to-redis-protocol/
* https://lethain.com/redis-protocol/
* https://redis.io/commands/command
* https://github.com/tokio-rs/tokio/blob/master/examples/chat.rs