Implement Transaction Support in a Redis Clone

// 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),
        }
    }
}

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")),
         }
     }
 }

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.

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