diff --git a/Cargo.lock b/Cargo.lock index 42c4c20..72f0c38 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -308,6 +308,12 @@ dependencies = [ "time", ] +[[package]] +name = "float-ord" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ce81f49ae8a0482e4c55ea62ebbd7e5a686af544c00b9d090bba3ff9be97b3d" + [[package]] name = "fnv" version = "1.0.7" @@ -628,6 +634,7 @@ dependencies = [ "bytes", "crc32fast", "flexi_logger", + "float-ord", "futures", "git-version", "glob", diff --git a/Cargo.toml b/Cargo.toml index 6f3a257..bab33c9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,6 +15,7 @@ tokio={version="1", features = ["full", "tracing"] } parking_lot="0.11.2" tokio-util={version="^0.6", features = ["full"] } crc32fast="1.3.2" +float-ord="^0.3" futures = { version = "0.3.0", features = ["thread-pool"]} hex = "0.4.3" git-version = "0.3.5" diff --git a/src/cmd/mod.rs b/src/cmd/mod.rs index 6ac4dea..acf868d 100644 --- a/src/cmd/mod.rs +++ b/src/cmd/mod.rs @@ -11,6 +11,7 @@ pub mod metrics; pub mod pubsub; pub mod server; pub mod set; +pub mod sorted_set; pub mod string; pub mod transaction; diff --git a/src/cmd/sorted_set.rs b/src/cmd/sorted_set.rs new file mode 100644 index 0000000..ead1ce4 --- /dev/null +++ b/src/cmd/sorted_set.rs @@ -0,0 +1,230 @@ +//! # Sorted Set command handlers +use crate::{ + connection::Connection, + error::Error, + value::{ + bytes_to_number, bytes_to_range_floatord, + sorted_set::{IOption, IResult}, + }, + value::{sorted_set::SortedSet, Value}, +}; +use bytes::Bytes; +use float_ord::FloatOrd; +use std::collections::VecDeque; + +/// Adds all the specified members with the specified scores to the sorted set +/// stored at key. It is possible to specify multiple score / member pairs. If a +/// specified member is already a member of the sorted set, the score is updated +/// and the element reinserted at the right position to ensure the correct +/// ordering. +/// +/// If key does not exist, a new sorted set with the specified members as sole +/// members is created, like if the sorted set was empty. If the key exists but +/// does not hold a sorted set, an error is returned. +/// +/// The score values should be the string representation of a double precision +/// floating point number. +inf and -inf values are valid values as well. +pub async fn zadd(conn: &Connection, mut args: VecDeque) -> Result { + let key = args.pop_front().ok_or(Error::Syntax)?; + let option = IOption::new(&mut args)?; + if args.is_empty() { + return Err(Error::InvalidArgsCount("ZADD".to_owned())); + } + if args.len() % 2 != 0 { + return Err(Error::Syntax); + } + if args.len() != 2 && option.incr { + return Err(Error::OptsNotCompatible( + "INCR option supports a single increment-element pair".to_owned(), + )); + } + let result = conn + .db() + .get(&key) + .map_mut(|v| match v { + Value::SortedSet(x) => { + let mut insert: usize = 0; + let mut updated: usize = 0; + + loop { + let score = match args.pop_front() { + Some(x) => bytes_to_number::(&x)?, + None => break, + }; + let value = args.pop_front().ok_or(Error::Syntax)?; + match x.insert(FloatOrd(score), value, &option) { + IResult::Inserted => insert += 1, + IResult::Updated => updated += 1, + _ => {} + } + } + + Ok(if option.return_change { + updated + } else { + insert + } + .into()) + } + _ => Err(Error::WrongType), + }) + .unwrap_or_else(|| { + let mut x = SortedSet::new(); + let mut insert: usize = 0; + let mut updated: usize = 0; + + loop { + let score = match args.pop_front() { + Some(x) => bytes_to_number::(&x)?, + None => break, + }; + let value = args.pop_front().ok_or(Error::Syntax)?; + match x.insert(FloatOrd(score), value, &option) { + IResult::Inserted => insert += 1, + IResult::Updated => updated += 1, + _ => {} + } + } + + conn.db().set(key.clone(), x.into(), None); + + Ok(if option.return_change { + updated + } else { + insert + } + .into()) + })?; + + conn.db().bump_version(&key); + + Ok(result) +} + +/// Increments the score of member in the sorted set stored at key by increment. +/// If member does not exist in the sorted set, it is added with increment as +/// its score (as if its previous score was 0.0). If key does not exist, a new +/// sorted set with the specified member as its sole member is created. +pub async fn zincr_by(conn: &Connection, mut args: VecDeque) -> Result { + let key = args.pop_front().ok_or(Error::Syntax)?; + let score = bytes_to_number::(&args.pop_front().ok_or(Error::Syntax)?)?; + let value = args.pop_front().ok_or(Error::Syntax)?; + let option = IOption::incr(); + let result = conn + .db() + .get(&key) + .map_mut(|v| match v { + Value::SortedSet(x) => { + let _ = x.insert(FloatOrd(score), value.clone(), &option); + Ok(x.get_score(&value).unwrap_or_default().into()) + } + _ => Err(Error::WrongType), + }) + .unwrap_or_else(|| { + #[allow(clippy::mutable_key_type)] + let mut x = SortedSet::new(); + let _ = x.insert(FloatOrd(score), value.clone(), &option); + Ok(x.get_score(&value).unwrap_or_default().into()) + })?; + + conn.db().bump_version(&key); + Ok(result) +} + +/// Returns the sorted set cardinality (number of elements) of the sorted set +/// stored at key. +pub async fn zcard(conn: &Connection, args: VecDeque) -> Result { + conn.db() + .get(&args[0]) + .map(|v| match v { + Value::SortedSet(x) => Ok(x.len().into()), + _ => Err(Error::WrongType), + }) + .unwrap_or(Ok(0.into())) +} + +/// Returns the number of elements in the sorted set at key with a score between +/// min and max. +/// +/// The min and max arguments have the same semantic as described for +/// ZRANGEBYSCORE. +pub async fn zcount(conn: &Connection, args: VecDeque) -> Result { + let min = bytes_to_range_floatord(&args[1])?; + let max = bytes_to_range_floatord(&args[2])?; + conn.db() + .get(&args[0]) + .map(|v| match v { + Value::SortedSet(x) => Ok(x.count_values_by_score_range(min, max).into()), + _ => Err(Error::WrongType), + }) + .unwrap_or(Ok(0.into())) +} + +#[cfg(test)] +mod test { + use crate::{ + cmd::test::{create_connection, run_command}, + error::Error, + }; + + #[tokio::test] + async fn test_set_wrong_type() { + let c = create_connection(); + + let _ = run_command(&c, &["set", "foo", "1"]).await; + + assert_eq!( + Err(Error::WrongType), + run_command(&c, &["zadd", "foo", "5", "bar", "1", "foo"]).await, + ); + } + + #[tokio::test] + async fn test_zadd() { + let c = create_connection(); + + assert_eq!( + Ok(2.into()), + run_command(&c, &["zadd", "foo", "5", "bar", "1", "foo"]).await, + ); + assert_eq!( + Ok(0.into()), + run_command(&c, &["zadd", "foo", "5", "bar", "1", "foo"]).await, + ); + assert_eq!(Ok(2.into()), run_command(&c, &["zcard", "foo"]).await,); + } + + #[tokio::test] + async fn test_zcount() { + let c = create_connection(); + + assert_eq!( + Ok(3.into()), + run_command( + &c, + &["zadd", "foo", "5", "bar", "1", "foo", "5.9", "foobar"] + ) + .await, + ); + assert_eq!( + Ok(0.into()), + run_command(&c, &["zadd", "foo", "nx", "511", "bar", "10", "foo"]).await, + ); + assert_eq!( + Ok(2.into()), + run_command(&c, &["zcount", "foo", "1", "5"]).await, + ); + assert_eq!( + Ok(1.into()), + run_command(&c, &["zcount", "foo", "1", "(5"]).await, + ); + assert_eq!( + Ok(0.into()), + run_command(&c, &["zcount", "foo", "(1", "(5"]).await, + ); + assert_eq!( + Ok(3.into()), + run_command(&c, &["zcount", "foo", "-inf", "+inf"]).await, + ); + } +} diff --git a/src/dispatcher/mod.rs b/src/dispatcher/mod.rs index c0282d3..38409c7 100644 --- a/src/dispatcher/mod.rs +++ b/src/dispatcher/mod.rs @@ -19,6 +19,44 @@ pub mod command; // Returns the server time dispatcher! { + sorted_set { + ZADD { + cmd::sorted_set::zadd, + [Flag::Write Flag::DenyOom Flag::Fast], + -4, + 1, + 1, + 1, + true, + }, + ZCARD { + cmd::sorted_set::zcard, + [Flag::ReadOnly Flag::Fast], + 2, + 1, + 1, + 1, + true, + }, + ZCOUNT { + cmd::sorted_set::zcount, + [Flag::ReadOnly Flag::Fast], + 4, + 1, + 1, + 1, + true, + }, + ZINCRBY { + cmd::sorted_set::zincr_by, + [Flag::Write Flag::DenyOom Flag::Fast], + 4, + 1, + 1, + 1, + true, + } + }, set { SADD { cmd::set::sadd, diff --git a/src/value/mod.rs b/src/value/mod.rs index cac93d6..b45f265 100644 --- a/src/value/mod.rs +++ b/src/value/mod.rs @@ -5,14 +5,16 @@ pub mod checksum; pub mod cursor; pub mod expiration; pub mod float; +pub mod sorted_set; pub mod typ; use crate::{error::Error, value_try_from, value_vec_try_from}; use bytes::{Bytes, BytesMut}; +use float_ord::FloatOrd; use redis_zero_protocol_parser::Value as ParsedValue; use sha2::{Digest, Sha256}; use std::{ - collections::{HashMap, HashSet, VecDeque}, + collections::{Bound, HashMap, HashSet, VecDeque}, convert::{TryFrom, TryInto}, str::FromStr, }; @@ -30,6 +32,8 @@ pub enum Value { List(VecDeque), /// Set. This type cannot be serialized Set(HashSet), + /// Sorted set + SortedSet(sorted_set::SortedSet), /// Vector/Array of values Array(Vec), /// Bytes/Strings/Binary data @@ -225,6 +229,29 @@ pub fn bytes_to_int(bytes: &[u8]) -> Result { .map_err(|_| Error::NotANumberType("an integer".to_owned())) } +/// Converts bytes to a Range number +pub fn bytes_to_range(bytes: &[u8]) -> Result, Error> { + match bytes { + b"-inf" | b"+inf" | b"inf" => Ok(Bound::Unbounded), + _ => { + if bytes[0] == b'(' { + Ok(Bound::Excluded(bytes_to_number::(&(bytes[1..]))?)) + } else { + Ok(Bound::Included(bytes_to_number::(bytes)?)) + } + } + } +} + +/// Converts bytes to a Range of float FloatOrd numbers +pub fn bytes_to_range_floatord(bytes: &[u8]) -> Result, Error> { + match bytes_to_range(bytes)? { + Bound::Included(n) => Ok(Bound::Included(FloatOrd(n))), + Bound::Excluded(n) => Ok(Bound::Excluded(FloatOrd(n))), + Bound::Unbounded => Ok(Bound::Unbounded), + } +} + impl<'a> From<&ParsedValue<'a>> for Value { fn from(value: &ParsedValue) -> Self { match value { @@ -281,6 +308,12 @@ impl From<&str> for Value { } } +impl From for Value { + fn from(value: sorted_set::SortedSet) -> Self { + Value::SortedSet(value) + } +} + impl From> for Value { fn from(value: HashMap) -> Value { Value::Hash(value) diff --git a/src/value/sorted_set/insert.rs b/src/value/sorted_set/insert.rs new file mode 100644 index 0000000..599607b --- /dev/null +++ b/src/value/sorted_set/insert.rs @@ -0,0 +1,112 @@ +use crate::error::Error; +use bytes::Bytes; +use std::{collections::VecDeque, fmt::Debug}; + +#[derive(Debug, Copy, Clone, Eq, PartialEq)] +pub(crate) enum IPolicy { + NX, + XX, +} + +#[derive(Debug, Copy, Clone, Eq, PartialEq)] +pub(crate) enum UPolicyScore { + LT, + GT, +} + +#[derive(Debug, Default, Clone)] +/// Insert option +pub struct IOption { + pub(crate) insert_policy: Option, + pub(crate) update_policy_score: Option, + /// Modify the return value from the number of new elements added, to the + /// total number of elements changed (CH is an abbreviation of changed). + /// Changed elements are new elements added and elements already existing + /// for which the score was updated. So elements specified in the command + /// line having the same score as they had in the past are not counted. + /// Note: normally the return value of ZADD only counts the number of new + /// elements added. + pub return_change: bool, + /// Increments instead of adding + pub incr: bool, +} + +impl IOption { + /// Creates a new instance where only incr is set to true + pub fn incr() -> Self { + Self { + incr: true, + ..Default::default() + } + } + /// Creates a new instance + pub fn new(args: &mut VecDeque) -> Result { + let mut update_policy = None; + let mut update_policy_score = None; + let mut return_change = false; + let mut incr = false; + while let Some(t) = args.front() { + let command = String::from_utf8_lossy(t); + match command.to_uppercase().as_str() { + "NX" => { + if update_policy == Some(IPolicy::XX) { + return Err(Error::OptsNotCompatible("XX AND NX".to_owned())); + } + update_policy = Some(IPolicy::NX); + args.pop_front(); + } + "XX" => { + if update_policy == Some(IPolicy::NX) { + return Err(Error::OptsNotCompatible("XX AND NX".to_owned())); + } + update_policy = Some(IPolicy::XX); + args.pop_front(); + } + "LT" => { + if update_policy == Some(IPolicy::NX) + || update_policy_score == Some(UPolicyScore::GT) + { + return Err(Error::OptsNotCompatible("GT, LT, and/or NX".to_owned())); + } + update_policy_score = Some(UPolicyScore::LT); + args.pop_front(); + } + "GT" => { + if update_policy == Some(IPolicy::NX) + || update_policy_score == Some(UPolicyScore::LT) + { + return Err(Error::OptsNotCompatible("GT, LT, and/or NX".to_owned())); + } + update_policy_score = Some(UPolicyScore::GT); + args.pop_front(); + } + "CH" => { + return_change = true; + args.pop_front(); + } + "INCR" => { + incr = true; + args.pop_front(); + } + _ => break, + } + } + Ok(Self { + insert_policy: update_policy, + update_policy_score, + return_change, + incr, + }) + } +} + +/// Insert result +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum IResult { + /// No operation has taken place + NoOp, + /// A new element has been added + Inserted, + /// An element has been updated + Updated, +} diff --git a/src/value/sorted_set/mod.rs b/src/value/sorted_set/mod.rs new file mode 100644 index 0000000..f7dc212 --- /dev/null +++ b/src/value/sorted_set/mod.rs @@ -0,0 +1,205 @@ +//! # Sorted Set module +use bytes::Bytes; +use float_ord::FloatOrd; +use std::{ + collections::{BTreeMap, HashMap}, + fmt::Debug, + ops::Bound, +}; + +mod insert; + +pub use insert::{IOption, IResult}; +use insert::{IPolicy, UPolicyScore}; + +/// Sorted set score +pub type Score = FloatOrd; +/// Sorted set data with score +pub type DataWithScore = (Score, Bytes); + +/// Sorted set structure +#[derive(Debug, Clone)] +pub struct SortedSet { + set: HashMap, + order: BTreeMap, +} + +impl PartialEq for SortedSet { + fn eq(&self, other: &SortedSet) -> bool { + self.order == other.order + } +} + +impl Default for SortedSet { + fn default() -> Self { + Self::new() + } +} + +impl SortedSet { + /// Creates a new instance + pub fn new() -> Self { + Self { + set: HashMap::new(), + order: BTreeMap::new(), + } + } + + /// Clears the map, removing all elements. + pub fn clear(&mut self) { + self.set.clear(); + self.order.clear(); + } + + /// Returns `true` if the map contains no elements. + pub fn is_empty(&self) -> bool { + self.set.is_empty() + } + + /// Returns the number of elements in the set + pub fn len(&self) -> usize { + self.set.len() + } + + /// Adds a value to the set. + /// If the set did not have this value present, true is returned. + /// + /// If the set did have this value present, false is returned. + pub fn insert(&mut self, score: Score, value: Bytes, option: &IOption) -> IResult { + if let Some((current_score, _)) = self.set.get(&value).cloned() { + if option.insert_policy == Some(IPolicy::NX) { + return IResult::NoOp; + } + let cmp = current_score.cmp(&score); + let update_based_on_score = + option + .update_policy_score + .map_or(true, |policy| match policy { + UPolicyScore::LT => cmp == std::cmp::Ordering::Less, + UPolicyScore::GT => cmp == std::cmp::Ordering::Greater, + }); + + if !update_based_on_score { + return IResult::NoOp; + } + // remove the previous order entry + self.order.remove(&(current_score, value.clone())); + + let score = if option.incr { + FloatOrd(current_score.0 + score.0) + } else { + score + }; + + // update and insert the new order entry + self.set.insert(value.clone(), (score, 0)); + self.order.insert((score, value), 0); + + self.update_value_position(); + IResult::Updated + } else { + if option.insert_policy == Some(IPolicy::XX) { + return IResult::NoOp; + } + self.set.insert(value.clone(), (score, 0)); + self.order.insert((score, value), 0); + self.update_value_position(); + IResult::Inserted + } + } + + /// Returns a reference to the score in the set, if any, that is equal to the given value. + pub fn get_score(&self, value: &Bytes) -> Option { + self.set.get(value).map(|(value, _)| value.0) + } + + /// Returns all the values sorted by their score + pub fn get_values(&self) -> Vec { + self.order.keys().map(|(_, value)| value.clone()).collect() + } + + #[inline] + fn convert_to_range( + min: Bound, + max: Bound, + ) -> (Bound, Bound) { + let min_bytes = Bytes::new(); + let max_bytes = Bytes::copy_from_slice(&vec![255u8; 4096]); + + ( + match min { + Bound::Included(value) => Bound::Included((value, min_bytes.clone())), + Bound::Excluded(value) => Bound::Excluded((value, max_bytes.clone())), + Bound::Unbounded => Bound::Unbounded, + }, + match max { + Bound::Included(value) => Bound::Included((value, max_bytes)), + Bound::Excluded(value) => Bound::Excluded((value, min_bytes)), + Bound::Unbounded => Bound::Unbounded, + }, + ) + } + + /// Get total number of values in a score range + pub fn count_values_by_score_range(&self, min: Bound, max: Bound) -> usize { + self.order.range(Self::convert_to_range(min, max)).count() + } + + /// Get values in a score range + pub fn get_values_by_score_range(&self, min: Bound, max: Bound) -> Vec { + self.order + .range(Self::convert_to_range(min, max)) + .map(|(k, _)| k.1.clone()) + .collect() + } + + /// Adds the position in the set to each value based on their score + #[inline] + fn update_value_position(&mut self) { + for (i, ((_, key), value)) in self.order.iter_mut().enumerate() { + *value = i; + if let Some(value) = self.set.get_mut(key) { + value.1 = i; + } + } + } + + /// Return the position into the set based on their score + pub fn get_value_pos(&self, value: &Bytes) -> Option { + Some(self.set.get(value)?.1) + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn basic_usage() { + let mut set: SortedSet = SortedSet::new(); + let mut op = IOption { + insert_policy: Some(IPolicy::NX), + ..Default::default() + }; + + assert_eq!( + set.insert(FloatOrd(1.0), "2".into(), &op), + IResult::Inserted + ); + assert_eq!( + set.insert(FloatOrd(0.0), "3".into(), &op), + IResult::Inserted + ); + assert_eq!(set.insert(FloatOrd(33.1), "3".into(), &op), IResult::NoOp); + + op.insert_policy = None; + op.incr = true; + assert_eq!(set.insert(FloatOrd(2.0), "2".into(), &op), IResult::Updated); + + //assert_eq!(vec![3, 2], set.get_values()); + assert_eq!(Some(3.0), set.get_score(&"2".into())); + assert_eq!(Some(1), set.get_value_pos(&"2".into())); + assert_eq!(Some(0), set.get_value_pos(&"3".into())); + assert_eq!(None, set.get_value_pos(&"5".into())); + } +}