DEV Community: Ashish Singh

Building Redis Server in Rust: Part 2

Ashish Singh — Mon, 15 Aug 2022 06:45:28 +0000

In the previous blog post, we covered single threaded client server communication. In this blog post, we will be covering multi-threaded server client communication. Every new client request will be handled in its own separate thread. We will be using tokio::spawn to spawn a new thread to handle the incoming request from the client.

We will also be implementing server shutdown using ctrl + c. It will close the server but the clients will be gracefully shutdown without panicking. For this we need to communication from server to client. Let's use tokio::sync::broadcast channel broadcast messages to client. But this is only possible when all the clients subscribe to the broadcast sender. This means our client will be holding at least fields socket and receiver.

Also, since multiple clients can connect simultaneously, they also read and write data simultaneously. This means, we also need to supply a copy of Db to each thread. As of now, we just have one instance variable of Db, this object needs to be shared across the threads. For this purpose, 2 things need to happen.

We will use Arc<Mutex> and wrap it around the HashMap<String, Bytes> in our Db to create shared references. We could use Rc<RefCell> for this purpose, but since we are talking about multi-threading Rc is not a great fit as it is not thread-safe.
Since, we are looking at many values associated with a client connection, we should create a struct Handler to encapsulate all the fields together.

Let's try to breakdown our code such that hitting ctrl + c while the server is running triggers the shutdown process for the server. In the main function of bin/server.rs let's capture the shutdown signal using tokio::signal.

use tokio::signal;

// in main function
let shutdown = signal::ctrl_c();

Whenever we hit ctrl+ c this shutdown future completes on the .await or Future::poll.

We want to run 2 branches at this point, one to handle the incoming requests from clients and the other to listen on this shutdown. So whenever ctrl+c is received the shutdown process starts dropping the client handle branch. This is where tokio::select comes into picture. It takes any number of async branches in form of futures(similar to promise in js) and runs them concurrently, waiting for value from any branch. Upon receiving value from a branch, it execuates the handler function and drops the rest of the branches.

tokio::select! {
    res = server::run(&mut listener) => {
            if let Err(_err) = res {
            println!("failed to accept connection");
        }
    }
    _ = shutdown => {
        println!("inside shutdown loop");
    }
}
// execution reaches here only when shutdown future returns value or server encounters and error.

In the code above, both the branches will run and wait for any one to return a value. To keep the server running, we will not return a value from server::run unless there is an error. If we hit ctrl+c at this point, then shutdown branch will run stopping our server::run() future.

At this point, lets create our Handler and Listener structs. The Listener struct is the server object holding TcpListener and broadcast channel.

Handler is also responsible for following in addition to manage socket and db clone:-

Receiving shutdown notification from server, in case server start shutting down.
Notifying the server when the handler is going out of scope.

Let's combine the shutdown related functionality into a separate field into a Shutdown struct and move the socket into a new field which is Connection struct.

File src/handler.rs

pub struct Handler {
    pub connection: Connection,
    pub db: Db,
    pub shutdown: Shutdown,
    // when the handler object is dropped, this sends a message to the receiver
    _shutdown_complete: mpsc::Sender<()>,
}

pub struct Connection {
    pub stream: TcpStream,
}

pub struct Shutdown {
    shutdown: bool,
    notify: broadcast::Receiver<()>,
}

The shutdown struct holds notify which is a broadcast receiver which subscribes to the server broadcasting shutdown. It's time to define the Listener struct and connect the broadcast from listener to handler subscriber.

Let's create a new file src/listener.rs.

use tokio::{
    net::{TcpListener, TcpStream},
    sync::{broadcast, mpsc},
};

use crate::Db;

pub struct Listener {
    pub db: Db,
    pub listener: TcpListener,
    pub notify_shutdown: broadcast::Sender<()>,
    pub shutdown_complete_rx: mpsc::Receiver<()>,
    pub shutdown_complete_tx: mpsc::Sender<()>,
}

In listener struct, notify_shutdown is a broadcast channel sender. It will broadcast a message when the listener notify_shutdown is dropped.

The other 2 fields shutdown_complete_rx and shutdown_complete_tx are mpsc objects used to communicate between the multiple handlers to the server. Imagine, there are many clients connected to the server and the server decides to shutdown. In that case, the mpsc::Receiver<()> will only await till all the senders have closed before the mpsc::Receiver<()> stops blocking for receiving the messages. In short, all senders should be completed/dropped for the receiver to allow for further execution of code.

The flow of messages is follows:-

Server start shutting down. It sends a message to all handlers.
All the handlers get this message from the subscription of broadcast::Sender<()> subscription.
All the handlers decide to wrap functioning and drop themselves.
Once the handler object goes out of scope (drops) they send message via mpsc::Sender<()> clone(explained later).
The server receives message in shutdown_complete_rx and knows for sure that handler object is dropped.
shutdown_complete_rx keeps receiving messages till all the handlers go out of scope.
Server closes gracefully having shutdown all the other handlers.

Note that we are using mpsc channel which stands for multiple producers single consumer. It fits our use case well as we are interested in sending messages from multiple handlers to a single server thread.

Now let's look at how we will create objects:-

Upon server start, we will create a listener object, which initiates the Db and channel receiver and senders.
Then we wait for client socket connections in a loop.
When a new socket connection starts, we will create a handler object which will subscribe to the broadcast::sender into a receiver. It will also initiate _shutdown_complete field using mpsc::sender clone.
This handler object then becomes the lifetime of the client connection.

listener object create

In bin/server.rs

use blog_redis::server;
use blog_redis::Listener;
use tokio::signal;
use tokio::{
    net::TcpListener,
    sync::{broadcast, mpsc},
};

#[tokio::main]
pub async fn main() -> Result<(), std::io::Error> {
    let listener = TcpListener::bind("127.0.0.1:8081").await?;
    let shutdown = signal::ctrl_c();
    let (notify_shutdown, _) = broadcast::channel(1);
    let (shutdown_complete_tx, shutdown_complete_rx) = mpsc::channel(1);

    let mut listener = Listener::new(
        listener,
        notify_shutdown,
        shutdown_complete_tx,
        shutdown_complete_rx,
    );

    tokio::select! {
        res = server::run(&mut listener) => {
             if let Err(_err) = res {
               println!("failed to accept connection");
            }
        }
        _ = shutdown => {
            println!("inside shutdown loop");
        }
    }

    // graceful shutdown code will go here. Refer Graceful shutdown section

    Ok(())
}

Next, we will change our src/server.rs file

use crate::Handler;
use crate::Listener;

pub async fn run(listener: &Listener) -> std::io::Result<()> {
    loop {
        let socket = listener.accept().await?;
        let mut handler = Handler::new(listener, socket);

        tokio::spawn(async move {
            if let Err(_err) = process_method(&mut handler).await {
                println!("Connection Error");
            }
        });
    }
}

async fn process_method(handler: &mut Handler) -> Result<(), std::io::Error> {
    while !handler.shutdown.is_shutdown() {
        let result = tokio::select! {
            _ = handler.shutdown.listen_recv() => {
                return Ok(());
            },
            res = handler.connection.read_buf_data() => res,
        };

        let (cmd, vec) = match result {
            Some((cmd, vec)) => (cmd, vec),
            None => return Ok(()),
        };

        handler.process_query(cmd, vec).await?;
    }
    Ok(())
}

Let's now implement all the methods on the Handler, Shutdown and Connection structs.

impl Connection {
    fn new(stream: TcpStream) -> Connection {
        Connection { stream: stream }
    }

    pub async fn read_buf_data(&mut self) -> Option<(Command, Vec<String>)> {
        let mut buf = BytesMut::with_capacity(1024);
        match self.stream.read_buf(&mut buf).await {
            Ok(size) => {
                if size == 0 {
                    // returning from empty buffer
                    return None;
                }
            }
            Err(err) => {
                println!("error {:?}", err);
                return None;
            }
        };
        let attrs = buffer_to_array(&mut buf);
        Some((Command::get_command(&attrs[0]), attrs))
    }
}

impl Shutdown {
    fn new(shutdown: bool, notify: broadcast::Receiver<()>) -> Shutdown {
        Shutdown { shutdown, notify }
    }

    pub async fn listen_recv(&mut self) -> Result<(),     tokio::sync::broadcast::error::RecvError> {
        self.notify.recv().await?; // returns error of type `tokio::sync::broadcast::error::RecvError`
        self.shutdown = true;
        Ok(())
    }

    pub fn is_shutdown(&self) -> bool {
        self.shutdown
    }
}

impl Handler {
    pub fn new(listener: &Listener, socket: TcpStream) -> Handler {
        Handler {
            connection: Connection::new(socket),
            db: listener.db.clone(),
            shutdown: Shutdown::new(false, listener.notify_shutdown.subscribe()),
            _shutdown_complete: listener.shutdown_complete_tx.clone(),
        }
    }

    pub async fn process_query(
        &mut self,
        command: Command,
        attrs: Vec<String>,
    ) -> Result<(), std::io::Error> {
        let connection = &mut self.connection;
        let db = &self.db;

        match command {
            Command::Get => {
                let result = db.read(&attrs);
                match result {
                    Ok(result) => {
                        connection.stream.write_all(&result).await?;
                    }
                    Err(_err) => {
                        connection.stream.write_all(b"").await?;
                    }
                }
                return Ok(());
            }
            Command::Set => {
                let resp = db.write(&attrs);
                match resp {
                    Ok(result) => {
                        connection.stream.write_all(&result.as_bytes()).await?;
                    }
                    Err(_err) => {
                        connection.stream.write_all(b"").await?;
                    }
                }
                return Ok(());
            }
            Command::Invalid => {
                connection.stream.write_all(b"invalid command").await?;
                Err(std::io::Error::from(ErrorKind::InvalidData))
            }
        }
    }
}

Db changes to Arc

In the previous blog post, we had defined Db entries as HashMap<String, Bytes> it worked fine because we just had one-one server client communication. But in case of multi-threading we need to pass shared mutable references to all the handler objects. One way to achieve this is using Arc<Mutex> wrapper on our HashMap. Arc<T> provides the shared ownership of the value it encapsulates T. It is a thread-safe way of updating the shared data. To pass mutable references, we will clone the Arc which produces a new Arc instance which points to the same memory location. Mutex provides a way to get the lock on the data such that only one acquirer can update the data at a time. Once the acquirer is done updating the data it releases the lock for other acquirers to do the same.

File src/db.rs

use bytes::Bytes;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};

#[derive(Clone, Debug)]
pub struct Db {
    pub entries: Arc<Mutex<HashMap<String, Bytes>>>,
}

impl Db {
    pub fn new() -> Db {
        Db {
            entries: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    pub fn write(&self, arr: &[String]) -> Result<&str, &'static str> {
        let key = &arr[1];
        let value = &arr[2];

        // we need to clone the referenced value since Bytes::from() function expects a 'static lifetime
        // variable but `value` has unknown lifetime in this function context
        let val = value.clone();
        let p = self.entries.lock().unwrap()
                    .insert(String::from(key), Bytes::from(val));

        match p {
            Some(_p) => Ok("r Ok"), // if they key was already present
            None => Ok("Ok"),       // if the key was not present
        }
    }

    /// Reads data from the database
    pub fn read(&self, arr: &[String]) -> Result<Bytes, &'static str> {
        let key = &arr[1];
        let query_result = self.entries.lock().unwrap();
        let res = query_result.get(key);
        match res {
            Some(value) => Ok(value.clone()),
            None => Err("no such key found"),
        }
    }
}

You might have noticed that we are calling lock() method on &self.entries where self is Db. There is no such method as lock on entries which is of type Arc. This method is present in Mutex module. The reason why this lock method is accessible on Arc is because both Arc and Mutex implement Deref trait. It stands for derefrencing of the smart pointer and Deref returns the reference to the inner data.

Deref & Deref Coercion

Deref does something known as Deref Coercion. Meaning it will modify the type of the reference into what is expected from the method being called on the reference in this case lock().

Here the entries field of type Arc<Mutex<HashMap<..>>> is coerced to &Mutex<Hashmap<..>>.

It does this by implementing the deref method which borrows the self (entries: Arc<>) and returns a reference to the inner data (&Mutex<>). It happens automatically when we pass a reference to a particular type’s value as an argument to a function or method that doesn’t match the parameter type in the function or method definition. A sequence of calls to the deref method converts the type we provided into the type the parameter needs.

Arc and Mutex implement both Deref and DerefMut. In case of get, we only need a reference to the &HashMap but in case of set we need a mutable reference to the &mut HashMap. .lock().unwrap() returns a MutexGuard<HashMap<..>> which gets coerced to &mut HashMap<..> since the insert method expects a mutable reference to the HashMap.

One thing to notice is we are calling the insert method in a single line call whereas, the get method is called in 2 separate lines. There is a reason for it.

The insert method returns an Option<Bytes> whereas the get method returns an Option<&Bytes>. In case of get we need to persist the reference of the Bytes until the function returns which depends on the lifetime of the hashmap under the hood. But in case of set, the return value is owned by the function itself. Let's try to understand the lifetime of the HashMap in both the cases.

The lock() method returns a &MutexGuard<HashMap<..>>. MutexGuard also implements Deref hence it will return the reference of HashMap using the deref function from Deref trait.

trait Deref {
    //...
    fn deref<'a>(&'a self) -> &'a T
    // T is the data encapsulated inside self eg. MutexGuard<HashMap<..>>
}

When we call query_result.get(key) under the hood, MutexGuard pointer derefrencing takes place using the above deref instance method. It means that the lifetime of the wrapped data HashMap will be the same as MutexGuard.

Since we do not store the value of MutexGuard anywhere it gets dropped right after the .get call if used in a single line. Since the hashmap has the same lifetime it also gets dropped and along with it any references it has returned as a result of get method also gets dropped.

However, when we are storing the mutexguard in a local variable of the function the lifetime of the mutexguard becomes the lifetime of the function scope. The lifetime of hashMap also becomes the lifetime of the function and hence the lifetime of the value returned from the hashMap also becomes the lifetime of the function scope.

We can use the query_result inside the function. The reason why we had to clone the value when returning from the function is because, query_result gets dropped after the scope of this function. The value is &Bytes type it also gets dropped after the function scope. Hence, using clone() to make a new copy is needed.

Handler object create

Let's move the main server logic to a new file src/server.rs. This file will hold main execution of server. We will move our code from bin/server.rs to src/server.rs.

File src/server.rs

use crate::Handler;
use crate::Listener;

pub async fn run(listener: &Listener) -> std::io::Result<()> {
    loop {
        let socket = listener.accept().await?;
        let mut handler = Handler::new(listener, socket);

        tokio::spawn(async move {
            if let Err(_err) = process_method(&mut handler).await {
                println!("Connection Error");
            }
        });
    }
}

async fn process_method(handler: &mut Handler) -> Result<(), std::io::Error> {
    while !handler.shutdown.is_shutdown() {
        let result = tokio::select! {
            _ = handler.shutdown.listen_recv() => {
                return Ok(());
            },
            res = handler.connection.read_buf_data() => res,
        };

        let (cmd, vec) = match result {
            Some((cmd, vec)) => (cmd, vec),
            None => return Ok(()),
        };

        handler.process_query(cmd, vec).await?;
    }
    Ok(())
}

One very important conideration in the code above is that we are invoking handler.shutdown.listen_recv() and handler.connection.read_buf_data() in the tokio::select! macro block. Since the handler object is &mut i.e. mutable reference one might think that this violates the borrow rule that there can't be more than one borrow mutable reference of one object. However, we have 2 different fields being passed down and there is no overlap between them hence the compiler doesn't complain of the borrow rule being voilated.

Initially, to keep things simple, I had not created Connection and Shutdown structs. So I was getting the error of borrow rules. Later I decided to encapsulate the respective fields in different structs and the compiler stopped complaining of the borrow error.

Graceful Shutdown

The handler objects are listening for messages from notify_shutdown of listener object on the notify field in the shutdown object inside handler.

Let's explore the process_method function inside src/server.rs. This method exits in following cases:

a successful completion of the process_query method at the bottom of the function.
When the buffer is empty and handler.read_buf_data returns None.
When handler.shutdown.listen_recv() returns a value.

In 3rd case, the value is only returned from handler.shutdown.listen_recv() when the notify_shutdown from listener sends a message.

So let's send a message from notify_shutdown using ctrl+c from bin/server.rs. This will result in tokio::select! going into the _ = shutdown => {} branch. The code moves out of the tokio::select! block. One way is to drop the notify_shutdown. All this message passing happens on the broadcast channel.

File bin/server.rs

// tokio::select! { ... }

drop(listener.notify_shutdown);

File src/handler.rs. This is where the message is received.

pub async fn listen_recv(&mut self) -> Result<(), tokio::sync::broadcast::error::RecvError> {
    self.notify.recv().await?; // the message from drop(notify_shutdown) is received here and await unblock the control flow
    self.shutdown = true;
    Ok(())
}

When the receiver notify receives the message, it immediately returns from the while loop into the run method of server.rs. Taking the control flow to the tokio::spawn(). This is where the handler object is dropped from the scope. At this point, the handler sends a message from _shutdown_complete to the receiver. This happens on the mpsc channel. Now we need to listen to these messages from the handler objects getting dropped.

// tokio::select! { ... }

drop(listener.notify_shutdown);
drop(listener.shutdown_complete_tx);

let _ = listener.shutdown_complete_rx.await;

// this point is only reached when all the handlers are dropped.

In the code above, we are first dropping the shutdown_complete_tx. This needs to be done because, if we don't drop this sender then the shutdown_complete_rx.await will wait forever for a message from the sender going into an infinite wait time. The listener.shutdown_complete_rx.await; will wait for all the handlers to go out of scope and then return None when all handlers(cloned senders) have been dropped.

Using .await we are guaranteed that all the handlers are dropped at this point.

Congratualtions! you have successfully created a multi-threaded redis server and client using tokio.

Here is the latest code repository found on github.

If you have any queries of questions feel free to ask me in the comments section or discord channel.

Building Redis Server in Rust: Part 1

Ashish Singh — Mon, 08 Aug 2022 05:21:00 +0000

In this blog post, we will be making a Redis server in Rust Programming Language. If you are new to rust and looking to learn the language this is a good starting point. In this post, no prior knowledge of the language is assumed. However, if you have read the rust book or some starter tutorial it would be great. If you are someone who has done some basic rust, then this is perfect to dive deeper by building a full scale software program.

The second part of this blog series.

The final implementation will have a fully functioning Redis server with get and set commands. You can get the code for this blog here.

This is going to be a 2 blog post series

In this first blog post we will create a basic single threaded server client interaction and data persistence and retrieval in server.
In the second post, we will implement multi-threaded client server interaction and implement shutdown notification to client using channels.

Prerequisite

You must have rust installed on your computer, if not follow instructions on the official installation guide.

Project Setup

Here we will create the application boilerplate, install all the necessary dependencies, and configure the project.

To get started, create a new rust project with the command.

$ cargo new my_redis
$ cd my_redis

Open the my_redis project in visual studio code or any other editor of your choice to start coding.

You will get 2 main files in the code. first is Cargo.toml and other is /src/main.rs. Cargo.toml file contains all the metadata of the project like project name and most importantly, the dependencies used in the project. Later, we will add libraries under the dependecies section of this file.

Dependencies used

Clap: Used to read command line arguments
Tokio: Used for async implementation.

main.rs has main function which is the entry point of our rust program whenever we run the command cargo run from terminal.

Redis consists of 2 parts, first is server and second is client. Client makes a request to the server to save or retrieve data. Server holds the responsibility to persist data in the server and retrieving it. So, in command line/terminal we will have 2 windows running server and client each.

By default, main.rs is the default executable file. However, since we need 2 executables each for server and client, we will put the server and client executable in src/bin folder which we will create soon. This folder holds as many executable files as we want. This is perfect for our use since we need 2 separate executable files for server and client to run in the terminal.

So, let's create a new folder inside src called bin. Next, create server.rs and client.rs inside the bin folder. The folder structure at this point should be:

src/
│   ├── main.rs
│   └── bin/
│       └── server.rs
│       └── client.rs

Add main() function in both files to get rid of the error which complains about missing main functions. Since, these are executable files, they must contain main functions which is the entry point for a rust program.

To run these executable files simply run the command like so:

$ cargo run --bin server 
$ cargo run --bin client // in new terminal window

--bin argument tell Cargo which executable to run in case there are multiple executables.

After this step, we don't need the main.rs file since we have our separate binaries ready to be executed. So go ahead and delete the main.rs file.

Server and client interaction will happen with a socket connection. Server will listen for incoming socket connection requests. Client will try to connect to server.

Intro to Tokio

Tokio provides asynchronous capabilities and utilities to work in rust. In this blog post, the major usage of tokio will be in form of TcpListener and TcpStream implementations from the tokio::net crate. Tokio has exposed .await functionality on these modules making them truly asynchronous. In short, whenever we use await keyword on the async function calls, the control flow waits till the execution of the function has completed.

To use tokio, add the tokio dependency in Cargo.toml file.
File Cargo.toml

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

Server implementation

Now, add #[tokio::main] macro right before the async fn main() line. This transforms the async main function into synchronous main function and wraps the code in the main function in an async block. The reason to put the code in the async block is because using .await on tokio functions will require the main function to be async type. But, since rust runtime expects the entry point main function to be synchronous in nature hence #[tokio::main] does that transformation for us.

#[tokio::main]
async fn main() {
  println!("in main");
} 

//gets transformed to 
fn main() {
  let rt = tokio::runtime::Runtime::new().unwrap()
  rt.block_on(async { // block_on takes a block of code which is async in nature, so any `.await` inside the block will be fine.
    println!("in main");
  })
}

Always remember, if you want to use .await inside a function or block of code then the function or block of code should be async type.

Next, we want to listen to incoming socket connections in the server.rs file. Import the TcpListener from tokio crate using use tokio::net::TcpListener.

Now in main function using TcpListen::bind() method and supplying it an address with a port. Since the bind method is an async method, we will use await to hold the execution till bind returns the value. The ? operator is a special operator which returns the value wrapped inside Result or propagate the error from the function one level up. You can read more about it here.

File bin/server.rs

pub async fn main() -> Result<(), std::io::Error> {
    let listener = TcpListener::bind("127.0.0.1:8081").await?;
    loop {
      let (mut socket, _) = listener.accept().await?;
      println!("connection accepted {:?}", socket);
    }
    Ok(())
}

At this point, please note we have changed the return type of main function to Result this is coming from core::Result<T, E>. The error gets propagated from the ? operator as explained above.

If we run the server with command $ cargo run --bin server It will run the program and wait at listener.accept() for new connections. With this our server is listening for socket connections on a single thread.

Now, let's move to the client for establishing a socket connection with server.

Since, the client initiates the socket connection, we will be using TcpStream from tokio::net crate. We will make the main function async and annotate with the macro #[tokio::main], which is the same thing we did in the the server setup.

To create a new connection use the same host address and port as server.

pub async fn main() -> Result<(), std::io::Error> {
    let mut stream = TcpStream::connect("127.0.0.1:8081").await?;

    Ok(())
}

At this point, note we have changed the return type of main function to Result this is coming from core::Result<T, E>. The error gets propagated from the ? operator as explained above.

Let's check this minimal server and client interaction. We expect the server to stop after accepting one client connection. It will also print the connection accepted print message.

Run the server and client in the separate terminal windows with the commands shared above and you will notice that server will print something like this.

connection accepted PollEvented { io: Some(TcpStream { addr: 127.0.0.1:8081, peer: 127.0.0.1:61643, fd: 10 }) }

It means, our connection is established and socket is ready to start sending and receiving data.

Writing to socket from client

The first thing we want to do is send the command from client to server. For the sake of simplicity we will send data in form of a string (Bytes in actuality) on socket.

Let's use write_all() on socket function which expects bytes to write data on the socket.

stream.write_all(b"set foo bar").await?;

In rust, any string literal prefixed with b in front of it makes it bytes sequence. If you want to understand the difference in between string literal, string and byte sequence checkout this awesome post.

That's all we need to write data to the socket connection.

Reading from socket in server

To read data from the socket on the server side, we will use read_buf method. This method accepts on argument as a buffer. It copies data from socket to the provided buffer.

We shall create a new buffer from BytesMut. It comes from the bytes crate. First add bytes dependency under the [dependencies] in Cargo.toml file.

[dependencies]
bytes = { version = "1" }

Then import BytesMut from bytes crate in the server.rs file. use bytes::BytesMut.

Let's create a mutable BytesMut with a capacity. This buf object is mutable because it will allow writing bytes into it. Next read the data from socket into the created buffer using read_buf method.
Bytes::with_capacity() takes in number of bytes to allocate to the buffer. This buffer with automatically grow in size as needed but it's more efficient to start with a generous buffer size.
Take for example "foo" will take 3 bytes but "ƒoo" will take 4 bytes. let's start by 1024 bytes as a standard staring size for our buffer.

let mut buf = BytesMut::with_capacity(1024); 
socket.read_buf(&mut buf).await?;
println!("buffer {:?}", buf); // printing the data in buffer

This is all we need to read data from socket into the server. Let's fire the server and client to see if our program is working as expected.

You will get buffer b"set foo bar" as output on the server terminal.

If you reached to this point, congratulations are in order as you have successfully sent from client to server that too in rust.

Command matching

Now, that we have read the data from client which is a command like "set foo bar". The first element is the command name and the next 2 are key and value.
Depending on what is the command we will either perform a get or set.

So, let's get started by converting the buffer to a Vector(sort of Array) of Strings. Then we can take the first string from the vector and fetch the command from Command Enum. We will create Command enum later in this post.

Let's create a method for this conversion called buffer_to_array. Create a new file src/helper.rs and put this function in that file. Since, this file is in src folder, we need to create another file called lib.rs and import the module Helper which gets created as a result of new file cmd.rs. Add pub mod helper; in the src/lib.rs file. This exposes the helper module to the library code. In rust, main.rs is the binary executable and lib.rs is the shared code. If you are creating a library lib.rs is where you'll import all your files.

File src/helper.rs

use bytes::Buf; // get_u8 

fn buffer_to_array(buf: &mut BytesMut) -> Vec<String> {
    let mut vec = vec![];
    let length = buf.len();
    let mut word = "".to_string();

    for i in 0..length {
        match buf.get_u8() {
            b' ' => { // match for space
                vec.push(word);
                word = "".to_string();
            }
            other => {
                // increase the word
                word.push(other as char);
                let new = word.clone();
                if i == length - 1 {
                    vec.push(new);
                }
            }
        }
    }
    vec
}

get_u8 method comes from the use bytes::Buf trait. Since BytesMut implements Buf trait so we can get_u8 on BytesMut.

buffer_at_array method essentially breaks the string at every space encountered (b' ') and returns a vector of strings vec<String>. Let's call this method in server.rs file in the loop to get a list of strings eg. ["set foo bar"].

use blog_redis::helper::buffer_to_array;

main() {
   loop {
      ...// previous code 
      let attrs = buffer_to_array(&mut buf);
   }
}

Next step, would be to find out the command. For this, take the first string from the vector and compare it with the list of known commands. In our case, Get or Set.

Let's create another file src/cmd.rs. We will create an enum for the commands inside this file. These are the set of commands we will be implementing.

pub enum Command {
  Get,
  Set,
  Invalid,
}

Now, we will try to match the first string and depending on the match return a valid command. Let's create a function inside the src/cmd.rs as implementation of Command to return the command(enum) based on the first string matching.

impl Command {
    pub fn get_command(str: &String) -> Command {
        match str.as_bytes() {
            b"set" => Command::Set,
            b"get" => Command::Get,
            _ => Command::Invalid,
        }
    }
}

Since this file is in src folder, let's import the code in lib.rs.

File lib.rs

pub mod cmd;
pub use cmd::Command;

Let's use this function to get the command.

loop {
    .. // previous code
    let attrs = buffer_to_array(&mut buf);
    let command = Command::get_command(&attrs[0]);
}

Persisting key and value

For our use case, we need to store key value pairs in the database. A data structure which is good for storing key, value pairs is HashMap. Any new connection request will take the reference to the db created and either update (set/write) data into it or get(read) data from it. We will create an instance of db first time when the server starts and for any new connections we will use the db instance to read or write data to it.

Let's create a struct for Database which will hold our key value pairs inside Hashmap. This struct should be added in its own file src/db.rs.

Let's import this module in our src/lib.rs like we did previously for cmd.rs and helper.rs.

File lib.rs

pub mod db;
pub use db::Db;

File src/db.rs

use bytes::Bytes;
use std::collections::HashMap;

pub struct Db {
    entries: HashMap<String, Bytes>
}

Next, let's create a new instance of this Db struct in the server.rs main function. This will be done before the server starts accepting the socket connection requests in the loop {}. Since, we want to persist the database object throughout the scope of the main function of the server.rs file. If we create db in the loop, then at the end of the first loop the object will be dropped from the scope. That's not what we want.

File server.rs

let db = Db::new();
loop {
    let (socket, _) = listener.accept().await?;
}

Let's create a new function process_query() in the server.rs file. It will accept the command, attrs, socket, db parameters and write the attributes to the db object.

pub async fn main() -> Result<(), std::io::Error> { 
    // previous main fn code...
    process_query(command, attrs, &mut socket, &mut db).await?;
}

async fn process_query(
    command: Command,
    attrs: Vec<String>,
    socket: &mut TcpStream,
    db: &mut Db,
) -> std::io::Result<()> {
    match command {
        Command::Get => {
            Ok(())
        }
        Command::Set => {
            let resp = db.write(&attrs);

            match resp {
                Ok(result) => {
                    println!("set result: {}", result);
                    socket.write_all(&result.as_bytes()).await?;
                }
                Err(_err) => {
                    socket.write_all(b"").await?;
                }
            }

            Ok(())
        }
        Command::Invalid => Ok(()),
    }

In process_query() method, we will match the command for Command::Get, Command::Set and Command::Invalid.

In Command::Set we will call db.write(&attrs). As of now, this write method on db is not implemented, so let's go ahead and do that.

File src/db.rs

impl Db {
    pub fn write(&mut self, arr: &[String]) -> Result<&str, &'static str> {
        let key = &arr[1];
        let value = &arr[2];

        let val = value.clone();
        let res: &Option<Bytes> = &self.entries.insert(String::from(key), Bytes::from(val));

        match res {
            Some(_res) => Ok("r Ok"),
            None => Ok("Ok"),
        }
    }

We need to clone the referenced value since Bytes::from() function expects a 'static lifetime variable but value has unknown lifetime in this function context, since it's being defined outside the scope of this function.

insert() function returns None if they key was not present in the hashmap, it returns Some(old_value) if the key was already present in the hashmap. We are returning "r Ok" in case of key already present and "Ok" when the key was not already present in the hash. This helps on the client side to give appropriate message to the client/consumer.

Some and None are types on the Option type. Option is used when we are expecting optional values, meaning either there will be a value with Some(value) or None.

Now, coming back to the bin/server.rs notice that we are getting the result of db.write(&attrs) into a result object. We are using match to capture the response from db.write(...) and writing the value directly into the socket using socket.write_all(&result).await?;. This will send bytes on the socket to be further read by the client socket connection.

client side result

File bin/client.rs

// previous main function code...
let mut buf = BytesMut::with_capacity(1024);
let _length = stream.read_buf(&mut buf).await?;
match std::str::from_utf8(&mut buf) {
    Ok(resp) => {
        if resp == "r Ok" {
            println!("key updated");
        } else if resp == "Ok" {
            println!("key set");
        }
    }
    Err(err: Utf8Error) => {
        // failed to convert bytes into string slice
        println!("error: {}", err);
    }
}

We will create buf of type BytesMut and read the data from socket into the buffer. Why we are using BytesMut is explained above.

Next, we will match the string data from socket and return response based on whether its fresh key write or key updated. We are also capturing Error from std::str::from_utf8(&mut buf) as it will return Utf8Error if it fails to convert the byte slice into string slice.

At this point, our implementation of set command is somewhat complete. We can not test this from the command line.

Let's go ahead and test this by firing the client and hope to see a success message. It will be either "key set" or "key updated".

$ cargo run --bin client
// response should be "key set" in terminal.

So far, we have set a key value pair in our server. Let's try to implement the get command and fetch the corresponding value. Let's implement this in next section.

Read data from server

To fetch the value associated with the key, we will create another socket connection in main() of client.rs. On this socket we will write the data get foo.

File bin/client.rs

let mut stream = TcpStream::connect("127.0.0.1:8081").await?;
stream.write_all(b"get foo").await?;

Next, we will change the bin/server.rs to accept the Command::Get in the match function in the process_query method.
File bin/server.rs

match command {
    Command::Get => {
        let result = db.read(&attrs);
        match result {
            Ok(result) => {
                socket.write_all(&result).await?;
            }
            Err(_err) => {
                println!("no key found {:?}", _err);
                socket.write_all(b"").await?;
            }
        }
        Ok(())
    }
}
Command::Set => { 
  // already implemented above
}

The server has returned the response by writing to the socket with the response from the db. Next, we will update the src/db.rs file to implement the read() method.

pub fn read(&mut self, arr: &[String]) -> Result<&Bytes, &'static str> {
    let key = &arr[1];
    let query_result = self.entries.get(key);

    if let Some(value) = query_result {
        return Ok(value);
    } else {
        return Err("no such key found");
    }
}

The error returned is a string slice str. The &'static lifetime generally means, that the value will persist till the end of the program. The reason for adding &'static lifetime to the str is because, the string is created in this function scope and will be dropped as soon as the function call ends. Hence, making it &'static helps convey to the compiler that this string lives forever. Rust lifetimes are not a simple concept, checkout this awesome tutorial for better understanding.

Final step would be to read the response in client.rs file and print the value in the command line.

let mut buf = BytesMut::with_capacity(1024);
let _length = stream.read_buf(&mut buf).await?;
println!("buffer: {:?}", &buf);
match std::str::from_utf8(&mut buf) {
    Ok(resp) => {
        if resp == "" {
            println!("no such key found");
        } else {
            println!("value: {}", resp);
        }
    }
    Err(_err) => {
        println!("in errr");
    }
}

This is all we need to implement a basic single-threaded redis server with set and get commands.

If you run the server and then the client, you will get response value: bar in the terminal. Hurray!!!

CLAP implementation

So far we have hardcoded our Redis commands into the client.rs file, that's not how a client should be sending the data. So, let's try to accept the command from the command line/terminal.

Clap is a command line parser crate in rust. It makes accepting command line arguments super easy.

First things first, add the dependency in Cargo.toml

clap = { version = "3.1.18", features = ["derive"] }

Now, let's import this in client.rs file and start adding needed code. We are changing way we were sending command from client file to server file. Previously, we created to socket connections from client file. Now, we will only create on socket connection per command line run of client with the command line arguments.

File bin/client.rs

use clap::{Parser, Subcommand};

#[derive(Parser, Debug)]
struct Cli {
    #[clap(subcommand)]
    command: Command,
}

#[derive(Subcommand, Debug)]
enum Command {
    Get {
        key: String,
    },
    Set {
        key: String,
        value: String,
    },
}

main() {
    let args = Cli::parse();

    let mut stream = TcpStream::connect("127.0.0.1:8081").await.unwrap();
    match args.command {
        Command::Set { key, value } => {
            stream.write_all(b"set").await?;
            stream.write_all(b" ").await?;

            stream.write_all(&key.as_bytes()).await?;
            stream.write_all(b" ").await?;

            stream.write_all(&value.as_bytes()).await?;
            let mut buf = BytesMut::with_capacity(1024);
            let _length = stream.read_buf(&mut buf).await?;
            match std::str::from_utf8(&mut buf) {
                Ok(resp) => {
                    if resp == "r Ok" {
                        println!("updated key");
                    } else if resp == "Ok" {
                        println!("key set");
                    }
                }
                Err(err) => {
                    // failed to convert bytes into string slice
                    println!("error: {}", err);
                }
            }
        }
        Command::Get { key } => {
            stream.write_all(b"get").await?;
            stream.write_all(b" ").await?;

            stream.write_all(&key.as_bytes()).await?;

            let mut buf = BytesMut::with_capacity(1024);
            let _length = stream.read_buf(&mut buf).await?;
            match std::str::from_utf8(&mut buf) {
                Ok(resp) => {
                    if resp == "" {
                        println!("no such key found");
                    } else {
                        println!("key: {} => value: {}", key, resp);
                    }
                }
                Err(_err) => {
                    println!("in errr");
                }
            }
            return Ok(());
        }
    }

}

In the code above, we are defining clap with subcommand, for this we are creating an Enum of commands with additional parameters and types.

We are doing simple matching on command supplied from the command line and writing the appropriate command on the socket connection for the server to read and execute.

What next?

The current implementation is a single threaded program. In reality, this is not what we want. Imagine that there are 5000 requests coming to our Redis server implementation, the current program will wait for the first request to complete to process the next one. However, we want all the requests to be handled without any wait time. This can be achieved using multi-threading. In the next post, we will change this single threaded implementation to a multi-threaded one. That is where Tokio shines the most with its advanced tokio::spawn and tokio::select functions.

We will also implement a shutdown mechanism powered by tokio channels like broadcast and mpsc strategies. Imagine there are 5000 active connections, what happens if the server is to shutdown. Graceful implementation of shutdown for all threads will be covered in the next post.

Happy Programming!