DEV Community: Dheeraj Gopinath

Implement Transaction Support in a Redis Clone

Dheeraj Gopinath — Tue, 17 Sep 2024 15:06:02 +0000

Redis transactions provide a mechanism to execute a group of commands as a single atomic operation. When a transaction is initiated, subsequent commands are queued rather than executed immediately. Once all desired commands are added to the transaction, it can be executed or aborted. During execution, Redis guarantees that all commands in the transaction will be processed sequentially and atomically, without interruption from other clients. However, Redis transactions differ from traditional database transactions in important ways. They don't support rollbacks - if a command fails during execution, other commands in the transaction will still be processed.

MULTI, EXEC and DISCARD

A transaction begins with the MULTI command, which tells Redis that subsequent commands should be queued rather than executed immediately. After MULTI, each command is validated but not executed, and Redis responds with "QUEUED". Once all desired commands are queued, the EXEC command is used to execute the entire transaction. The DISCARD command can be used after MULTI but before EXEC to cancel the transaction and clear the queued commands.

Let’s start with the code changes now. The transaction logic is enclosed in a Transaction struct, which has the below methods:

init - Initializes a transaction.
add_command - Queues a command into the transaction.
exec - Executes all commands inside the transaction sequentially.
discard - Discards all commands inside the transaction.

// src/command/transactions.rs

use crate::{resp::types::RespType, storage::db::DB};

use super::Command;

/// Represents a Redis transaction that can be executed atomically (MULTI and EXEC).
pub struct Transaction {
    /// The queue of commands to be executed.
    commands: Vec<Command>,
    /// Indicates whether a transaction is currently active.
    is_active: bool,
}

impl Transaction {
    /// Creates a new `Transaction` instance.
    pub fn new() -> Transaction {
        Transaction {
            commands: vec![],
            is_active: false,
        }
    }

    /// Initializes a new transaction (MULTI command).
    ///
    /// # Returns
    ///
    /// * `Ok(())` if the transaction was successfully initialized.
    /// * `Err(TransactionError::CannotNestMulti)` if a transaction is already active.
    pub fn init(&mut self) -> Result<(), TransactionError> {
        if self.is_active {
            return Err(TransactionError::CannotNestMulti);
        }
        self.is_active = true;

        Ok(())
    }

    /// Adds a command to the transaction.
    ///
    /// # Arguments
    ///
    /// * `cmd` - The command to be added to the transaction.
    pub fn add_command(&mut self, cmd: Command) {
        self.commands.push(cmd);
    }

    /// Checks if a transaction is currently active.
    pub fn is_active(&self) -> bool {
        self.is_active
    }

    /// Executes the commands in the transaction and returns the array of responses.
    ///
    /// This method will execute all the commands in the transaction and return the
    /// responses as a `RespType::Array`. After the execution, the transaction is
    /// automatically discarded.
    ///
    /// # Arguments
    ///
    /// * `db` - The database where the key and values are stored.
    ///
    /// # Returns
    ///
    /// A `RespType::Array` containing the responses for each command in the transaction.
    pub async fn exec(&mut self, db: &DB) -> RespType {
        let mut responses: Vec<RespType> = vec![];

        for cmd in self.commands.iter() {
            // execute the command
            let res = cmd.execute(db);

            responses.push(res);
        }

        // discard txn after executing all commands
        self.discard();

        RespType::Array(responses)
    }

    /// Discards the current transaction.
    ///
    /// This method clears the queue of commands and resets the `is_active` flag.
    pub fn discard(&mut self) {
        self.commands = vec![];
        self.is_active = false;
    }
}

/// Represents errors that can occur during transaction operations.
#[derive(Debug)]
pub enum TransactionError {
    /// Indicates that a MULTI command cannot be nested within another active transaction.
    CannotNestMulti,
}

impl std::error::Error for TransactionError {}

impl std::fmt::Display for TransactionError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            TransactionError::CannotNestMulti => "MULTI calls cannot be nested".fmt(f),
        }
    }
}

Now lets add the MULTI, EXEC and DISCARD commands to our Command enum.

diff --git a/src/command/mod.rs b/src/command/mod.rs
index 71196fc..2804775 100644
--- a/src/command/mod.rs
+++ b/src/command/mod.rs
@@ -17,6 +17,7 @@ mod lrange;
 mod ping;
 mod rpush;
 mod set;
+pub mod transactions;

 /// Represents the supported Nimblecache commands.
 #[derive(Debug, Clone)]
@@ -33,6 +34,12 @@ pub enum Command {
     RPush(RPush),
     /// The LRANGE command.
     LRange(LRange),
+    /// The MULTI command.
+    Multi,
+    /// The EXEC command.
+    Exec,
+    /// The DISCARD command.
+    Discard,
 }

 impl Command {
@@ -91,6 +98,9 @@ impl Command {
                     Err(e) => return Err(e),
                 }
             }
+            "multi" => Command::Multi,
+            "exec" => Command::Exec,
+            "discard" => Command::Discard,
             _ => {
                 return Err(CommandError::UnknownCommand(ErrUnknownCommand {
                     cmd: cmd_name,
@@ -118,6 +128,12 @@ impl Command {
             Command::LPush(lpush) => lpush.apply(db),
             Command::RPush(rpush) => rpush.apply(db),
             Command::LRange(lrange) => lrange.apply(db),
+            // MULTI calls are handled inside FrameHandler.handle since it involves command queueing.
+            Command::Multi => RespType::SimpleString(String::from("OK")),
+            // EXEC calls are handled inside FrameHandler.handle too, since it involves executing queued commands.
+            Command::Exec => RespType::NullBulkString,
+            // DISCARD calls are handled inside FrameHandler.handle too, since it involves discarding queued commands.
+            Command::Discard => RespType::SimpleString(String::from("OK")),
         }
     }
 }

You’ll notice that the commands simply return an OK or nil, and there’s no handler attached to each command. This is because once a transaction is started using the MULTI command, all following commands should be queued. Therefore, MULTI, EXEC, and DISCARD are handled as special cases in our FrameHandler.handle method.

diff --git a/src/handler.rs b/src/handler.rs
index 5ee044c..9bbe529 100644
--- a/src/handler.rs
+++ b/src/handler.rs
@@ -7,7 +7,7 @@ use tokio::net::TcpStream;
 use tokio_util::codec::Framed;

 use crate::{
-    command::Command,
+    command::{transactions::Transaction, Command},
     resp::{frame::RespCommandFrame, types::RespType},
     storage::db::DB,
 };
@@ -29,6 +29,20 @@ impl FrameHandler {
     /// processes them, and sends back the responses. It continues until
     /// an error occurs or the connection is closed.
     ///
+    /// The server's behavior depends on whether a `MULTI` command has been issued.
+    ///
+    /// ## No MULTI Command Issued
+    ///
+    /// If no `MULTI` command has been issued, each command is executed
+    /// immediately and its response is sent back.
+    ///
+    /// ## MULTI Command Issued
+    ///
+    /// If a `MULTI` command has been issued, the method will enter a transaction
+    /// mode. In this mode, all subsequent commands will be queued until an
+    /// `EXEC` command is received. When `EXEC` is called, all the queued
+    /// commands are executed, and the array of responses is sent back.
+    ///
     /// # Arguments
     ///
     /// * `db` - Reference to the database where the key-value pairs are stored.
@@ -42,17 +56,59 @@ impl FrameHandler {
     /// This method will return an error if there's an issue with reading
     /// from or writing to the connection.
     pub async fn handle(mut self, db: &DB) -> Result<()> {
+        // commands are queued here if MULTI command was issued
+        let mut multicommand = Transaction::new();
+
         while let Some(resp_cmd) = self.conn.next().await {
             match resp_cmd {
                 Ok(cmd_frame) => {
                     // Read the command from the frame.
                     let resp_cmd = Command::from_resp_command_frame(cmd_frame);

-                    // Execute the command and get the RESP response.
-                    // If command fails, return RESP SimpleError as response.
+                    // If command is parsed successfully, execute it and get the RESP response,
+                    // otherwise set a SimpleError RESP value as the response.
                     let response = match resp_cmd {
-                        Ok(cmd) => cmd.execute(db),
-                        Err(e) => RespType::SimpleError(format!("{}", e)),
+                        Ok(cmd) => match cmd {
+                            // Initialize pipeline if MULTI command is issued
+                            Command::Multi => {
+                                let init_multicommand = &mut multicommand.init();
+                                match init_multicommand {
+                                    Ok(_) => cmd.execute(db),
+                                    Err(e) => RespType::SimpleError(format!("{}", e)),
+                                }
+                            }
+                            // Execute all commands in pipeline if EXEC command is issued
+                            Command::Exec => {
+                                if multicommand.is_active() {
+                                    multicommand.exec(db).await
+                                } else {
+                                    RespType::SimpleError(String::from("EXEC without MULTI"))
+                                }
+                            }
+                            Command::Discard => {
+                                if multicommand.is_active() {
+                                    multicommand.discard();
+                                    cmd.execute(db)
+                                } else {
+                                    RespType::SimpleError(String::from("DISCARD without MULTI"))
+                                }
+                            }
+                            _ => {
+                                // Queue commands if pipeline is active, else execute the command
+                                if multicommand.is_active() {
+                                    multicommand.add_command(cmd);
+                                    RespType::SimpleString(String::from("QUEUED"))
+                                } else {
+                                    cmd.execute(db)
+                                }
+                            }
+                        },
+                        Err(e) => {
+                            if multicommand.is_active() {
+                                multicommand.discard();
+                            }
+                            RespType::SimpleError(format!("{}", e))
+                        }
                     };

                     // Write the RESP response into the TCP stream.

Testing

Start Nimblecache using the below command and use Redis client to connect to it.

# --port parameter is optional
RUST_LOG=info cargo run -- --port 6380

Now run the below commands:

MULTI
SET a 1 - At this point, if you connect to Nimblecache using redis-cli in a different terminal and execute GET a, you’ll get nil as the response.
EXEC - All queued commands will be executed now. If you execute GET a now, you’ll get 1 as the response.

Source code

The source code for this specific part is available at https://github.com/dheerajgopi/nimblecache/tree/blog-7.

If you are interested in seeing the git-diff between this part and the previous part of this series, have a look at this commit: https://github.com/dheerajgopi/nimblecache/commit/42a4f164ec1ba683e6fc7dd5d652e315e149e63a

Feel free to check the main branch of the Nimblecache repository to see the latest code.

How to Run Multiple Tokio Runtimes in a Rust Application

Dheeraj Gopinath — Wed, 11 Sep 2024 14:02:50 +0000

Getting started with Tokio is straightforward. By simply adding the #[tokio::main] macro to your entry point and using tokio::spawn for task management, you can quickly build an asynchronous application that handles typical use cases effectively.

However, as your application grows in complexity and you need finer control over thread allocation and performance, you'll want to move beyond the default setup and configure the Tokio runtime manually.

In this article, we'll explore a technique where multiple Tokio runtimes are used within a single application. Each runtime is configured with its own parameters and dedicated to specific types of tasks. For example, you might use one runtime to manage lightweight tasks, while another handles more resource-intensive operations.

To illustrate this approach, we'll walk through a practical example: a TCP echo server. In this setup, one Tokio runtime will be responsible for accepting incoming connections, while a separate runtime processes the TCP messages, demonstrating how this technique can help you balance workloads and improve overall application performance.

Here’s the code for a very basic echo server built using the #[tokio::main] macro. We’ll use this as a starting point.

// rustc version 1.77.2
// tokio version 1.38.0, features: full 

use std::error::Error;

use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::TcpListener,
};

/// This macro sets up a default Tokio runtime, and will execute the main method
/// within it.
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let listener = TcpListener::bind("127.0.0.1:8000").await?;

    // loop for accepting incoming connections
    loop {
        let (mut sock, _) = listener.accept().await?;

        // Spawn a async task for each connection.
        // The task handles both read and write operations on the connection.
        tokio::spawn(async move {
            let mut buf = vec![0; 512];

            loop {
                let bytes_read = sock.read(&mut buf).await.expect("failed to read request");

                if bytes_read == 0 {
                    return;
                }

                buf.push(b'\n');

                sock.write_all(&buf[0..bytes_read + 1])
                    .await
                    .expect("failed to write response");

                sock.flush().await.expect("failed to flush response")
            }
        });
    }
}

Echo Server With Multiple Tokio Runtimes

Instead of setting up a Tokio runtime using the macro, we will use tokio::runtime::Builder to do the same, but with much more control over its configuration.

This is an overview of what we’ll be doing:

Create 2 Tokio runtimes: acceptor_runtime and echo_runtime.
Create a tokio::sync::mpsc channel.
acceptor_runtime tasks accept new TCP connections and send the TcpStream to the mpsc channel.
echo_runtime tasks receive the TcpStream from the receiver half of the mpsc channel, and write the echo response.

Here’s the code for the above logic:

// rustc version 1.77.2
// tokio version 1.38.0, features: rt-multi-thread, net, io-util 

use std::error::Error;

use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::{TcpListener, TcpStream},
    sync::mpsc,
};

fn main() -> Result<(), Box<dyn Error>> {
    // Create a tokio runtime whose job is to simply accept new incoming TCP connections.
    let acceptor_runtime = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(2)
        .thread_name("acceptor-pool")
        .enable_all()
        .build()?;

    // Create another tokio runtime whose job is only to write the response bytes to the outgoing TCP message.
    let echo_runtime = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(8)
        .thread_name("echo-handler-pool")
        .enable_all()
        .build()?;

    // this channel is used to pass the TcpStream from acceptor_runtime task to
    // to echo_runtime task where the request handling is done.
    let (tx, mut rx) = mpsc::channel::<TcpStream>(4000);

    // The receiver part of the channel is moved inside a echo_runtime task.
    // This task simply writes the echo response to the TcpStreams coming through the
    // channel receiver.
    echo_runtime.spawn(async move {
        while let Some(mut sock) = rx.recv().await {
            tokio::spawn(async move {
                loop {
                    let mut buf = vec![0; 512];

                    loop {
                        let bytes_read = sock.read(&mut buf).await.expect("failed to read request");

                        if bytes_read == 0 {
                            return;
                        }

                        buf.push(b'\n');

                        sock.write_all(&buf[0..bytes_read + 1])
                            .await
                            .expect("failed to write response");
                    }
                }
            });
        }
    });

    // acceptor_runtime task is run in a blocking manner, so that our server starts
    // accepting new TCP connections. This task just accepts the incoming TcpStreams
    // and are sent to the sender half of the channel.
    acceptor_runtime.block_on(async move {
        let listener = match TcpListener::bind("127.0.0.1:8000").await {
            Ok(l) => l,
            Err(e) => panic!("error binding TCP listener: {}", e),
        };

        loop {
            let sock = match accept_conn(&listener).await {
                Ok(stream) => stream,
                Err(e) => panic!("error reading TCP stream: {}", e),
            };
            let _ = tx.send(sock).await;
        }
    });

    Ok(())
}

async fn accept_conn(listener: &TcpListener) -> Result<TcpStream, Box<dyn Error>> {
    loop {
        match listener.accept().await {
            Ok((sock, _)) => return Ok(sock),
            Err(e) => panic!("error accepting connection: {}", e),
        }
    }
}

Advantages Over the Default Setup

Even though our code is slightly more complex, it has some advantages over the default macro setup.

We can control the number of worker threads using the worker_threads option in tokio::runtime::Builder to better optimize the Tokio runtime based on the type of tasks it performs.
Using a channel gives us better control over concurrency. We can easily build backpressure mechanisms with channels. Our default setup does not provide any control over the number of worker tasks spawned.
The modified echo server has a better separation of concerns. It clearly separates the connection acceptance logic from the connection handling logic.

Implement List Data Type for a Redis Clone

Dheeraj Gopinath — Mon, 09 Sep 2024 11:16:20 +0000

In our ongoing series on building a Redis clone using Rust and Tokio, we've already set up a key-value storage and implemented the essential GET and SET commands. Now, it's time to introduce a more complex data structure: the LIST.

Before proceeding further, take a look at the official Redis documentation: https://redis.io/docs/latest/develop/data-types/lists/.

Please note that not all List commands will be implemented in this exercise. We'll cover only LPUSH, RPUSH, and LRANGE.

Add Support for List Data Type

Building upon our existing key-value storage system, implementing the LIST data type in our Redis clone is surprisingly straightforward. We'll leverage Rust's standard library, specifically the VecDeque (double-ended queue) data structure, to efficiently handle operations at both ends of the list. VecDeque is an ideal choice for our implementation as it provides O(1) time complexity for insertions and deletions at both the front and back of the list. This aligns perfectly with Redis' LIST operations like LPUSH, RPUSH, LPOP, and RPOP, which primarily manipulate the head and tail of the list. By using VecDeque, we can ensure that our LIST commands will perform optimally, mirroring the efficiency users expect from Redis.

Below are the required code changes:

 diff --git a/src/storage/db.rs b/src/storage/db.rs
index 187965a..c567b48 100644
--- a/src/storage/db.rs
+++ b/src/storage/db.rs
@@ -1,7 +1,7 @@
 // src/storage/db.rs

 use std::{
-    collections::HashMap,
+    collections::{HashMap, VecDeque},
     sync::{Arc, RwLock},
 };

@@ -30,11 +30,12 @@ pub struct Entry {
 }

 /// The `Value` enum allows for storing various types of data associated with a key.
-/// Currently, it supports only String data type. But it can be expanded in the future
-/// to support more data types as needed (like List, Hash etc).
+/// Currently, it supports only String and List data type. But it can be expanded in the future
+/// to support more data types as needed (like Hash, SortedSet etc).
 #[derive(Debug, Clone)]
 pub enum Value {
     String(String),
+    List(VecDeque<String>),
 }

 impl Storage {
@@ -122,6 +123,184 @@ impl DB {

         return Ok(());
     }
+
+    /// Add new elements to the head of a list.
+    /// If the key is not present in the DB, and empty list is initialized
+    /// against the key before adding the elements to the head.
+    ///
+    /// # Arguments
+    ///
+    /// * `k` - The key on which list is stored.
+    ///
+    /// * `v` - The values to be added to the head of the list.
+    ///
+    /// # Returns
+    ///
+    /// * `Ok(())` - If values are added successfully to the head of the list.
+    /// * `Err(DBError)` - if key already exists and has non-list data.
+    pub fn lpush(&self, k: String, v: Vec<String>) -> Result<usize, DBError> {
+        let mut data = match self.data.write() {
+            Ok(data) => data,
+            Err(e) => return Err(DBError::Other(format!("{}", e))),
+        };
+
+        let entry = match data.get_mut(k.as_str()) {
+            Some(entry) => Some(entry),
+            None => None,
+        };
+
+        match entry {
+            Some(e) => {
+                let val = &mut e.value;
+                match val {
+                    Value::List(l) => {
+                        for each in v.iter().cloned() {
+                            l.push_front(each);
+                        }
+                        Ok(l.len())
+                    }
+                    _ => Err(DBError::WrongType),
+                }
+            }
+            None => {
+                let list = VecDeque::from(v);
+                let l_len = list.len();
+                data.insert(k.to_string(), Entry::new(Value::List(list)));
+
+                Ok(l_len)
+            }
+        }
+    }
+
+    /// Adds new elements to the tail of a list.
+    /// If the key is not present in the DB, and empty list is initialized
+    /// against the key before adding the elements to the tail.
+    ///
+    /// # Arguments
+    ///
+    /// * `k` - The key on which list is stored.
+    ///
+    /// * `v` - The values to be added to the tail of the list.
+    ///
+    /// # Returns
+    ///
+    /// * `Ok(())` - If value are added successfully to the tail of the list.
+    /// * `Err(DBError)` - if key already exists and has non-list data.
+    pub fn rpush(&self, k: String, v: Vec<String>) -> Result<usize, DBError> {
+        let mut data = match self.data.write() {
+            Ok(data) => data,
+            Err(e) => return Err(DBError::Other(format!("{}", e))),
+        };
+
+        let entry = match data.get_mut(k.as_str()) {
+            Some(entry) => Some(entry),
+            None => None,
+        };
+
+        match entry {
+            Some(e) => {
+                let val = &mut e.value;
+                match val {
+                    Value::List(l) => {
+                        for each in v.iter().cloned() {
+                            l.push_back(each);
+                        }
+                        Ok(l.len())
+                    }
+                    _ => Err(DBError::WrongType),
+                }
+            }
+            None => {
+                let list = VecDeque::from(v);
+                let l_len = list.len();
+                data.insert(k.to_string(), Entry::new(Value::List(list)));
+
+                Ok(l_len)
+            }
+        }
+    }
+
+    /// Returns the specified number of elements of the list stored at key, based on the start and stop indices.
+    /// These offsets can also be negative numbers indicating offsets starting at the end of the list.
+    /// For example, -1 is the last element of the list, -2 the penultimate, and so on.
+    /// Please note that the item at stop index is also included in the result.
+    ///
+    /// If the specified key is not found, an empty list is returned.
+    ///
+    /// # Arguments
+    ///
+    /// * `k` - The key on which list is stored.
+    ///
+    /// * `start_idx` - The start index.
+    ///
+    /// * `stop_idx` - The end index.
+    ///
+    /// # Returns
+    ///
+    /// * `Ok(Vec<String>)` - If values are retrieved successfully from the list.
+    /// * `Err(DBError)` - if key already exists and has non-list data.
+    pub fn lrange(&self, k: String, start_idx: i64, stop_idx: i64) -> Result<Vec<String>, DBError> {
+        let data = match self.data.read() {
+            Ok(data) => data,
+            Err(e) => return Err(DBError::Other(format!("{}", e))),
+        };
+
+        let entry = match data.get(k.as_str()) {
+            Some(entry) => entry,
+            None => return Ok(vec![]),
+        };
+
+        match &entry.value {
+            Value::List(l) => {
+                let l_len = l.len() as i64;
+                let (rounded_start_idx, rounded_stop_idx) =
+                    Self::round_list_indices(l_len, start_idx, stop_idx);
+                Ok(l.range(rounded_start_idx..rounded_stop_idx)
+                    .cloned()
+                    .collect())
+            }
+            _ => Err(DBError::WrongType),
+        }
+    }
+
+    /// Round index to 0, if the given index value is less than zero.
+    /// Round index to list length, if the given index value is greater then the list length.
+    fn round_list_index(list_len: i64, idx: i64) -> usize {
+        if idx < 0 {
+            let idx = list_len - idx.abs();
+            if idx < 0 {
+                return 0;
+            } else {
+                return idx as usize;
+            }
+        }
+
+        if idx >= list_len {
+            return (list_len - 1) as usize;
+        }
+
+        return idx as usize;
+    }
+
+    /// Round the start and stop indices using `Self::round_list_index` method and return them as
+    /// a tuple.
+    /// Special condition: If stop index is lower than start index, return (0, 0).
+    fn round_list_indices(list_len: i64, start_idx: i64, stop_idx: i64) -> (usize, usize) {
+        if stop_idx < start_idx {
+            return (0, 0);
+        }
+
+        let rounded_start_idx = Self::round_list_index(list_len, start_idx);
+        let rounded_stop_idx = Self::round_list_index(list_len, stop_idx);
+
+        if rounded_start_idx < rounded_stop_idx {
+            (rounded_start_idx, rounded_stop_idx + 1)
+        } else if rounded_stop_idx < rounded_start_idx {
+            (0, 0)
+        } else {
+            (rounded_start_idx, rounded_start_idx + 1)
+        }
+    }
 }

 impl Entry {

Handle LPUSH, RPUSH and LRANGE

Implementing these commands is similar to how we did the same for PING, SET and GET. Before proceeding further, have a look at the documentation for the above commands:

LPUSH - https://redis.io/docs/latest/commands/lpush/
RPUSH - https://redis.io/docs/latest/commands/rpush/
LRANGE - https://redis.io/docs/latest/commands/lrange/

If you read the above documentation, you'll see that the LPUSH and RPUSH commands return an integer in RESP format. Currently, our RespType doesn't support integer values, so let's add that feature first.

Add Integer Support in RespType

diff --git a/src/resp/types.rs b/src/resp/types.rs
index 9dd2930..051a21a 100644
--- a/src/resp/types.rs
+++ b/src/resp/types.rs
@@ -11,6 +11,8 @@ pub enum RespType {
     SimpleError(String),
     /// Refer <https://redis.io/docs/latest/develop/reference/protocol-spec/#arrays>
     Array(Vec<RespType>),
+    /// Refer <https://redis.io/docs/latest/develop/reference/protocol-spec/#integers>
+    Integer(i64),
 }

 use super::RespError;
@@ -139,6 +141,7 @@ impl RespType {
                 Bytes::from_iter(arr_bytes)
             }
             RespType::SimpleError(es) => Bytes::from_iter(format!("-{}\r\n", es).into_bytes()),
+            RespType::Integer(i) => Bytes::from_iter(format!(":{}\r\n", i).into_bytes()),
         };
     }

Add Command Handlers for LPUSH, RPUSH and LRANGE

// src/command/lpush.rs

use crate::{resp::types::RespType, storage::db::DB};

use super::CommandError;

/// Represents the LPUSH command in Nimblecache.
#[derive(Debug, Clone)]
pub struct LPush {
    key: String,
    values: Vec<String>,
}

impl LPush {
    /// Creates a new `LPUSH` instance from the given arguments.
    ///
    /// # Arguments
    ///
    /// * `args` - A vector of `RespType` representing the arguments to the SET command.
    ///
    /// # Returns
    ///
    /// * `Ok(LPush)` if parsing succeeds.
    /// * `Err(CommandError)` if parsing fails.
    pub fn with_args(args: Vec<RespType>) -> Result<LPush, CommandError> {
        if args.len() < 2 {
            return Err(CommandError::Other(String::from(
                "Wrong number of arguments specified for 'LPUSH' command",
            )));
        }

        // parse key
        let key = &args[0];
        let key = match key {
            RespType::BulkString(k) => k,
            _ => {
                return Err(CommandError::Other(String::from(
                    "Invalid argument. Key must be a bulk string",
                )));
            }
        };

        // parse values
        let mut values: Vec<String> = vec![];
        for arg in args[1..].iter() {
            match arg {
                RespType::BulkString(v) => values.push(v.to_string()),
                _ => {
                    return Err(CommandError::Other(String::from(
                        "Invalid argument. Value must be a bulk string",
                    )));
                }
            }
        }

        Ok(LPush {
            key: key.to_string(),
            values,
        })
    }

    /// Executes the LPUSH command.
    ///
    /// # Arguments
    ///
    /// * `db` - The database where the key and values are stored.
    ///
    /// # Returns
    ///
    /// It returns the length of the list if value is successfully written.
    pub fn apply(&self, db: &DB) -> RespType {
        match db.lpush(self.key.clone(), self.values.clone()) {
            Ok(len) => RespType::Integer(len as i64),
            Err(e) => RespType::SimpleError(format!("{}", e)),
        }
    }

    pub fn build_command(&self) -> RespType {
        let mut args: Vec<RespType> = vec![
            RespType::BulkString(String::from("LPUSH")),
            RespType::BulkString(self.key.clone()),
        ];

        let arg_vals = self.values.clone();
        for arg in arg_vals.iter() {
            args.push(RespType::BulkString(arg.to_string()));
        }

        RespType::Array(args)
    }
}

// src/command/rpush.rs

use crate::{resp::types::RespType, storage::db::DB};

use super::CommandError;

/// Represents the RPUSH command in Nimblecache.
#[derive(Debug, Clone)]
pub struct RPush {
    key: String,
    values: Vec<String>,
}

impl RPush {
    /// Creates a new `RPUSH` instance from the given arguments.
    ///
    /// # Arguments
    ///
    /// * `args` - A vector of `RespType` representing the arguments to the SET command.
    ///
    /// # Returns
    ///
    /// * `Ok(RPush)` if parsing succeeds.
    /// * `Err(CommandError)` if parsing fails.
    pub fn with_args(args: Vec<RespType>) -> Result<RPush, CommandError> {
        if args.len() < 2 {
            return Err(CommandError::Other(String::from(
                "Wrong number of arguments specified for 'RPUSH' command",
            )));
        }

        // parse key
        let key = &args[0];
        let key = match key {
            RespType::BulkString(k) => k,
            _ => {
                return Err(CommandError::Other(String::from(
                    "Invalid argument. Key must be a bulk string",
                )));
            }
        };

        // parse values
        let mut values: Vec<String> = vec![];
        for arg in args[1..].iter() {
            match arg {
                RespType::BulkString(v) => values.push(v.to_string()),
                _ => {
                    return Err(CommandError::Other(String::from(
                        "Invalid argument. Value must be a bulk string",
                    )));
                }
            }
        }

        Ok(RPush {
            key: key.to_string(),
            values,
        })
    }

    /// Executes the RPUSH command.
    ///
    /// # Arguments
    ///
    /// * `db` - The database where the key and values are stored.
    ///
    /// # Returns
    ///
    /// It returns the length of the list if value is successfully written.
    pub fn apply(&self, db: &DB) -> RespType {
        match db.rpush(self.key.clone(), self.values.clone()) {
            Ok(len) => RespType::Integer(len as i64),
            Err(e) => RespType::SimpleError(format!("{}", e)),
        }
    }

    pub fn build_command(&self) -> RespType {
        let mut args: Vec<RespType> = vec![
            RespType::BulkString(String::from("RPUSH")),
            RespType::BulkString(self.key.clone()),
        ];

        let arg_vals = self.values.clone();
        for arg in arg_vals.iter() {
            args.push(RespType::BulkString(arg.to_string()));
        }

        RespType::Array(args)
    }
}

// src/command/lrange.rs

use crate::{resp::types::RespType, storage::db::DB};

use super::CommandError;

/// Represents the LRANGE command in Nimblecache.
#[derive(Debug, Clone)]
pub struct LRange {
    key: String,
    start_idx: i64,
    end_idx: i64,
}

impl LRange {
    /// Creates a new `LRANGE` instance from the given arguments.
    ///
    /// # Arguments
    ///
    /// * `args` - A vector of `RespType` representing the arguments to the SET command.
    ///
    /// # Returns
    ///
    /// * `Ok(LRange)` if parsing succeeds.
    /// * `Err(CommandError)` if parsing fails.
    pub fn with_args(args: Vec<RespType>) -> Result<LRange, CommandError> {
        if args.len() < 3 {
            return Err(CommandError::Other(String::from(
                "Wrong number of arguments specified for 'LRANGE' command",
            )));
        }

        // parse key
        let key = &args[0];
        let key = match key {
            RespType::BulkString(k) => k,
            _ => {
                return Err(CommandError::Other(String::from(
                    "Invalid argument. Key must be a bulk string",
                )));
            }
        };

        // parse start index
        let value = &args[1];
        let start_idx = match value {
            RespType::BulkString(v) => {
                let start_idx = v.parse::<i64>();
                match start_idx {
                    Ok(i) => i,
                    Err(_) => {
                        return Err(CommandError::Other(String::from(
                            "Start index should be an integer",
                        )))
                    }
                }
            }
            _ => {
                return Err(CommandError::Other(String::from(
                    "Invalid argument. Value must be an integer in bulk string format",
                )));
            }
        };

        // parse end index
        let value = &args[2];
        let end_idx = match value {
            RespType::BulkString(v) => {
                let end_idx = v.parse::<i64>();
                match end_idx {
                    Ok(i) => i,
                    Err(_) => {
                        return Err(CommandError::Other(String::from(
                            "End index should be an integer",
                        )))
                    }
                }
            }
            _ => {
                return Err(CommandError::Other(String::from(
                    "Invalid argument. Value must be an integer in bulk string format",
                )));
            }
        };

        Ok(LRange {
            key: key.to_string(),
            start_idx,
            end_idx,
        })
    }

    /// Executes the LRANGE command.
    ///
    /// # Arguments
    ///
    /// * `db` - The database where the key and values are stored.
    ///
    /// # Returns
    ///
    /// It returns the specified number of elements in the list stored at key, based on start and stop indices.
    pub fn apply(&self, db: &DB) -> RespType {
        match db.lrange(self.key.clone(), self.start_idx, self.end_idx) {
            Ok(elems) => {
                let sub_list = elems
                    .iter()
                    .cloned()
                    .map(|e| RespType::BulkString(e))
                    .collect();
                RespType::Array(sub_list)
            }
            Err(e) => RespType::SimpleError(format!("{}", e)),
        }
    }
}

diff --git a/src/command/mod.rs b/src/command/mod.rs
index ba4b8b4..71196fc 100644
--- a/src/command/mod.rs
+++ b/src/command/mod.rs
@@ -3,13 +3,19 @@
 use core::fmt;

 use get::Get;
+use lpush::LPush;
+use lrange::LRange;
 use ping::Ping;
+use rpush::RPush;
 use set::Set;

 use crate::{resp::types::RespType, storage::db::DB};

 mod get;
+mod lpush;
+mod lrange;
 mod ping;
+mod rpush;
 mod set;

 /// Represents the supported Nimblecache commands.
@@ -21,6 +27,12 @@ pub enum Command {
     Set(Set),
     /// The GET command.
     Get(Get),
+    /// The LPUSH command.
+    LPush(LPush),
+    /// The RPUSH command.
+    RPush(RPush),
+    /// The LRANGE command.
+    LRange(LRange),
 }

 impl Command {
@@ -58,6 +70,27 @@ impl Command {
                     Err(e) => return Err(e),
                 }
             }
+            "lpush" => {
+                let cmd = LPush::with_args(Vec::from(args));
+                match cmd {
+                    Ok(cmd) => Command::LPush(cmd),
+                    Err(e) => return Err(e),
+                }
+            }
+            "rpush" => {
+                let cmd = RPush::with_args(Vec::from(args));
+                match cmd {
+                    Ok(cmd) => Command::RPush(cmd),
+                    Err(e) => return Err(e),
+                }
+            }
+            "lrange" => {
+                let cmd = LRange::with_args(Vec::from(args));
+                match cmd {
+                    Ok(cmd) => Command::LRange(cmd),
+                    Err(e) => return Err(e),
+                }
+            }
             _ => {
                 return Err(CommandError::UnknownCommand(ErrUnknownCommand {
                     cmd: cmd_name,
@@ -82,6 +115,9 @@ impl Command {
             Command::Ping(ping) => ping.apply(),
             Command::Set(set) => set.apply(db),
             Command::Get(get) => get.apply(db),
+            Command::LPush(lpush) => lpush.apply(db),
+            Command::RPush(rpush) => rpush.apply(db),
+            Command::LRange(lrange) => lrange.apply(db),
         }
     }
 }

Testing

Start Nimblecache using the below command and use Redis client to connect to it.

# --port parameter is optional
RUST_LOG=info cargo run -- --port 6380

Now, try issuing LPUSH, RPUSH and LRANGE commands to see if they work properly.

Source code

The source code for this specific part is available at https://github.com/dheerajgopi/nimblecache/tree/blog-6.

Feel free to check the main branch of the Nimblecache repository to see the latest code.

Conclusion and what's next

In this article, we successfully implemented the LIST data type in our Redis clone, covering the essential LPUSH, RPUSH, and LRANGE commands.

In the next part of this series, we will implement transactions using the MULTI and EXEC commands. This will allow us to execute multiple commands at once.

Stay tuned!

Building key-value storage for a Redis clone

Dheeraj Gopinath — Mon, 02 Sep 2024 14:10:25 +0000

In the previous part of this series, we explored how to handle a basic PING command in our Redis clone. In this part, we will implement the Redis SET and GET commands. Additionally, we will build a simple HashMap-based storage system to support these operations.

Hash-map based storage

The storage system we're building for our Redis clone is conceptually very simple—essentially, it's just a global HashMap that employs thread-safe mechanisms to manage concurrent access. However, what might be straightforward to implement in other programming languages becomes slightly complex in Rust due to its strict ownership and concurrency model. In Rust, we have to handle shared state across threads using tools like Arc and Mutex or RwLock. These ensure that our data remains safe and consistent across multiple threads.

Check the code below for the HashMap storage implementation:

The rustdocs will help you understand each part of the code.

// src/storage/db.rs

use std::{
    collections::HashMap,
    sync::{Arc, RwLock},
};

use super::DBError;

/// The Storage struct is designed to act as a wrapper around the core database,
/// allowing it to be shared across multiple connections. The database is encapsulated within an Arc,
/// to enable concurrent access.
#[derive(Debug, Clone)]
pub struct Storage {
    db: Arc<DB>,
}

/// The DB struct is the component that houses the actual data,
/// which is stored in a RwLock wrapped around a HashMap. This ensures thread-safe read and write operations.
#[derive(Debug)]
pub struct DB {
    data: RwLock<HashMap<String, Entry>>,
}

/// The Entry struct represents the value associated with a particular key in the database.
/// This struct encapsulates the Value enum, which allows for different types of data to be stored.
#[derive(Debug, Clone)]
pub struct Entry {
    value: Value,
}

/// The `Value` enum allows for storing various types of data associated with a key.
/// Currently, it supports only String data type. But it can be expanded in the future
/// to support more data types as needed (like List, Hash etc).
#[derive(Debug, Clone)]
pub enum Value {
    String(String),
}

impl Storage {
    /// Create a new instance of `Storage` which contains the DB.
    pub fn new(db: DB) -> Storage {
        Storage { db: Arc::new(db) }
    }

    /// Returns a clone of the shared database (`Arc<DB>`).
    ///
    /// This method provides access to the underlying database, which is shared across all
    /// connections. The database is wrapped in an `Arc` to ensure concurrent access by multiple threads.
    pub fn db(&self) -> Arc<DB> {
        self.db.clone()
    }
}

impl DB {
    /// Create a new instance of DB.
    pub fn new() -> DB {
        DB {
            data: RwLock::new(HashMap::new()),
        }
    }

    /// Get the string value stored against a key.
    ///
    /// # Arguments
    ///
    /// * `k` - The key on which lookup is performed.
    ///
    /// # Returns
    ///
    /// * `Ok(Option<String>)` - `Some(String)` if key is found in DB, else `None`
    /// * `Err(DBError)` - if key already exists and has non-string data.
    pub fn get(&self, k: &str) -> Result<Option<String>, DBError> {
        let data = match self.data.read() {
            Ok(data) => data,
            Err(e) => return Err(DBError::Other(format!("{}", e))),
        };

        let entry = match data.get(k) {
            Some(entry) => entry,
            None => return Ok(None),
        };

        if let Value::String(s) = &entry.value {
            return Ok(Some(s.to_string()));
        }

        Err(DBError::WrongType)
    }

    /// Set a string value against a key.
    ///
    /// # Arguments
    ///
    /// * `k` - The key on which value is to be set.
    ///
    /// * `v` - The value to be set against the key.
    ///
    /// # Returns
    ///
    /// * `Ok(())` - If value is successfully added against the key.
    /// * `Err(DBError)` - if key already exists and has non-string data.
    pub fn set(&self, k: String, v: Value) -> Result<(), DBError> {
        let mut data = match self.data.write() {
            Ok(data) => data,
            Err(e) => return Err(DBError::Other(format!("{}", e))),
        };

        let entry = match data.get(k.as_str()) {
            Some(entry) => Some(entry),
            None => None,
        };

        if entry.is_some() {
            match entry.unwrap().value {
                Value::String(_) => {}
                _ => return Err(DBError::WrongType),
            }
        }

        data.insert(k.to_string(), Entry::new(v));

        return Ok(());
    }
}

impl Entry {
    pub fn new(value: Value) -> Entry {
        Entry { value }
    }
}

// src/storage/mod.rs

pub mod db;

/// Represents errors that can occur during DB operations.
#[derive(Debug)]
pub enum DBError {
    /// Represents an error where wrong data type is encountered against a key.
    /// For e.g. If you try to perform list related operation (such as lpush, rpush) on a key
    /// which stores a string value.
    WrongType,
    /// Represents any other error with a descriptive message.
    Other(String),
}

impl std::fmt::Display for DBError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            DBError::WrongType => {
                "WRONGTYPE Operation against a key holding the wrong kind of value".fmt(f)
            }
            DBError::Other(msg) => msg.as_str().fmt(f),
        }
    }
}

Implementing SET and GET commands

Now that we've implemented a basic database, let's perform simple read and write operations using the SET and GET commands. The structure for implementing these commands is similar to how the PING command was implemented. However, the key difference is that the apply method for the SET and GET commands requires a reference to the storage layer, allowing them to interact directly with the database to store and retrieve data.

SET command

Our implementation of SET command is very basic. Advanced options like TTL are ignored for now.

Here's the Redis documentation for the SET command: https://redis.io/docs/latest/commands/set/

// src/command/set.rs

use crate::{
    resp::types::RespType,
    storage::db::{Value, DB},
};

use super::CommandError;

/// Represents the SET command in Nimblecache.
///
/// The `Set` struct encapsulates the key-value pair for the SET command, which is used
/// to store a string value against a key in the database.
#[derive(Debug, Clone)]
pub struct Set {
    key: String,
    value: String,
}

impl Set {
    /// Creates a new `Set` instance from the given arguments.
    ///
    /// This function parses the arguments provided in the form of a `RespType` vector.
    /// It validates and extracts the key and value for the SET command.
    ///
    /// # Arguments
    ///
    /// * `args` - A vector of `RespType` representing the arguments to the SET command.
    ///
    /// # Returns
    ///
    /// * `Ok(Set)` - If parsing succeeds and the key-value pair is valid.
    /// * `Err(CommandError)` - if parsing fails due to validation errors.
    pub fn with_args(args: Vec<RespType>) -> Result<Set, CommandError> {
        if args.len() < 2 {
            return Err(CommandError::Other(String::from(
                "Wrong number of arguments specified for 'SET' command",
            )));
        }

        // parse key
        let key = &args[0];
        let key = match key {
            RespType::BulkString(k) => k,
            _ => {
                return Err(CommandError::Other(String::from(
                    "Invalid argument. Key must be a bulk string",
                )));
            }
        };

        // parse value
        let value = &args[1];
        let value = match value {
            RespType::BulkString(v) => v.to_string(),
            _ => {
                return Err(CommandError::Other(String::from(
                    "Invalid argument. Value must be a bulk string",
                )));
            }
        };

        Ok(Set {
            key: key.to_string(),
            value,
        })
    }

    /// Executes the SET command.
    ///
    /// This method writes the string value to the database under the specified key.
    /// If the operation is successful, it returns an "OK" response as a `BulkString`.
    /// If the operation fails, it returns an error response.
    ///
    /// # Arguments
    ///
    /// * `db` - A reference to the `DB` instance where the key-value pair should be stored.
    ///
    /// # Returns
    ///
    /// * `BulkString("OK")` - If the value is successfully written.
    /// * `SimpleError` - If the operation fails due to some error.
    pub fn apply(&self, db: &DB) -> RespType {
        match db.set(self.key.clone(), Value::String(self.value.clone())) {
            Ok(_) => RespType::BulkString("OK".to_string()),
            Err(e) => RespType::SimpleError(format!("{}", e)),
        }
    }
}

GET command

Here's the Redis documentation for GET command: https://redis.io/docs/latest/commands/get/

If you read the above documentation, you'll find that Redis responds with a nil if the specified key does not exist. But as of now, our Redis clone does not have a corresponding RespType value for nil. Let's implement that too as part of building the GET command handler.

nil values are represented in RESP using NullBulkString. It's simply a BulkString with length of negative one (-1).

diff --git a/src/resp/types.rs b/src/resp/types.rs
index ebedb2c..9dd2930 100644
--- a/src/resp/types.rs
+++ b/src/resp/types.rs
@@ -5,6 +5,8 @@ pub enum RespType {
     SimpleString(String),
     /// Refer <https://redis.io/docs/latest/develop/reference/protocol-spec/#bulk-strings>
     BulkString(String),
+    /// Null representation in RESP2. It's simply a BulkString with length of negative one (-1).
+    NullBulkString,
     /// Refer <https://redis.io/docs/latest/develop/reference/protocol-spec/#simple-errors>
     SimpleError(String),
     /// Refer <https://redis.io/docs/latest/develop/reference/protocol-spec/#arrays>
@@ -127,6 +129,7 @@ impl RespType {
                 let bulkstr_bytes = format!("${}\r\n{}\r\n", bs.chars().count(), bs).into_bytes();
                 Bytes::from_iter(bulkstr_bytes)
             }
+            RespType::NullBulkString => Bytes::from("$-1\r\n"),
             RespType::Array(arr) => {
                 let mut arr_bytes = format!("*{}\r\n", arr.len()).into_bytes();
                 arr.iter()

Now lets implement the GET command handler.

// src/command/get.rs

use crate::{resp::types::RespType, storage::db::DB};

use super::CommandError;

/// Represents the GET command in Nimblecache.
///
/// The `Get` struct is used to retrieve the value associated with a specified key
/// from the database.
#[derive(Debug, Clone)]
pub struct Get {
    /// Key to be searched in the database
    key: String,
}

impl Get {
    /// Creates a new `Get` instance from the given arguments.
    ///
    /// This function parses the arguments provided in the form of a `RespType` vector.
    /// It validates and extracts the key for the GET command.
    ///
    /// # Arguments
    ///
    /// * `args` - A vector of `RespType` representing the arguments to the GET command.
    ///
    /// # Returns
    ///
    /// * `Ok(Get)` - If parsing succeeds and the key is valid.
    /// * `Err(CommandError)` - if parsing fails due to validation errors.
    pub fn with_args(args: Vec<RespType>) -> Result<Get, CommandError> {
        if args.len() < 1 {
            return Err(CommandError::Other(String::from(
                "Wrong number of arguments specified for 'GET' command",
            )));
        }

        // parse key
        let key = &args[0];
        let key = match key {
            RespType::BulkString(k) => k.to_string(),
            _ => {
                return Err(CommandError::Other(String::from(
                    "Invalid argument. Key must be a bulk string",
                )));
            }
        };

        Ok(Get { key })
    }

    /// Executes the GET command.
    ///
    /// # Arguments
    ///
    /// * `db` - The database where the key and values are stored.
    ///
    /// # Returns
    ///
    /// - If key is present in DB - Value of the key as a `BulkString`
    /// - If key is not found in DB - A `NullBulkString`
    /// - If an error is encountered - A `SimpleError` with an error message
    pub fn apply(&self, db: &DB) -> RespType {
        match db.get(self.key.as_str()) {
            Ok(val) => match val {
                Some(s) => RespType::BulkString(s),
                None => RespType::NullBulkString,
            },
            Err(e) => RespType::SimpleError(format!("{}", e)),
        }
    }
}

Now, let's add the SET and GET command handlers to the Command enum.

diff --git a/src/command/mod.rs b/src/command/mod.rs
index eb914ac..ba4b8b4 100644
--- a/src/command/mod.rs
+++ b/src/command/mod.rs
@@ -2,17 +2,25 @@

 use core::fmt;

+use get::Get;
 use ping::Ping;
+use set::Set;

-use crate::resp::types::RespType;
+use crate::{resp::types::RespType, storage::db::DB};

+mod get;
 mod ping;
+mod set;

 /// Represents the supported Nimblecache commands.
 #[derive(Debug, Clone)]
 pub enum Command {
     /// The PING command.
     Ping(Ping),
+    /// The SET command.
+    Set(Set),
+    /// The GET command.
+    Get(Get),
 }

 impl Command {
@@ -36,6 +44,20 @@ impl Command {

         let cmd = match cmd_name.to_lowercase().as_str() {
             "ping" => Command::Ping(Ping::with_args(Vec::from(args))?),
+            "set" => {
+                let cmd = Set::with_args(Vec::from(args));
+                match cmd {
+                    Ok(cmd) => Command::Set(cmd),
+                    Err(e) => return Err(e),
+                }
+            }
+            "get" => {
+                let cmd = Get::with_args(Vec::from(args));
+                match cmd {
+                    Ok(cmd) => Command::Get(cmd),
+                    Err(e) => return Err(e),
+                }
+            }
             _ => {
                 return Err(CommandError::UnknownCommand(ErrUnknownCommand {
                     cmd: cmd_name,
@@ -48,12 +70,18 @@ impl Command {

     /// Executes the Nimblecache command.
     ///
+    /// # Arguments
+    ///
+    /// * `db` - Reference to the database where the key-value pairs are stored.
+    ///
     /// # Returns
     ///
     /// The result of the command execution as a `RespType`.
-    pub fn execute(&self) -> RespType {
+    pub fn execute(&self, db: &DB) -> RespType {
         match self {
             Command::Ping(ping) => ping.apply(),
+            Command::Set(set) => set.apply(db),
+            Command::Get(get) => get.apply(db),
         }
     }
 }

Other code changes

Below are some additional code changes required to pass the DB reference to the command handlers:

diff --git a/src/handler.rs b/src/handler.rs
index 2f6ae69..5ee044c 100644
--- a/src/handler.rs
+++ b/src/handler.rs
@@ -9,6 +9,7 @@ use tokio_util::codec::Framed;
 use crate::{
     command::Command,
     resp::{frame::RespCommandFrame, types::RespType},
+    storage::db::DB,
 };

 /// Handles RESP command frames over a single TCP connection.
@@ -28,6 +29,10 @@ impl FrameHandler {
     /// processes them, and sends back the responses. It continues until
     /// an error occurs or the connection is closed.
     ///
+    /// # Arguments
+    ///
+    /// * `db` - Reference to the database where the key-value pairs are stored.
+    ///
     /// # Returns
     ///
     /// A `Result` indicating whether the operation succeeded or failed.
@@ -36,7 +41,7 @@ impl FrameHandler {
     ///
     /// This method will return an error if there's an issue with reading
     /// from or writing to the connection.
-    pub async fn handle(mut self) -> Result<()> {
+    pub async fn handle(mut self, db: &DB) -> Result<()> {
         while let Some(resp_cmd) = self.conn.next().await {
             match resp_cmd {
                 Ok(cmd_frame) => {
@@ -46,7 +51,7 @@ impl FrameHandler {
                     // Execute the command and get the RESP response.
                     // If command fails, return RESP SimpleError as response.
                     let response = match resp_cmd {
-                        Ok(cmd) => cmd.execute(),
+                        Ok(cmd) => cmd.execute(db),
                         Err(e) => RespType::SimpleError(format!("{}", e)),
                     };

diff --git a/src/server.rs b/src/server.rs
index 9e45bee..5fe3af2 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -1,5 +1,7 @@
 // src/server.rs

+use std::sync::Arc;
+
 // anyhow provides the Error and Result types for convenient error handling
 use anyhow::{Error, Result};

@@ -9,24 +11,33 @@ use log::error;
 use tokio::net::{TcpListener, TcpStream};
 use tokio_util::codec::Framed;

-use crate::{handler::FrameHandler, resp::frame::RespCommandFrame};
+use crate::{handler::FrameHandler, resp::frame::RespCommandFrame, storage::db::Storage};

-/// The Server struct holds the tokio TcpListener which listens for
-/// incoming TCP connections.
+/// The Server struct holds:
+///
+/// * the tokio TcpListener which listens for incoming TCP connections.
+///
+/// * Shared storage
+///
 #[derive(Debug)]
 pub struct Server {
+    /// The TCP listener for accepting incoming connections.
     listener: TcpListener,
+    /// Contains the shared storage.
+    storage: Storage,
 }

 impl Server {
-    /// Creates a new Server instance with the given TcpListener.
-    pub fn new(listener: TcpListener) -> Server {
-        Server { listener }
+    /// Creates a new Server instance with the given TcpListener and shared storage.
+    pub fn new(listener: TcpListener, storage: Storage) -> Server {
+        Server { listener, storage }
     }

     /// Runs the server in an infinite loop, continuously accepting and handling
     /// incoming connections.
     pub async fn run(&mut self) -> Result<()> {
+        let db = self.storage.db().clone();
+
         loop {
             // accept a new TCP connection.
             // If successful the corresponding TcpStream is stored
@@ -44,11 +55,14 @@ impl Server {
             // and to write RespType values into outgoing TCP messages.
             let resp_command_frame = Framed::with_capacity(sock, RespCommandFrame::new(), 8 * 1024);

+            // Clone the Arc of DB for passing it to the tokio task.
+            let db = Arc::clone(&db);
+
             // Spawn a new asynchronous task to handle the connection.
             // This allows the server to handle multiple connections concurrently.
             tokio::spawn(async move {
                 let handler = FrameHandler::new(resp_command_frame);
-                if let Err(e) = handler.handle().await {
+                if let Err(e) = handler.handle(db.as_ref()).await {
                     error!("Failed to handle command: {}", e);
                 }
             });

diff --git a/src/main.rs b/src/main.rs
index 59926d6..ca4e16b 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -4,6 +4,7 @@ mod command;
 mod handler;
 mod resp;
 mod server;
+mod storage;

 use crate::server::Server;
 use anyhow::Result;
@@ -52,8 +53,11 @@ async fn main() -> Result<()> {
         Err(e) => panic!("Could not bind the TCP listener to {}. Err: {}", &addr, e),
     };

+    // initialize shared storage
+    let shared_storage = storage::db::Storage::new(storage::db::DB::new());
+
     // Create a new instance of the Server with the bound TcpListener
-    let mut server = Server::new(listener);
+    let mut server = Server::new(listener, shared_storage);

     // Run the server to start accepting and handling connections
     // This will run indefinitely until the program is terminated

Testing SET and GET

Start Nimblecache using the below command and use Redis client to connect to it.

# --port parameter is optional
RUST_LOG=info cargo run -- --port 6380

Now, try issuing SET and GET commands to see if they work properly.

Source code

The source code for this specific part is available at https://github.com/dheerajgopi/nimblecache/tree/blog-5.

Feel free to check the main branch of the Nimblecache repository to see the latest code.

Conclusion and what's next

In this part, we successfully implemented the SET and GET commands for our Redis clone, along with a simple HashMap-based storage system. This allows us to perform basic key-value operations, laying the foundation for more complex features.

In the next part of this series, we will implement the Redis LIST data type and its related operations, enabling us to store and manipulate lists within our Redis clone.

Stay tuned!

Handle Redis PING command

Dheeraj Gopinath — Thu, 29 Aug 2024 12:19:24 +0000

In our journey to build a Redis clone using Rust and Tokio, we've laid the groundwork with a TCP server, a RESP parser and Redis command parsing capability. Let's now take it forward by implementing our first real Redis command: PING.

The PING command is a simple operation in Redis, mainly used to test the connection between the client and the server. It responds with PONG if no argument is provided; otherwise, it echoes the argument back to the client.

Implement PING command handler

Let's create a struct called Ping which handles the Redis PING command and returns the output RESP value. We will place the Ping struct inside a command Rust module.

Here's the general idea for whats being done in the below code:

Ping::with_args accepts a vector of RespType values as arguments (remember that a Redis command is simply a list of RESP BulkString values), validates for the correct RESP data type, and creates a new instance of the Ping struct, which holds the optional message value.
apply returns the response for the PING command. If an optional message is provided, it returns the same as a BulkString; otherwise, it returns the string PONG as a SimpleString.
Command enum represents all the supported Redis commands. The Ping struct will be accessed as a variant of the Command enum.
Command::from_resp_command_frame tries to parse the exact Redis command from a Vec<RespType> value. If parsing fails, it will return CommandError.

// src/command/ping.rs

use crate::resp::types::RespType;

use super::CommandError;

/// Represents the PING command in Nimblecache.
#[derive(Debug, Clone)]
pub struct Ping {
    /// Custom message
    msg: Option<String>,
}

impl Ping {
    /// Creates a new `Ping` instance from the given arguments.
    ///
    /// # Arguments
    ///
    /// * `args` - A vector of `RespType` representing the arguments to the PING command.
    ///
    /// # Returns
    ///
    /// * `Ok(Ping)` if parsing succeeds.
    /// * `Err(CommandError)` if parsing fails.
    pub fn with_args(args: Vec<RespType>) -> Result<Ping, CommandError> {
        if args.len() == 0 {
            return Ok(Ping { msg: None });
        }

        let msg = match &args[0] {
            RespType::BulkString(s) => s.clone(),
            _ => return Err(CommandError::Other(String::from("Invalid message"))),
        };

        Ok(Ping { msg: Some(msg) })
    }

    /// Executes the PING command.
    ///
    /// # Returns
    ///
    /// A `RespType` representing the response:
    /// - If no message was provided, it returns "PONG" as a `SimpleString`.
    /// - If a message was provided, it returns that message as a `BulkString`.
    pub fn apply(&self) -> RespType {
        if let Some(msg) = &self.msg {
            RespType::BulkString(msg.to_string())
        } else {
            RespType::SimpleString(String::from("PONG"))
        }
    }
}

// src/command/mod.rs

use core::fmt;

use ping::Ping;

use crate::resp::types::RespType;

mod ping;

/// Represents the supported Nimblecache commands.
#[derive(Debug, Clone)]
pub enum Command {
    /// The PING command.
    Ping(Ping),
}

impl Command {
    /// Attempts to parse a Nimblecache command from a RESP command frame.
    ///
    /// # Arguments
    ///
    /// * `frame` - A vector of `RespType` representing the command and its arguments.
    /// The first item is always the command name, and the rest are its arguments.
    ///
    /// # Returns
    ///
    /// * `Ok(Command)` if parsing succeeds.
    /// * `Err(CommandError)` if parsing fails.
    pub fn from_resp_command_frame(frame: Vec<RespType>) -> Result<Command, CommandError> {
        let (cmd_name, args) = frame.split_at(1);
        let cmd_name = match &cmd_name[0] {
            RespType::BulkString(s) => s.clone(),
            _ => return Err(CommandError::InvalidFormat),
        };

        let cmd = match cmd_name.to_lowercase().as_str() {
            "ping" => Command::Ping(Ping::with_args(Vec::from(args))?),
            _ => {
                return Err(CommandError::UnknownCommand(ErrUnknownCommand {
                    cmd: cmd_name,
                }));
            }
        };

        Ok(cmd)
    }

    /// Executes the Nimblecache command.
    ///
    /// # Returns
    ///
    /// The result of the command execution as a `RespType`.
    pub fn execute(&self) -> RespType {
        match self {
            Command::Ping(ping) => ping.apply(),
        }
    }
}

/// Represents all possible errors that can occur during command parsing and execution.
#[derive(Debug)]
pub enum CommandError {
    /// Indicates that the command format is invalid.
    InvalidFormat,
    /// Indicates that the command is unknown.
    UnknownCommand(ErrUnknownCommand),
    /// Represents any other error with a descriptive message.
    Other(String),
}

/// Represents an error for an unknown command.
#[derive(Debug)]
pub struct ErrUnknownCommand {
    /// The name of the unknown command.
    pub cmd: String,
}

impl std::error::Error for CommandError {}

impl fmt::Display for CommandError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            CommandError::InvalidFormat => "Invalid command format".fmt(f),
            CommandError::UnknownCommand(e) => write!(f, "Unknown command: {}", e.cmd),
            CommandError::Other(msg) => msg.as_str().fmt(f),
        }
    }
}

Now that we have the Ping command handler ready, let's update src/handler.rs to properly handle incoming Redis commands instead of just echoing them back to the client.

The below code changes do the following:

Parse the exact Command type from the incoming Redis command.
Execute the command and write the output RESP value into the outgoing TCP stream.
If command execution fails, a RESP value of SimpleError type will be written into the outgoing TCP stream.

diff --git a/src/handler.rs b/src/handler.rs
index 8f2a002..2f6ae69 100644
--- a/src/handler.rs
+++ b/src/handler.rs
@@ -6,7 +6,10 @@ use log::error;
 use tokio::net::TcpStream;
 use tokio_util::codec::Framed;

-use crate::resp::{frame::RespCommandFrame, types::RespType};
+use crate::{
+    command::Command,
+    resp::{frame::RespCommandFrame, types::RespType},
+};

 /// Handles RESP command frames over a single TCP connection.
 pub struct FrameHandler {
@@ -22,7 +25,7 @@ impl FrameHandler {
     /// Handles incoming RESP command frames.
     ///
     /// This method continuously reads command frames from the connection,
-    /// and echo it back to the client. It continues until
+    /// processes them, and sends back the responses. It continues until
     /// an error occurs or the connection is closed.
     ///
     /// # Returns
@@ -37,8 +40,18 @@ impl FrameHandler {
         while let Some(resp_cmd) = self.conn.next().await {
             match resp_cmd {
                 Ok(cmd_frame) => {
+                    // Read the command from the frame.
+                    let resp_cmd = Command::from_resp_command_frame(cmd_frame);
+
+                    // Execute the command and get the RESP response.
+                    // If command fails, return RESP SimpleError as response.
+                    let response = match resp_cmd {
+                        Ok(cmd) => cmd.execute(),
+                        Err(e) => RespType::SimpleError(format!("{}", e)),
+                    };
+
                     // Write the RESP response into the TCP stream.
-                    if let Err(e) = self.conn.send(RespType::Array(cmd_frame)).await {
+                    if let Err(e) = self.conn.send(response).await {
                         error!("Error sending response: {}", e);
                         break;
                     }

Setting server port using CLI parameter

This time the updated code will be tested using Redis CLI tool. But if you have Redis CLI tool in your machine, there's a good chance that Redis itself is running on port 6379. So, instead of instead of manually changing the port in code, let's add a CLI parameter to easily specify the port on which our Redis clone should run.

CLI parameters can be easily added to our application using the clap crate.

The code changes for adding CLI parameter support are given below:

diff --git a/Cargo.toml b/Cargo.toml
index b50bb57..e0ac3a2 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -8,6 +8,7 @@ edition = "2021"
 [dependencies]
 anyhow = "1.0.86"
 bytes = "1.6.0"
+clap = { version = "4.5.8", features = ["derive"] }
 tokio = { version = "1.38.0", features = [
     "rt-multi-thread",
     "macros",

diff --git a/src/main.rs b/src/main.rs
index de0058f..59926d6 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -7,24 +7,44 @@ mod server;

 use crate::server::Server;
 use anyhow::Result;
+use clap::Parser;
 use log::info;
 use tokio::net::TcpListener;

+const DEFAULT_PORT: u16 = 6379;
+
+#[derive(Debug, Parser)]
+#[command(
+    name = "nimblecache-server",
+    version,
+    author,
+    about = "A RESP based in-memory cache"
+)]
+struct Cli {
+    /// Port to be bound to Nimblecache server
+    #[arg(long)]
+    port: Option<u16>,
+}
+
 #[tokio::main]
 async fn main() -> Result<()> {
     // Initialize the logger.
     // This sets up logging based on the RUST_LOG environment variable
     env_logger::init();

+    // Get port from --port CLI parameter. Defaults to 6379
+    let cli = Cli::parse();
+    let port = cli.port.unwrap_or(DEFAULT_PORT);
+
     // Define the address and port for the TCP server to listen on
     // Here we're using localhost (127.0.0.1) and port 6379 (commonly used for Redis)
-    let addr = format!("127.0.0.1:{}", 6379);
+    let addr = format!("127.0.0.1:{}", port);

     // Attempt to bind the TCP listener to the specified address and port
     let listener = match TcpListener::bind(&addr).await {
         // If successful, return the TcpListener
         Ok(tcp_listener) => {
-            info!("TCP listener started on port 6379");
+            info!("TCP listener started on port {}", port);
             tcp_listener
         }
         // If there is an error, panic and print the error message

Now you can use the below command to run the Redis clone on port 6380. If the --port parameter is not specified, it will default to 6379.

RUST_LOG=info cargo run -- --port 6380

Testing PING using Redis CLI

Testing will be pretty simple now, since we are going to use the Redis CLI tool. Start the server in the port of your choice, and start the Redis CLI tool by specifying the same port as the server.

Next, issue a PING command in the CLI tool and check that our server responds with a PONG. Also, try sending a custom message with the PING command to see if the server responds with your custom message instead of PONG.

Source Code

The source code for this specific part is available at https://github.com/dheerajgopi/nimblecache/tree/blog-4.

Feel free to check the main branch of the Nimblecache repository to see the latest code.

Conclusion and what's next

In this part, we successfully implemented the PING command handler for our Redis clone, allowing us to test the connection and echo custom messages. We also added CLI parameter support to easily specify the server port, making testing more flexible.

The next part will focus on implementing the storage mechanism and the basic SET and GET commands. This will enable our Redis clone to store and retrieve data, bringing us closer to a fully functional Redis-like server.

Stay tuned!

Parsing Redis Commands with tokio-util

Dheeraj Gopinath — Wed, 28 Aug 2024 14:52:18 +0000

In our ongoing series on building a Redis clone using Rust, we've covered the basics of setting up a TCP server and parsing basic RESP data types. Now, it's time to advance our project by understanding the structure of Redis commands and how to parse them efficiently.

In this article, we'll explore how Redis commands are structured and sent to the server, learn to parse incoming TCP messages as Redis commands using the Framed utility from the tokio-util library, and understand the RESP Array type as part of this process.

Structure of a Redis command

A client sends a request to the Redis server as an array of strings, where the first element is the command name, followed by its arguments. In RESP terms, a Redis command can be represented by an Array of BulkStrings. At this point, our RESP parser supports only SimpleString, Bulk String and SimpleError. Now its time to add support for RESP array.

RESP Array

Here's how an Array data type is represented in RESP.

*<number-of-elements>\r\n<element-1>...<element-n>

The asterisk symbol (*) as the first byte indicates it's an Array.
One or more decimal digits (0..9) as the string's length, in bytes, as an unsigned, base-10 value.
The CRLF terminator.
An additional RESP type for every element of the array.

Here's an example for a RESP array which contains 2 BulkStrings:

*2\r\n$4\r\nPING\r\n$5\r\nhello\r\n,

which represents the Redis PING command in RESP format.

Now, let's add support for RESP Arrays in our Redis clone. Please note that the functions parse_array_len and parse_bulk_string_len in the code below will be used later in this article. You don't need to worry about the intent of those functions at this point :)

diff --git a/src/resp/types.rs b/src/resp/types.rs
index cae0dcb..ebedb2c 100644
--- a/src/resp/types.rs
+++ b/src/resp/types.rs
@@ -7,6 +7,8 @@ pub enum RespType {
     BulkString(String),
     /// Refer <https://redis.io/docs/latest/develop/reference/protocol-spec/#simple-errors>
     SimpleError(String),
+    /// Refer <https://redis.io/docs/latest/develop/reference/protocol-spec/#arrays>
+    Array(Vec<RespType>),
 }

 use super::RespError;
@@ -125,10 +127,87 @@ impl RespType {
                 let bulkstr_bytes = format!("${}\r\n{}\r\n", bs.chars().count(), bs).into_bytes();
                 Bytes::from_iter(bulkstr_bytes)
             }
+            RespType::Array(arr) => {
+                let mut arr_bytes = format!("*{}\r\n", arr.len()).into_bytes();
+                arr.iter()
+                    .map(|v| v.to_bytes())
+                    .for_each(|b| arr_bytes.extend(b));
+
+                Bytes::from_iter(arr_bytes)
+            }
             RespType::SimpleError(es) => Bytes::from_iter(format!("-{}\r\n", es).into_bytes()),
         };
     }

+    /// Parses the length of a RESP array from the given byte buffer.
+    ///
+    /// This function attempts to read the first few bytes of a RESP array to determine its length.
+    /// It expects the input to start with a '*' character followed by the length and terminated by CRLF.
+    ///
+    /// # Arguments
+    ///
+    /// * `src` - A `BytesMut` containing the bytes to parse.
+    ///
+    /// # Returns
+    ///
+    /// * `Ok(Some((usize, usize)))` - If successful, returns a tuple containing:
+    ///   - The parsed length of the array
+    ///   - The number of bytes read from the input
+    /// * `Ok(None)` - If there's not enough data in the buffer to parse the length
+    /// * `Err(RespError)` - If the input is not a valid RESP array prefix or if parsing fails
+    pub fn parse_array_len(src: BytesMut) -> Result<Option<(usize, usize)>, RespError> {
+        let (array_prefix_bytes, bytes_read) = match Self::read_till_crlf(&src[..]) {
+            Some((b, size)) => (b, size),
+            None => return Ok(None),
+        };
+
+        if bytes_read < 4 || array_prefix_bytes[0] as char != '*' {
+            return Err(RespError::InvalidArray(String::from(
+                "Not a valid RESP array",
+            )));
+        }
+
+        match Self::parse_usize_from_buf(&array_prefix_bytes[1..]) {
+            Ok(len) => Ok(Some((len, bytes_read))),
+            Err(e) => Err(e),
+        }
+    }
+
+    /// Parses the length of a RESP bulk string from the given byte buffer.
+    ///
+    /// This function attempts to read the first few bytes of a RESP bulk string to determine its length.
+    /// It expects the input to start with a '$' character followed by the length and terminated by CRLF.
+    ///
+    /// # Arguments
+    ///
+    /// * `src` - A `BytesMut` containing the bytes to parse.
+    ///
+    /// # Returns
+    ///
+    /// * `Ok(Some((usize, usize)))` - If successful, returns a tuple containing:
+    ///   - The parsed length of the bulk string
+    ///   - The number of bytes read from the input
+    /// * `Ok(None)` - If there's not enough data in the buffer to parse the length
+    /// * `Err(RespError)` - If the input is not a valid RESP bulk string prefix or if parsing fails
+    ///
+    pub fn parse_bulk_string_len(src: BytesMut) -> Result<Option<(usize, usize)>, RespError> {
+        let (bulkstr_prefix_bytes, bytes_read) = match Self::read_till_crlf(&src[..]) {
+            Some((b, size)) => (b, size),
+            None => return Ok(None),
+        };
+
+        if bytes_read < 4 || bulkstr_prefix_bytes[0] as char != '$' {
+            return Err(RespError::InvalidBulkString(String::from(
+                "Not a valid RESP bulk string",
+            )));
+        }
+
+        match Self::parse_usize_from_buf(&bulkstr_prefix_bytes[1..]) {
+            Ok(len) => Ok(Some((len, bytes_read))),
+            Err(e) => Err(e),
+        }
+    }
+
     // Read the bytes till reaching CRLF ("\r\n")
     fn read_till_crlf(buf: &[u8]) -> Option<(&[u8], usize)> {
         for i in 1..buf.len() {

diff --git a/src/resp/mod.rs b/src/resp/mod.rs
index 9cdde64..24a806d 100644
--- a/src/resp/mod.rs
+++ b/src/resp/mod.rs
@@ -1,3 +1,4 @@
+pub mod frame;
 pub mod types;

 /// Represents errors that can occur during RESP parsing.
@@ -7,6 +8,8 @@ pub enum RespError {
     InvalidBulkString(String),
     /// Represents an error in parsing a simple string, with an error message.
     InvalidSimpleString(String),
+    /// Represents an error in parsing an array, with an error message.
+    InvalidArray(String),
     /// Represents any other error with a descriptive message.
     Other(String),
 }
@@ -17,6 +20,7 @@ impl std::fmt::Display for RespError {
             RespError::Other(msg) => msg.as_str().fmt(f),
             RespError::InvalidBulkString(msg) => msg.as_str().fmt(f),
             RespError::InvalidSimpleString(msg) => msg.as_str().fmt(f),
+            RespError::InvalidArray(msg) => msg.as_str().fmt(f),
         }
     }
 }

Parsing Redis commands using tokio-util Framed

Before we handle Redis command parsing, it's important to revisit the existing TCP message parsing logic our TCP server. Specifically, the below lines:

// src/server.rs

...
...

// read the TCP message and move the raw bytes into a buffer
let mut buffer = BytesMut::with_capacity(512);
if let Err(e) = sock.read_buf(&mut buffer).await {
    panic!("Error reading request: {}", e);
}

// Try parsing the RESP data from the bytes in the buffer.
// If parsing fails return the error message as a RESP SimpleError data type.
let resp_data = match RespType::parse(buffer) {
    Ok((data, _)) => data,
    Err(e) => RespType::SimpleError(format!("{}", e)),
};

...
...

Some of the issues with the above code are:

Potential for incomplete reads: The code reads into the buffer once, but doesn't account for partial reads where the full command might not have arrived yet.
Conversely, if multiple commands arrive in one read, the code might not handle it correctly.
Crude error handling: There's no proper error handling. The code just panics if it fails to read into the buffer.

The root cause of these issues is that we are using a buffer with a fixed size. While it is possible to fix these problems, manually managing the buffer can be error-prone. In such cases, it is better to use existing utilities like Framed in tokio-util.

What is tokio-util Framed?

To put it simply, the Framed utility is a tool for converting a blob of bytes into a domain-specific type and vice versa. In our case, Framed is going to operate on the socket, converting the raw bytes of the incoming TCP messages into RespType values, and also converting RespType values back into raw bytes for outgoing TCP messages. The best part is that Framed handles the buffering internally.

Framed relies on two simple traits: Decode and Encode. Together, these form a Codec, which is necessary for Framed to work on a socket. So, as part of implementing Framed this is what we need to do:

Implement Decoder trait which will convert the incoming TCP messages into Vec<RespType>.
Implement Encoder trait which will convert RespType values into raw bytes which will be sent via the outgoing TCP message.

Before implementing the Codec, let's add the tokio-util dependency in Cargo.toml.

diff --git a/Cargo.toml b/Cargo.toml
index 0e36917..b50bb57 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -15,5 +15,7 @@ tokio = { version = "1.38.0", features = [
     "io-util",
     "sync",
 ] }
+tokio-util = { version = "0.7.11", features = ["codec"] }
 log = "0.4.22"
 env_logger = "0.11.3"
+futures = { version = "0.3", default-features = true }

Implementing Decoder

The job of the decoder is to look into a data buffer and determine if it contains enough data to produce a complete item (Vec<RespType>) in the stream.

A decoder should work based on the below logic:

If the buffer contains enough data to construct the Redis command (Vec<RespType>), the decoder should remove those bytes from the buffer and return Ok(Some(Vec<RespType>)).
If the buffer does not contain enough data to construct the Redis command, leave the buffer unmodified and return Ok(None). The Framed utility will then read more data from the TCP stream and append those bytes to the buffer. After that, the decoder is called again to check if it has enough data to produce the command.
In case of any decoding errors, return Err(std::io::Error).

The code for the decoder is given below. Here are few points to notice:

The struct RespCommandFrame, implements the Decoder trait. It utilizes a CommandBuilder struct internally to accumulate the parts of a command.
The decoder examines the first few bytes to find the RESP array length. If it finds the length, it initializes the CommandBuilder struct. If the array length is n, the decoder expects n BulkStrings as part of the incoming TCP message.
The decoder then tries to read n BulkStrings and add each item to the CommandBuilder. Once all n items are added to the CommandBuilder, the command (Vec<RespType>) is returned successfully.
The functions RespType::parse_array_len and RespType::parse_bulk_string_len, mentioned earlier in this article, are utilized in the decoder.

// src/resp/frame.rs

use bytes::{Buf, BufMut};
use core::fmt;
use std::io::Error;
use tokio_util::codec::{Decoder, Encoder};

use crate::resp::types::RespType;

use super::RespError;

/// This codec handles Nimblecache commands, which are always represented
/// as array of bulk strings in the RESP (REdis Serialization Protocol) protocol.
///
/// The codec uses a `CommandBuilder` internally to construct the array of bulk strings
/// that make up a Nimblecache command.
pub struct RespCommandFrame {
    /// Builder for appending the bulk strings in the command array.
    cmd_builder: Option<CommandBuilder>,
}

impl RespCommandFrame {
    /// Creates a new `RespCommandFrame`.
    ///
    /// # Returns
    ///
    /// A new instance of `RespCommandFrame` with no command builder initialized.
    pub fn new() -> RespCommandFrame {
        RespCommandFrame { cmd_builder: None }
    }
}

impl Decoder for RespCommandFrame {
    type Item = Vec<RespType>;

    type Error = std::io::Error;

    /// Decodes bytes from the input stream into a `Vec<RespType>` representing a Nimblecache command.
    ///
    /// This method implements the RESP protocol decoding logic, specifically handling
    /// arrays of bulk strings which represent Nimblecache commands. It uses a `CommandBuilder`
    /// to accumulate the parts of the command as they are received.
    ///
    /// # Arguments
    ///
    /// * `src` - A mutable reference to the input buffer containing bytes to decode.
    ///
    /// # Returns
    ///
    /// * `Ok(Some(Vec<RespType>))` if a complete command (array of bulk strings) was successfully decoded.
    /// * `Ok(None)` if more data is needed to complete the command.
    /// * `Err(std::io::Error)` if an error occurred during decoding.
    fn decode(
        &mut self,
        src: &mut bytes::BytesMut,
    ) -> std::result::Result<Option<Self::Item>, Self::Error> {
        // A command in RESP protocol should always be an array of Bulk Strings.
        // Check the first 2 bytes to validate if its a RESP array.
        if self.cmd_builder.is_none() {
            let (cmd_len, bytes_read) = match RespType::parse_array_len(src.clone()) {
                Ok(arr_len) => match arr_len {
                    Some((len, bytes_read)) => (len, bytes_read),
                    None => return Ok(None),
                },
                Err(e) => {
                    return Err(Error::new(
                        std::io::ErrorKind::InvalidData,
                        FrameError::from(e),
                    ));
                }
            };

            // initilize command builder, if its a valid RESP array.
            self.cmd_builder = Some(CommandBuilder::new(cmd_len));

            // advance buffer
            src.advance(bytes_read);
        }

        // Read all bytes in buffer
        while src.len() != 0 {
            // Validate and check the length of next bulk string
            let (bulkstr_len, bytes_read) = match RespType::parse_bulk_string_len(src.clone()) {
                Ok(bulkstr_len) => match bulkstr_len {
                    Some((len, bytes_read)) => (len, bytes_read),
                    None => return Ok(None),
                },
                Err(e) => {
                    return Err(Error::new(
                        std::io::ErrorKind::InvalidData,
                        FrameError::from(e),
                    ));
                }
            };

            // A bulk string has the below format
            //
            // `${string length in bytes }\r\n{string value}\r\n`
            //
            // Check if the buffer contains the required number of bytes to parse
            // the bulk string (including the CRLF at the end)
            let bulkstr_bytes = bulkstr_len + bytes_read + 2;
            if src.len() < bulkstr_bytes {
                return Ok(None);
            }

            // now that its sure the buffer has all the bytes required to parse the bulk string, parse it.
            let (bulkstr, bytes_read) = match RespType::parse_bulk_string(src.clone()) {
                Ok((resp_type, bytes_read)) => (resp_type, bytes_read),
                Err(e) => {
                    return Err(Error::new(
                        std::io::ErrorKind::InvalidData,
                        FrameError::from(e),
                    ));
                }
            };

            // append the bulk string to the command builder
            self.cmd_builder.as_mut().unwrap().add_part(bulkstr);

            // advance buffer
            src.advance(bytes_read);

            // if the command builder has all the parts, return it, else check buffer again
            let cmd_builder = self.cmd_builder.as_ref().unwrap();
            if cmd_builder.all_parts_received() {
                let cmd = cmd_builder.build();
                self.cmd_builder = None;
                return Ok(Some(cmd));
            }
        }

        Ok(None)
    }
}

/// This struct is used to accumulate the parts of a Nimblecache command, which are
/// typically represented as an array of bulk strings in the RESP protocol.
struct CommandBuilder {
    parts: Vec<RespType>,
    num_parts: usize,
    parts_parsed: usize,
}

impl CommandBuilder {
    /// Creates a new `CommandBuilder` with the specified number of parts.
    pub fn new(num_parts: usize) -> CommandBuilder {
        CommandBuilder {
            parts: vec![],
            num_parts,
            parts_parsed: 0,
        }
    }

    /// Adds a part to the command being built and increments the count of parsed parts.
    ///
    /// # Arguments
    ///
    /// * `part` - A `RespType` representing a part of the command.
    pub fn add_part(&mut self, part: RespType) {
        self.parts.push(part);
        self.parts_parsed += 1;
    }

    /// Checks if all expected parts of the command have been received.
    ///
    /// # Returns
    ///
    /// `true` if the number of parsed parts equals the expected number of parts,
    /// `false` otherwise.
    pub fn all_parts_received(&self) -> bool {
        self.num_parts == self.parts_parsed
    }

    /// Builds and returns the complete command as a vector of RESP values.
    ///
    /// # Returns
    ///
    /// A vector of `RespType` containing all the parts of the command.
    pub fn build(&self) -> Vec<RespType> {
        self.parts.clone()
    }
}

/// Represents error that can occur during RESP command frame parsing.
#[derive(Debug)]
pub struct FrameError {
    err: RespError,
}

impl FrameError {
    pub fn from(err: RespError) -> FrameError {
        FrameError { err }
    }
}

impl std::error::Error for FrameError {}

impl fmt::Display for FrameError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.err.fmt(f)
    }
}

Implementing Encoder

The job of the encoder is to convert RespType values into raw bytes and write the same into the buffer.

Compared to the decoder, this is straight-forward to implement. Add the below code to src/resp/frame.rs for implementing Encoder trait.

impl Encoder<RespType> for RespCommandFrame {
    type Error = std::io::Error;

    /// Encodes a `RespType` into bytes and writes them to the output buffer.
    ///
    /// It's primarily used for sending responses to Nimblecache commands.
    ///
    /// # Arguments
    ///
    /// * `item` - The `RespType` to encode.
    /// * `dst` - The output buffer to write the encoded bytes to.
    ///
    /// # Returns
    ///
    /// * `Ok(())` if the encoding was successful.
    /// * `Err(std::io::Error)` if an error occurred during encoding.
    fn encode(&mut self, item: RespType, dst: &mut bytes::BytesMut) -> Result<(), Self::Error> {
        dst.put_slice(&item.to_bytes());

        Ok(())
    }
}

Using Framed

Now let's use Framed with RespCommandFrame codec in our server loop. Make the modifications to the code as given below:

// src/handler.rs

use anyhow::Result;
use futures::{SinkExt, StreamExt};
use log::error;
use tokio::net::TcpStream;
use tokio_util::codec::Framed;

use crate::resp::{frame::RespCommandFrame, types::RespType};

/// Handles RESP command frames over a single TCP connection.
pub struct FrameHandler {
    /// The framed connection using `RespCommandFrame` as the codec.
    conn: Framed<TcpStream, RespCommandFrame>,
}
impl FrameHandler {
    /// Creates a new `FrameHandler` instance.
    pub fn new(conn: Framed<TcpStream, RespCommandFrame>) -> FrameHandler {
        FrameHandler { conn }
    }

    /// Handles incoming RESP command frames.
    ///
    /// This method continuously reads command frames from the connection,
    /// and echo it back to the client. It continues until
    /// an error occurs or the connection is closed.
    ///
    /// # Returns
    ///
    /// A `Result` indicating whether the operation succeeded or failed.
    ///
    /// # Errors
    ///
    /// This method will return an error if there's an issue with reading
    /// from or writing to the connection.
    pub async fn handle(mut self) -> Result<()> {
        while let Some(resp_cmd) = self.conn.next().await {
            match resp_cmd {
                Ok(cmd_frame) => {
                    // Write the RESP response into the TCP stream.
                    if let Err(e) = self.conn.send(RespType::Array(cmd_frame)).await {
                        error!("Error sending response: {}", e);
                        break;
                    }
                }
                Err(e) => {
                    error!("Error reading the request: {}", e);
                    break;
                }
            };

            // flush the buffer into the TCP stream.
            self.conn.flush().await?;
        }

        Ok(())
    }
}

diff --git a/src/server.rs b/src/server.rs
index a249175..9e45bee 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -3,17 +3,13 @@
 // anyhow provides the Error and Result types for convenient error handling
 use anyhow::{Error, Result};

-use bytes::BytesMut;
 // log crate provides macros for logging at various levels (error, warn, info, debug, trace)
 use log::error;

-use tokio::{
-    // AsyncWriteExt trait provides asynchronous write methods like write_all
-    io::{AsyncReadExt, AsyncWriteExt},
-    net::{TcpListener, TcpStream},
-};
+use tokio::net::{TcpListener, TcpStream};
+use tokio_util::codec::Framed;

-use crate::resp::types::RespType;
+use crate::{handler::FrameHandler, resp::frame::RespCommandFrame};

 /// The Server struct holds the tokio TcpListener which listens for
 /// incoming TCP connections.
@@ -35,7 +31,7 @@ impl Server {
             // accept a new TCP connection.
             // If successful the corresponding TcpStream is stored
             // in the variable `sock`, else a panic will occur.
-            let mut sock = match self.accept_conn().await {
+            let sock = match self.accept_conn().await {
                 Ok(stream) => stream,
                 // Log the error and panic if there is an issue accepting a connection.
                 Err(e) => {
@@ -44,29 +40,17 @@ impl Server {
                 }
             };

+            // Use RespCommandFrame codec to read incoming TCP messages as Redis command frames,
+            // and to write RespType values into outgoing TCP messages.
+            let resp_command_frame = Framed::with_capacity(sock, RespCommandFrame::new(), 8 * 1024);
+
             // Spawn a new asynchronous task to handle the connection.
             // This allows the server to handle multiple connections concurrently.
             tokio::spawn(async move {
-                // read the TCP message and move the raw bytes into a buffer
-                let mut buffer = BytesMut::with_capacity(512);
-                if let Err(e) = sock.read_buf(&mut buffer).await {
-                    panic!("Error reading request: {}", e);
-                }
-
-                // Try parsing the RESP data from the bytes in the buffer.
-                // If parsing fails return the error message as a RESP SimpleError data type.
-                let resp_data = match RespType::parse(buffer) {
-                    Ok((data, _)) => data,
-                    Err(e) => RespType::SimpleError(format!("{}", e)),
-                };
-
-                // Echo the RESP message back to the client.
-                if let Err(e) = &mut sock.write_all(&resp_data.to_bytes()[..]).await {
-                    // Log the error and panic if there is an issue writing the response.
-                    error!("{}", e);
-                    panic!("Error writing response")
+                let handler = FrameHandler::new(resp_command_frame);
+                if let Err(e) = handler.handle().await {
+                    error!("Failed to handle command: {}", e);
                 }
-                // The connection is closed automatically when `sock` goes out of scope.
             });
         }
     }

diff --git a/src/main.rs b/src/main.rs
index 7b13f63..dee666f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,6 +1,7 @@
 // src/main.rs

 // Include the server module defined in server.rs
+mod handler;
 mod resp;
 mod server;

Running the TCP server

Run the following command to run the application.

RUST_LOG=info cargo run

RUST_LOG=info will set the log level to info.

Use the below command to send Redis command using nc.

# send a RESP array. This will echo the same bulk string back.
{ echo -e '*2\r\n$4\r\nPING\r\n$5\r\nhello\r\n'; sleep 1; } | nc localhost 6379

With the new changes, you might notice that our server won't drop the TCP connection after sending the response back to the client. This is a good thing. However, it also means that testing our Redis clone using nc will become even more clunkier. We'll try to eliminate nc and use an actual Redis client in the following parts.

Source Code

The source code for this specific part is available at https://github.com/dheerajgopi/nimblecache/tree/blog-3.

Feel free to check the main branch of the Nimblecache repository to see the latest code.

Conclusion and what's next

In this article, we've learned about the structure of Redis commands and how to parse them efficiently using the Framed utility from the tokio-util library. It's worth noting that Framed can be applied in any server application which involves parsing and writing of protocol-specific messages.

In the next part of this series, we will focus on handling a real Redis command. Stay tuned!

Build a Bare-Bones RESP parser

Dheeraj Gopinath — Tue, 27 Aug 2024 07:13:56 +0000

In this article, we explore the basics of RESP (REdis Serialization Protocol), a protocol used by Redis for client-server communication. We begin by understanding the structure and representation of basic RESP data types, such as BulkString, SimpleString, and SimpleError. Next, we build a simple RESP parser that can handle these data types, and then, we modify our existing TCP server to be RESP-aware, allowing it to process RESP-formatted TCP messages. This prepares us for implementing Redis commands and handling more complex RESP types in future articles.

What is RESP?

RESP (REdis Serialization Protocol) is a binary-safe serialization protocol that supports several data types (Strings, Integers, Arrays, etc.). Redis uses RESP for its client-server communication. The general idea is that the first byte determines the type, and the subsequent bytes constitute the type's content. The type's content might be prefixed by its length so that the parser knows how many bytes to read. The \r\n (CRLF) is the protocol's terminator, which separates the parts in the type's content.

Here's how a Redis BulkString data type is represented in RESP.

$<length>\r\n<data>\r\n

The dollar sign ($) as the first byte indicates it's a BulkString.
One or more decimal digits (0..9) as the string's length, in bytes, as an unsigned, base-10 value.
The CRLF terminator.
The data.
A final CRLF.

And here's an example of a BulkString in RESP.

$5\r\nhello\r\n

Please note that the RESP format differs for other data types. Redis has provided excellent documentation on RESP and its data types. I highly recommend you go through the documentation before building the RESP parser.

Here's the link to the documentation: https://redis.io/docs/latest/develop/reference/protocol-spec/

Redis data types in Rust

Lets start with a RespType enum which will be a wrapper around all supported data types. For now, we will support SimpleString, SimpleError and BulkString.

// src/resp/types.rs

/// This enum is a wrapper for the different data types in RESP.
#[derive(Clone, Debug)]
pub enum RespType {
    /// Refer <https://redis.io/docs/latest/develop/reference/protocol-spec/#simple-strings>
    SimpleString(String),
    /// Refer <https://redis.io/docs/latest/develop/reference/protocol-spec/#bulk-strings>
    BulkString(String),
    /// Refer <https://redis.io/docs/latest/develop/reference/protocol-spec/#simple-errors>
    SimpleError(String),
}

Parsing RESP values from a buffer of raw bytes is pretty straightforward. Before we start with the code for parsing, lets create some Error types which are used in case of parsing failures.

// src/resp/mod.rs

pub mod types;

/// Represents errors that can occur during RESP parsing.
#[derive(Debug)]
pub enum RespError {
    /// Represents an error in parsing a bulk string, with an error message.
    InvalidBulkString(String),
    /// Represents an error in parsing a simple string, with an error message.
    InvalidSimpleString(String),
    /// Represents any other error with a descriptive message.
    Other(String),
}

impl std::fmt::Display for RespError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            RespError::Other(msg) => msg.as_str().fmt(f),
            RespError::InvalidBulkString(msg) => msg.as_str().fmt(f),
            RespError::InvalidSimpleString(msg) => msg.as_str().fmt(f),
        }
    }
}

Don't forget to add mod resp; in src/main.rs, or else the files inside the resp module wont be considered during compilation process.

Since we are dealing with buffers of raw bytes, lets add the bytes crate into the list of dependencies.

[dependencies]
...
bytes = "1.6.0"
...

Let's now start with the actual parsing logic.

Parsing BulkString

Parsing BulkString involves the following steps:

Check if the first byte is $.
Read the subsequent bytes until the CRLF terminator (\r\n) is encountered. This contains the length (number of bytes) of the actual string data. Once the length is parsed, the parser knows how many bytes to read for extracting the complete string data.
Read the subsequent bytes based on the length parsed in the previous step. These bytes contain the actual string data.
The extracted bytes are then converted to a UTF-8 string.

Here's the code for the above logic (Add below code in src/resp/types.rs file).

use bytes::{Bytes, BytesMut};
use super::RespError;

impl RespType {
    /// Parse the given bytes into its respective RESP type and return the parsed RESP value and
    /// the number of bytes read from the buffer.
    ///
    /// More details on the parsing logic is available at
    /// <https://redis.io/docs/latest/develop/reference/protocol-spec/#resp-protocol-description>.
    ///
    /// # Errors
    /// Error will be returned in the following scenarios:
    /// - If first byte is an invalid character.
    /// - If the parsing fails due to encoding issues etc.
    pub fn parse(buffer: BytesMut) -> Result<(RespType, usize)> {
        let c = buffer[0] as char;
        return match c {
            '$' => Self::parse_bulk_string(buffer),
            _ => Err(RespError::Other(String::from(
                "Invalid RESP data type",
            ))),
        };
    }

    /// Parse the given bytes into a BulkString RESP value. This will return the parsed RESP
    /// value and the number of bytes read from the buffer.
    ///
    /// Example BulkString: `$5\r\nhello\r\n`
    ///
    /// # BulkString Parts:
    /// ``
    ///     $      |            5           | \r\n |    hello     | \r\n
    /// identifier | string length in bytes | CRLF | string value | CRLF
    /// ``
    ///
    /// # Parsing Logic:
    /// - The buffer is read until CRLF characters ("\r\n") are encountered.
    /// - That slice of bytes are then parsed into an int. That will be the string length in bytes (let's say `bulkstr_len`)
    /// - `bulkstr_len` number of bytes are read from the buffer again from where it was stopped previously.
    /// - This 2nd slice of bytes is then parsed into an UTF-8 string.
    ///
    /// Note: The first byte in the buffer is skipped since it's just an identifier for the
    /// RESP type and is not the part of the actual value itself.
    pub fn parse_bulk_string(buffer: BytesMut) -> Result<(RespType, usize), RespError> {
        // read until CRLF and parse length
        let (bulkstr_len, bytes_consumed) =
            if let Some((buf_data, len)) = Self::read_till_crlf(&buffer[1..]) {
                let bulkstr_len = Self::parse_usize_from_buf(buf_data)?;
                (bulkstr_len, len + 1)
            } else {
                return Err(RespError::InvalidBulkString(String::from(
                    "Invalid value for bulk string",
                )));
            };

        // validate if buffer contains the complete string data based on
        // the length parsed in the previous step.
        let bulkstr_end_idx = bytes_consumed + bulkstr_len as usize;
        if bulkstr_end_idx >= buffer.len() {
            return Err(RespError::InvalidBulkString(String::from(
                "Invalid value for bulk string length",
            )));
        }

        // convert raw bytes into UTF-8 string.
        let bulkstr = String::from_utf8(buffer[bytes_consumed..bulkstr_end_idx].to_vec());

        match bulkstr {
            Ok(bs) => Ok((RespType::BulkString(bs), bulkstr_end_idx + 2)),
            Err(_) => Err(RespError::InvalidBulkString(String::from(
                "Bulk string value is not a valid UTF-8 string",
            ))),
        }
    }

    // Read the bytes till reaching CRLF ("\r\n")
    fn read_till_crlf(buf: &[u8]) -> Option<(&[u8], usize)> {
        for i in 1..buf.len() {
            if buf[i - 1] == b'\r' && buf[i] == b'\n' {
                return Some((&buf[0..(i - 1)], i + 1));
            }
        }

        None
    }

    // Parse usize from bytes. The number is provided in string format.
    // So convert raw bytes into UTF-8 string and then convert the string
    // into usize.
    fn parse_usize_from_buf(buf: &[u8]) -> Result<usize, RespError> {
        let utf8_str = String::from_utf8(buf.to_vec());
        let parsed_int = match utf8_str {
            Ok(s) => {
                let int = s.parse::<usize>();
                match int {
                    Ok(n) => Ok(n),
                    Err(_) => Err(RespError::Other(String::from(
                        "Invalid value for an integer",
                    ))),
                }
            }
            Err(_) => Err(RespError::Other(String::from("Invalid UTF-8 string"))),
        };

        parsed_int
    }
}

Parsing SimpleString

Parsing SimpleString is even simpler. It does not have the length prefix as found in BulkString. The logic is as follows:

Check if the first byte is +.
Read the subsequent bytes until the CRLF terminator (\r\n) is encountered. These bytes contain the actual string data.
The extracted bytes are then converted to a UTF-8 string.

Here's the code for the above logic (Modify src/resp/types.rs with the code given below).

impl RespType {

    ...
    ...

    pub fn parse(buffer: BytesMut) -> Result<(RespType, usize), RespError> {
        ...
        ...
        return match c {
            '$' => Self::parse_bulk_string(buffer),
            '+' => Self::parse_simple_string(buffer),
            _ => Err(RespError::Other(String::from(
                "Invalid RESP data type",
            ))),
        };
    }

    ...    
    ...

    /// Parse the given bytes into a SimpleString RESP value. This will return the parsed RESP
    /// value and the number of bytes read from the buffer.
    ///
    /// Example SimpleString: `+OK\r\n`
    ///
    /// # SimpleString Parts:
    /// ``
    ///      +      |      OK      | \r\n
    ///  identifier | string value | CRLF
    /// ``
    ///
    /// # Parsing Logic:
    /// - The buffer is read until CRLF characters ("\r\n") are encountered. That slice of bytes are then
    /// parsed into an UTF-8 string.
    pub fn parse_simple_string(buffer: BytesMut) -> Result<(RespType, usize), RespError> {
        // read until CRLF and parse the bytes into an UTF-8 string.
        if let Some((buf_data, len)) = Self::read_till_crlf(&buffer[1..]) {
            let utf8_str = String::from_utf8(buf_data.to_vec());

            return match utf8_str {
                Ok(simple_str) => Ok((RespType::SimpleString(simple_str), len + 1)),
                Err(_) => {
                    return Err(RespError::InvalidSimpleString(String::from(
                        "Simple string value is not a valid UTF-8 string",
                    )))
                }
            };
        }

        Err(RespError::InvalidSimpleString(String::from(
            "Invalid value for simple string",
        )))
    }

    ...    
    ...
}

Parsing SimpleError is the same as parsing SimpleString, except the first byte should be - instead of +. However, a separate parsing function for SimpleError is not needed right now because errors are usually sent back to the client as a response and are not part of the request message.

RESP-aware TCP server loop

Now that we have a basic RESP parser, let's convert our existing TCP server into a RESP-aware echo server.

For this we need to make modifications on src/resp/types.rs and /src/server.rs.

Below are the modifications made to src/resp/types.rs. This is for adding a method which converts the RESP type back into raw bytes.

impl RespType {
    ...
    ...

    /// Convert the RESP value into its byte values.
    pub fn to_bytes(&self) -> Bytes {
        return match self {
            RespType::SimpleString(ss) => Bytes::from_iter(format!("+{}\r\n", ss).into_bytes()),
            RespType::BulkString(bs) => {
                let bulkstr_bytes = format!("${}\r\n{}\r\n", bs.chars().count(), bs).into_bytes();
                Bytes::from_iter(bulkstr_bytes)
            }
            RespType::SimpleError(es) => Bytes::from_iter(format!("-{}\r\n", es).into_bytes()),
        };
    }

    ...
    ...
}

The modifications made to /src/server.rs are as given below.

diff --git a/src/server.rs b/src/server.rs
index ae3fef9..a249175 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -3,15 +3,18 @@
 // anyhow provides the Error and Result types for convenient error handling
 use anyhow::{Error, Result};

+use bytes::BytesMut;
 // log crate provides macros for logging at various levels (error, warn, info, debug, trace)
 use log::error;

 use tokio::{
     // AsyncWriteExt trait provides asynchronous write methods like write_all
-    io::AsyncWriteExt,
+    io::{AsyncReadExt, AsyncWriteExt},
     net::{TcpListener, TcpStream},
 };

+use crate::resp::types::RespType;
+
 /// The Server struct holds the tokio TcpListener which listens for
 /// incoming TCP connections.
 #[derive(Debug)]
@@ -44,8 +47,21 @@ impl Server {
             // Spawn a new asynchronous task to handle the connection.
             // This allows the server to handle multiple connections concurrently.
             tokio::spawn(async move {
-                // Write a "Hello!" message to the client.
-                if let Err(e) = &mut sock.write_all("Hello!".as_bytes()).await {
+                // read the TCP message and move the raw bytes into a buffer
+                let mut buffer = BytesMut::with_capacity(512);
+                if let Err(e) = sock.read_buf(&mut buffer).await {
+                    panic!("Error reading request: {}", e);
+                }
+
+                // Try parsing the RESP data from the bytes in the buffer.
+                // If parsing fails return the error message as a RESP SimpleError data type.
+                let resp_data = match RespType::parse(buffer) {
+                    Ok((data, _)) => data,
+                    Err(e) => RespType::SimpleError(format!("{}", e)),
+                };
+
+                // Echo the RESP message back to the client.
+                if let Err(e) = &mut sock.write_all(&resp_data.to_bytes()[..]).await {
                     // Log the error and panic if there is an issue writing the response.
                     error!("{}", e);
                     panic!("Error writing response")
@@ -70,4 +86,4 @@ impl Server {
             }
         }
     }
-}
\ No newline at end of file
+}

Running the TCP server

Run the following command to run the application.

RUST_LOG=info cargo run

RUST_LOG=info will set the log level to info.

Use the below commands to send RESP data using nc.

# send bulk string. This will echo the same bulk string back.
{ echo -e '$5\r\nhello\r\n'; sleep 1; } | nc localhost 6379

# send simple string. This will echo the same bulk string back.
{ echo -e '+OK\r\n'; sleep 1; } | nc localhost 6379

# send invalid resp data. This will respond with a simple error.
{ echo -e '+OK'; sleep 1; } | nc localhost 6379

Source Code

The source code for this specific part is available at https://github.com/dheerajgopi/nimblecache/tree/blog-2.

Feel free to check the main branch of the Nimblecache repository to see the latest code.

Conclusion and what's next

In this article, we explored the basics of RESP and built a simple RESP parser that can handle RESP data types like BulkString, SimpleString and SimpleError. We also modified our TCP server to be RESP-aware, allowing it to process RESP-formatted messages.

In the next part of this series, we will delve into implementing Redis command parsing and handling more RESP types. Stay tuned!

Setting Up a Simple TCP Server

Dheeraj Gopinath — Sat, 24 Aug 2024 09:13:26 +0000

In this tutorial, you'll start the journey towards building Nimblecache (Redis clone) by first creating a simple TCP server using Tokio, a powerful asynchronous runtime for the Rust programming language. By the end of this tutorial, you will have a basic understanding of how to handle TCP connections and manage asynchronous tasks using Tokio. Let's dive in!

Why Tokio?

We could use standard threads to build the TCP server, as it is simpler than using something like Tokio. However, we are choosing Tokio because its async model is perfect for TCP servers. It handles concurrent connections efficiently and uses fewer system resources than a threaded model. It excels in I/O efficiency by allowing the server to perform other tasks while waiting for network operations, unlike standard threads, which require a separate thread for each connection and can quickly exhaust resources.

What we're building

Before we dive into the code, let's briefly cover what we're building:

A TCP server that listens for incoming connections
Handles multiple clients concurrently
Responds with a simple "Hello!" message to each client

This TCP server will serve as a foundation on which we can build the Redis clone.

Setting Up the Project

First, create a new Rust project:

cargo new nimblecache
cd nimblecache

Add the following dependencies to your Cargo.toml:

[package]
name = "nimblecache"
version = "0.1.0"
edition = "2021"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

[dependencies]
anyhow = "1.0.86"
tokio = { version = "1.38.0", features = [
    "rt-multi-thread",
    "macros",
    "net",
    "io-util",
    "sync",
] }
log = "0.4.22"
env_logger = "0.11.3"

anyhow: Anyhow is an error-handling library. It simplifies error management by providing anyhow::Error and anyhow::Result types. These can be used to handle various error types without excessive boilerplate.

tokio: Tokio is the asynchronous runtime for Rust. The features enabled are:

"rt-multi-thread": Enables the multi-threaded runtime.
"macros": Includes helpful macros like #[tokio::main].
"net": Provides networking types like TcpListener and TcpStream.
"io-util": Includes I/O utilities for working with streams.
"sync": Offers synchronization primitives (like tokio Mutex) for concurrent programming.

The Tokio dependency is more specific in its feature selection, enabling only the necessary features instead of using features = ["full"]. This approach helps achieve faster compile times and a smaller binary size.

log: This is the logging facade for Rust. It provides macros for logging at different levels (info, warn, error etc.).

env_logger: This is a logger implementation that can be configured via environment variables. It works in conjunction with the log crate.

The server

The code below is the core of the TCP server. It defines a Server struct, which encapsulates the main functionality of our TCP server:

It holds a TcpListener which listens for incoming connections.
The run method enters an infinite loop, accepting new connections.
For each connection, it spawns a new asynchronous task to handle the client.
The accept_conn method wraps the process of accepting a new connection.

It's a basic implementation with detailed comments to make it easy to understand.

// src/server.rs

// anyhow provides the Error and Result types for convenient error handling
use anyhow::{Error, Result};

// log crate provides macros for logging at various levels (error, warn, info, debug, trace)
use log::error;

use tokio::{
    // AsyncWriteExt trait provides asynchronous write methods like write_all
    io::AsyncWriteExt,
    net::{TcpListener, TcpStream},
};

/// The Server struct holds the tokio TcpListener which listens for
/// incoming TCP connections.
#[derive(Debug)]
pub struct Server {
    listener: TcpListener,
}

impl Server {
    /// Creates a new Server instance with the given TcpListener.
    pub fn new(listener: TcpListener) -> Server {
        Server { listener }
    }

    /// Runs the server in an infinite loop, continuously accepting and handling
    /// incoming connections.
    pub async fn run(&mut self) -> Result<()> {
        loop {
            // accept a new TCP connection.
            // If successful the corresponding TcpStream is stored
            // in the variable `sock`, else a panic will occur.
            let mut sock = match self.accept_conn().await {
                Ok(stream) => stream,
                // Log the error and panic if there is an issue accepting a connection.
                Err(e) => {
                    error!("{}", e);
                    panic!("Error accepting connection");
                }
            };

            // Spawn a new asynchronous task to handle the connection.
            // This allows the server to handle multiple connections concurrently.
            tokio::spawn(async move {
                // Write a "Hello!" message to the client.
                if let Err(e) = &mut sock.write_all("Hello!".as_bytes()).await {
                    // Log the error and panic if there is an issue writing the response.
                    error!("{}", e);
                    panic!("Error writing response")
                }
                // The connection is closed automatically when `sock` goes out of scope.
            });
        }
    }

    /// Accepts a new incoming TCP connection and returns the corresponding
    /// tokio TcpStream.
    async fn accept_conn(&mut self) -> Result<TcpStream> {
        loop {
            // Wait for an incoming connection.
            // The `accept()` method returns a tuple of (TcpStream, SocketAddr),
            // but we only need the TcpStream.
            match self.listener.accept().await {
                // Return the TcpStream if a connection is successfully accepted.
                Ok((sock, _)) => return Ok(sock),
                // Return an error if there is an issue accepting a connection.
                Err(e) => return Err(Error::from(e)),
            }
        }
    }
}

The Main Function

Now, let's set up our main.rs to use the above server. This main function does the following:

Initializes the logger.
Sets up the address and port for the server (127.0.0.1:6379).
Binds a TcpListener to this address.
Creates a new Server instance with this listener.
Runs the server.

// src/main.rs

// Include the server module defined in server.rs
mod server;

use crate::server::Server;
use anyhow::Result;
use log::info;
use tokio::net::TcpListener;

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize the logger.
    // This sets up logging based on the RUST_LOG environment variable
    env_logger::init();

    // Define the address and port for the TCP server to listen on
    // Here we're using localhost (127.0.0.1) and port 6379 (commonly used for Redis)
    let addr = format!("127.0.0.1:{}", 6379);

    // Attempt to bind the TCP listener to the specified address and port
    let listener = match TcpListener::bind(&addr).await {
        // If successful, return the TcpListener
        Ok(tcp_listener) => {
            info!("TCP listener started on port 6379");
            tcp_listener
        },
        // If there is an error, panic and print the error message
        // This could happen if the port is already in use, for example
        Err(e) => panic!("Could not bind the TCP listener to {}. Err: {}", &addr, e),
    };

    // Create a new instance of the Server with the bound TcpListener
    let mut server = Server::new(listener);

    // Run the server to start accepting and handling connections
    // This will run indefinitely until the program is terminated
    server.run().await?;

    // This Ok(()) is technically unreachable as server.run() loops infinitely,
    // but it's needed to satisfy the Result return type of main()
    Ok(())
}

Running the TCP server

Run the following command to run the application.

RUST_LOG=info cargo run

RUST_LOG=info will set the log level to info.

Now you can send TCP messages to the server using command-line utilities like nc.

nc -v localhost 6379

The above command will be responded with a Hello! from the TCP server.

Source Code

The source code for this specific part is available at https://github.com/dheerajgopi/nimblecache/tree/blog-1. Feel free to check the main branch of the Nimblecache repository to see the latest code.

Conclusion and what's next

This simple TCP server demonstrates the basics of network programming with Rust and Tokio. The next step is to expand this TCP server to handle the Redis Serialization Protocol (or RESP), enabling it to process and respond to Redis commands effectively. This will lay the groundwork for implementing more advanced features and functionalities of our Redis clone (Nimblecache). We'll cover this in the following part of the series.

Introduction

Dheeraj Gopinath — Sat, 24 Aug 2024 09:03:43 +0000

In this series of articles, we will embark on a journey to build a Redis clone using Rust. This project, named Nimblecache, is a personal project to deepen my understanding of Rust. Inspired by the CodeCrafter's Build Your Own Redis course, which I highly recommend, I have used it as a foundational guide for creating Nimblecache.

Parts

This series is a work in progress. Each new part will build on the previous ones, gradually adding more features to Nimblecache. Stay tuned for future updates!

RAG using LLMSmith and FastAPI

Dheeraj Gopinath — Tue, 28 May 2024 12:00:36 +0000

What is LLMSmith?

LLMSmith is a lightweight Python library designed for developing functionalities powered by Large Language Models (LLMs). It allows developers to integrate generative AI capabilities into all sorts of applications. In this case, we will be creating an RAG based chatbot using LLMSmith and expose it as an API endpoint in a FastAPI app.

LLMSmith repo: https://github.com/dheerajgopi/llmsmith

FYI.. I’m the author of this library :)

Now, lets get to the interesting part.

Implement the RAG functionality using LLMSmith

This is what we will be doing here.

Pre-process the user’s query for stripping out info that is irrelevant for retrieval using OpenAI LLM.
Retrieve relevant documents from Qdrant vector DB.
Rerank the retrieved documents so that they are ordered based on semantic relevance.
Generate answer using OpenAI LLM. The reranked documents are passed as context in the prompt.

The below piece of code uses LLMSmith library for implementing the above mentioned RAG flow.

from textwrap import dedent
import cohere
from fastembed import TextEmbedding
import openai
from qdrant_client import AsyncQdrantClient
from llmsmith.task.retrieval.vector.qdrant import QdrantRetriever
from llmsmith.reranker.cohere import CohereReranker
from llmsmith.task.textgen.openai import OpenAITextGenTask, OpenAITextGenOptions
from llmsmith.job.job import SequentialJob

from rag_llmsmith_fastapi.config import settings


preprocess_prompt = (
    dedent("""
    Convert the natural language query from a user into a query for a vectorstore.
    In this process, you strip out information that is not relevant for the retrieval task.
    Return only the query converted for retrieval and nothing else.
    Here is the user query: {{root}}""")
    .strip("\n")
    .replace("\n", " ")
)


class RAGService:
    def __init__(
        self,
        llm_client: openai.AsyncOpenAI,
        vectordb_client: AsyncQdrantClient,
        reranker_client: cohere.AsyncClient,
        embedder: TextEmbedding,
        **_,
    ) -> None:
        self.llm_client = llm_client
        self.vectordb_client = vectordb_client
        self.reranker_client = reranker_client
        self.embedder = embedder

    async def chat(self, user_prompt):
        # Create Cohere reranker
        reranker = CohereReranker(client=self.reranker_client)

        # Embedding function to be passed into the Qdrant retriever
        def embedding_func(x):
            return list(self.embedder.query_embed(x))

        # Define the Qdrant retriever task. The embedding function and reranker are passed as parameters.
        retrieval_task = QdrantRetriever(
            name="qdrant-retriever",
            client=self.vectordb_client,
            collection_name=settings.QDRANT.COLLECTION_NAME,
            embedding_func=embedding_func,
            embedded_field_name="description",  # name of the field in the document on which embeddedings are created while uploading data to the Qdrant collection
            reranker=reranker,
        )

        # Define the OpenAI LLM task for rephrasing the query
        preprocess_task = OpenAITextGenTask(
            name="openai-preprocessor",
            llm=self.llm_client,
            llm_options=OpenAITextGenOptions(model="gpt-4-turbo", temperature=0),
        )

        # Define the OpenAI LLM task for answering the query
        answer_generate_task = OpenAITextGenTask(
            name="openai-answer-generator",
            llm=self.llm_client,
            llm_options=OpenAITextGenOptions(model="gpt-4-turbo", temperature=0),
        )

        # define the sequence of tasks
        # {{root}} is a special placeholer in `input_template` which will be replaced with the prompt entered by the user (`user_prompt`).
        # The placeholder {{qdrant-retriever.output}} will be replaced with the output from Qdrant DB retriever task.
        # The placeholder {{openai-preprocessor.output}} will be replaced with the output from the query preprocessing task done by OpenAI LLM.
        job: SequentialJob[str, str] = (
            SequentialJob()
            .add_task(
                preprocess_task,
                input_template=preprocess_prompt,
            )
            .add_task(retrieval_task, input_template="{{openai-preprocessor.output}}")
            .add_task(
                answer_generate_task,
                input_template="Answer the question based on the context: \n\n QUESTION:\n{{root}}\n\nCONTEXT:\n{{qdrant-retriever.output}}",
            )
        )

        # Now, run the job
        await job.run(user_prompt)

        return job.task_output("openai-answer-generator")

Using LLMSmith makes it quite easy to implement such LLM based functionalities. In the above code, we first create the following tasks:

QdrantRetriever — To retrieve documents from Qdrant.
CohereReranker — Rerank documents based on sematic relevance (passed into QdrantRetriever)
OpenAITextGenTask — To execute LLM calls. Used for both pre-processing and answer generation in our case.

After the task definitions, all these tasks are run sequentially using SequentialJob. In an LLMSmith job, the output of previous tasks can be easily passed into the next task via placeholders using input_template parameter.

Integrate the RAGService into a FastAPI endpoint

This is easy! Just call the chat method from the RAGService instance in your route handler function.

from fastapi import APIRouter

from llmsmith.task.models import TaskOutput

from rag_llmsmith_fastapi.chat.model import ChatRequest, ChatResponse
from rag_llmsmith_fastapi.chat.service import RAGService


class ChatController:
    def __init__(self, rag_svc: RAGService) -> None:
        self.rag_svc = rag_svc
        self.router: APIRouter = APIRouter(tags=["Chat endpoint"], prefix="/api")

        self.router.add_api_route(
            path="/chat",
            endpoint=self.chat,
            methods=["POST"],
        )

    async def chat(self, req_body: ChatRequest):
        rag_response: TaskOutput = await self.rag_svc.chat(req_body.content)
        return ChatResponse(content=rag_response.content)

The complete code can be found in this repo.

For running the application, clone the repository from Github and follow the README.md file.

Contributors Welcome

All contributions (no matter if small) to LLMSmith are always welcome. To see how you can help and where to start see CONTRIBUTING.md.

Function Calling Agent using OpenAI Assistant

Dheeraj Gopinath — Thu, 23 May 2024 13:34:20 +0000

What is “function calling” in LLMs?

“Function calling” in LLMs refer to its capability to accept a list of user defined functions (aka tools) and to intelligently choose which “tool” to use based on the prompt provided by the user.

For example, let’s say we have a function get_weather(location: str) which provides the current weather based on the location argument. If we pass the get_weather function to the LLM as a tool and ask “What’s the current weather in Budapest?”, the LLM can decide to use the get_weather tool instead of responding with incorrect data. In this case, the LLM will simply return the tool name, which is get_weather and the arguments to be passed which is location = "Budapest". Now, we can call the get_weather function using the arguments provided by the LLM (note that LLM cannot call the function for us, we should do that ourselves). After executing the function, pass the return value of the function back to the LLM, and the LLM will provide us with a sensible response.

So what is a function calling agent?

The general idea of a function calling agent is that, we pass the user query and a list of tools to the LLM, and then call the LLM in a loop until we get the desired response. I know its a vague description :) So, let’s consider an example for better understanding.

Let’s consider a mathematical query for example, since LLMs are generally bad at math.

Tools provided to the LLM

add(a: float, b: float)
subtract(a: float, b: float)
multiply(a: float, b: float)

Query:

calculate sum of 1 and 5 and multiply it with the difference of 6 and 3

The Agent Loop

The below image will do a better job in explaining the agent loop than writing a long explanation.

You can see that a “conversation” is going on between the app and the LLM inside the agent loop until the LLM can find the result for the user’s query.

Implementation using OpenAI assistant

Below is the complete code for the OpenAI function calling agent. The explanation for each step has been provided as comments in the code.

Before executing this script, install the required libraries using the below command.

pip install openai python-dotenv

Also, store your OpenAI API key in a .env file (OPENAI_API_KEY=your-api-key).

import os
import sys
from dotenv import load_dotenv
import openai
from openai.types.beta import Assistant, Thread
from openai.types.beta.threads import Run

import json
import logging

# Load env variables (OPENAI_API_KEY)
load_dotenv()

# set log handlers
log_handler = logging.StreamHandler(sys.stdout)
log = logging.getLogger(__name__)
log.addHandler(log_handler)
log.setLevel(logging.INFO)

# initialize OpenAI client
client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

# Define the functions which will be the part of the LLM toolkit
def add(a: float, b: float) -> float:
    return a + b

def subtract(a: float, b: float) -> float:
    return a - b

def multiply(a: float, b: float) -> float:
    return a * b

tool_callables = {
    "add": add,
    "subtract": subtract,
    "multiply": multiply
}

# declaration of tools (functions) to be passed into the OpenAI assistant
math_tools = [
    {
        "function": {
            "name": "add",
            "description": "Returns the sum of two numbers.",
            "parameters": {
                "type": "object",
                "properties": {"a": {"type": "number"}, "b": {"type": "number"}},
                "required": ["a", "b"],
            },
        },
        "type": "function",
    },
    {
        "function": {
            "name": "subtract",
            "description": "Returns the difference of two numbers.",
            "parameters": {
                "type": "object",
                "properties": {"a": {"type": "number"}, "b": {"type": "number"}},
                "required": ["a", "b"],
            },
        },
        "type": "function",
    },
    {
        "function": {
            "name": "multiply",
            "description": "Returns the product of two numbers.",
            "parameters": {
                "type": "object",
                "properties": {"a": {"type": "number"}, "b": {"type": "number"}},
                "required": ["a", "b"],
            },
        },
        "type": "function",
    }
]

openai_assistant: Assistant = client.beta.assistants.create(
    model="gpt-4-turbo",
    instructions="you are a math tutor, who explains the solutions to math problems",
    name="math-tutor",
    tools=math_tools
)


def run_math_agent(query: str, max_turns: int = 5) -> str:
    # Initialize the OpenAI assistant. The assistant will have its own unique id.
    openai_assistant = client.beta.assistants.create(
        model="gpt-4-turbo",
        instructions="you are a math tutor, who explains the solutions to math problems",
        name="math-tutor",
        tools=math_tools
    )

    # Create a Thread. In OpenAI lingo, a `Thread` can be considered like a conversation thread (not the multithreading one)
    # The to and fro communication between the script and the LLM will be stored against this thread id.
    thread: Thread = client.beta.threads.create()

    # Send the user query as part of the newly created thread
    client.beta.threads.messages.create(
        thread_id=thread.id, role="user", content=query
    )

    # The user query is now part of the thread. Now call the LLM (or "run" the thread in OpenAI lingo).
    # `create_and_poll` is just a helper method which polls the LLM until a terminal state is reached.
    # 
    # The terminal states are given below:
    # "requires_action" - A function call is required. Execute the function and submit the response back.
    #                     The results should be submitted back before the `expires_at` timestamp.
    # "completed"       - The Run is completed successfully.
    # "cancelled"       - The run was cancelled (its possible to cancel an in-progress Run).
    # "failed"          - Failed due to some error.
    # "expired"         - Run can get expired if we fail to submit function call results before `expires_at` timestamp.
    run: Run = client.beta.threads.runs.create_and_poll(
        thread_id=thread.id,
        assistant_id=openai_assistant.id,
    )

    # The agent loop. `max_turns` will set a limit on the number of LLM calls made inside the agent loop.
    # Its better to set a limit since LLM calls are costly.
    for turn in range(max_turns):

        # Fetch the last message from the thread
        messages = client.beta.threads.messages.list(
            thread_id=thread.id,
            run_id=run.id,
            order="desc",
            limit=1,
        )

        # Check for the terminal state of the Run.
        # If state is "completed", exit agent loop and return the LLM response.
        if run.status == "completed":
            assistant_res: str = next(
                (
                    content.text.value
                    for content in messages.data[0].content
                    if content.type == "text"
                ),
                None,
            )

            return assistant_res

        # If state is "requires_action", function calls are required. Execute the functions and send their outputs to the LLM.
        if run.status == "requires_action":
            func_tool_outputs = []

            # LLM can ask for multiple functions to be executed. Execute all function calls in loop and
            # append the results into `func_tool_outputs` list.
            for tool in run.required_action.submit_tool_outputs.tool_calls:
                # parse the arguments required for the function call from the LLM response
                args = (
                    json.loads(tool.function.arguments)
                    if tool.function.arguments
                    else {}
                )
                func_output = tool_callables[tool.function.name](**args)

                # OpenAI needs the output of the function call against the tool_call_id
                func_tool_outputs.append(
                    {"tool_call_id": tool.id, "output": str(func_output)}
                )

            # Submit the function call outputs back to OpenAI
            run = client.beta.threads.runs.submit_tool_outputs_and_poll(
                thread_id=thread.id, run_id=run.id, tool_outputs=func_tool_outputs
            )

            # Continue the agent loop.
            # Agent will check the output of the function output submission as part of next iteration.
            continue

        # Handle errors if terminal state is "failed"
        else:
            if run.status == "failed":
                log.error(
                    f"OpenAIFunctionAgent turn-{turn+1} | Run failure reason: {run.last_error}"
                )

            raise Exception(
                f"Failed to generate text due to: {run.last_error}"
            )

    # Raise error if turn-limit is reached.
    raise MaxTurnsReachedException()


class MaxTurnsReachedException(Exception):
    def __init__(self):
        super().__init__("Reached maximum number of turns")

if __name__ == "__main__":
    log.info(run_math_agent("calculate sum of 1 and 5 and multiply it with difference of 6 and 3"))

On executing the above script the LLM will respond with the correct result, which is 18. The LLM will also explain (like a math tutor) each step it took to find the result.

Shameless plug :)

You can use function calling agents for your application in a much easier way by just adding a python library called LLMSmith to your list of dependencies. As you might have guessed, I’m the author of that library :).

You can install the library using the following command.

pip install "llmsmith[openai]"

Here’s the code for using a function calling agent using llmsmith.

import asyncio
import os

from dotenv import load_dotenv
import openai
from llmsmith.agent.function.openai import OpenAIFunctionAgent
from llmsmith.agent.function.options.openai import OpenAIAssistantOptions
from llmsmith.agent.tool.openai import OpenAIAssistantTool

from llmsmith.task.models import TaskInput


# load env vars for getting OPENAI_API_KEY
load_dotenv()

# Define the functions which will be the part of the LLM toolkit
def add(a: float, b: float) -> float:
    return a + b


async def run():
    # initialize OpenAI client
    llm = openai.AsyncOpenAI(api_key=os.getenv("OPENAI_API_KEY"))

    # declaration of tools (functions) to be passed into the OpenAIFunctionAgent
    add_tool = OpenAIAssistantTool(
        declaration={
            "function": {
                "name": "add",
                "description": "Returns the sum of two numbers.",
                "parameters": {
                    "type": "object",
                    "properties": {"a": {"type": "number"}, "b": {"type": "number"}},
                    "required": ["a", "b"],
                },
            },
            "type": "function",
        },
        callable=add,
    )

    # create the agent
    task: OpenAIFunctionAgent = await OpenAIFunctionAgent.create(
        name="testfunc",
        llm=llm,
        assistant_options=OpenAIAssistantOptions(
            model="gpt-4-turbo",
            system_prompt="you are a math tutor, who explains the solutions to math problems"
        ),
        tools=[add_tool],
        max_turns=5,
    )

    # run the agent
    res = await task.execute(TaskInput("Add sum of 1 and 2 to the sum of 5 and 6"))

    print(f"\n\nAgent response: {res.content}")


if __name__ == "__main__":
    loop = asyncio.get_event_loop()
    loop.run_until_complete(run())

llmsmith can be used for building all sorts of LLM based functionalities (not just agents). Refer the documentation for getting a better idea about the library. The Examples section has sample codes for implementing RAG using llmsmith and more.