use std::{
cmp::Reverse,
collections::BinaryHeap,
io::{self, Read, Write},
};
use fastrand::Rng;
use gxhash::{GxBuildHasher, HashMap, HashSet};
use wbase::{glob::glob_match, time::now_ticks};
use wresp::{
ExpireOption,
cmd_strings::RESP_ERR_GENERIC_UNSUPPORTED_OPERATION as RESP_ERR_UNSUPPORTED_OPERATION,
};
use wval::GarnetObjectType;
use crate::{
object_store_utils::GarnetObjectPayload,
types::{
ObjectInput,
expiration_queue::{ExpirationQueue, ExpirationQueueEntry},
garnet_object_base::read_scan_input,
object_output::{ObjectOutput, ObjectOutputFlags},
},
};
#[derive(
Debug, Clone, Copy, PartialEq, Eq, num_enum::TryFromPrimitive, num_enum::IntoPrimitive,
)]
#[repr(u8)]
pub enum HashOperation {
Hcollect = 0,
Hexpire = 1,
Httl = 2,
Hpersist = 3,
Hget = 4,
Hmget = 5,
Hset = 6,
Hmset = 7,
Hsetnx = 8,
Hlen = 9,
Hdel = 10,
Hexists = 11,
Hgetall = 12,
Hkeys = 13,
Hvals = 14,
Hincrby = 15,
Hincrbyfloat = 16,
Hrandfield = 17,
Hscan = 18,
Hstrlen = 19,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(i32)]
pub enum HashExpireResult {
KeyNotFound = -2,
ExpireConditionNotMet = 0,
ExpireUpdated = 1,
KeyAlreadyExpired = 2,
}
#[derive(Debug, Clone, Default)]
pub struct HashObject {
pub hash: HashMap<Vec<u8>, Vec<u8>>,
pub expiration_times: Option<HashMap<Vec<u8>, i64>>,
pub expiration_queue: Option<ExpirationQueue>,
pub heap_memory_size: i64,
mutated_by_ttl: bool,
}
#[derive(Debug, Clone, bitcode::Encode, bitcode::Decode)]
struct HashWire {
entries: Vec<(Vec<u8>, Vec<u8>)>,
expirations: Option<Vec<(Vec<u8>, i64)>>,
}
impl HashObject {
pub fn new() -> Self {
Self::default()
}
#[inline]
pub fn len(&self) -> usize {
self.hash.len()
}
#[inline]
pub fn is_empty(&self) -> bool {
self.hash.is_empty()
}
pub fn deserialize_from_slice(slice: &[u8]) -> io::Result<Self> {
let wire: HashWire =
bitcode::decode(slice).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
let mut obj = Self::new();
obj.hash.reserve(wire.entries.len());
let now = now_ticks();
for (item, value) in wire.entries {
obj.update_size(&item, &value, true);
obj.hash.insert(item, value);
}
if let Some(expirations) = wire.expirations {
for (item, expiration) in expirations {
if expiration < now {
if let Some(val) = obj.hash.remove(&item) {
obj.update_size(&item, &val, false);
}
} else if obj.hash.contains_key(&item) {
obj.initialize_expiration_structures();
if let Some(times) = obj.expiration_times.as_mut() {
times.insert(item.clone(), expiration);
}
if let Some(queue) = obj.expiration_queue.as_mut() {
queue.push(Reverse(ExpirationQueueEntry {
expiration,
key: item,
}));
}
obj.update_expiration_size(true, true);
}
}
obj.cleanup_expiration_structures_if_empty();
}
Ok(obj)
}
pub fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
let mut buf = Vec::new();
reader.read_to_end(&mut buf)?;
Self::deserialize_from_slice(&buf)
}
pub fn serialize_to_vec(&self) -> Vec<u8> {
let now = now_ticks();
let has_expirations = self.expiration_times.is_some();
let mut entries: Vec<(Vec<u8>, Vec<u8>)> = Vec::with_capacity(self.hash.len());
for (k, v) in &self.hash {
if !has_expirations || !self.is_expired_at(k, now) {
entries.push((k.clone(), v.clone()));
}
}
let expirations = self.expiration_times.as_ref().and_then(|times| {
let mut active: Vec<(Vec<u8>, i64)> = Vec::with_capacity(times.len());
for (k, exp) in times {
if *exp >= now && self.hash.contains_key(k) {
active.push((k.clone(), *exp));
}
}
(!active.is_empty()).then_some(active)
});
bitcode::encode(&HashWire {
entries,
expirations,
})
}
pub fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
writer.write_all(&self.serialize_to_vec())
}
pub fn operate(
&mut self,
input: &ObjectInput,
output: &mut ObjectOutput,
resp_protocol_version: u8,
) -> bool {
if input.header.data[0] != GarnetObjectType::Hash as u8 {
output.output_flags |= ObjectOutputFlags::WRONG_TYPE;
output.payload.clear();
return true;
}
let Some(op) = hash_op_from_header(input) else {
output.write_error(RESP_ERR_UNSUPPORTED_OPERATION.as_bytes());
return true;
};
match op {
HashOperation::Hset | HashOperation::Hmset | HashOperation::Hsetnx => {
self.hash_set(input, output);
}
HashOperation::Hget => self.hash_get(input, output, resp_protocol_version),
HashOperation::Hmget => self.hash_multiple_get(input, output, resp_protocol_version),
HashOperation::Hgetall => self.hash_get_all(output, resp_protocol_version),
HashOperation::Hdel => self.hash_delete(input, output),
HashOperation::Hlen => self.hash_length(output),
HashOperation::Hstrlen => self.hash_str_length(input, output),
HashOperation::Hexists => self.hash_exists(input, output),
HashOperation::Hexpire => self.hash_expire(input, output),
HashOperation::Httl => self.hash_time_to_live(input, output),
HashOperation::Hpersist => self.hash_persist(input, output),
HashOperation::Hkeys | HashOperation::Hvals => self.hash_get_keys_or_values(input, output),
HashOperation::Hincrby => self.hash_increment(input, output),
HashOperation::Hincrbyfloat => self.hash_increment_float(input, output),
HashOperation::Hrandfield => self.hash_random_field(input, output, resp_protocol_version),
HashOperation::Hcollect => self.hash_collect(output),
HashOperation::Hscan => self.scan_operate(input, output),
}
if self.hash.is_empty() {
output.output_flags |= ObjectOutputFlags::REMOVE_KEY;
}
true
}
#[inline]
pub fn update_size(&mut self, key: &[u8], value: &[u8], add: bool) {
let memory_size =
(key.len().div_ceil(8) * 8 + value.len().div_ceil(8) * 8 + 16 + 16 + 16) as i64;
if add {
self.heap_memory_size += memory_size;
} else {
self.heap_memory_size -= memory_size;
}
}
pub fn initialize_expiration_structures(&mut self) {
if self.expiration_times.is_none() {
self.expiration_times = Some(HashMap::with_hasher(GxBuildHasher::default()));
self.expiration_queue = Some(BinaryHeap::new());
self.heap_memory_size += 16; }
}
#[inline]
pub fn update_expiration_size(&mut self, add: bool, include_pq: bool) {
let mut memory_size = 16 + 16;
if include_pq {
memory_size += 16 + 16;
}
if add {
self.heap_memory_size += memory_size;
} else {
self.heap_memory_size -= memory_size;
}
}
pub fn cleanup_expiration_structures_if_empty(&mut self) {
let Some(times) = self.expiration_times.as_ref() else {
return;
};
if !times.is_empty() {
return;
}
if let Some(queue) = self.expiration_queue.as_ref() {
self.heap_memory_size -= (16 + 16) * queue.len() as i64;
}
self.heap_memory_size -= 16; self.expiration_times = None;
self.expiration_queue = None;
}
pub fn scan(
&self,
start: i64,
count: i64,
pattern: &[u8],
is_no_value: bool,
) -> (Vec<Vec<u8>>, i64) {
let mut items: Vec<Vec<u8>> = Vec::new();
let mut cursor = start;
if (self.hash.len() as i64) < start {
cursor = 0;
return (items, cursor);
}
let count = if is_no_value { count } else { count * 2 };
let mut index = 0_i64;
let mut expired_keys_count = 0_i64;
for (key, value) in self.hash.iter() {
if self.is_expired(key) {
expired_keys_count += 1;
continue;
}
if index < start {
index += 1;
continue;
}
if pattern.is_empty() || glob_match(pattern, key) {
items.push(key.clone());
if !is_no_value {
items.push(value.clone());
}
}
cursor += 1;
if items.len() as i64 == count {
break;
}
}
if cursor + expired_keys_count == self.hash.len() as i64 {
cursor = 0;
}
(items, cursor)
}
#[inline]
pub fn is_expired_at(&self, key: &[u8], now: i64) -> bool {
self
.expiration_times
.as_ref()
.and_then(|t| t.get(key))
.is_some_and(|&expiration| expiration < now)
}
#[inline]
pub fn is_expired(&self, key: &[u8]) -> bool {
self.is_expired_at(key, now_ticks())
}
#[inline]
pub fn has_expirable_items(&self) -> bool {
self.expiration_times.is_some()
}
#[inline]
pub const fn mutated_by_ttl(&self) -> bool {
self.mutated_by_ttl
}
pub fn delete_expired_items(&mut self) {
if self.expiration_times.is_none() {
return;
}
self.delete_expired_items_worker();
}
fn delete_expired_items_worker(&mut self) {
let now = now_ticks();
while let Some(queue) = self.expiration_queue.as_mut() {
let Some(Reverse(head)) = queue.peek() else {
break;
};
if head.expiration >= now {
break;
}
let key = head.key.clone();
let expiration = head.expiration;
let in_times = self
.expiration_times
.as_ref()
.and_then(|t| t.get(&key))
.is_some_and(|&actual| actual == expiration);
if in_times {
self.expiration_times.as_mut().unwrap().remove(&key);
queue.pop();
self.update_expiration_size(false, true);
if let Some(value) = self.hash.get(&key).cloned() {
self.hash.remove(&key);
self.update_size(&key, &value, false);
self.mutated_by_ttl = true;
}
} else {
queue.pop();
self.heap_memory_size -= 16 + 16;
}
}
self.cleanup_expiration_structures_if_empty();
}
#[inline]
pub fn try_get_value(&self, key: &[u8]) -> Option<&Vec<u8>> {
if self.is_expired(key) {
return None;
}
self.hash.get(key)
}
pub fn remove(&mut self, key: &[u8]) -> Option<Vec<u8>> {
self.delete_expired_items();
let value = self.hash.remove(key)?;
if self.has_expirable_items() {
self.expiration_times.as_mut().unwrap().remove(key);
self.update_expiration_size(false, false);
}
self.update_size(key, &value, false);
Some(value)
}
pub fn count(&self) -> usize {
let Some(times) = self.expiration_times.as_ref() else {
return self.hash.len();
};
let expired_keys_count = times.keys().filter(|k| self.is_expired(k)).count();
self.hash.len() - expired_keys_count
}
pub fn contains_key(&self, key: &[u8]) -> bool {
self.hash.contains_key(key) && !self.is_expired(key)
}
pub(crate) fn add(&mut self, key: &[u8], value: Vec<u8>) {
self.delete_expired_items();
self.update_size(key, &value, true);
self.hash.insert(key.to_vec(), value);
}
pub fn set_expiration(
&mut self,
key: &[u8],
expiration: i64,
expire_option: ExpireOption,
) -> HashExpireResult {
if !self.contains_key(key) {
return HashExpireResult::KeyNotFound;
}
if expiration <= now_ticks() {
self.remove(key);
return HashExpireResult::KeyAlreadyExpired;
}
self.initialize_expiration_structures();
let current_expiration = self
.expiration_times
.as_ref()
.and_then(|t| t.get(key))
.copied();
let denied = match current_expiration {
Some(current) => {
expire_option.contains(ExpireOption::NX)
|| (expire_option.contains(ExpireOption::GT) && expiration <= current)
|| (expire_option.contains(ExpireOption::LT) && expiration >= current)
}
None => expire_option.contains(ExpireOption::XX) || expire_option.contains(ExpireOption::GT),
};
if denied {
return HashExpireResult::ExpireConditionNotMet;
}
let key_vec = key.to_vec();
if current_expiration.is_some() {
if let Some(slot) = self.expiration_times.as_mut().unwrap().get_mut(key) {
*slot = expiration;
}
self.heap_memory_size += 16 + 16;
} else {
self
.expiration_times
.as_mut()
.unwrap()
.insert(key_vec.clone(), expiration);
self.update_expiration_size(true, true);
}
self
.expiration_queue
.as_mut()
.unwrap()
.push(Reverse(ExpirationQueueEntry {
expiration,
key: key_vec,
}));
HashExpireResult::ExpireUpdated
}
pub fn persist(&mut self, key: &[u8]) -> i32 {
if !self.contains_key(key) {
return HashExpireResult::KeyNotFound as i32;
}
if self.has_expirable_items()
&& let Some(_) = self.expiration_times.as_mut().unwrap().remove(key)
{
self.heap_memory_size -= 16 + 16;
self.cleanup_expiration_structures_if_empty();
return HashExpireResult::ExpireUpdated as i32;
}
-1
}
pub fn get_expiration(&self, key: &[u8]) -> i64 {
if !self.contains_key(key) {
return i64::from(HashExpireResult::KeyNotFound as i32);
}
if let Some(&expiration) = self.expiration_times.as_ref().and_then(|t| t.get(key)) {
return expiration;
}
-1
}
pub fn element_at(&self, index: usize) -> Option<(Vec<u8>, Vec<u8>)> {
if self.has_expirable_items() {
return self
.hash
.iter()
.filter(|(k, _)| !self.is_expired(k))
.nth(index)
.map(|(k, v)| (k.clone(), v.clone()));
}
self
.hash
.iter()
.nth(index)
.map(|(k, v)| (k.clone(), v.clone()))
}
}
#[inline]
pub fn hash_op_from_header(input: &ObjectInput) -> Option<HashOperation> {
HashOperation::try_from(input.header.sub_id()).ok()
}
pub(crate) fn pick_k_random_indexes(n: usize, k: usize, seed: i32, distinct: bool) -> Vec<usize> {
const K_OVER_N_THRESHOLD: f64 = 0.1;
let mut rng = Rng::with_seed(u64::from(seed as u32));
if n == 0 || k == 0 {
return Vec::new();
}
if !distinct || (k as f64) / (n as f64) < K_OVER_N_THRESHOLD {
let mut indexes = Vec::with_capacity(k);
if !distinct {
indexes.extend((0..k).map(|_| rng.usize(..n)));
} else {
let mut picked = HashSet::with_capacity_and_hasher(k, GxBuildHasher::default());
while indexes.len() < k {
let idx = rng.usize(..n);
if picked.insert(idx) {
indexes.push(idx);
}
}
}
indexes
} else {
let mut perm: Vec<usize> = (0..n).collect();
for i in 0..k.min(n) {
let j = rng.usize(i..perm.len());
perm.swap(i, j);
}
perm.truncate(k);
perm
}
}
#[inline]
pub(crate) fn pick_random_index(n: usize, rand: i32) -> usize {
(rand as u32 as usize) % n
}
pub(crate) fn scan_operate_shared(
input: &ObjectInput,
output: &mut ObjectOutput,
do_scan: impl FnOnce(i64, i64, &[u8], bool) -> (Vec<Vec<u8>>, i64),
) {
let params = match read_scan_input(input, input.arg2) {
Ok(params) => params,
Err(msg) => {
output.write_error(msg);
return;
}
};
let (items, cursor_output) = do_scan(
params.cursor,
params.count,
params.pattern,
params.is_no_value,
);
let items_len = items.len();
output.write_array_length(2);
output.write_int64_as_bulk_string(cursor_output);
if items.is_empty() {
output.write_empty_array();
} else {
output.write_array_length(items.len());
for item in items {
output.write_bulk_string(&item);
}
}
output.result1 = items_len as i64;
}
impl GarnetObjectPayload for HashObject {
const OBJECT_TAG: GarnetObjectType = GarnetObjectType::Hash;
#[inline]
fn from_blob(raw: &[u8]) -> Self {
Self::deserialize_from_slice(raw).unwrap_or_default()
}
#[inline]
fn to_blob(&self) -> Vec<u8> {
self.serialize_to_vec()
}
#[inline]
fn is_empty(&self) -> bool {
self.hash.is_empty()
}
}