From dc7837949c969fb7776bd0e673e370afd5859555 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Thu, 5 Feb 2026 22:30:33 -0500 Subject: [PATCH] fix: add AOF persistence for hash and set commands critical bug fix: HSET, HDEL, HINCRBY, SADD, and SREM mutations were not being logged to the append-only file, causing data loss on restart. changes: - add AofRecord variants for HSet, HDel, HIncrBy, SAdd, SRem - implement serialization and deserialization for new record types - update recovery to replay hash and set mutations - extend to_aof_record() in shard.rs to emit new records - add round-trip tests for all new AOF record types - improve memory.rs documentation with platform notes - clarify MSET fan-out logic in connection.rs --- crates/ember-core/src/memory.rs | 10 ++ crates/ember-core/src/shard.rs | 156 +++++++++++++++++++++ crates/ember-persistence/src/aof.rs | 164 +++++++++++++++++++++++ crates/ember-persistence/src/recovery.rs | 62 +++++++++ crates/ember-server/src/connection.rs | 6 +- 5 files changed, 397 insertions(+), 1 deletion(-) diff --git a/crates/ember-core/src/memory.rs b/crates/ember-core/src/memory.rs index 3341d162..7de82691 100644 --- a/crates/ember-core/src/memory.rs +++ b/crates/ember-core/src/memory.rs @@ -3,6 +3,16 @@ //! Provides byte-level accounting of memory used by entries. Updated //! on every mutation so the engine can enforce memory limits and //! report stats without scanning the entire keyspace. +//! +//! # Platform notes +//! +//! Overhead constants are empirical estimates for 64-bit platforms (x86-64, +//! aarch64). On 32-bit systems these would be smaller; the effect is that +//! we'd overestimate memory usage, which triggers eviction earlier than +//! necessary but doesn't cause correctness issues. +//! +//! The constants assume Rust's standard library allocator. Custom allocators +//! (jemalloc, mimalloc) may have different per-allocation overhead. use crate::types::Value; diff --git a/crates/ember-core/src/shard.rs b/crates/ember-core/src/shard.rs index db531ce9..300e0975 100644 --- a/crates/ember-core/src/shard.rs +++ b/crates/ember-core/src/shard.rs @@ -764,6 +764,39 @@ fn to_aof_record(req: &ShardRequest, resp: &ShardResponse) -> Option milliseconds: *milliseconds, }) } + // Hash commands + (ShardRequest::HSet { key, fields }, ShardResponse::Len(_)) => Some(AofRecord::HSet { + key: key.clone(), + fields: fields.clone(), + }), + (ShardRequest::HDel { key, .. }, ShardResponse::HDelLen { removed, .. }) + if !removed.is_empty() => + { + Some(AofRecord::HDel { + key: key.clone(), + fields: removed.clone(), + }) + } + (ShardRequest::HIncrBy { key, field, delta }, ShardResponse::Integer(_)) => { + Some(AofRecord::HIncrBy { + key: key.clone(), + field: field.clone(), + delta: *delta, + }) + } + // Set commands + (ShardRequest::SAdd { key, members }, ShardResponse::Len(count)) if *count > 0 => { + Some(AofRecord::SAdd { + key: key.clone(), + members: members.clone(), + }) + } + (ShardRequest::SRem { key, members }, ShardResponse::Len(count)) if *count > 0 => { + Some(AofRecord::SRem { + key: key.clone(), + members: members.clone(), + }) + } _ => None, } } @@ -1501,4 +1534,127 @@ mod tests { _ => panic!("expected Scan response"), } } + + #[test] + fn to_aof_record_for_hset() { + let req = ShardRequest::HSet { + key: "h".into(), + fields: vec![("f1".into(), Bytes::from("v1"))], + }; + let resp = ShardResponse::Len(1); + let record = to_aof_record(&req, &resp).unwrap(); + match record { + AofRecord::HSet { key, fields } => { + assert_eq!(key, "h"); + assert_eq!(fields.len(), 1); + } + _ => panic!("expected HSet record"), + } + } + + #[test] + fn to_aof_record_for_hdel() { + let req = ShardRequest::HDel { + key: "h".into(), + fields: vec!["f1".into(), "f2".into()], + }; + let resp = ShardResponse::HDelLen { + count: 2, + removed: vec!["f1".into(), "f2".into()], + }; + let record = to_aof_record(&req, &resp).unwrap(); + match record { + AofRecord::HDel { key, fields } => { + assert_eq!(key, "h"); + assert_eq!(fields.len(), 2); + } + _ => panic!("expected HDel record"), + } + } + + #[test] + fn to_aof_record_skips_hdel_when_none_removed() { + let req = ShardRequest::HDel { + key: "h".into(), + fields: vec!["f1".into()], + }; + let resp = ShardResponse::HDelLen { + count: 0, + removed: vec![], + }; + assert!(to_aof_record(&req, &resp).is_none()); + } + + #[test] + fn to_aof_record_for_hincrby() { + let req = ShardRequest::HIncrBy { + key: "h".into(), + field: "counter".into(), + delta: 5, + }; + let resp = ShardResponse::Integer(10); + let record = to_aof_record(&req, &resp).unwrap(); + match record { + AofRecord::HIncrBy { key, field, delta } => { + assert_eq!(key, "h"); + assert_eq!(field, "counter"); + assert_eq!(delta, 5); + } + _ => panic!("expected HIncrBy record"), + } + } + + #[test] + fn to_aof_record_for_sadd() { + let req = ShardRequest::SAdd { + key: "s".into(), + members: vec!["m1".into(), "m2".into()], + }; + let resp = ShardResponse::Len(2); + let record = to_aof_record(&req, &resp).unwrap(); + match record { + AofRecord::SAdd { key, members } => { + assert_eq!(key, "s"); + assert_eq!(members.len(), 2); + } + _ => panic!("expected SAdd record"), + } + } + + #[test] + fn to_aof_record_skips_sadd_when_none_added() { + let req = ShardRequest::SAdd { + key: "s".into(), + members: vec!["m1".into()], + }; + let resp = ShardResponse::Len(0); + assert!(to_aof_record(&req, &resp).is_none()); + } + + #[test] + fn to_aof_record_for_srem() { + let req = ShardRequest::SRem { + key: "s".into(), + members: vec!["m1".into()], + }; + let resp = ShardResponse::Len(1); + let record = to_aof_record(&req, &resp).unwrap(); + match record { + AofRecord::SRem { key, members } => { + assert_eq!(key, "s"); + assert_eq!(members.len(), 1); + } + _ => panic!("expected SRem record"), + } + } + + #[test] + fn to_aof_record_skips_srem_when_none_removed() { + let req = ShardRequest::SRem { + key: "s".into(), + members: vec!["m1".into()], + }; + let resp = ShardResponse::Len(0); + assert!(to_aof_record(&req, &resp).is_none()); + } } diff --git a/crates/ember-persistence/src/aof.rs b/crates/ember-persistence/src/aof.rs index 2d5bfde1..11aeb62e 100644 --- a/crates/ember-persistence/src/aof.rs +++ b/crates/ember-persistence/src/aof.rs @@ -49,6 +49,11 @@ const TAG_PERSIST: u8 = 10; const TAG_PEXPIRE: u8 = 11; const TAG_INCR: u8 = 12; const TAG_DECR: u8 = 13; +const TAG_HSET: u8 = 14; +const TAG_HDEL: u8 = 15; +const TAG_HINCRBY: u8 = 16; +const TAG_SADD: u8 = 17; +const TAG_SREM: u8 = 18; /// A single mutation record stored in the AOF. #[derive(Debug, Clone, PartialEq)] @@ -86,6 +91,23 @@ pub enum AofRecord { Incr { key: String }, /// DECR key. Decr { key: String }, + /// HSET key field value [field value ...]. + HSet { + key: String, + fields: Vec<(String, Bytes)>, + }, + /// HDEL key field [field ...]. + HDel { key: String, fields: Vec }, + /// HINCRBY key field delta. + HIncrBy { + key: String, + field: String, + delta: i64, + }, + /// SADD key member [member ...]. + SAdd { key: String, members: Vec }, + /// SREM key member [member ...]. + SRem { key: String, members: Vec }, } impl AofRecord { @@ -170,6 +192,45 @@ impl AofRecord { format::write_u8(&mut buf, TAG_DECR).expect("vec write"); format::write_bytes(&mut buf, key.as_bytes()).expect("vec write"); } + AofRecord::HSet { key, fields } => { + format::write_u8(&mut buf, TAG_HSET).expect("vec write"); + format::write_bytes(&mut buf, key.as_bytes()).expect("vec write"); + format::write_u32(&mut buf, fields.len() as u32).expect("vec write"); + for (field, value) in fields { + format::write_bytes(&mut buf, field.as_bytes()).expect("vec write"); + format::write_bytes(&mut buf, value).expect("vec write"); + } + } + AofRecord::HDel { key, fields } => { + format::write_u8(&mut buf, TAG_HDEL).expect("vec write"); + format::write_bytes(&mut buf, key.as_bytes()).expect("vec write"); + format::write_u32(&mut buf, fields.len() as u32).expect("vec write"); + for field in fields { + format::write_bytes(&mut buf, field.as_bytes()).expect("vec write"); + } + } + AofRecord::HIncrBy { key, field, delta } => { + format::write_u8(&mut buf, TAG_HINCRBY).expect("vec write"); + format::write_bytes(&mut buf, key.as_bytes()).expect("vec write"); + format::write_bytes(&mut buf, field.as_bytes()).expect("vec write"); + format::write_i64(&mut buf, *delta).expect("vec write"); + } + AofRecord::SAdd { key, members } => { + format::write_u8(&mut buf, TAG_SADD).expect("vec write"); + format::write_bytes(&mut buf, key.as_bytes()).expect("vec write"); + format::write_u32(&mut buf, members.len() as u32).expect("vec write"); + for member in members { + format::write_bytes(&mut buf, member.as_bytes()).expect("vec write"); + } + } + AofRecord::SRem { key, members } => { + format::write_u8(&mut buf, TAG_SREM).expect("vec write"); + format::write_bytes(&mut buf, key.as_bytes()).expect("vec write"); + format::write_u32(&mut buf, members.len() as u32).expect("vec write"); + for member in members { + format::write_bytes(&mut buf, member.as_bytes()).expect("vec write"); + } + } } buf } @@ -256,6 +317,50 @@ impl AofRecord { let key = read_string(&mut cursor, "key")?; Ok(AofRecord::Decr { key }) } + TAG_HSET => { + let key = read_string(&mut cursor, "key")?; + let count = format::read_u32(&mut cursor)?; + let mut fields = Vec::with_capacity(count as usize); + for _ in 0..count { + let field = read_string(&mut cursor, "field")?; + let value = Bytes::from(format::read_bytes(&mut cursor)?); + fields.push((field, value)); + } + Ok(AofRecord::HSet { key, fields }) + } + TAG_HDEL => { + let key = read_string(&mut cursor, "key")?; + let count = format::read_u32(&mut cursor)?; + let mut fields = Vec::with_capacity(count as usize); + for _ in 0..count { + fields.push(read_string(&mut cursor, "field")?); + } + Ok(AofRecord::HDel { key, fields }) + } + TAG_HINCRBY => { + let key = read_string(&mut cursor, "key")?; + let field = read_string(&mut cursor, "field")?; + let delta = format::read_i64(&mut cursor)?; + Ok(AofRecord::HIncrBy { key, field, delta }) + } + TAG_SADD => { + let key = read_string(&mut cursor, "key")?; + let count = format::read_u32(&mut cursor)?; + let mut members = Vec::with_capacity(count as usize); + for _ in 0..count { + members.push(read_string(&mut cursor, "member")?); + } + Ok(AofRecord::SAdd { key, members }) + } + TAG_SREM => { + let key = read_string(&mut cursor, "key")?; + let count = format::read_u32(&mut cursor)?; + let mut members = Vec::with_capacity(count as usize); + for _ in 0..count { + members.push(read_string(&mut cursor, "member")?); + } + Ok(AofRecord::SRem { key, members }) + } _ => Err(FormatError::UnknownTag(tag)), } } @@ -893,4 +998,63 @@ mod tests { let p = aof_path(Path::new("/data"), 3); assert_eq!(p, PathBuf::from("/data/shard-3.aof")); } + + #[test] + fn record_round_trip_hset() { + let rec = AofRecord::HSet { + key: "hash".into(), + fields: vec![ + ("f1".into(), Bytes::from("v1")), + ("f2".into(), Bytes::from("v2")), + ], + }; + let bytes = rec.to_bytes(); + let decoded = AofRecord::from_bytes(&bytes).unwrap(); + assert_eq!(rec, decoded); + } + + #[test] + fn record_round_trip_hdel() { + let rec = AofRecord::HDel { + key: "hash".into(), + fields: vec!["f1".into(), "f2".into()], + }; + let bytes = rec.to_bytes(); + let decoded = AofRecord::from_bytes(&bytes).unwrap(); + assert_eq!(rec, decoded); + } + + #[test] + fn record_round_trip_hincrby() { + let rec = AofRecord::HIncrBy { + key: "hash".into(), + field: "counter".into(), + delta: -42, + }; + let bytes = rec.to_bytes(); + let decoded = AofRecord::from_bytes(&bytes).unwrap(); + assert_eq!(rec, decoded); + } + + #[test] + fn record_round_trip_sadd() { + let rec = AofRecord::SAdd { + key: "set".into(), + members: vec!["m1".into(), "m2".into(), "m3".into()], + }; + let bytes = rec.to_bytes(); + let decoded = AofRecord::from_bytes(&bytes).unwrap(); + assert_eq!(rec, decoded); + } + + #[test] + fn record_round_trip_srem() { + let rec = AofRecord::SRem { + key: "set".into(), + members: vec!["m1".into()], + }; + let bytes = rec.to_bytes(); + let decoded = AofRecord::from_bytes(&bytes).unwrap(); + assert_eq!(rec, decoded); + } } diff --git a/crates/ember-persistence/src/recovery.rs b/crates/ember-persistence/src/recovery.rs index 14f7bdb6..1d4e6641 100644 --- a/crates/ember-persistence/src/recovery.rs +++ b/crates/ember-persistence/src/recovery.rs @@ -299,6 +299,68 @@ fn replay_aof( AofRecord::Decr { key } => { apply_incr(map, key, -1); } + AofRecord::HSet { key, fields } => { + let entry = map + .entry(key) + .or_insert_with(|| (RecoveredValue::Hash(HashMap::new()), None)); + if let RecoveredValue::Hash(ref mut hash) = entry.0 { + for (field, value) in fields { + hash.insert(field, value); + } + } + } + AofRecord::HDel { key, fields } => { + if let Some(entry) = map.get_mut(&key) { + if let RecoveredValue::Hash(ref mut hash) = entry.0 { + for field in fields { + hash.remove(&field); + } + if hash.is_empty() { + map.remove(&key); + count += 1; + continue; + } + } + } + } + AofRecord::HIncrBy { key, field, delta } => { + let entry = map + .entry(key) + .or_insert_with(|| (RecoveredValue::Hash(HashMap::new()), None)); + if let RecoveredValue::Hash(ref mut hash) = entry.0 { + let current: i64 = hash + .get(&field) + .and_then(|v| std::str::from_utf8(v).ok()) + .and_then(|s| s.parse().ok()) + .unwrap_or(0); + let new_val = current.saturating_add(delta); + hash.insert(field, Bytes::from(new_val.to_string())); + } + } + AofRecord::SAdd { key, members } => { + let entry = map + .entry(key) + .or_insert_with(|| (RecoveredValue::Set(HashSet::new()), None)); + if let RecoveredValue::Set(ref mut set) = entry.0 { + for member in members { + set.insert(member); + } + } + } + AofRecord::SRem { key, members } => { + if let Some(entry) = map.get_mut(&key) { + if let RecoveredValue::Set(ref mut set) = entry.0 { + for member in members { + set.remove(&member); + } + if set.is_empty() { + map.remove(&key); + count += 1; + continue; + } + } + } + } } count += 1; } diff --git a/crates/ember-server/src/connection.rs b/crates/ember-server/src/connection.rs index e50b358e..a9a00990 100644 --- a/crates/ember-server/src/connection.rs +++ b/crates/ember-server/src/connection.rs @@ -260,12 +260,16 @@ async fn execute(cmd: Command, engine: &Engine) -> Frame { } Command::MSet { pairs } => { - // fan out individual SET requests — MSET always succeeds (or OOMs) + // Fan out individual SET requests — MSET always succeeds (or OOMs). + // We build a HashMap for O(1) value lookups during routing. If there + // are duplicate keys in pairs, the HashMap keeps the last value, which + // matches Redis semantics (last write wins). let keys: Vec = pairs.iter().map(|(k, _)| k.clone()).collect(); let values: std::collections::HashMap = pairs.into_iter().collect(); match engine .route_multi(&keys, |k| { + // Safe: k comes from keys, which came from pairs, so it exists in values. let value = values.get(&k).cloned().unwrap_or_default(); ShardRequest::Set { key: k,