From 0e9c83578a8a109c94e703a684e92b8730c26558 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Tue, 10 Feb 2026 21:29:59 -0500 Subject: [PATCH 1/6] thread schema registry into shards and recover schemas from AOF schemas were lost on restart because ProtoRegister records were skipped during AOF replay. now recovery collects them (last-wins dedup) and the shard restores them into the shared registry on startup. changes: - RecoveryResult gains a `schemas` vec (behind protobuf feature) - replay_aof collects ProtoRegister records into a schema map - spawn_shard/run_shard accept an optional SharedSchemaRegistry - after recovery, schemas are restored via registry.restore() - engine passes its schema_registry clone to each shard --- crates/ember-core/src/engine.rs | 2 + crates/ember-core/src/shard.rs | 74 ++++++++++++++++++++++-- crates/ember-persistence/src/recovery.rs | 68 ++++++++++++++++++++-- 3 files changed, 134 insertions(+), 10 deletions(-) diff --git a/crates/ember-core/src/engine.rs b/crates/ember-core/src/engine.rs index 480e3302..400c496a 100644 --- a/crates/ember-core/src/engine.rs +++ b/crates/ember-core/src/engine.rs @@ -71,6 +71,8 @@ impl Engine { shard_config, config.persistence.clone(), Some(drop_handle.clone()), + #[cfg(feature = "protobuf")] + config.schema_registry.clone(), ) }) .collect(); diff --git a/crates/ember-core/src/shard.rs b/crates/ember-core/src/shard.rs index 8354a193..eca1ff30 100644 --- a/crates/ember-core/src/shard.rs +++ b/crates/ember-core/src/shard.rs @@ -407,9 +407,17 @@ pub fn spawn_shard( config: ShardConfig, persistence: Option, drop_handle: Option, + #[cfg(feature = "protobuf")] schema_registry: Option, ) -> ShardHandle { let (tx, rx) = mpsc::channel(buffer); - tokio::spawn(run_shard(rx, config, persistence, drop_handle)); + tokio::spawn(run_shard( + rx, + config, + persistence, + drop_handle, + #[cfg(feature = "protobuf")] + schema_registry, + )); ShardHandle { tx } } @@ -420,6 +428,7 @@ async fn run_shard( config: ShardConfig, persistence: Option, drop_handle: Option, + #[cfg(feature = "protobuf")] schema_registry: Option, ) { let shard_id = config.shard_id; let mut keyspace = Keyspace::with_config(config); @@ -466,6 +475,24 @@ async fn run_shard( "recovered shard state" ); } + + // restore schemas found in the AOF into the shared registry + #[cfg(feature = "protobuf")] + if let Some(ref registry) = schema_registry { + if !result.schemas.is_empty() { + if let Ok(mut reg) = registry.write() { + let schema_count = result.schemas.len(); + for (name, descriptor) in result.schemas { + reg.restore(name, descriptor); + } + info!( + shard_id, + schemas = schema_count, + "restored schemas from AOF" + ); + } + } + } } // -- AOF writer -- @@ -1269,7 +1296,14 @@ mod tests { #[tokio::test] async fn shard_round_trip() { - let handle = spawn_shard(16, ShardConfig::default(), None, None); + let handle = spawn_shard( + 16, + ShardConfig::default(), + None, + None, + #[cfg(feature = "protobuf")] + None, + ); let resp = handle .send(ShardRequest::Set { @@ -1299,7 +1333,14 @@ mod tests { #[tokio::test] async fn expired_key_through_shard() { - let handle = spawn_shard(16, ShardConfig::default(), None, None); + let handle = spawn_shard( + 16, + ShardConfig::default(), + None, + None, + #[cfg(feature = "protobuf")] + None, + ); handle .send(ShardRequest::Set { @@ -1323,7 +1364,14 @@ mod tests { #[tokio::test] async fn active_expiration_cleans_up_without_access() { - let handle = spawn_shard(16, ShardConfig::default(), None, None); + let handle = spawn_shard( + 16, + ShardConfig::default(), + None, + None, + #[cfg(feature = "protobuf")] + None, + ); // set a key with a short TTL handle @@ -1389,7 +1437,14 @@ mod tests { // write some keys then trigger a snapshot { - let handle = spawn_shard(16, config.clone(), Some(pcfg.clone()), None); + let handle = spawn_shard( + 16, + config.clone(), + Some(pcfg.clone()), + None, + #[cfg(feature = "protobuf")] + None, + ); handle .send(ShardRequest::Set { key: "a".into(), @@ -1430,7 +1485,14 @@ mod tests { // start a new shard with the same config — should recover { - let handle = spawn_shard(16, config, Some(pcfg), None); + let handle = spawn_shard( + 16, + config, + Some(pcfg), + None, + #[cfg(feature = "protobuf")] + None, + ); // give it a moment to recover tokio::time::sleep(Duration::from_millis(50)).await; diff --git a/crates/ember-persistence/src/recovery.rs b/crates/ember-persistence/src/recovery.rs index c106e0cf..97e325b0 100644 --- a/crates/ember-persistence/src/recovery.rs +++ b/crates/ember-persistence/src/recovery.rs @@ -77,6 +77,10 @@ pub struct RecoveryResult { pub loaded_snapshot: bool, /// Whether an AOF was replayed. pub replayed_aof: bool, + /// Schemas found in the AOF, deduplicated by name (last wins). + /// Each entry is `(schema_name, descriptor_bytes)`. + #[cfg(feature = "protobuf")] + pub schemas: Vec<(String, Bytes)>, } /// Recovers a shard's state from snapshot and/or AOF files. @@ -110,6 +114,8 @@ fn recover_shard_impl( let mut map: HashMap = HashMap::new(); let mut loaded_snapshot = false; let mut replayed_aof = false; + #[cfg(feature = "protobuf")] + let mut schema_map: HashMap = HashMap::new(); // step 1: load snapshot let snap_path = snapshot::snapshot_path(data_dir, shard_id); @@ -130,7 +136,13 @@ fn recover_shard_impl( // step 2: replay AOF let aof_path = aof::aof_path(data_dir, shard_id); if aof_path.exists() { - match replay_aof(&aof_path, &mut map, encryption_key) { + match replay_aof( + &aof_path, + &mut map, + encryption_key, + #[cfg(feature = "protobuf")] + &mut schema_map, + ) { Ok(count) => { if count > 0 { replayed_aof = true; @@ -164,6 +176,8 @@ fn recover_shard_impl( entries, loaded_snapshot, replayed_aof, + #[cfg(feature = "protobuf")] + schemas: schema_map.into_iter().collect(), } } @@ -218,6 +232,7 @@ fn replay_aof( path: &Path, map: &mut HashMap, #[allow(unused_variables)] encryption_key: Option>, + #[cfg(feature = "protobuf")] schema_map: &mut HashMap, ) -> Result { #[cfg(feature = "encryption")] let mut reader = if let Some(key) = encryption_key { @@ -436,9 +451,10 @@ fn replay_aof( map.insert(key, (RecoveredValue::Proto { type_name, data }, expire_ms)); } #[cfg(feature = "protobuf")] - AofRecord::ProtoRegister { .. } => { - // schema registration is handled separately by the engine, - // not in the per-shard recovery map + AofRecord::ProtoRegister { name, descriptor } => { + // last-wins: if the same schema name appears multiple times + // in the AOF, the final registration is the one we keep. + schema_map.insert(name, descriptor); } } count += 1; @@ -869,4 +885,48 @@ mod tests { assert_eq!(result.entries.len(), 1); assert!(result.entries[0].ttl.is_some()); } + + #[cfg(feature = "protobuf")] + #[test] + fn proto_schemas_recovered_from_aof() { + let dir = temp_dir(); + let path = aof::aof_path(dir.path(), 0); + + { + let mut writer = AofWriter::open(&path).unwrap(); + writer + .write_record(&AofRecord::ProtoRegister { + name: "users".into(), + descriptor: Bytes::from("fake-descriptor-a"), + }) + .unwrap(); + // a proto value that depends on the schema + writer + .write_record(&AofRecord::ProtoSet { + key: "user:1".into(), + type_name: "test.User".into(), + data: Bytes::from("some-proto-data"), + expire_ms: -1, + }) + .unwrap(); + // re-registration of same schema (last wins) + writer + .write_record(&AofRecord::ProtoRegister { + name: "users".into(), + descriptor: Bytes::from("fake-descriptor-b"), + }) + .unwrap(); + writer.sync().unwrap(); + } + + let result = recover_shard(dir.path(), 0); + assert!(result.replayed_aof); + assert_eq!(result.entries.len(), 1); + + // schemas should be collected with last-wins dedup + assert_eq!(result.schemas.len(), 1); + let (name, desc) = &result.schemas[0]; + assert_eq!(name, "users"); + assert_eq!(desc, &Bytes::from("fake-descriptor-b")); + } } From 26ed962067b687ac6d1fa3ef6b87214d5b1e1a58 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Tue, 10 Feb 2026 21:35:28 -0500 Subject: [PATCH 2/6] atomic SETFIELD/DELFIELD via shard-local mutation the previous read-modify-write pattern for PROTO.SETFIELD and PROTO.DELFIELD spanned two shard roundtrips, leaving a window where another client could overwrite the key between the read and write. moving the mutation into the shard's single-threaded dispatch makes it inherently atomic. changes: - add ProtoSetField/ProtoDelField ShardRequest variants - add ProtoFieldUpdated ShardResponse variant - dispatch_proto_field_op helper reads, mutates, and writes back within a single dispatch call - to_aof_record maps field ops to full ProtoSet records - connection.rs and concurrent_handler.rs simplified to just route the request and match on the response --- crates/ember-core/src/shard.rs | 218 ++++++++++++++---- crates/ember-server/src/concurrent_handler.rs | 93 ++------ crates/ember-server/src/connection.rs | 85 ++----- 3 files changed, 218 insertions(+), 178 deletions(-) diff --git a/crates/ember-core/src/shard.rs b/crates/ember-core/src/shard.rs index eca1ff30..628d41dc 100644 --- a/crates/ember-core/src/shard.rs +++ b/crates/ember-core/src/shard.rs @@ -288,6 +288,21 @@ pub enum ShardRequest { name: String, descriptor: Bytes, }, + /// Atomically reads a proto value, sets a field, and writes it back. + /// Runs entirely within the shard's single-threaded dispatch. + #[cfg(feature = "protobuf")] + ProtoSetField { + key: String, + field_path: String, + value: String, + }, + /// Atomically reads a proto value, clears a field, and writes it back. + /// Runs entirely within the shard's single-threaded dispatch. + #[cfg(feature = "protobuf")] + ProtoDelField { + key: String, + field_path: String, + }, } /// The shard's response to a request. @@ -350,6 +365,14 @@ pub enum ShardResponse { /// PROTO.TYPE result: message type name or None. #[cfg(feature = "protobuf")] ProtoTypeName(Option), + /// Result of an atomic SETFIELD/DELFIELD: carries the updated value + /// for AOF persistence. + #[cfg(feature = "protobuf")] + ProtoFieldUpdated { + type_name: String, + data: Bytes, + expire: Option, + }, } /// A request bundled with its reply channel. @@ -536,7 +559,12 @@ async fn run_shard( match msg { Some(msg) => { let request_kind = describe_request(&msg.request); - let response = dispatch(&mut keyspace, &msg.request); + let response = dispatch( + &mut keyspace, + &msg.request, + #[cfg(feature = "protobuf")] + &schema_registry, + ); // write AOF record for successful mutations if let Some(ref mut writer) = aof_writer { @@ -627,7 +655,11 @@ fn describe_request(req: &ShardRequest) -> RequestKind { } /// Executes a single request against the keyspace. -fn dispatch(ks: &mut Keyspace, req: &ShardRequest) -> ShardResponse { +fn dispatch( + ks: &mut Keyspace, + req: &ShardRequest, + #[cfg(feature = "protobuf")] schema_registry: &Option, +) -> ShardResponse { match req { ShardRequest::Get { key } => match ks.get(key) { Ok(val) => ShardResponse::Value(val), @@ -914,6 +946,30 @@ fn dispatch(ks: &mut Keyspace, req: &ShardRequest) -> ShardResponse { // is written by the to_aof_record path after dispatch returns Ok. #[cfg(feature = "protobuf")] ShardRequest::ProtoRegisterAof { .. } => ShardResponse::Ok, + #[cfg(feature = "protobuf")] + ShardRequest::ProtoSetField { + key, + field_path, + value, + } => dispatch_proto_field_op(ks, schema_registry, key, |reg, type_name, data, ttl| { + let new_data = reg.set_field(type_name, data, field_path, value)?; + Ok(ShardResponse::ProtoFieldUpdated { + type_name: type_name.to_owned(), + data: new_data, + expire: ttl, + }) + }), + #[cfg(feature = "protobuf")] + ShardRequest::ProtoDelField { key, field_path } => { + dispatch_proto_field_op(ks, schema_registry, key, |reg, type_name, data, ttl| { + let new_data = reg.clear_field(type_name, data, field_path)?; + Ok(ShardResponse::ProtoFieldUpdated { + type_name: type_name.to_owned(), + data: new_data, + expire: ttl, + }) + }) + } // snapshot/rewrite/flush_async are handled in the main loop, not here ShardRequest::Snapshot | ShardRequest::RewriteAof | ShardRequest::FlushDbAsync => { ShardResponse::Ok @@ -921,6 +977,60 @@ fn dispatch(ks: &mut Keyspace, req: &ShardRequest) -> ShardResponse { } } +/// Shared logic for atomic proto field operations (SETFIELD/DELFIELD). +/// +/// Reads the proto value, acquires the schema registry, calls the +/// provided mutation closure, then writes the result back to the keyspace +/// — all within the single-threaded shard dispatch. +#[cfg(feature = "protobuf")] +fn dispatch_proto_field_op( + ks: &mut Keyspace, + schema_registry: &Option, + key: &str, + mutate: F, +) -> ShardResponse +where + F: FnOnce( + &crate::schema::SchemaRegistry, + &str, + &[u8], + Option, + ) -> Result, +{ + let registry = match schema_registry { + Some(r) => r, + None => return ShardResponse::Err("protobuf support is not enabled".into()), + }; + + let (type_name, data, remaining_ttl) = match ks.proto_get(key) { + Ok(Some(tuple)) => tuple, + Ok(None) => return ShardResponse::Value(None), + Err(_) => return ShardResponse::WrongType, + }; + + let reg = match registry.read() { + Ok(r) => r, + Err(_) => return ShardResponse::Err("schema registry lock poisoned".into()), + }; + + let resp = match mutate(®, &type_name, &data, remaining_ttl) { + Ok(r) => r, + Err(e) => return ShardResponse::Err(e.to_string()), + }; + + // write the updated value back, preserving the original TTL + if let ShardResponse::ProtoFieldUpdated { + ref type_name, + ref data, + expire, + } = resp + { + ks.proto_set(key.to_owned(), type_name.clone(), data.clone(), expire); + } + + resp +} + /// Converts a successful mutation request+response pair into an AOF record. /// Returns None for non-mutation requests or failed mutations. fn to_aof_record(req: &ShardRequest, resp: &ShardResponse) -> Option { @@ -1083,6 +1193,24 @@ fn to_aof_record(req: &ShardRequest, resp: &ShardResponse) -> Option descriptor: descriptor.clone(), }) } + // atomic field ops persist as a full ProtoSet (the whole re-encoded value) + #[cfg(feature = "protobuf")] + ( + ShardRequest::ProtoSetField { key, .. } | ShardRequest::ProtoDelField { key, .. }, + ShardResponse::ProtoFieldUpdated { + type_name, + data, + expire, + }, + ) => { + let expire_ms = expire.map(|d| d.as_millis() as i64).unwrap_or(-1); + Some(AofRecord::ProtoSet { + key: key.clone(), + type_name: type_name.clone(), + data: data.clone(), + expire_ms, + }) + } _ => None, } } @@ -1210,11 +1338,21 @@ fn write_snapshot( mod tests { use super::*; + /// Test helper: dispatch without a schema registry. + fn test_dispatch(ks: &mut Keyspace, req: &ShardRequest) -> ShardResponse { + dispatch( + ks, + req, + #[cfg(feature = "protobuf")] + &None, + ) + } + #[test] fn dispatch_set_and_get() { let mut ks = Keyspace::new(); - let resp = dispatch( + let resp = test_dispatch( &mut ks, &ShardRequest::Set { key: "k".into(), @@ -1226,7 +1364,7 @@ mod tests { ); assert!(matches!(resp, ShardResponse::Ok)); - let resp = dispatch(&mut ks, &ShardRequest::Get { key: "k".into() }); + let resp = test_dispatch(&mut ks, &ShardRequest::Get { key: "k".into() }); match resp { ShardResponse::Value(Some(Value::String(data))) => { assert_eq!(data, Bytes::from("v")); @@ -1238,7 +1376,7 @@ mod tests { #[test] fn dispatch_get_missing() { let mut ks = Keyspace::new(); - let resp = dispatch(&mut ks, &ShardRequest::Get { key: "nope".into() }); + let resp = test_dispatch(&mut ks, &ShardRequest::Get { key: "nope".into() }); assert!(matches!(resp, ShardResponse::Value(None))); } @@ -1247,10 +1385,10 @@ mod tests { let mut ks = Keyspace::new(); ks.set("key".into(), Bytes::from("val"), None); - let resp = dispatch(&mut ks, &ShardRequest::Del { key: "key".into() }); + let resp = test_dispatch(&mut ks, &ShardRequest::Del { key: "key".into() }); assert!(matches!(resp, ShardResponse::Bool(true))); - let resp = dispatch(&mut ks, &ShardRequest::Del { key: "key".into() }); + let resp = test_dispatch(&mut ks, &ShardRequest::Del { key: "key".into() }); assert!(matches!(resp, ShardResponse::Bool(false))); } @@ -1259,10 +1397,10 @@ mod tests { let mut ks = Keyspace::new(); ks.set("yes".into(), Bytes::from("here"), None); - let resp = dispatch(&mut ks, &ShardRequest::Exists { key: "yes".into() }); + let resp = test_dispatch(&mut ks, &ShardRequest::Exists { key: "yes".into() }); assert!(matches!(resp, ShardResponse::Bool(true))); - let resp = dispatch(&mut ks, &ShardRequest::Exists { key: "no".into() }); + let resp = test_dispatch(&mut ks, &ShardRequest::Exists { key: "no".into() }); assert!(matches!(resp, ShardResponse::Bool(false))); } @@ -1271,7 +1409,7 @@ mod tests { let mut ks = Keyspace::new(); ks.set("key".into(), Bytes::from("val"), None); - let resp = dispatch( + let resp = test_dispatch( &mut ks, &ShardRequest::Expire { key: "key".into(), @@ -1280,7 +1418,7 @@ mod tests { ); assert!(matches!(resp, ShardResponse::Bool(true))); - let resp = dispatch(&mut ks, &ShardRequest::Ttl { key: "key".into() }); + let resp = test_dispatch(&mut ks, &ShardRequest::Ttl { key: "key".into() }); match resp { ShardResponse::Ttl(TtlResult::Seconds(s)) => assert!((58..=60).contains(&s)), other => panic!("expected Ttl(Seconds), got {other:?}"), @@ -1290,7 +1428,7 @@ mod tests { #[test] fn dispatch_ttl_missing() { let mut ks = Keyspace::new(); - let resp = dispatch(&mut ks, &ShardRequest::Ttl { key: "gone".into() }); + let resp = test_dispatch(&mut ks, &ShardRequest::Ttl { key: "gone".into() }); assert!(matches!(resp, ShardResponse::Ttl(TtlResult::NotFound))); } @@ -1577,7 +1715,7 @@ mod tests { #[test] fn dispatch_incr_new_key() { let mut ks = Keyspace::new(); - let resp = dispatch(&mut ks, &ShardRequest::Incr { key: "c".into() }); + let resp = test_dispatch(&mut ks, &ShardRequest::Incr { key: "c".into() }); assert!(matches!(resp, ShardResponse::Integer(1))); } @@ -1585,7 +1723,7 @@ mod tests { fn dispatch_decr_existing() { let mut ks = Keyspace::new(); ks.set("n".into(), Bytes::from("10"), None); - let resp = dispatch(&mut ks, &ShardRequest::Decr { key: "n".into() }); + let resp = test_dispatch(&mut ks, &ShardRequest::Decr { key: "n".into() }); assert!(matches!(resp, ShardResponse::Integer(9))); } @@ -1593,7 +1731,7 @@ mod tests { fn dispatch_incr_non_integer() { let mut ks = Keyspace::new(); ks.set("s".into(), Bytes::from("hello"), None); - let resp = dispatch(&mut ks, &ShardRequest::Incr { key: "s".into() }); + let resp = test_dispatch(&mut ks, &ShardRequest::Incr { key: "s".into() }); assert!(matches!(resp, ShardResponse::Err(_))); } @@ -1601,7 +1739,7 @@ mod tests { fn dispatch_incrby() { let mut ks = Keyspace::new(); ks.set("n".into(), Bytes::from("10"), None); - let resp = dispatch( + let resp = test_dispatch( &mut ks, &ShardRequest::IncrBy { key: "n".into(), @@ -1615,7 +1753,7 @@ mod tests { fn dispatch_decrby() { let mut ks = Keyspace::new(); ks.set("n".into(), Bytes::from("10"), None); - let resp = dispatch( + let resp = test_dispatch( &mut ks, &ShardRequest::DecrBy { key: "n".into(), @@ -1628,7 +1766,7 @@ mod tests { #[test] fn dispatch_incrby_new_key() { let mut ks = Keyspace::new(); - let resp = dispatch( + let resp = test_dispatch( &mut ks, &ShardRequest::IncrBy { key: "new".into(), @@ -1642,7 +1780,7 @@ mod tests { fn dispatch_incrbyfloat() { let mut ks = Keyspace::new(); ks.set("n".into(), Bytes::from("10.5"), None); - let resp = dispatch( + let resp = test_dispatch( &mut ks, &ShardRequest::IncrByFloat { key: "n".into(), @@ -1662,7 +1800,7 @@ mod tests { fn dispatch_append() { let mut ks = Keyspace::new(); ks.set("k".into(), Bytes::from("hello"), None); - let resp = dispatch( + let resp = test_dispatch( &mut ks, &ShardRequest::Append { key: "k".into(), @@ -1676,14 +1814,14 @@ mod tests { fn dispatch_strlen() { let mut ks = Keyspace::new(); ks.set("k".into(), Bytes::from("hello"), None); - let resp = dispatch(&mut ks, &ShardRequest::Strlen { key: "k".into() }); + let resp = test_dispatch(&mut ks, &ShardRequest::Strlen { key: "k".into() }); assert!(matches!(resp, ShardResponse::Len(5))); } #[test] fn dispatch_strlen_missing() { let mut ks = Keyspace::new(); - let resp = dispatch(&mut ks, &ShardRequest::Strlen { key: "nope".into() }); + let resp = test_dispatch(&mut ks, &ShardRequest::Strlen { key: "nope".into() }); assert!(matches!(resp, ShardResponse::Len(0))); } @@ -1707,7 +1845,7 @@ mod tests { #[test] fn dispatch_incrbyfloat_new_key() { let mut ks = Keyspace::new(); - let resp = dispatch( + let resp = test_dispatch( &mut ks, &ShardRequest::IncrByFloat { key: "new".into(), @@ -1782,17 +1920,17 @@ mod tests { Some(Duration::from_secs(60)), ); - let resp = dispatch(&mut ks, &ShardRequest::Persist { key: "key".into() }); + let resp = test_dispatch(&mut ks, &ShardRequest::Persist { key: "key".into() }); assert!(matches!(resp, ShardResponse::Bool(true))); - let resp = dispatch(&mut ks, &ShardRequest::Ttl { key: "key".into() }); + let resp = test_dispatch(&mut ks, &ShardRequest::Ttl { key: "key".into() }); assert!(matches!(resp, ShardResponse::Ttl(TtlResult::NoExpiry))); } #[test] fn dispatch_persist_missing_key() { let mut ks = Keyspace::new(); - let resp = dispatch(&mut ks, &ShardRequest::Persist { key: "nope".into() }); + let resp = test_dispatch(&mut ks, &ShardRequest::Persist { key: "nope".into() }); assert!(matches!(resp, ShardResponse::Bool(false))); } @@ -1805,7 +1943,7 @@ mod tests { Some(Duration::from_secs(60)), ); - let resp = dispatch(&mut ks, &ShardRequest::Pttl { key: "key".into() }); + let resp = test_dispatch(&mut ks, &ShardRequest::Pttl { key: "key".into() }); match resp { ShardResponse::Ttl(TtlResult::Milliseconds(ms)) => { assert!(ms > 59_000 && ms <= 60_000); @@ -1817,7 +1955,7 @@ mod tests { #[test] fn dispatch_pttl_missing() { let mut ks = Keyspace::new(); - let resp = dispatch(&mut ks, &ShardRequest::Pttl { key: "nope".into() }); + let resp = test_dispatch(&mut ks, &ShardRequest::Pttl { key: "nope".into() }); assert!(matches!(resp, ShardResponse::Ttl(TtlResult::NotFound))); } @@ -1826,7 +1964,7 @@ mod tests { let mut ks = Keyspace::new(); ks.set("key".into(), Bytes::from("val"), None); - let resp = dispatch( + let resp = test_dispatch( &mut ks, &ShardRequest::Pexpire { key: "key".into(), @@ -1835,7 +1973,7 @@ mod tests { ); assert!(matches!(resp, ShardResponse::Bool(true))); - let resp = dispatch(&mut ks, &ShardRequest::Pttl { key: "key".into() }); + let resp = test_dispatch(&mut ks, &ShardRequest::Pttl { key: "key".into() }); match resp { ShardResponse::Ttl(TtlResult::Milliseconds(ms)) => { assert!(ms > 4000 && ms <= 5000); @@ -1889,7 +2027,7 @@ mod tests { #[test] fn dispatch_set_nx_when_key_missing() { let mut ks = Keyspace::new(); - let resp = dispatch( + let resp = test_dispatch( &mut ks, &ShardRequest::Set { key: "k".into(), @@ -1908,7 +2046,7 @@ mod tests { let mut ks = Keyspace::new(); ks.set("k".into(), Bytes::from("old"), None); - let resp = dispatch( + let resp = test_dispatch( &mut ks, &ShardRequest::Set { key: "k".into(), @@ -1932,7 +2070,7 @@ mod tests { let mut ks = Keyspace::new(); ks.set("k".into(), Bytes::from("old"), None); - let resp = dispatch( + let resp = test_dispatch( &mut ks, &ShardRequest::Set { key: "k".into(), @@ -1952,7 +2090,7 @@ mod tests { #[test] fn dispatch_set_xx_when_key_missing() { let mut ks = Keyspace::new(); - let resp = dispatch( + let resp = test_dispatch( &mut ks, &ShardRequest::Set { key: "k".into(), @@ -1989,7 +2127,7 @@ mod tests { assert_eq!(ks.len(), 2); - let resp = dispatch(&mut ks, &ShardRequest::FlushDb); + let resp = test_dispatch(&mut ks, &ShardRequest::FlushDb); assert!(matches!(resp, ShardResponse::Ok)); assert_eq!(ks.len(), 0); } @@ -2001,7 +2139,7 @@ mod tests { ks.set("user:2".into(), Bytes::from("b"), None); ks.set("item:1".into(), Bytes::from("c"), None); - let resp = dispatch( + let resp = test_dispatch( &mut ks, &ShardRequest::Scan { cursor: 0, @@ -2026,7 +2164,7 @@ mod tests { ks.set("user:2".into(), Bytes::from("b"), None); ks.set("item:1".into(), Bytes::from("c"), None); - let resp = dispatch( + let resp = test_dispatch( &mut ks, &ShardRequest::Scan { cursor: 0, @@ -2176,7 +2314,7 @@ mod tests { ks.set("user:1".into(), Bytes::from("a"), None); ks.set("user:2".into(), Bytes::from("b"), None); ks.set("item:1".into(), Bytes::from("c"), None); - let resp = dispatch( + let resp = test_dispatch( &mut ks, &ShardRequest::Keys { pattern: "user:*".into(), @@ -2195,7 +2333,7 @@ mod tests { fn dispatch_rename() { let mut ks = Keyspace::new(); ks.set("old".into(), Bytes::from("value"), None); - let resp = dispatch( + let resp = test_dispatch( &mut ks, &ShardRequest::Rename { key: "old".into(), @@ -2210,7 +2348,7 @@ mod tests { #[test] fn dispatch_rename_missing_key() { let mut ks = Keyspace::new(); - let resp = dispatch( + let resp = test_dispatch( &mut ks, &ShardRequest::Rename { key: "missing".into(), diff --git a/crates/ember-server/src/concurrent_handler.rs b/crates/ember-server/src/concurrent_handler.rs index 2267e720..115420fd 100644 --- a/crates/ember-server/src/concurrent_handler.rs +++ b/crates/ember-server/src/concurrent_handler.rs @@ -481,48 +481,26 @@ async fn execute_concurrent( field_path, value, } => { - let registry = match _engine.schema_registry() { - Some(r) => r, - None => return Frame::Error("ERR protobuf support is not enabled".into()), - }; - let req = ember_core::ShardRequest::ProtoGet { key: key.clone() }; - let (type_name, data, existing_ttl) = match _engine.route(&key, req).await { - Ok(ember_core::ShardResponse::ProtoValue(Some(tuple))) => tuple, - Ok(ember_core::ShardResponse::ProtoValue(None)) => return Frame::Null, - Ok(ember_core::ShardResponse::WrongType) => { - return Frame::Error( - "WRONGTYPE Operation against a key holding the wrong kind of value".into(), - ) - } - Ok(other) => { - return Frame::Error(format!("ERR unexpected shard response: {other:?}")) - } - Err(e) => return Frame::Error(format!("ERR {e}")), - }; - let new_data = { - let reg = match registry.read() { - Ok(r) => r, - Err(_) => return Frame::Error("ERR schema registry lock poisoned".into()), - }; - match reg.set_field(&type_name, &data, &field_path, &value) { - Ok(d) => d, - Err(e) => return Frame::Error(format!("ERR {e}")), - } - }; - let req = ember_core::ShardRequest::ProtoSet { + if _engine.schema_registry().is_none() { + return Frame::Error("ERR protobuf support is not enabled".into()); + } + let req = ember_core::ShardRequest::ProtoSetField { key: key.clone(), - type_name, - data: new_data, - expire: existing_ttl, - nx: false, - xx: true, + field_path, + value, }; match _engine.route(&key, req).await { - Ok(ember_core::ShardResponse::Ok) => Frame::Simple("OK".into()), + Ok(ember_core::ShardResponse::ProtoFieldUpdated { .. }) => { + Frame::Simple("OK".into()) + } Ok(ember_core::ShardResponse::Value(None)) => Frame::Null, + Ok(ember_core::ShardResponse::WrongType) => Frame::Error( + "WRONGTYPE Operation against a key holding the wrong kind of value".into(), + ), Ok(ember_core::ShardResponse::OutOfMemory) => { Frame::Error("OOM command not allowed when used memory > 'maxmemory'".into()) } + Ok(ember_core::ShardResponse::Err(msg)) => Frame::Error(format!("ERR {msg}")), Ok(other) => Frame::Error(format!("ERR unexpected shard response: {other:?}")), Err(e) => Frame::Error(format!("ERR {e}")), } @@ -530,48 +508,23 @@ async fn execute_concurrent( #[cfg(feature = "protobuf")] Command::ProtoDelField { key, field_path } => { - let registry = match _engine.schema_registry() { - Some(r) => r, - None => return Frame::Error("ERR protobuf support is not enabled".into()), - }; - let req = ember_core::ShardRequest::ProtoGet { key: key.clone() }; - let (type_name, data, existing_ttl) = match _engine.route(&key, req).await { - Ok(ember_core::ShardResponse::ProtoValue(Some(tuple))) => tuple, - Ok(ember_core::ShardResponse::ProtoValue(None)) => return Frame::Null, - Ok(ember_core::ShardResponse::WrongType) => { - return Frame::Error( - "WRONGTYPE Operation against a key holding the wrong kind of value".into(), - ) - } - Ok(other) => { - return Frame::Error(format!("ERR unexpected shard response: {other:?}")) - } - Err(e) => return Frame::Error(format!("ERR {e}")), - }; - let new_data = { - let reg = match registry.read() { - Ok(r) => r, - Err(_) => return Frame::Error("ERR schema registry lock poisoned".into()), - }; - match reg.clear_field(&type_name, &data, &field_path) { - Ok(d) => d, - Err(e) => return Frame::Error(format!("ERR {e}")), - } - }; - let req = ember_core::ShardRequest::ProtoSet { + if _engine.schema_registry().is_none() { + return Frame::Error("ERR protobuf support is not enabled".into()); + } + let req = ember_core::ShardRequest::ProtoDelField { key: key.clone(), - type_name, - data: new_data, - expire: existing_ttl, - nx: false, - xx: true, + field_path, }; match _engine.route(&key, req).await { - Ok(ember_core::ShardResponse::Ok) => Frame::Integer(1), + Ok(ember_core::ShardResponse::ProtoFieldUpdated { .. }) => Frame::Integer(1), Ok(ember_core::ShardResponse::Value(None)) => Frame::Null, + Ok(ember_core::ShardResponse::WrongType) => Frame::Error( + "WRONGTYPE Operation against a key holding the wrong kind of value".into(), + ), Ok(ember_core::ShardResponse::OutOfMemory) => { Frame::Error("OOM command not allowed when used memory > 'maxmemory'".into()) } + Ok(ember_core::ShardResponse::Err(msg)) => Frame::Error(format!("ERR {msg}")), Ok(other) => Frame::Error(format!("ERR unexpected shard response: {other:?}")), Err(e) => Frame::Error(format!("ERR {e}")), } diff --git a/crates/ember-server/src/connection.rs b/crates/ember-server/src/connection.rs index 20ae9849..4d06ab8b 100644 --- a/crates/ember-server/src/connection.rs +++ b/crates/ember-server/src/connection.rs @@ -1808,45 +1808,20 @@ async fn execute( field_path, value, } => { - let registry = match engine.schema_registry() { - Some(r) => r, - None => return Frame::Error("ERR protobuf support is not enabled".into()), - }; - // step 1: fetch current value - let req = ShardRequest::ProtoGet { key: key.clone() }; - let (type_name, data, existing_ttl) = match engine.route(&key, req).await { - Ok(ShardResponse::ProtoValue(Some(tuple))) => tuple, - Ok(ShardResponse::ProtoValue(None)) => return Frame::Null, - Ok(ShardResponse::WrongType) => return wrongtype_error(), - Ok(other) => { - return Frame::Error(format!("ERR unexpected shard response: {other:?}")) - } - Err(e) => return Frame::Error(format!("ERR {e}")), - }; - // step 2: decode, mutate, re-encode - let new_data = { - let reg = match registry.read() { - Ok(r) => r, - Err(_) => return Frame::Error("ERR schema registry lock poisoned".into()), - }; - match reg.set_field(&type_name, &data, &field_path, &value) { - Ok(d) => d, - Err(e) => return Frame::Error(format!("ERR {e}")), - } - }; - // step 3: store back (XX = only if key still exists), preserving TTL - let req = ShardRequest::ProtoSet { + if engine.schema_registry().is_none() { + return Frame::Error("ERR protobuf support is not enabled".into()); + } + let req = ShardRequest::ProtoSetField { key: key.clone(), - type_name, - data: new_data, - expire: existing_ttl, - nx: false, - xx: true, + field_path, + value, }; match engine.route(&key, req).await { - Ok(ShardResponse::Ok) => Frame::Simple("OK".into()), + Ok(ShardResponse::ProtoFieldUpdated { .. }) => Frame::Simple("OK".into()), Ok(ShardResponse::Value(None)) => Frame::Null, + Ok(ShardResponse::WrongType) => wrongtype_error(), Ok(ShardResponse::OutOfMemory) => oom_error(), + Ok(ShardResponse::Err(msg)) => Frame::Error(format!("ERR {msg}")), Ok(other) => Frame::Error(format!("ERR unexpected shard response: {other:?}")), Err(e) => Frame::Error(format!("ERR {e}")), } @@ -1854,45 +1829,19 @@ async fn execute( #[cfg(feature = "protobuf")] Command::ProtoDelField { key, field_path } => { - let registry = match engine.schema_registry() { - Some(r) => r, - None => return Frame::Error("ERR protobuf support is not enabled".into()), - }; - // step 1: fetch current value - let req = ShardRequest::ProtoGet { key: key.clone() }; - let (type_name, data, existing_ttl) = match engine.route(&key, req).await { - Ok(ShardResponse::ProtoValue(Some(tuple))) => tuple, - Ok(ShardResponse::ProtoValue(None)) => return Frame::Null, - Ok(ShardResponse::WrongType) => return wrongtype_error(), - Ok(other) => { - return Frame::Error(format!("ERR unexpected shard response: {other:?}")) - } - Err(e) => return Frame::Error(format!("ERR {e}")), - }; - // step 2: decode, clear field, re-encode - let new_data = { - let reg = match registry.read() { - Ok(r) => r, - Err(_) => return Frame::Error("ERR schema registry lock poisoned".into()), - }; - match reg.clear_field(&type_name, &data, &field_path) { - Ok(d) => d, - Err(e) => return Frame::Error(format!("ERR {e}")), - } - }; - // step 3: store back (XX = only if key still exists), preserving TTL - let req = ShardRequest::ProtoSet { + if engine.schema_registry().is_none() { + return Frame::Error("ERR protobuf support is not enabled".into()); + } + let req = ShardRequest::ProtoDelField { key: key.clone(), - type_name, - data: new_data, - expire: existing_ttl, - nx: false, - xx: true, + field_path, }; match engine.route(&key, req).await { - Ok(ShardResponse::Ok) => Frame::Integer(1), + Ok(ShardResponse::ProtoFieldUpdated { .. }) => Frame::Integer(1), Ok(ShardResponse::Value(None)) => Frame::Null, + Ok(ShardResponse::WrongType) => wrongtype_error(), Ok(ShardResponse::OutOfMemory) => oom_error(), + Ok(ShardResponse::Err(msg)) => Frame::Error(format!("ERR {msg}")), Ok(other) => Frame::Error(format!("ERR unexpected shard response: {other:?}")), Err(e) => Frame::Error(format!("ERR {e}")), } From d71f1b1885c29b6999e198b3c476027708267d5d Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Tue, 10 Feb 2026 21:37:14 -0500 Subject: [PATCH 3/6] persist schemas through AOF rewrite MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit after a BGREWRITEAOF, the snapshot captures proto values but the AOF is truncated — losing any ProtoRegister records. on the next restart, recovery loads the values but has no schemas, breaking all field-level operations. now handle_rewrite writes ProtoRegister records for all registered schemas into the AOF immediately after truncation. --- crates/ember-core/src/shard.rs | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/crates/ember-core/src/shard.rs b/crates/ember-core/src/shard.rs index 628d41dc..f4e18c03 100644 --- a/crates/ember-core/src/shard.rs +++ b/crates/ember-core/src/shard.rs @@ -596,6 +596,8 @@ async fn run_shard( &persistence, &mut aof_writer, shard_id, + #[cfg(feature = "protobuf")] + &schema_registry, ); let _ = msg.reply.send(resp); continue; @@ -1247,11 +1249,15 @@ fn handle_snapshot( } /// Writes a snapshot and then truncates the AOF. +/// +/// When protobuf is enabled, re-persists all registered schemas to the +/// AOF after truncation so they survive the next restart. fn handle_rewrite( keyspace: &Keyspace, persistence: &Option, aof_writer: &mut Option, shard_id: u16, + #[cfg(feature = "protobuf")] schema_registry: &Option, ) -> ShardResponse { let pcfg = match persistence { Some(p) => p, @@ -1273,6 +1279,22 @@ fn handle_rewrite( if let Err(e) = writer.truncate() { warn!(shard_id, "aof truncate after rewrite failed: {e}"); } + + // re-persist schemas so they survive the next recovery + #[cfg(feature = "protobuf")] + if let Some(ref registry) = schema_registry { + if let Ok(reg) = registry.read() { + for (name, descriptor) in reg.iter_schemas() { + let record = AofRecord::ProtoRegister { + name: name.to_owned(), + descriptor: descriptor.clone(), + }; + if let Err(e) = writer.write_record(&record) { + warn!(shard_id, "failed to re-persist schema after rewrite: {e}"); + } + } + } + } } info!(shard_id, entries = count, "aof rewrite complete"); ShardResponse::Ok From 0303fc6dab5f5e9f8082e231233152aafca0b369 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Tue, 10 Feb 2026 21:39:19 -0500 Subject: [PATCH 4/6] broadcast error logging and concurrent mode warning - log a warning when broadcasting ProtoRegisterAof fails instead of silently discarding the error with `let _ =` - add a startup warning when both concurrent mode and protobuf are enabled, since generic commands (DEL, TTL, EXPIRE) don't affect proto keys in concurrent mode --- crates/ember-server/src/concurrent_handler.rs | 7 +++++-- crates/ember-server/src/connection.rs | 7 +++++-- crates/ember-server/src/main.rs | 9 +++++++++ 3 files changed, 19 insertions(+), 4 deletions(-) diff --git a/crates/ember-server/src/concurrent_handler.rs b/crates/ember-server/src/concurrent_handler.rs index 115420fd..a84c4acf 100644 --- a/crates/ember-server/src/concurrent_handler.rs +++ b/crates/ember-server/src/concurrent_handler.rs @@ -308,12 +308,15 @@ async fn execute_concurrent( }; match result { Ok(types) => { - let _ = _engine + if let Err(e) = _engine .broadcast(|| ember_core::ShardRequest::ProtoRegisterAof { name: name.clone(), descriptor: descriptor.clone(), }) - .await; + .await + { + tracing::warn!("failed to persist proto registration to AOF: {e}"); + } Frame::Array( types .into_iter() diff --git a/crates/ember-server/src/connection.rs b/crates/ember-server/src/connection.rs index 4d06ab8b..31928590 100644 --- a/crates/ember-server/src/connection.rs +++ b/crates/ember-server/src/connection.rs @@ -1644,12 +1644,15 @@ async fn execute( match result { Ok(types) => { // persist the registration to all shards' AOF - let _ = engine + if let Err(e) = engine .broadcast(|| ShardRequest::ProtoRegisterAof { name: name.clone(), descriptor: descriptor.clone(), }) - .await; + .await + { + tracing::warn!("failed to persist proto registration to AOF: {e}"); + } Frame::Array( types .into_iter() diff --git a/crates/ember-server/src/main.rs b/crates/ember-server/src/main.rs index 2c6dba6c..52c93900 100644 --- a/crates/ember-server/src/main.rs +++ b/crates/ember-server/src/main.rs @@ -23,6 +23,8 @@ use clap::Parser; use ember_cluster::{GossipConfig, NodeId}; use ember_core::ShardPersistenceConfig; use tracing::info; +#[cfg(feature = "protobuf")] +use tracing::warn; use crate::cluster::ClusterCoordinator; use crate::config::{ @@ -266,6 +268,13 @@ async fn main() { if args.protobuf { engine_config.schema_registry = Some(ember_core::schema::SchemaRegistry::shared()); info!("protobuf value storage enabled"); + + if args.concurrent { + warn!( + "concurrent mode with protobuf: DEL, TTL, EXPIRE, and other generic \ + commands do not affect proto keys. use sharded mode for full proto support." + ); + } } if let Some(limit) = max_memory { From 0002dfffbea687dc6681e7533d2d01d51719f984 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Tue, 10 Feb 2026 21:40:57 -0500 Subject: [PATCH 5/6] proto value size limit, memory correction, schema count cap - add MAX_PROTO_VALUE_BYTES (64MB) check in validate() to prevent a single PROTO.SET from exhausting shard memory - add MAX_SCHEMAS (1024) check in register() to cap the number of registered schemas - fix proto memory overhead from 24 to 48 bytes to properly account for both String (24 bytes) and Bytes (24 bytes) structs --- crates/ember-core/src/memory.rs | 5 +++-- crates/ember-core/src/schema.rs | 26 ++++++++++++++++++++++++++ 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/crates/ember-core/src/memory.rs b/crates/ember-core/src/memory.rs index 77018d9c..c2a56b40 100644 --- a/crates/ember-core/src/memory.rs +++ b/crates/ember-core/src/memory.rs @@ -223,9 +223,10 @@ pub fn value_size(value: &Value) -> usize { let member_bytes: usize = set.iter().map(|m| m.len() + HASHSET_MEMBER_OVERHEAD).sum(); HASHSET_BASE_OVERHEAD + member_bytes } - // type_name String (24 bytes ptr+len+cap on heap) + data Bytes (24 bytes). + // type_name: String struct = 24 bytes (ptr+len+cap) on 64-bit. + // data: Bytes struct = ~24 bytes (ptr+len+vtable/arc). #[cfg(feature = "protobuf")] - Value::Proto { type_name, data } => type_name.len() + data.len() + 24, + Value::Proto { type_name, data } => type_name.len() + data.len() + 48, } } diff --git a/crates/ember-core/src/schema.rs b/crates/ember-core/src/schema.rs index 2567592f..21f88b22 100644 --- a/crates/ember-core/src/schema.rs +++ b/crates/ember-core/src/schema.rs @@ -25,6 +25,14 @@ const MAX_DESCRIPTOR_BYTES: usize = 10 * 1024 * 1024; /// Deep nesting beyond this is almost certainly a bug or an abuse vector. const MAX_FIELD_PATH_DEPTH: usize = 16; +/// Maximum allowed size for a single proto value in bytes (64 MB). +/// Prevents a single PROTO.SET from exhausting shard memory. +const MAX_PROTO_VALUE_BYTES: usize = 64 * 1024 * 1024; + +/// Maximum number of schemas that can be registered. +/// Prevents unbounded growth of the schema registry. +const MAX_SCHEMAS: usize = 1024; + /// Errors that can occur during schema operations. #[derive(Debug, Error)] pub enum SchemaError { @@ -48,6 +56,12 @@ pub enum SchemaError { #[error("field path too deep: {0} segments (max {1})")] PathTooDeep(usize, usize), + + #[error("proto value too large: {0} bytes (max {1})")] + ValueTooLarge(usize, usize), + + #[error("schema limit reached: {0} schemas (max {1})")] + TooManySchemas(usize, usize), } /// A registered schema: the raw descriptor bytes and the parsed pool. @@ -106,6 +120,10 @@ impl SchemaRegistry { return Err(SchemaError::AlreadyExists(name)); } + if self.schemas.len() >= MAX_SCHEMAS { + return Err(SchemaError::TooManySchemas(self.schemas.len(), MAX_SCHEMAS)); + } + if descriptor_bytes.len() > MAX_DESCRIPTOR_BYTES { return Err(SchemaError::DescriptorTooLarge( descriptor_bytes.len(), @@ -145,8 +163,16 @@ impl SchemaRegistry { /// Validates that `data` is a valid encoding of `message_type`. /// + /// Checks the value size limit before decoding. /// Searches all registered schemas for the type name. pub fn validate(&self, message_type: &str, data: &[u8]) -> Result<(), SchemaError> { + if data.len() > MAX_PROTO_VALUE_BYTES { + return Err(SchemaError::ValueTooLarge( + data.len(), + MAX_PROTO_VALUE_BYTES, + )); + } + let descriptor = self.find_message(message_type)?; DynamicMessage::decode(descriptor, data) From f72aadff95d9e9e1766fb43390ebd376bec3191a Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Tue, 10 Feb 2026 21:47:57 -0500 Subject: [PATCH 6/6] tests and sync fix for recovery, rewrite, and limits MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit add unit tests for value size limit rejection, at-limit acceptance, and schema count cap. add integration tests for schema recovery after restart, schema survival through AOF rewrite, DEL on proto keys, and concurrent mode schema recovery. fix: sync AOF after re-persisting schemas during rewrite — without this, schema records could remain in the BufWriter buffer and be lost if the server is killed before the next fsync cycle. --- crates/ember-core/src/schema.rs | 48 ++++++++++ crates/ember-core/src/shard.rs | 5 + tests/integration/src/proto.rs | 160 ++++++++++++++++++++++++++++++++ 3 files changed, 213 insertions(+) diff --git a/crates/ember-core/src/schema.rs b/crates/ember-core/src/schema.rs index 21f88b22..639c45b7 100644 --- a/crates/ember-core/src/schema.rs +++ b/crates/ember-core/src/schema.rs @@ -1390,4 +1390,52 @@ mod tests { let frame = registry.get_field("test.BigNum", &buf, "val").unwrap(); assert_eq!(frame, Frame::Integer(42)); } + + // --- value size limit / schema count limit tests --- + + #[test] + fn value_too_large_rejected() { + let mut registry = SchemaRegistry::new(); + let desc = make_descriptor("test", "User", "name"); + registry.register("users".into(), desc).unwrap(); + + let oversized = vec![0u8; MAX_PROTO_VALUE_BYTES + 1]; + let err = registry.validate("test.User", &oversized).unwrap_err(); + assert!(matches!(err, SchemaError::ValueTooLarge(_, _))); + } + + #[test] + fn value_at_limit_allowed() { + let mut registry = SchemaRegistry::new(); + let desc = make_descriptor("test", "User", "name"); + registry.register("users".into(), desc).unwrap(); + + // a value exactly at the limit should pass the size check + // (it will fail validation since it's not valid protobuf, but + // it should NOT fail with ValueTooLarge) + let at_limit = vec![0u8; MAX_PROTO_VALUE_BYTES]; + let err = registry.validate("test.User", &at_limit).unwrap_err(); + assert!( + !matches!(err, SchemaError::ValueTooLarge(_, _)), + "expected validation error, not size limit" + ); + } + + #[test] + fn schema_count_limit() { + let mut registry = SchemaRegistry::new(); + + // register up to the limit + for i in 0..MAX_SCHEMAS { + let desc = make_descriptor(&format!("pkg{i}"), &format!("Msg{i}"), "val"); + registry + .register(format!("schema-{i}"), desc) + .unwrap_or_else(|e| panic!("failed to register schema {i}: {e}")); + } + + // the next one should fail + let desc = make_descriptor("overflow", "Overflow", "val"); + let err = registry.register("overflow".into(), desc).unwrap_err(); + assert!(matches!(err, SchemaError::TooManySchemas(_, _))); + } } diff --git a/crates/ember-core/src/shard.rs b/crates/ember-core/src/shard.rs index f4e18c03..3aa36552 100644 --- a/crates/ember-core/src/shard.rs +++ b/crates/ember-core/src/shard.rs @@ -1295,6 +1295,11 @@ fn handle_rewrite( } } } + + // flush so schemas are durable before we report success + if let Err(e) = writer.sync() { + warn!(shard_id, "aof sync after rewrite failed: {e}"); + } } info!(shard_id, entries = count, "aof rewrite complete"); ShardResponse::Ok diff --git a/tests/integration/src/proto.rs b/tests/integration/src/proto.rs index 2701ff0d..6c1222dd 100644 --- a/tests/integration/src/proto.rs +++ b/tests/integration/src/proto.rs @@ -969,6 +969,121 @@ async fn duplicate_registration_returns_error() { assert!(matches!(resp, Frame::Error(ref s) if s.contains("already registered"))); } +// ---- schema recovery tests ---- + +#[tokio::test] +async fn schema_recovery_after_restart() { + let data_dir = tempfile::tempdir().unwrap(); + let path = data_dir.path().to_path_buf(); + + let desc = make_multi_field_descriptor(); + let data = encode_profile(&desc, "alice", 25, true); + + // start server, register schema, set proto value + { + let server = TestServer::start_with(ServerOptions { + protobuf: true, + appendonly: true, + data_dir_path: Some(path.clone()), + ..Default::default() + }); + let mut c = server.connect().await; + + c.cmd_raw(&[b"PROTO.REGISTER", b"profiles", &desc]).await; + c.cmd_raw(&[b"PROTO.SET", b"p:1", b"test.Profile", &data]) + .await; + + // wait for fsync + tokio::time::sleep(std::time::Duration::from_millis(300)).await; + } + + // restart — schemas should be recovered from AOF + let server = TestServer::start_with(ServerOptions { + protobuf: true, + appendonly: true, + data_dir_path: Some(path), + ..Default::default() + }); + let mut c = server.connect().await; + + // GETFIELD requires the schema to be loaded — this proves recovery works + let resp = c.cmd(&["PROTO.GETFIELD", "p:1", "name"]).await; + assert_eq!(resp, Frame::Bulk(Bytes::from("alice"))); + + let resp = c.cmd(&["PROTO.GETFIELD", "p:1", "age"]).await; + assert_eq!(resp, Frame::Integer(25)); + + drop(data_dir); +} + +#[tokio::test] +async fn schema_survives_aof_rewrite() { + let data_dir = tempfile::tempdir().unwrap(); + let path = data_dir.path().to_path_buf(); + + let desc = make_multi_field_descriptor(); + let data = encode_profile(&desc, "alice", 30, true); + + // start server, register, set, then trigger AOF rewrite + { + let server = TestServer::start_with(ServerOptions { + protobuf: true, + appendonly: true, + data_dir_path: Some(path.clone()), + ..Default::default() + }); + let mut c = server.connect().await; + + c.cmd_raw(&[b"PROTO.REGISTER", b"profiles", &desc]).await; + c.cmd_raw(&[b"PROTO.SET", b"p:1", b"test.Profile", &data]) + .await; + + // trigger AOF rewrite — this truncates the AOF and writes a snapshot + let resp = c.cmd(&["BGREWRITEAOF"]).await; + assert!(matches!(resp, Frame::Simple(_))); + // give the rewrite time to complete + tokio::time::sleep(std::time::Duration::from_millis(500)).await; + } + + // restart — schema must survive the rewrite + let server = TestServer::start_with(ServerOptions { + protobuf: true, + appendonly: true, + data_dir_path: Some(path), + ..Default::default() + }); + let mut c = server.connect().await; + + let resp = c.cmd(&["PROTO.GETFIELD", "p:1", "name"]).await; + assert_eq!(resp, Frame::Bulk(Bytes::from("alice"))); + + let resp = c.cmd(&["PROTO.GETFIELD", "p:1", "age"]).await; + assert_eq!(resp, Frame::Integer(30)); + + drop(data_dir); +} + +#[tokio::test] +async fn del_removes_proto_key() { + let server = start_proto_server(false); + let mut c = server.connect().await; + + let desc = make_descriptor("test", "User", "name"); + c.cmd_raw(&[b"PROTO.REGISTER", b"users", &desc]).await; + + let data = encode_message(&desc, "test.User", "name", "alice"); + c.cmd_raw(&[b"PROTO.SET", b"user:del", b"test.User", &data]) + .await; + + // DEL should remove the proto key + let resp = c.get_int(&["DEL", "user:del"]).await; + assert_eq!(resp, 1); + + // PROTO.GET should now return null + let resp = c.cmd(&["PROTO.GET", "user:del"]).await; + assert!(matches!(resp, Frame::Null)); +} + // ---- concurrent mode nested path and misc tests ---- // note: TTL preservation tests are sharded-only because concurrent mode // routes proto commands through engine shards while TTL checks the @@ -1013,3 +1128,48 @@ async fn concurrent_delfield_nested() { let resp = c.cmd(&["PROTO.GETFIELD", "o:1", "inner.value"]).await; assert_eq!(resp, Frame::Bulk(Bytes::from(""))); } + +// ---- concurrent mode recovery and DEL tests ---- + +#[tokio::test] +async fn concurrent_schema_recovery_after_restart() { + let data_dir = tempfile::tempdir().unwrap(); + let path = data_dir.path().to_path_buf(); + + let desc = make_multi_field_descriptor(); + let data = encode_profile(&desc, "bob", 42, false); + + { + let server = TestServer::start_with(ServerOptions { + protobuf: true, + concurrent: true, + appendonly: true, + data_dir_path: Some(path.clone()), + ..Default::default() + }); + let mut c = server.connect().await; + + c.cmd_raw(&[b"PROTO.REGISTER", b"profiles", &desc]).await; + c.cmd_raw(&[b"PROTO.SET", b"p:1", b"test.Profile", &data]) + .await; + + tokio::time::sleep(std::time::Duration::from_millis(300)).await; + } + + let server = TestServer::start_with(ServerOptions { + protobuf: true, + concurrent: true, + appendonly: true, + data_dir_path: Some(path), + ..Default::default() + }); + let mut c = server.connect().await; + + let resp = c.cmd(&["PROTO.GETFIELD", "p:1", "name"]).await; + assert_eq!(resp, Frame::Bulk(Bytes::from("bob"))); + + let resp = c.cmd(&["PROTO.GETFIELD", "p:1", "age"]).await; + assert_eq!(resp, Frame::Integer(42)); + + drop(data_dir); +}