Parsing Redis Commands with tokio-util

#redis #rust #tokio #database

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