Skip to content

Commit

Permalink
feat: support Get command (#6)
Browse files Browse the repository at this point in the history
* refactor: split the runtime part and static part when implementing Command

* feat: support get command

* feat: command id will convert to uppercase at comptime when making command_type_stub

* refactor: remove redundant code while registering mod
  • Loading branch information
J0HN50N133 authored Apr 13, 2024
1 parent 51bd770 commit 73514da
Show file tree
Hide file tree
Showing 7 changed files with 254 additions and 101 deletions.
140 changes: 113 additions & 27 deletions src/commands/src/commands.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,49 +23,135 @@ use crate::error::Result;

mod string;

pub static COMMANDS_TABLE: Lazy<HashMap<CommandId, Command>> = Lazy::new(|| {
let mut table = HashMap::new();
table.insert(CommandId::SET, Command::new(CommandId::SET));
table
});

/// ```plaintext
/// COMMAND: string => Command ID => Command --create--> CommandInstance
/// ```
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, strum::Display, strum::EnumString)]
#[strum(serialize_all = "lowercase")]
#[strum(ascii_case_insensitive)]
pub enum CommandId {
SET,
// pub type GlobalCommandTable = HashMap<CommandId, Command>;
#[derive(Default)]
pub struct GlobalCommandTable {
table: HashMap<CommandIdRef<'static>, Command>,
}

impl GlobalCommandTable {
pub fn register<T: CommandTypeInfo>(&mut self) {
let cmd = Command::new::<T>();
self.table.insert(cmd.id, cmd);
}

pub fn get(&self, cmd_id: &CommandId) -> Option<&Command> {
self.table.get(cmd_id.to_ascii_uppercase().as_str())
}

pub fn new() -> GlobalCommandTable {
GlobalCommandTable {
table: HashMap::new(),
}
}
}

macro_rules! register_mod {
($($mod:ident),*) => {
Lazy::new(|| {
let mut table = Default::default();
$(
$mod::register(&mut table);
);*
table
})
};
}

#[macro_export]
macro_rules! register {
($($cmd:ident),*) => {
pub(in $crate::commands) fn register(table: &mut GlobalCommandTable) {
static START: Once = Once::new();
START.call_once(move || {
$(
table.register::<$cmd>();
)*
});
}
};
}

/// create a command type stub for a command
///
/// [`cmd_id`] will be converted to uppercase
#[macro_export]
macro_rules! command_type_stub {
(id: $cmd_id:literal) => {
paste::paste! {
fn command() -> &'static Command {
static STUB: Lazy<Command> =
Lazy::new(|| Command::new_stub(
stringify!([< $cmd_id:upper >])
));
&STUB
}
}
};
}

pub static GLOBAL_COMMANDS_TABLE: Lazy<GlobalCommandTable> = register_mod! {string};

/// Mapping relationship between command id and command instance
/// [`CommandId`] => [`Command`] --create--> [`CommandInstance`]
pub type CommandId = str;
pub type CommandIdRef<'a> = &'a CommandId;

type CreateInstanceFn = fn() -> Box<dyn CommandInstance>;

#[derive(Clone)]
pub struct Command {
/// command id i.e. command name
id: CommandId,
id: CommandIdRef<'static>,
create_instance_fn: CreateInstanceFn,
}
impl Command {
pub fn create_instance(&self) -> Box<dyn CommandInstance> {
(self.create_instance_fn)()
}
}

impl Command {
pub fn new(cmd_id: CommandId) -> Self {
Self { id: cmd_id }
pub(crate) fn new<F: CommandTypeInfo>() -> Self {
let mut cmd = F::command().clone();
cmd.create_instance_fn = F::boxed;
cmd
}

pub fn new_instance(&self) -> impl CommandInstance {
match self.id {
CommandId::SET => string::Set::new(),
pub(crate) fn new_stub(cmd_id: CommandIdRef<'static>) -> Self {
/// [`dummy_create_inst`] is a dummy function to satisfy the type of `CommandInstance`
/// and prevent cyclic dependency between [`create_inst`] and [`new`] of CommandInst
fn dummy_create_inst() -> Box<dyn CommandInstance> {
unreachable!()
}
Self {
id: cmd_id,
create_instance_fn: dummy_create_inst,
}
}
}

pub trait CommandInstance {
fn get_attr(&self) -> &Command;
pub trait CommandInstance: Send {
/// Parse command arguments representing in an array of Bytes, since client can only send RESP3 Array frames
fn parse(&mut self, input: &[Bytes]) -> Result<()>;
fn execute(&mut self, storage: &Storage, namespace: Bytes) -> Result<BytesFrame>;
}

pub trait CommandTypeInfo: CommandInstance + Sized + 'static {
/// Tell the system how to create instance
fn new() -> Self;

/// Static typing infomation of command, which is used to register the command
fn command() -> &'static Command;

fn id(&self) -> CommandId {
self.get_attr().id
/// Boxed version of `new`
fn boxed() -> Box<dyn CommandInstance> {
Box::new(Self::new())
}

/// Parse an array of Bytes, since client can only send RESP3 Array frames
fn parse(&mut self, input: Vec<Bytes>) -> Result<()>;
fn execute(self, storage: &Storage, namespace: Bytes) -> BytesFrame;
fn id() -> CommandIdRef<'static> {
Self::command().id
}
}

#[cfg(test)]
Expand Down
97 changes: 61 additions & 36 deletions src/commands/src/commands/string.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,18 +12,22 @@
// See the License for the specific language governing permissions and
// limitations under the License.

use std::sync::Once;

use bytes::Bytes;
use common_base::bytes::StringBytes;
use once_cell::sync::Lazy;
use redis_protocol::resp3::types::{BytesFrame, FrameKind};
use roxy::datatypes::string::{RedisString, StringSetArgs, StringSetArgsBuilder, StringSetType};
use roxy::storage::Storage;
use snafu::ResultExt;

use super::{Command, CommandId, CommandInstance};
use crate::error::{InvalidCmdSyntaxSnafu, Result};
use super::{Command, CommandInstance, CommandTypeInfo, GlobalCommandTable};
use crate::error::{FailInStorageSnafu, InvalidCmdSyntaxSnafu, Result};
use crate::parser::{
chain, key, keyword, optional, string, ttl, value_of_type, Parser, TTLOption, Tokens,
};
use crate::{command_type_stub, register};

pub struct SetArgs {
key: Bytes,
Expand All @@ -33,33 +37,17 @@ pub struct SetArgs {

/// [`SET`] is an instance of `SET` command
pub struct Set {
cmd: Command,
args: Option<SetArgs>,
}

impl Set {
pub fn new() -> Self {
Self {
cmd: Command { id: CommandId::SET },
args: None,
}
}
}

impl CommandInstance for Set {
fn get_attr(&self) -> &Command {
&self.cmd
}

fn parse(&mut self, input: Vec<Bytes>) -> Result<()> {
fn parse(&mut self, input: &[Bytes]) -> Result<()> {
let mut set_args_builder = StringSetArgsBuilder::default();
let mut tokens = Tokens::new(&input[..]);
// skip the command name which is already ensured by strum
tokens.advance();
let mut tokens = Tokens::new(input);

let (key, value) = chain(key(), string())
.parse(&mut tokens)
.context(InvalidCmdSyntaxSnafu { cmd_id: self.id() })?;
.context(InvalidCmdSyntaxSnafu { cmd_id: Self::id() })?;

// TODO: Add permutation parser
let set_type = optional(value_of_type::<StringSetType>())
Expand All @@ -84,7 +72,7 @@ impl CommandInstance for Set {
let set_args = set_args_builder.build().unwrap();
tokens
.expect_eot()
.context(InvalidCmdSyntaxSnafu { cmd_id: self.id() })?;
.context(InvalidCmdSyntaxSnafu { cmd_id: Self::id() })?;
self.args = Some(SetArgs {
key: key.into(),
value,
Expand All @@ -93,27 +81,64 @@ impl CommandInstance for Set {
Ok(())
}

fn execute(mut self, storage: &Storage, namespace: Bytes) -> BytesFrame {
fn execute(&mut self, storage: &Storage, namespace: Bytes) -> Result<BytesFrame> {
let db = RedisString::new(storage, namespace.into());
let args = self.args.take().unwrap();
// TODO: handle ttl

let (Ok(res) | Err(res)) = db
.set(args.key, args.value, &args.set_args)
db.set(args.key, args.value, &args.set_args)
.map(|opt_old| match opt_old {
Some(old) if args.set_args.get => (FrameKind::BlobString, old).try_into().unwrap(),
Some(_) => (FrameKind::SimpleString, "OK").try_into().unwrap(),
Some(old) => (FrameKind::BlobString, old).try_into().unwrap(),
None if args.set_args.get => (FrameKind::SimpleString, "OK").try_into().unwrap(),
None => BytesFrame::Null,
})
.map_err(|err| {
(FrameKind::SimpleError, err.to_string())
.try_into()
.unwrap()
});
res
.context(FailInStorageSnafu { cmd_id: Self::id() })
}
}

impl CommandTypeInfo for Set {
fn new() -> Self {
Self { args: None }
}

command_type_stub! { id: "Set" }
}

pub struct Get {
key: Bytes,
}

impl CommandInstance for Get {
fn parse(&mut self, input: &[Bytes]) -> Result<()> {
let key = key()
.parse(&mut Tokens::new(input))
.context(InvalidCmdSyntaxSnafu { cmd_id: Self::id() })?;
self.key = key.into();
Ok(())
}

fn execute(&mut self, storage: &Storage, namespace: Bytes) -> Result<BytesFrame> {
RedisString::new(storage, namespace.into())
.get(self.key.clone())
.map(|opt_value| {
opt_value
.map(|value| (FrameKind::BlobString, value).try_into().unwrap())
.unwrap_or_else(|| BytesFrame::Null)
})
.context(FailInStorageSnafu { cmd_id: Self::id() })
}
}

impl CommandTypeInfo for Get {
fn new() -> Self {
Self { key: Bytes::new() }
}

command_type_stub! { id: "Get" }
}

register! {Set, Get}

#[cfg(test)]
mod tests {
use std::error::Error;
Expand All @@ -125,7 +150,7 @@ mod tests {
fn test_valid_set_cmd_parse() {
let input = resp3_encode_command("SeT key value nX");
let mut set_cmd = Set::new();
assert!(set_cmd.parse(input).is_ok());
assert!(set_cmd.parse(&input[1..]).is_ok());
let args = set_cmd.args.as_ref().unwrap();
assert_eq!(args.key, &b"key"[..]);
assert_eq!(args.value.as_utf8(), "value");
Expand All @@ -138,8 +163,8 @@ mod tests {
let mut set_cmd = Set::new();
println!(
"{}",
set_cmd.parse(input.clone()).unwrap_err().source().unwrap()
set_cmd.parse(&input[1..]).unwrap_err().source().unwrap()
);
assert!(set_cmd.parse(input).is_err());
assert!(set_cmd.parse(&input[..]).is_err());
}
}
10 changes: 5 additions & 5 deletions src/commands/src/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,22 +14,22 @@

use snafu::Snafu;

use crate::commands::CommandId;
use crate::commands::CommandIdRef;
use crate::parser::ParseError;

#[derive(Debug, Snafu)]
#[snafu(visibility(pub))]
#[snafu(visibility(pub(crate)))]
pub enum Error {
#[snafu(display("Invalid '{}' command", cmd_id))]
InvalidCmdSyntax {
source: ParseError,
cmd_id: CommandId,
cmd_id: CommandIdRef<'static>,
},
#[snafu(display("Fail to execute command '{}' because of storage error", cmd_id))]
FailInStorage {
source: roxy::error::Error,
cmd_id: CommandId,
cmd_id: CommandIdRef<'static>,
},
}

pub type Result<T> = std::result::Result<T, Error>;
pub(crate) type Result<T> = std::result::Result<T, Error>;
2 changes: 1 addition & 1 deletion src/commands/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,5 +13,5 @@
// limitations under the License.

pub mod commands;
pub(crate) mod error;
pub mod error;
pub mod parser;
Loading

0 comments on commit 73514da

Please sign in to comment.