From 6c20f6b0eec1f2c5bce4ebb237124410f5b4e150 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Sat, 14 Feb 2026 17:01:30 -0500 Subject: [PATCH 1/2] refactor: add AddResult::UNCHANGED, collapse sorted_set flag checks the add_with_flags method had 5 identical `AddResult { added: false, updated: false }` returns scattered across separate if-blocks. added a const UNCHANGED and collapsed nx/gt/lt/eq checks into a single guard. --- crates/ember-core/src/types/sorted_set.rs | 46 +++++++++-------------- 1 file changed, 17 insertions(+), 29 deletions(-) diff --git a/crates/ember-core/src/types/sorted_set.rs b/crates/ember-core/src/types/sorted_set.rs index 9cd5f610..8345b194 100644 --- a/crates/ember-core/src/types/sorted_set.rs +++ b/crates/ember-core/src/types/sorted_set.rs @@ -37,6 +37,14 @@ pub struct AddResult { pub updated: bool, } +impl AddResult { + /// No change: member was neither added nor updated. + pub const UNCHANGED: Self = Self { + added: false, + updated: false, + }; +} + /// A sorted set of unique string members, each with a floating-point score. /// /// Members are ordered by `(score, member_name)`. Rank is determined by @@ -69,30 +77,13 @@ impl SortedSet { let new_score = OrderedFloat(score); if let Some(&old_score) = self.scores.get(&member) { - // member exists — check flags - if flags.nx { - return AddResult { - added: false, - updated: false, - }; - } - if flags.gt && new_score <= old_score { - return AddResult { - added: false, - updated: false, - }; - } - if flags.lt && new_score >= old_score { - return AddResult { - added: false, - updated: false, - }; - } - if new_score == old_score { - return AddResult { - added: false, - updated: false, - }; + // member exists — skip if any flag condition prevents the update + if flags.nx + || (flags.gt && new_score <= old_score) + || (flags.lt && new_score >= old_score) + || new_score == old_score + { + return AddResult::UNCHANGED; } // update: remove old entry, insert new self.tree.remove(&(old_score, member.clone())); @@ -103,12 +94,9 @@ impl SortedSet { updated: true, } } else { - // new member — check XX flag + // new member — XX means only update, so skip if flags.xx { - return AddResult { - added: false, - updated: false, - }; + return AddResult::UNCHANGED; } self.scores.insert(member.clone(), new_score); self.tree.insert((new_score, member), ()); From 10ab1fb93c61152bafa3c1e138e71cb809c9736d Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Sat, 14 Feb 2026 17:04:19 -0500 Subject: [PATCH 2/2] refactor: extract collection-write helpers in keyspace list_push, zadd, hset, and sadd all followed the same 10-step pattern: expire check, type guard, memory estimation, limit enforcement, empty entry creation, old size measurement, mutation, new size measurement, memory adjustment, touch. extracted four helpers that capture the shared steps: - ensure_collection_type(): is-new check + type guard - reserve_memory(): estimate increase + enforce limit - insert_empty(): create entry with empty collection - track_size(): measure before/after + adjust tracker also simplified zrem with let-else and added a comment on lrange's empty-range check. each collection-write function is now ~15 lines instead of ~50. --- crates/ember-core/src/keyspace.rs | 252 ++++++++++++++---------------- 1 file changed, 116 insertions(+), 136 deletions(-) diff --git a/crates/ember-core/src/keyspace.rs b/crates/ember-core/src/keyspace.rs index f50a612d..88aa86a9 100644 --- a/crates/ember-core/src/keyspace.rs +++ b/crates/ember-core/src/keyspace.rs @@ -338,6 +338,63 @@ impl Keyspace { } } + /// Checks whether a key either doesn't exist or holds the expected + /// collection type. Returns `Ok(true)` if the key is new, + /// `Ok(false)` if it exists with the right type, or `Err(WrongType)` + /// if the key exists with a different type. + fn ensure_collection_type( + &self, + key: &str, + type_check: fn(&Value) -> bool, + ) -> Result { + match self.entries.get(key) { + None => Ok(true), + Some(e) if type_check(&e.value) => Ok(false), + Some(_) => Err(WriteError::WrongType), + } + } + + /// Estimates the memory cost of a collection write and enforces the + /// limit. `base_overhead` is the fixed cost of a new collection + /// (e.g. VECDEQUE_BASE_OVERHEAD). Returns `Ok(())` on success. + fn reserve_memory( + &mut self, + is_new: bool, + key: &str, + base_overhead: usize, + element_increase: usize, + ) -> Result<(), WriteError> { + let estimated_increase = if is_new { + memory::ENTRY_OVERHEAD + key.len() + base_overhead + element_increase + } else { + element_increase + }; + if self.enforce_memory_limit(estimated_increase) { + Ok(()) + } else { + Err(WriteError::OutOfMemory) + } + } + + /// Inserts a new key with an empty collection value. Used by + /// collection-write methods after type-checking and memory reservation. + fn insert_empty(&mut self, key: &str, value: Value) { + self.memory.add(key, &value); + self.entries.insert(key.to_owned(), Entry::new(value, None)); + } + + /// Measures entry size before and after a mutation, adjusting the + /// memory tracker for the difference. Touches the entry afterwards. + fn track_size(&mut self, key: &str, f: impl FnOnce(&mut Entry) -> T) -> T { + let entry = self.entries.get_mut(key).expect("caller verified key exists"); + let old_size = memory::entry_size(key, &entry.value); + let result = f(entry); + let entry = self.entries.get(key).expect("mutation should not remove key"); + let new_size = memory::entry_size(key, &entry.value); + self.memory.adjust(old_size, new_size); + result + } + /// Adjusts the expiry count when replacing an entry whose TTL status /// may have changed (e.g. SET overwriting an existing key). fn adjust_expiry_count(&mut self, had_expiry: bool, has_expiry: bool) { @@ -1024,6 +1081,7 @@ impl Keyspace { Value::List(deque) => { let len = deque.len() as i64; let (s, e) = normalize_range(start, stop, len); + // empty range: inverted indices or both out of bounds if s > e { return Ok(vec![]); } @@ -1064,53 +1122,33 @@ impl Keyspace { fn list_push(&mut self, key: &str, values: &[Bytes], left: bool) -> Result { self.remove_if_expired(key); - let is_new = !self.entries.contains_key(key); - - if !is_new && !matches!(self.entries[key].value, Value::List(_)) { - return Err(WriteError::WrongType); - } + let is_new = self.ensure_collection_type(key, |v| matches!(v, Value::List(_)))?; - // estimate the memory increase and enforce the limit before mutating let element_increase: usize = values .iter() .map(|v| memory::VECDEQUE_ELEMENT_OVERHEAD + v.len()) .sum(); - let estimated_increase = if is_new { - memory::ENTRY_OVERHEAD + key.len() + memory::VECDEQUE_BASE_OVERHEAD + element_increase - } else { - element_increase - }; - if !self.enforce_memory_limit(estimated_increase) { - return Err(WriteError::OutOfMemory); - } + self.reserve_memory(is_new, key, memory::VECDEQUE_BASE_OVERHEAD, element_increase)?; if is_new { - let value = Value::List(VecDeque::new()); - self.memory.add(key, &value); - self.entries.insert(key.to_owned(), Entry::new(value, None)); + self.insert_empty(key, Value::List(VecDeque::new())); } - let entry = self - .entries - .get_mut(key) - .expect("just inserted or verified"); - let old_entry_size = memory::entry_size(key, &entry.value); - - let Value::List(ref mut deque) = entry.value else { - return Err(WriteError::WrongType); - }; - for val in values { - if left { - deque.push_front(val.clone()); - } else { - deque.push_back(val.clone()); + let len = self.track_size(key, |entry| { + let Value::List(ref mut deque) = entry.value else { + unreachable!("type verified by ensure_collection_type"); + }; + for val in values { + if left { + deque.push_front(val.clone()); + } else { + deque.push_back(val.clone()); + } } - } - let len = deque.len(); - entry.touch(); - - let new_entry_size = memory::entry_size(key, &entry.value); - self.memory.adjust(old_entry_size, new_entry_size); + let len = deque.len(); + entry.touch(); + len + }); Ok(len) } @@ -1163,40 +1201,26 @@ impl Keyspace { ) -> Result { self.remove_if_expired(key); - let is_new = !self.entries.contains_key(key); - if !is_new && !matches!(self.entries[key].value, Value::SortedSet(_)) { - return Err(WriteError::WrongType); - } + let is_new = + self.ensure_collection_type(key, |v| matches!(v, Value::SortedSet(_)))?; // worst-case estimate: assume all members are new let member_increase: usize = members .iter() .map(|(_, m)| SortedSet::estimated_member_cost(m)) .sum(); - let estimated_increase = if is_new { - memory::ENTRY_OVERHEAD + key.len() + SortedSet::BASE_OVERHEAD + member_increase - } else { - member_increase - }; - if !self.enforce_memory_limit(estimated_increase) { - return Err(WriteError::OutOfMemory); - } + self.reserve_memory(is_new, key, SortedSet::BASE_OVERHEAD, member_increase)?; if is_new { - let value = Value::SortedSet(SortedSet::new()); - self.memory.add(key, &value); - self.entries.insert(key.to_owned(), Entry::new(value, None)); + self.insert_empty(key, Value::SortedSet(SortedSet::new())); } - let entry = self - .entries - .get_mut(key) - .expect("just inserted or verified"); - let old_entry_size = memory::entry_size(key, &entry.value); - - let mut count = 0; - let mut applied = Vec::new(); - if let Value::SortedSet(ref mut ss) = entry.value { + let (count, applied) = self.track_size(key, |entry| { + let Value::SortedSet(ref mut ss) = entry.value else { + unreachable!("type verified by ensure_collection_type"); + }; + let mut count = 0; + let mut applied = Vec::new(); for (score, member) in members { let result = ss.add_with_flags(member.clone(), *score, flags); if result.added || result.updated { @@ -1210,15 +1234,13 @@ impl Keyspace { count += 1; } } - } - entry.touch(); - - let new_entry_size = memory::entry_size(key, &entry.value); - self.memory.adjust(old_entry_size, new_entry_size); + entry.touch(); + (count, applied) + }); - // clean up if the set is still empty (e.g. XX on a new key) - if let Value::SortedSet(ref ss) = entry.value { - if ss.is_empty() { + // clean up if the set is still empty (e.g. XX flag on a new key) + if let Some(entry) = self.entries.get(key) { + if matches!(&entry.value, Value::SortedSet(ss) if ss.is_empty()) { self.memory .remove_with_size(memory::entry_size(key, &entry.value)); self.entries.remove(key); @@ -1238,13 +1260,9 @@ impl Keyspace { return Ok(vec![]); } - match self.entries.get(key) { - None => return Ok(vec![]), - Some(e) => { - if !matches!(e.value, Value::SortedSet(_)) { - return Err(WrongType); - } - } + let Some(entry) = self.entries.get(key) else { return Ok(vec![]) }; + if !matches!(entry.value, Value::SortedSet(_)) { + return Err(WrongType); } let old_entry_size = memory::entry_size(key, &self.entries[key].value); @@ -1370,50 +1388,31 @@ impl Keyspace { self.remove_if_expired(key); - let is_new = !self.entries.contains_key(key); - - if !is_new && !matches!(self.entries[key].value, Value::Hash(_)) { - return Err(WriteError::WrongType); - } + let is_new = self.ensure_collection_type(key, |v| matches!(v, Value::Hash(_)))?; - // estimate memory increase before mutating let field_increase: usize = fields .iter() .map(|(f, v)| f.len() + v.len() + memory::HASHMAP_ENTRY_OVERHEAD) .sum(); - let estimated_increase = if is_new { - memory::ENTRY_OVERHEAD + key.len() + memory::HASHMAP_BASE_OVERHEAD + field_increase - } else { - field_increase - }; - if !self.enforce_memory_limit(estimated_increase) { - return Err(WriteError::OutOfMemory); - } + self.reserve_memory(is_new, key, memory::HASHMAP_BASE_OVERHEAD, field_increase)?; if is_new { - let value = Value::Hash(HashMap::new()); - self.memory.add(key, &value); - self.entries.insert(key.to_owned(), Entry::new(value, None)); + self.insert_empty(key, Value::Hash(HashMap::new())); } - let entry = self - .entries - .get_mut(key) - .expect("just inserted or verified"); - let old_entry_size = memory::entry_size(key, &entry.value); - - let mut added = 0; - if let Value::Hash(ref mut map) = entry.value { + let added = self.track_size(key, |entry| { + let Value::Hash(ref mut map) = entry.value else { + unreachable!("type verified by ensure_collection_type"); + }; + let mut added = 0; for (field, value) in fields { if map.insert(field.clone(), value.clone()).is_none() { added += 1; } } - } - entry.touch(); - - let new_entry_size = memory::entry_size(key, &entry.value); - self.memory.adjust(old_entry_size, new_entry_size); + entry.touch(); + added + }); Ok(added) } @@ -1662,50 +1661,31 @@ impl Keyspace { self.remove_if_expired(key); - let is_new = !self.entries.contains_key(key); - - if !is_new && !matches!(self.entries[key].value, Value::Set(_)) { - return Err(WriteError::WrongType); - } + let is_new = self.ensure_collection_type(key, |v| matches!(v, Value::Set(_)))?; - // estimate memory increase before mutating let member_increase: usize = members .iter() .map(|m| m.len() + memory::HASHSET_MEMBER_OVERHEAD) .sum(); - let estimated_increase = if is_new { - memory::ENTRY_OVERHEAD + key.len() + memory::HASHSET_BASE_OVERHEAD + member_increase - } else { - member_increase - }; - if !self.enforce_memory_limit(estimated_increase) { - return Err(WriteError::OutOfMemory); - } + self.reserve_memory(is_new, key, memory::HASHSET_BASE_OVERHEAD, member_increase)?; if is_new { - let value = Value::Set(std::collections::HashSet::new()); - self.memory.add(key, &value); - self.entries.insert(key.to_owned(), Entry::new(value, None)); + self.insert_empty(key, Value::Set(std::collections::HashSet::new())); } - let entry = self - .entries - .get_mut(key) - .expect("just inserted or verified"); - let old_entry_size = memory::entry_size(key, &entry.value); - - let mut added = 0; - if let Value::Set(ref mut set) = entry.value { + let added = self.track_size(key, |entry| { + let Value::Set(ref mut set) = entry.value else { + unreachable!("type verified by ensure_collection_type"); + }; + let mut added = 0; for member in members { if set.insert(member.clone()) { added += 1; } } - } - entry.touch(); - - let new_entry_size = memory::entry_size(key, &entry.value); - self.memory.adjust(old_entry_size, new_entry_size); + entry.touch(); + added + }); Ok(added) }