//! # in-memory database
//!
//! This database module is the core of the miniredis project. All other modules around this
//! database module.
use self::utils::{far_future, ExpirationOpts, Override};
use crate::{
error::Error,
value::{
bytes_to_number,
cursor::Cursor,
typ::{Typ, ValueTyp},
VDebug, Value,
},
};
use bytes::{BufMut, Bytes, BytesMut};
use core::num;
use entry::{new_version, Entry};
use expiration::ExpirationDb;
use glob::Pattern;
use log::trace;
use num_traits::CheckedAdd;
use parking_lot::{Mutex, RwLock};
use seahash::hash;
use std::{
collections::HashMap,
convert::{TryFrom, TryInto},
ops::{AddAssign, Deref},
str::FromStr,
sync::Arc,
thread,
};
use tokio::time::{Duration, Instant};
mod entry;
mod expiration;
pub mod pool;
pub mod scan;
pub(crate) mod utils;
/// Database structure
///
/// Each connection has their own clone of the database and the conn_id is stored in each instance.
/// The slots property is shared for all connections.
///
/// To avoid lock contention this database is *not* a single HashMap, instead it is a vector of
/// HashMaps. Each key is pre-sharded and a bucket is selected. By doing this pre-step instead of
/// locking the entire database, only a small portion is locked (shared or exclusively) at a time,
/// making this database implementation thread-friendly. The number of number_of_slots available cannot be
/// changed at runtime.
///
/// The database is also aware of other connections locking other keys exclusively (for
/// transactions).
///
/// Each entry is wrapped with an entry::Entry struct, which is aware of expirations and data
/// versioning (in practice the nanosecond of last modification).
#[derive(Debug)]
pub struct Db {
/// A vector of hashmaps.
///
/// Instead of having a single HashMap, and having all threads fighting for
/// blocking the single HashMap, we have a vector of N HashMap
/// (configurable), which in theory allow to have faster reads and writes.
///
/// Because all operations are always key specific, the key is used to hash
/// and select to which HashMap the data might be stored.
slots: Arc<Vec<RwLock<HashMap<Bytes, Entry>>>>,
/// Data structure to store all expiring keys
expirations: Arc<Mutex<ExpirationDb>>,
/// Number of HashMaps that are available.
number_of_slots: usize,
/// Databases unique ID. This is an internal identifier to avoid deadlocks
/// when copying and moving data between databases.
pub db_id: u128,
/// Current connection ID
///
/// A Database is attached to a conn_id. The slots and expiration data
/// structures are shared between all connections, regardless of conn_id.
///
/// This particular database instace is attached to a conn_id, which is used
/// to lock keys exclusively for transactions and other atomic operations.
conn_id: u128,
/// HashMap of all blocked keys by other connections. If a key appears in
/// here and it is not being hold by the current connection, current
/// connection must wait.
tx_key_locks: Arc<RwLock<HashMap<Bytes, u128>>>,
}
impl Db {
/// Creates a new database instance
pub fn new(number_of_slots: usize) -> Self {
let slots = (0..number_of_slots)
.map(|_| RwLock::new(HashMap::new()))
.collect();
Self {
slots: Arc::new(slots),
expirations: Arc::new(Mutex::new(ExpirationDb::new())),
conn_id: 0,
db_id: new_version(),
tx_key_locks: Arc::new(RwLock::new(HashMap::new())),
number_of_slots,
}
}
/// Creates a new Database instance bound to a connection.
///
/// This is particular useful when locking keys exclusively.
///
/// All the internal data are shjared through an Arc.
pub fn new_db_instance(self: Arc<Db>, conn_id: u128) -> Arc<Db> {
Arc::new(Self {
slots: self.slots.clone(),
tx_key_locks: self.tx_key_locks.clone(),
expirations: self.expirations.clone(),
conn_id,
db_id: self.db_id,
number_of_slots: self.number_of_slots,
})
}
#[inline]
/// Returns a slot where a key may be hosted.
///
/// In order to avoid too much locks, instead of having a single hash a
/// database instance is a set of hashes. Each key is pre-shared with a
/// quick hashing algorithm to select a 'slot' or HashMap where it may be
/// hosted.
fn get_slot(&self, key: &Bytes) -> usize {
let id = (hash(key) as usize) % self.number_of_slots;
trace!("selected slot {} for key {:?}", id, key);
let waiting = Duration::from_nanos(100);
while let Some(blocker) = self.tx_key_locks.read().get(key) {
// Loop while the key we are trying to access is being blocked by a
// connection in a transaction
if *blocker == self.conn_id {
// the key is being blocked by ourself, it is safe to break the
// waiting loop
break;
}
thread::sleep(waiting);
}
id
}
/// Locks keys exclusively
///
/// The locked keys are only accesible (read or write) by the connection
/// that locked them, any other connection must wait until the locking
/// connection releases them.
///
/// This is used to simulate redis transactions. Transaction in Redis are
/// atomic but pausing a multi threaded Redis just to keep the same promises
/// was a bit extreme, that's the reason why a transaction will lock
/// exclusively all keys involved.
pub fn lock_keys(&self, keys: &[Bytes]) {
let waiting = Duration::from_nanos(100);
loop {
let mut lock = self.tx_key_locks.write();
let mut i = 0;
for key in keys.iter() {
if let Some(blocker) = lock.get(key) {
if *blocker == self.conn_id {
// It is blocked by us already.
continue;
}
// It is blocked by another tx, we need to break
// and retry to gain the lock over this key
break;
}
lock.insert(key.clone(), self.conn_id);
i += 1;
}
if i == keys.len() {
// All the involved keys are successfully being blocked
// exclusively.
break;
}
// We need to sleep a bit and retry.
drop(lock);
thread::sleep(waiting);
}
}
/// Releases the lock on keys
pub fn unlock_keys(&self, keys: &[Bytes]) {
let mut lock = self.tx_key_locks.write();
for key in keys.iter() {
lock.remove(key);
}
}
/// Return debug info for a key
pub fn debug(&self, key: &Bytes) -> Result<VDebug, Error> {
let slot = self.slots[self.get_slot(key)].read();
Ok(slot
.get(key)
.filter(|x| x.is_valid())
.ok_or(Error::NotFound)?
.value
.debug())
}
/// Return the digest for each key. This used for testing only
pub fn digest(&self, keys: &[Bytes]) -> Result<Vec<Value>, Error> {
Ok(keys
.iter()
.map(|key| {
let slot = self.slots[self.get_slot(key)].read();
Value::Blob(
slot.get(key)
.filter(|v| v.is_valid())
.map(|v| hex::encode(&v.value.digest()))
.unwrap_or("00000".into())
.as_str()
.into(),
)
})
.collect::<Vec<Value>>())
}
/// Flushes the entire database
pub fn flushdb(&self) -> Result<Value, Error> {
self.expirations.lock().flush();
self.slots
.iter()
.map(|s| {
let mut s = s.write();
s.clear();
})
.for_each(drop);
Ok(Value::Ok)
}
/// Returns the number of elements in the database
pub fn len(&self) -> Result<usize, Error> {
self.purge();
Ok(self.slots.iter().map(|s| s.read().len()).sum())
}
/// Round numbers to store efficiently, specially float numbers. For instance `1.00` will be converted to `1`.
fn round_numbers<T>(number: T) -> BytesMut
where
T: ToString,
{
let number_to_str = number.to_string();
if number_to_str.find('.').is_none() {
return number_to_str.as_bytes().into();
}
let number_to_str = number_to_str
.trim_end_matches(|c| c == '0' || c == '.')
.to_string();
if number_to_str.is_empty() {
"0"
} else {
number_to_str.as_str()
}
.into()
}
// Converts a given number to a correct Value, it should be used with Self::round_numbers()
fn number_to_value(number: &[u8]) -> Result<Value, Error> {
if number.iter().find(|x| **x == b'.').is_some() {
Ok(Value::new(number))
} else {
Ok(Value::Integer(bytes_to_number(number)?))
}
}
/// Increment a sub-key in a hash
///
/// If the stored value cannot be converted into a number an error will be thrown
pub fn hincrby<T>(
&self,
key: &Bytes,
sub_key: &Bytes,
incr_by: &Bytes,
typ: &str,
) -> Result<Value, Error>
where
T: ToString
+ FromStr
+ CheckedAdd
+ for<'a> TryFrom<&'a Value, Error = Error>
+ Into<Value>
+ Copy,
{
let mut slot = self.slots[self.get_slot(key)].write();
let mut incr_by: T =
bytes_to_number(incr_by).map_err(|_| Error::NotANumberType(typ.to_owned()))?;
match slot.get_mut(key).filter(|x| x.is_valid()).map(|x| x.get()) {
Some(Value::Hash(h)) => {
let mut h = h.write();
if let Some(n) = h.get(sub_key) {
incr_by = incr_by
.checked_add(
&bytes_to_number(n)
.map_err(|_| Error::NotANumberType(typ.to_owned()))?,
)
.ok_or(Error::Overflow)?;
}
let incr_by_bytes = Self::round_numbers(incr_by).freeze();
h.insert(sub_key.clone(), incr_by_bytes.clone());
Self::number_to_value(&incr_by_bytes)
}
None => {
#[allow(clippy::mutable_key_type)]
let mut h = HashMap::new();
let incr_by_bytes = Self::round_numbers(incr_by).freeze();
h.insert(sub_key.clone(), incr_by_bytes.clone());
let _ = slot.insert(key.clone(), Entry::new(h.into(), None));
Self::number_to_value(&incr_by_bytes)
}
_ => Err(Error::WrongType),
}
}
/// Increments a key's value by a given number
///
/// If the stored value cannot be converted into a number an error will be
/// thrown.
pub fn incr<T>(&self, key: &Bytes, incr_by: T) -> Result<T, Error>
where
T: ToString + CheckedAdd + for<'a> TryFrom<&'a Value, Error = Error> + Into<Value> + Copy,
{
let mut slot = self.slots[self.get_slot(key)].write();
match slot.get_mut(key).filter(|x| x.is_valid()) {
Some(x) => {
if !x.is_clonable() {
return Err(Error::WrongType);
}
let value = x.get();
let mut number: T = value.try_into()?;
number = incr_by.checked_add(&number).ok_or(Error::Overflow)?;
x.change_value(Value::Blob(Self::round_numbers(number)));
Ok(number)
}
None => {
slot.insert(
key.clone(),
Entry::new(Value::Blob(Self::round_numbers(incr_by)), None),
);
Ok(incr_by)
}
}
}
/// Removes any expiration associated with a given key
pub fn persist(&self, key: &Bytes) -> Value {
let mut slot = self.slots[self.get_slot(key)].write();
slot.get_mut(key)
.filter(|x| x.is_valid())
.map_or(0.into(), |x| {
if x.has_ttl() {
self.expirations.lock().remove(key);
x.persist();
1.into()
} else {
0.into()
}
})
}
/// Set time to live for a given key
pub fn set_ttl(
&self,
key: &Bytes,
expires_in: Duration,
opts: ExpirationOpts,
) -> Result<Value, Error> {
if (opts.NX && opts.XX) || (opts.NX && opts.GT) || (opts.NX && opts.LT) {
return Err(Error::OptsNotCompatible("NX and XX, GT or LT".to_owned()));
}
if opts.GT && opts.LT {
return Err(Error::OptsNotCompatible("GT and LT".to_owned()));
}
let mut slot = self.slots[self.get_slot(key)].write();
let expires_at = Instant::now()
.checked_add(expires_in)
.unwrap_or_else(far_future);
Ok(slot
.get_mut(key)
.filter(|x| x.is_valid())
.map_or(0.into(), |x| {
let current_expire = x.get_ttl();
if opts.NX && current_expire.is_some() {
return 0.into();
}
if opts.XX && current_expire.is_none() {
return 0.into();
}
if opts.GT {
if let Some(current_expire) = current_expire {
if expires_at <= current_expire {
return 0.into();
}
} else {
return 0.into();
}
}
if opts.LT {
if let Some(current_expire) = current_expire {
if expires_at >= current_expire {
return 0.into();
}
}
}
self.expirations.lock().add(key, expires_at);
x.set_ttl(expires_at);
1.into()
}))
}
/// Overwrites part of the string stored at key, starting at the specified
/// offset, for the entire length of value. If the offset is larger than the
/// current length of the string at key, the string is padded with zero-bytes to
/// make offset fit. Non-existing keys are considered as empty strings, so this
/// command will make sure it holds a string large enough to be able to set
/// value at offset.
pub fn set_range(&self, key: &Bytes, offset: i128, data: &[u8]) -> Result<Value, Error> {
let mut slot = self.slots[self.get_slot(key)].write();
let value = slot.get_mut(key).map(|value| {
if !value.is_valid() {
self.expirations.lock().remove(key);
value.persist();
}
value.get_mut()
});
if offset < 0 {
return Err(Error::OutOfRange);
}
if offset >= 512 * 1024 * 1024 - 4 {
return Err(Error::MaxAllowedSize);
}
let length = offset as usize + data.len();
match value {
Some(Value::Blob(bytes)) => {
if bytes.capacity() < length {
bytes.resize(length, 0);
}
let writer = &mut bytes[offset as usize..length];
writer.copy_from_slice(data);
Ok(bytes.len().into())
}
None => {
if data.len() == 0 {
return Ok(0.into());
}
let mut bytes = BytesMut::new();
bytes.resize(length, 0);
let writer = &mut bytes[offset as usize..];
writer.copy_from_slice(data);
slot.insert(key.clone(), Entry::new(Value::new(&bytes), None));
Ok(bytes.len().into())
}
_ => Err(Error::WrongType),
}
}
/// Copies a key
pub fn copy(
&self,
source: &Bytes,
target: &Bytes,
replace: Override,
target_db: Option<Arc<Db>>,
) -> Result<bool, Error> {
let slot = self.slots[self.get_slot(source)].read();
let value = if let Some(value) = slot.get(source).filter(|x| x.is_valid()) {
value.clone()
} else {
return Ok(false);
};
drop(slot);
if let Some(db) = target_db {
if db.db_id == self.db_id && source == target {
return Err(Error::SameEntry);
}
if replace == Override::No && db.exists(&[target.clone()]) > 0 {
return Ok(false);
}
let _ = db.set_advanced(
target,
value.value.clone(),
value.get_ttl().map(|v| v - Instant::now()),
replace,
false,
false,
);
Ok(true)
} else {
if source == target {
return Err(Error::SameEntry);
}
if replace == Override::No && self.exists(&[target.clone()]) > 0 {
return Ok(false);
}
let mut slot = self.slots[self.get_slot(target)].write();
slot.insert(target.clone(), value);
Ok(true)
}
}
/// Moves a given key between databases
pub fn move_key(&self, source: &Bytes, target_db: Arc<Db>) -> Result<bool, Error> {
if self.db_id == target_db.db_id {
return Err(Error::SameEntry);
}
let mut slot = self.slots[self.get_slot(source)].write();
let (expires_in, value) = if let Some(value) = slot.get(source).filter(|v| v.is_valid()) {
(
value.get_ttl().map(|t| t - Instant::now()),
value.value.clone(),
)
} else {
return Ok(false);
};
if Value::Integer(1)
== target_db.set_advanced(&source, value, expires_in, Override::No, false, false)
{
slot.remove(source);
Ok(true)
} else {
Ok(false)
}
}
/// Renames a key
pub fn rename(
&self,
source: &Bytes,
target: &Bytes,
override_value: Override,
) -> Result<bool, Error> {
let slot1 = self.get_slot(source);
let slot2 = self.get_slot(target);
if slot1 == slot2 {
let mut slot = self.slots[slot1].write();
if override_value == Override::No && slot.get(target).is_some() {
return Ok(false);
}
if let Some(value) = slot.remove(source) {
slot.insert(target.clone(), value);
Ok(true)
} else {
Err(Error::NotFound)
}
} else {
let mut slot1 = self.slots[slot1].write();
let mut slot2 = self.slots[slot2].write();
if override_value == Override::No && slot2.get(target).is_some() {
return Ok(false);
}
if let Some(value) = slot1.remove(source) {
slot2.insert(target.clone(), value);
Ok(true)
} else {
Err(Error::NotFound)
}
}
}
/// Removes keys from the database
pub fn del(&self, keys: &[Bytes]) -> Value {
let mut expirations = self.expirations.lock();
keys.iter()
.filter_map(|key| {
expirations.remove(key);
self.slots[self.get_slot(key)].write().remove(key)
})
.filter(|key| key.is_valid())
.count()
.into()
}
/// Returns all keys that matches a given pattern. This is a very expensive command.
pub fn get_all_keys(&self, pattern: &Bytes) -> Result<Vec<Value>, Error> {
let pattern = String::from_utf8_lossy(pattern);
let pattern =
Pattern::new(&pattern).map_err(|_| Error::InvalidPattern(pattern.to_string()))?;
Ok(self
.slots
.iter()
.map(|slot| {
slot.read()
.keys()
.filter(|key| {
let str_key = String::from_utf8_lossy(key);
pattern.matches(&str_key)
})
.map(|key| Value::new(key))
.collect::<Vec<Value>>()
})
.flatten()
.collect())
}
/// Check if keys exists in the database
pub fn exists(&self, keys: &[Bytes]) -> usize {
let mut matches = 0;
keys.iter()
.map(|key| {
let slot = self.slots[self.get_slot(key)].read();
if let Some(key) = slot.get(key) {
matches += if key.is_valid() { 1 } else { 0 };
}
})
.for_each(drop);
matches
}
/// get_map_or
///
/// Instead of returning an entry of the database, to avoid clonning, this function will
/// execute a callback function with the entry as a parameter. If no record is found another
/// callback function is going to be executed, dropping the lock before doing so.
///
/// If an entry is found, the lock is not dropped before doing the callback. Avoid inserting
/// new entries. In this case the value is passed by reference, so it is possible to modify the
/// entry itself.
///
/// This function is useful to read non-scalar values from the database. Non-scalar values are
/// forbidden to clone, attempting cloning will endup in an error (Error::WrongType)
pub fn get_map_or<F1, F2>(&self, key: &Bytes, found: F1, not_found: F2) -> Result<Value, Error>
where
F1: FnOnce(&Value) -> Result<Value, Error>,
F2: FnOnce() -> Result<Value, Error>,
{
let slot = self.slots[self.get_slot(key)].read();
let entry = slot.get(key).filter(|x| x.is_valid()).map(|e| e.get());
if let Some(entry) = entry {
found(entry)
} else {
// drop lock
drop(slot);
not_found()
}
}
/// Updates the entry version of a given key
pub fn bump_version(&self, key: &Bytes) -> bool {
let mut slot = self.slots[self.get_slot(key)].write();
slot.get_mut(key)
.filter(|x| x.is_valid())
.map(|entry| {
entry.bump_version();
})
.is_some()
}
/// Returns the version of a given key
pub fn get_version(&self, key: &Bytes) -> u128 {
let slot = self.slots[self.get_slot(key)].read();
slot.get(key)
.filter(|x| x.is_valid())
.map(|entry| entry.version())
.unwrap_or_default()
}
/// Returns the name of the value type
pub fn get_data_type(&self, key: &Bytes) -> String {
let slot = self.slots[self.get_slot(key)].read();
slot.get(key)
.filter(|x| x.is_valid())
.map_or("none".to_owned(), |x| {
Typ::get_type(x.get()).to_string().to_lowercase()
})
}
/// Get a copy of an entry
pub fn get(&self, key: &Bytes) -> Value {
let slot = self.slots[self.get_slot(key)].read();
slot.get(key)
.filter(|x| x.is_valid())
.map_or(Value::Null, |x| x.clone_value())
}
/// Get a copy of an entry and modifies the expiration of the key
pub fn getex(&self, key: &Bytes, expires_in: Option<Duration>, make_persistent: bool) -> Value {
let mut slot = self.slots[self.get_slot(key)].write();
slot.get_mut(key)
.filter(|x| x.is_valid())
.map(|value| {
if make_persistent {
self.expirations.lock().remove(key);
value.persist();
} else if let Some(expires_in) = expires_in {
let expires_at = Instant::now()
.checked_add(expires_in)
.unwrap_or_else(far_future);
self.expirations.lock().add(key, expires_at);
value.set_ttl(expires_at);
}
value
})
.map_or(Value::Null, |x| x.clone_value())
}
/// Get multiple copies of entries
pub fn get_multi(&self, keys: &[Bytes]) -> Value {
keys.iter()
.map(|key| {
let slot = self.slots[self.get_slot(key)].read();
slot.get(key)
.filter(|x| x.is_valid() && x.is_clonable())
.map_or(Value::Null, |x| x.clone_value())
})
.collect::<Vec<Value>>()
.into()
}
/// Get a key or set a new value for the given key.
pub fn getset(&self, key: &Bytes, value: Value) -> Value {
let mut slot = self.slots[self.get_slot(key)].write();
self.expirations.lock().remove(key);
slot.insert(key.clone(), Entry::new(value, None))
.filter(|x| x.is_valid())
.map_or(Value::Null, |x| x.clone_value())
}
/// Takes an entry from the database.
pub fn getdel(&self, key: &Bytes) -> Value {
let mut slot = self.slots[self.get_slot(key)].write();
slot.remove(key).map_or(Value::Null, |x| {
self.expirations.lock().remove(key);
x.clone_value()
})
}
/// Set a key, value with an optional expiration time
pub fn append(&self, key: &Bytes, value_to_append: &Bytes) -> Result<Value, Error> {
let mut slot = self.slots[self.get_slot(key)].write();
let mut entry = slot.get_mut(key).filter(|x| x.is_valid());
if let Some(entry) = slot.get_mut(key).filter(|x| x.is_valid()) {
match entry.get_mut() {
Value::Blob(value) => {
value.put(value_to_append.as_ref());
Ok(value.len().into())
}
_ => Err(Error::WrongType),
}
} else {
slot.insert(key.clone(), Entry::new(Value::new(value_to_append), None));
Ok(value_to_append.len().into())
}
}
/// Set multiple key/value pairs. Are involved keys are locked exclusively
/// like a transaction.
///
/// If override_all is set to false, all entries must be new entries or the
/// entire operation fails, in this case 1 or is returned. Otherwise `Ok` is
/// returned.
pub fn multi_set(&self, key_values: &[Bytes], override_all: bool) -> Result<Value, Error> {
if key_values.len() % 2 == 1 {
return Err(Error::Syntax);
}
let keys = key_values
.iter()
.step_by(2)
.cloned()
.collect::<Vec<Bytes>>();
self.lock_keys(&keys);
if !override_all {
for key in keys.iter() {
let slot = self.slots[self.get_slot(key)].read();
if slot.get(key).is_some() {
self.unlock_keys(&keys);
return Ok(0.into());
}
}
}
for (i, _) in key_values.iter().enumerate().step_by(2) {
let mut slot = self.slots[self.get_slot(&key_values[i])].write();
slot.insert(
key_values[i].clone(),
Entry::new(Value::new(&key_values[i + 1]), None),
);
}
self.unlock_keys(&keys);
if override_all {
Ok(Value::Ok)
} else {
Ok(1.into())
}
}
/// Set a key, value with an optional expiration time
pub fn set(&self, key: &Bytes, value: Value, expires_in: Option<Duration>) -> Value {
self.set_advanced(key, value, expires_in, Default::default(), false, false)
}
/// Set a value in the database with various settings
pub fn set_advanced(
&self,
key: &Bytes,
value: Value,
expires_in: Option<Duration>,
override_value: Override,
keep_ttl: bool,
return_previous: bool,
) -> Value {
let mut slot = self.slots[self.get_slot(key)].write();
let expires_at = expires_in.map(|duration| {
Instant::now()
.checked_add(duration)
.unwrap_or_else(far_future)
});
let previous = slot.get(key).filter(|x| x.is_valid());
let expires_at = if keep_ttl {
if let Some(previous) = previous {
previous.get_ttl()
} else {
expires_at
}
} else {
expires_at
};
let to_return = if return_previous {
let previous = previous.map_or(Value::Null, |v| v.clone_value());
if previous.is_err() {
// Error while trying to clone the previous value to return, we
// must halt and return immediately.
return previous;
}
Some(previous)
} else {
None
};
match override_value {
Override::No => {
if previous.is_some() {
return if let Some(to_return) = to_return {
to_return
} else {
0.into()
};
}
}
Override::Only => {
if previous.is_none() {
return if let Some(to_return) = to_return {
to_return
} else {
0.into()
};
}
}
_ => {}
};
if let Some(expires_at) = expires_at {
self.expirations.lock().add(key, expires_at);
} else {
/// Make sure to remove the new key (or replaced) from the
/// expiration table (from any possible past value).
self.expirations.lock().remove(key);
}
slot.insert(key.clone(), Entry::new(value, expires_at));
if let Some(to_return) = to_return {
to_return
} else if override_value == Override::Yes {
Value::Ok
} else {
1.into()
}
}
/// Returns the TTL of a given key
pub fn ttl(&self, key: &Bytes) -> Option<Option<Instant>> {
let slot = self.slots[self.get_slot(key)].read();
slot.get(key).filter(|x| x.is_valid()).map(|x| x.get_ttl())
}
/// Check whether a given key is in the list of keys to be purged or not.
/// This function is mainly used for unit testing
pub fn is_key_in_expiration_list(&self, key: &Bytes) -> bool {
self.expirations.lock().has(key)
}
/// Remove expired entries from the database.
///
/// This function should be called from a background thread every few seconds. Calling it more
/// often is a waste of resources.
///
/// Expired keys are automatically hidden by the database, this process is just claiming back
/// the memory from those expired keys.
pub fn purge(&self) -> u64 {
let mut expirations = self.expirations.lock();
let mut removed = 0;
trace!("Watching {} keys for expirations", expirations.len());
let keys = expirations.get_expired_keys(None);
drop(expirations);
keys.iter()
.map(|key| {
let mut slot = self.slots[self.get_slot(key)].write();
if slot.remove(key).is_some() {
trace!("Removed key {:?} due timeout", key);
removed += 1;
}
})
.for_each(drop);
removed
}
}
impl scan::Scan for Db {
fn scan(
&self,
cursor: Cursor,
pattern: Option<&Bytes>,
count: Option<usize>,
typ: Option<Typ>,
) -> Result<scan::Result, Error> {
let mut keys = vec![];
let mut slot_id = cursor.bucket as usize;
let mut last_pos = cursor.last_position as usize;
let pattern = pattern
.map(|pattern| {
let pattern = String::from_utf8_lossy(pattern);
Pattern::new(&pattern).map_err(|_| Error::InvalidPattern(pattern.to_string()))
})
.transpose()?;
loop {
let slot = if let Some(value) = self.slots.get(slot_id) {
value.read()
} else {
// We iterated through all the entries, time to signal that to
// the client but returning a "0" cursor.
slot_id = 0;
last_pos = 0;
break;
};
for (key, value) in slot.iter().skip(last_pos) {
if !value.is_valid() {
// Entry still exists in memory but it is not longer valid
// and will soon be gargabe collected.
last_pos += 1;
continue;
}
if let Some(pattern) = &pattern {
let str_key = String::from_utf8_lossy(key);
if !pattern.matches(&str_key) {
last_pos += 1;
continue;
}
}
if let Some(typ) = &typ {
if !typ.is_value_type(value.get()) {
last_pos += 1;
continue;
}
}
keys.push(Value::new(key));
last_pos += 1;
if keys.len() == count.unwrap_or(10) {
break;
}
}
if keys.len() == count.unwrap_or(10) {
break;
}
last_pos = 0;
slot_id += 1;
}
Ok(scan::Result {
cursor: Cursor::new(slot_id as u16, last_pos as u64)?,
result: keys,
})
}
}
#[cfg(test)]
mod test {
use super::*;
use crate::{bytes, db::scan::Scan, value::float::Float};
use std::str::FromStr;
#[test]
fn incr_wrong_type() {
let db = Db::new(100);
db.set(&bytes!(b"num"), Value::Blob(bytes!("some string")), None);
let r = db.incr(&bytes!("num"), 1);
assert!(r.is_err());
assert_eq!(Error::NotANumber, r.expect_err("should fail"));
assert_eq!(Value::Blob(bytes!("some string")), db.get(&bytes!("num")));
}
#[test]
fn incr_blob_float() {
let db = Db::new(100);
db.set(&bytes!(b"num"), Value::Blob(bytes!("1.1")), None);
assert_eq!(Ok(2.2.into()), db.incr::<Float>(&bytes!("num"), 1.1.into()));
assert_eq!(Value::Blob(bytes!("2.2")), db.get(&bytes!("num")));
}
#[test]
fn incr_blob_int_float() {
let db = Db::new(100);
db.set(&bytes!(b"num"), Value::Blob(bytes!("1")), None);
assert_eq!(Ok(2.1.into()), db.incr::<Float>(&bytes!("num"), 1.1.into()));
assert_eq!(Value::Blob(bytes!("2.1")), db.get(&bytes!("num")));
}
#[test]
fn incr_blob_int() {
let db = Db::new(100);
db.set(&bytes!(b"num"), Value::Blob(bytes!("1")), None);
assert_eq!(Ok(2), db.incr(&bytes!("num"), 1));
assert_eq!(Value::Blob(bytes!("2")), db.get(&bytes!("num")));
}
#[test]
fn incr_blob_int_set() {
let db = Db::new(100);
assert_eq!(Ok(1), db.incr(&bytes!("num"), 1));
assert_eq!(Value::Blob(bytes!("1")), db.get(&bytes!("num")));
}
#[test]
fn incr_blob_float_set() {
let db = Db::new(100);
assert_eq!(Ok(1.1.into()), db.incr::<Float>(&bytes!("num"), 1.1.into()));
assert_eq!(Value::Blob(bytes!("1.1")), db.get(&bytes!("num")));
}
#[test]
fn del() {
let db = Db::new(100);
db.set(&bytes!(b"expired"), Value::Ok, Some(Duration::from_secs(0)));
db.set(&bytes!(b"valid"), Value::Ok, None);
db.set(
&bytes!(b"expiring"),
Value::Ok,
Some(Duration::from_secs(5)),
);
assert_eq!(
Value::Integer(2),
db.del(&[
bytes!(b"expired"),
bytes!(b"valid"),
bytes!(b"expiring"),
bytes!(b"not_existing_key")
])
);
}
#[test]
fn ttl() {
let db = Db::new(100);
db.set(&bytes!(b"expired"), Value::Ok, Some(Duration::from_secs(0)));
db.set(&bytes!(b"valid"), Value::Ok, None);
db.set(
&bytes!(b"expiring"),
Value::Ok,
Some(Duration::from_secs(5)),
);
assert_eq!(None, db.ttl(&bytes!(b"expired")));
assert_eq!(None, db.ttl(&bytes!(b"not_existing_key")));
assert_eq!(Some(None), db.ttl(&bytes!(b"valid")));
assert!(match db.ttl(&bytes!(b"expiring")) {
Some(Some(_)) => true,
_ => false,
});
}
#[test]
fn persist_bug() {
let db = Db::new(100);
db.set(&bytes!(b"one"), Value::Ok, Some(Duration::from_secs(1)));
assert_eq!(Value::Ok, db.get(&bytes!(b"one")));
assert!(db.is_key_in_expiration_list(&bytes!(b"one")));
db.persist(&bytes!(b"one"));
assert!(!db.is_key_in_expiration_list(&bytes!(b"one")));
}
#[test]
fn purge_keys() {
let db = Db::new(100);
db.set(&bytes!(b"one"), Value::Ok, Some(Duration::from_secs(0)));
// Expired keys should not be returned, even if they are not yet
// removed by the purge process.
assert_eq!(Value::Null, db.get(&bytes!(b"one")));
// Purge twice
assert_eq!(1, db.purge());
assert_eq!(0, db.purge());
assert_eq!(Value::Null, db.get(&bytes!(b"one")));
}
#[test]
fn replace_purge_keys() {
let db = Db::new(100);
db.set(&bytes!(b"one"), Value::Ok, Some(Duration::from_secs(0)));
// Expired keys should not be returned, even if they are not yet
// removed by the purge process.
assert_eq!(Value::Null, db.get(&bytes!(b"one")));
db.set(&bytes!(b"one"), Value::Ok, Some(Duration::from_secs(5)));
assert_eq!(Value::Ok, db.get(&bytes!(b"one")));
// Purge should return 0 as the expired key has been removed already
assert_eq!(0, db.purge());
}
#[test]
fn scan_skip_expired() {
let db = Db::new(100);
db.set(&bytes!(b"one"), Value::Ok, Some(Duration::from_secs(0)));
db.set(&bytes!(b"two"), Value::Ok, Some(Duration::from_secs(0)));
for i in 0u64..20u64 {
let key: Bytes = i.to_string().into();
db.set(&key, Value::Ok, None);
}
let result = db
.scan(Cursor::from_str("0").unwrap(), None, None, None)
.unwrap();
// first 10 records
assert_eq!(10, result.result.len());
// make sure the cursor is valid
assert_ne!("0", result.cursor.to_string());
let result = db.scan(result.cursor, None, None, None).unwrap();
// 10 more records
assert_eq!(10, result.result.len());
// make sure the cursor is valid
assert_ne!("0", result.cursor.to_string());
let result = db.scan(result.cursor, None, None, None).unwrap();
// No more results!
assert_eq!(0, result.result.len());
assert_eq!("0", result.cursor.to_string());
}
#[test]
fn scan_limit() {
let db = Db::new(10);
db.set(&bytes!(b"one"), Value::Ok, Some(Duration::from_secs(0)));
db.set(&bytes!(b"two"), Value::Ok, Some(Duration::from_secs(0)));
for i in 0u64..2000u64 {
let key: Bytes = i.to_string().into();
db.set(&key, Value::Ok, None);
}
let result = db
.scan(Cursor::from_str("0").unwrap(), None, Some(2), None)
.unwrap();
assert_eq!(2, result.result.len());
assert_ne!("0", result.cursor.to_string());
}
#[test]
fn scan_filter() {
let db = Db::new(100);
db.set(&bytes!(b"fone"), Value::Ok, None);
db.set(&bytes!(b"ftwo"), Value::Ok, None);
for i in 0u64..20u64 {
let key: Bytes = i.to_string().into();
db.set(&key, Value::Ok, None);
}
let result = db
.scan(
Cursor::from_str("0").unwrap(),
Some(&bytes!(b"f*")),
None,
None,
)
.unwrap();
assert_eq!(2, result.result.len());
assert_eq!("0", result.cursor.to_string());
}
}