From f09e61f364e6e596f3ca285f0d06a423d90173cc Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Tue, 10 Feb 2026 11:24:14 -0500 Subject: [PATCH 1/8] feat: add protobuf feature flag and prost-reflect dependency adds the `protobuf` feature gate across the workspace. when enabled, ember-core gains prost-reflect for dynamic protobuf message handling. the feature cascades: ember-server/protobuf enables emberkv-core/protobuf and ember-persistence/protobuf. --- Cargo.toml | 3 +++ crates/ember-core/Cargo.toml | 2 ++ crates/ember-persistence/Cargo.toml | 1 + crates/ember-server/Cargo.toml | 1 + 4 files changed, 7 insertions(+) diff --git a/Cargo.toml b/Cargo.toml index d2eac414..0136fc12 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -61,6 +61,9 @@ rustls = { version = "0.23", default-features = false, features = ["std", "tls12 rustls-pemfile = "2" rustls-native-certs = "0.8" +# dynamic protobuf messages +prost-reflect = "0.14" + # internal crates (version required for crates.io publishing) emberkv-core = { version = "0.4.2", path = "crates/ember-core" } ember-protocol = { version = "0.4.2", path = "crates/ember-protocol" } diff --git a/crates/ember-core/Cargo.toml b/crates/ember-core/Cargo.toml index 352b9e02..b61d1d07 100644 --- a/crates/ember-core/Cargo.toml +++ b/crates/ember-core/Cargo.toml @@ -11,6 +11,7 @@ readme = "README.md" [features] encryption = ["ember-persistence/encryption"] +protobuf = ["prost-reflect"] [lib] name = "ember_core" @@ -25,6 +26,7 @@ tokio = { workspace = true } tracing = { workspace = true } rand = { workspace = true } ordered-float = { workspace = true } +prost-reflect = { workspace = true, optional = true } dashmap = "6" parking_lot = "0.12" diff --git a/crates/ember-persistence/Cargo.toml b/crates/ember-persistence/Cargo.toml index eca0d0b0..0ca5c0d6 100644 --- a/crates/ember-persistence/Cargo.toml +++ b/crates/ember-persistence/Cargo.toml @@ -11,6 +11,7 @@ readme = "README.md" [features] encryption = ["aes-gcm", "rand"] +protobuf = [] [dependencies] thiserror = { workspace = true } diff --git a/crates/ember-server/Cargo.toml b/crates/ember-server/Cargo.toml index 14ee2857..a1472c3c 100644 --- a/crates/ember-server/Cargo.toml +++ b/crates/ember-server/Cargo.toml @@ -13,6 +13,7 @@ readme = "README.md" default = ["jemalloc"] jemalloc = ["tikv-jemallocator"] encryption = ["emberkv-core/encryption", "ember-persistence/encryption"] +protobuf = ["emberkv-core/protobuf", "ember-persistence/protobuf"] [dependencies] bytes = { workspace = true } From cb4b48e75d61186b35de928070b69e267550b7f9 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Tue, 10 Feb 2026 11:25:35 -0500 Subject: [PATCH 2/8] feat: add schema registry for protobuf message validation SchemaRegistry stores compiled FileDescriptorSet descriptors and validates protobuf values at write time. supports registration, validation, introspection (describe/list), and recovery (restore). gated behind the `protobuf` feature flag. --- crates/ember-core/src/lib.rs | 3 + crates/ember-core/src/schema.rs | 322 ++++++++++++++++++++++++++++++++ 2 files changed, 325 insertions(+) create mode 100644 crates/ember-core/src/schema.rs diff --git a/crates/ember-core/src/lib.rs b/crates/ember-core/src/lib.rs index a39fe166..b09d5df5 100644 --- a/crates/ember-core/src/lib.rs +++ b/crates/ember-core/src/lib.rs @@ -15,6 +15,9 @@ pub mod shard; pub mod time; pub mod types; +#[cfg(feature = "protobuf")] +pub mod schema; + pub use concurrent::ConcurrentKeyspace; pub use engine::{Engine, EngineConfig}; pub use error::ShardError; diff --git a/crates/ember-core/src/schema.rs b/crates/ember-core/src/schema.rs new file mode 100644 index 00000000..9b817c1b --- /dev/null +++ b/crates/ember-core/src/schema.rs @@ -0,0 +1,322 @@ +//! Schema registry for protobuf message validation. +//! +//! Stores compiled `FileDescriptorSet` descriptors so that ember can +//! validate protobuf values at write time and return typed metadata +//! on reads. Users register schemas via `PROTO.REGISTER` and the +//! registry is shared (behind an `Arc`) across all connections. +//! +//! Only compiled when the `protobuf` feature is enabled. + +use std::collections::HashMap; +use std::sync::{Arc, RwLock}; + +use bytes::Bytes; +use prost_reflect::{DescriptorPool, DynamicMessage, MessageDescriptor}; +use thiserror::Error; + +/// Errors that can occur during schema operations. +#[derive(Debug, Error)] +pub enum SchemaError { + #[error("invalid descriptor: {0}")] + InvalidDescriptor(String), + + #[error("unknown message type: {0}")] + UnknownMessageType(String), + + #[error("validation failed: {0}")] + ValidationFailed(String), + + #[error("schema already registered: {0}")] + AlreadyExists(String), +} + +/// A registered schema: the raw descriptor bytes and the parsed pool. +struct RegisteredSchema { + /// Raw `FileDescriptorSet` bytes, kept for persistence. + descriptor_bytes: Bytes, + /// Parsed descriptor pool for message lookup and validation. + pool: DescriptorPool, + /// All message type full names in this schema. + message_types: Vec, +} + +/// Registry of protobuf schemas. +/// +/// Each schema is identified by a user-chosen name (e.g. "users/v1") +/// and contains one or more message type definitions. +pub struct SchemaRegistry { + schemas: HashMap, +} + +/// Thread-safe handle to a shared schema registry. +pub type SharedSchemaRegistry = Arc>; + +impl SchemaRegistry { + /// Creates an empty registry. + pub fn new() -> Self { + Self { + schemas: HashMap::new(), + } + } + + /// Creates a new `SharedSchemaRegistry` wrapped in `Arc`. + pub fn shared() -> SharedSchemaRegistry { + Arc::new(RwLock::new(Self::new())) + } + + /// Registers a schema from compiled `FileDescriptorSet` bytes. + /// + /// Returns the list of message type names defined in the schema. + /// Fails if the name is already registered or the descriptor is invalid. + pub fn register(&mut self, name: String, descriptor_bytes: Bytes) -> Result, SchemaError> { + if self.schemas.contains_key(&name) { + return Err(SchemaError::AlreadyExists(name)); + } + + let pool = DescriptorPool::decode(descriptor_bytes.as_ref()) + .map_err(|e| SchemaError::InvalidDescriptor(e.to_string()))?; + + let message_types: Vec = pool + .all_messages() + .map(|m| m.full_name().to_owned()) + .collect(); + + if message_types.is_empty() { + return Err(SchemaError::InvalidDescriptor( + "no message types found in descriptor".into(), + )); + } + + self.schemas.insert( + name, + RegisteredSchema { + descriptor_bytes, + pool, + message_types: message_types.clone(), + }, + ); + + Ok(message_types) + } + + /// Validates that `data` is a valid encoding of `message_type`. + /// + /// Searches all registered schemas for the type name. + pub fn validate(&self, message_type: &str, data: &[u8]) -> Result<(), SchemaError> { + let descriptor = self.find_message(message_type)?; + + DynamicMessage::decode(descriptor, data) + .map_err(|e| SchemaError::ValidationFailed(e.to_string()))?; + + Ok(()) + } + + /// Returns the names of all registered schemas. + pub fn schema_names(&self) -> Vec { + let mut names: Vec = self.schemas.keys().cloned().collect(); + names.sort(); + names + } + + /// Returns the message type names defined in a schema, or `None` + /// if the schema isn't registered. + pub fn describe(&self, name: &str) -> Option> { + self.schemas.get(name).map(|s| s.message_types.clone()) + } + + /// Iterates over all schemas, yielding `(name, descriptor_bytes)`. + /// Used for persistence (snapshot/AOF). + pub fn iter_schemas(&self) -> impl Iterator { + self.schemas + .iter() + .map(|(name, schema)| (name.as_str(), &schema.descriptor_bytes)) + } + + /// Restores a schema during recovery. Skips duplicates silently + /// (idempotent — safe for AOF replay). + pub fn restore(&mut self, name: String, descriptor_bytes: Bytes) { + if self.schemas.contains_key(&name) { + return; + } + + let pool = match DescriptorPool::decode(descriptor_bytes.as_ref()) { + Ok(p) => p, + Err(e) => { + tracing::warn!(schema = %name, "failed to restore schema: {e}"); + return; + } + }; + + let message_types: Vec = pool + .all_messages() + .map(|m| m.full_name().to_owned()) + .collect(); + + self.schemas.insert( + name, + RegisteredSchema { + descriptor_bytes, + pool, + message_types, + }, + ); + } + + /// Looks up a message descriptor by full name across all schemas. + fn find_message(&self, message_type: &str) -> Result { + for schema in self.schemas.values() { + if let Some(desc) = schema.pool.get_message_by_name(message_type) { + return Ok(desc); + } + } + Err(SchemaError::UnknownMessageType(message_type.to_owned())) + } +} + +impl Default for SchemaRegistry { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Builds a minimal FileDescriptorSet containing a single message type. + /// Uses prost-reflect's own encoding rather than shelling out to protoc. + fn make_descriptor(package: &str, message_name: &str, field_name: &str) -> Bytes { + use prost_reflect::prost_types::{ + DescriptorProto, FieldDescriptorProto, FileDescriptorProto, FileDescriptorSet, + }; + use prost_reflect::prost::Message; + + let fds = FileDescriptorSet { + file: vec![FileDescriptorProto { + name: Some(format!("{package}.proto")), + package: Some(package.to_owned()), + message_type: vec![DescriptorProto { + name: Some(message_name.to_owned()), + field: vec![FieldDescriptorProto { + name: Some(field_name.to_owned()), + number: Some(1), + r#type: Some(9), // TYPE_STRING + label: Some(1), // LABEL_OPTIONAL + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }], + }; + + let mut buf = Vec::new(); + fds.encode(&mut buf).expect("encode descriptor"); + Bytes::from(buf) + } + + #[test] + fn register_and_describe() { + let mut registry = SchemaRegistry::new(); + let desc = make_descriptor("test", "User", "name"); + + let types = registry.register("users".into(), desc).unwrap(); + assert_eq!(types, vec!["test.User"]); + + let described = registry.describe("users").unwrap(); + assert_eq!(described, vec!["test.User"]); + } + + #[test] + fn double_registration_fails() { + let mut registry = SchemaRegistry::new(); + let desc = make_descriptor("test", "User", "name"); + + registry.register("users".into(), desc.clone()).unwrap(); + let err = registry.register("users".into(), desc).unwrap_err(); + assert!(matches!(err, SchemaError::AlreadyExists(_))); + } + + #[test] + fn invalid_descriptor_fails() { + let mut registry = SchemaRegistry::new(); + let err = registry + .register("bad".into(), Bytes::from("not a protobuf")) + .unwrap_err(); + assert!(matches!(err, SchemaError::InvalidDescriptor(_))); + } + + #[test] + fn validate_valid_message() { + let mut registry = SchemaRegistry::new(); + let desc = make_descriptor("test", "User", "name"); + registry.register("users".into(), desc).unwrap(); + + // encode a valid User message with name = "alice" + let pool = ®istry.schemas["users"].pool; + let msg_desc = pool.get_message_by_name("test.User").unwrap(); + let mut msg = DynamicMessage::new(msg_desc); + msg.set_field_by_name("name", prost_reflect::Value::String("alice".into())); + + let mut buf = Vec::new(); + use prost_reflect::prost::Message; + msg.encode(&mut buf).unwrap(); + + registry.validate("test.User", &buf).unwrap(); + } + + #[test] + fn validate_unknown_type_fails() { + let registry = SchemaRegistry::new(); + let err = registry.validate("no.Such.Type", &[]).unwrap_err(); + assert!(matches!(err, SchemaError::UnknownMessageType(_))); + } + + #[test] + fn schema_names_sorted() { + let mut registry = SchemaRegistry::new(); + registry + .register("z-schema".into(), make_descriptor("z", "Z", "val")) + .unwrap(); + registry + .register("a-schema".into(), make_descriptor("a", "A", "val")) + .unwrap(); + + let names = registry.schema_names(); + assert_eq!(names, vec!["a-schema", "z-schema"]); + } + + #[test] + fn describe_unknown_returns_none() { + let registry = SchemaRegistry::new(); + assert!(registry.describe("nope").is_none()); + } + + #[test] + fn restore_is_idempotent() { + let mut registry = SchemaRegistry::new(); + let desc = make_descriptor("test", "User", "name"); + + registry.restore("users".into(), desc.clone()); + registry.restore("users".into(), desc); + + assert_eq!(registry.schema_names(), vec!["users"]); + } + + #[test] + fn iter_schemas_returns_all() { + let mut registry = SchemaRegistry::new(); + let desc1 = make_descriptor("a", "A", "val"); + let desc2 = make_descriptor("b", "B", "val"); + + registry.register("alpha".into(), desc1).unwrap(); + registry.register("beta".into(), desc2).unwrap(); + + let mut pairs: Vec<_> = registry + .iter_schemas() + .map(|(name, _)| name.to_owned()) + .collect(); + pairs.sort(); + assert_eq!(pairs, vec!["alpha", "beta"]); + } +} From 35379d2cbbdc303f1bf8f20ada589d57c6d5a526 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Tue, 10 Feb 2026 11:28:02 -0500 Subject: [PATCH 3/8] feat: add Value::Proto variant with memory tracking adds a new Proto { type_name, data } variant to the Value enum, gated behind the protobuf feature. updates PartialEq, type_name(), value_size(), is_large_value(), and snapshot/recovery code to handle the new type. --- crates/ember-core/Cargo.toml | 2 +- crates/ember-core/src/memory.rs | 7 ++++ crates/ember-core/src/shard.rs | 7 ++++ crates/ember-core/src/types/mod.rs | 19 ++++++++++ crates/ember-persistence/src/recovery.rs | 5 +++ crates/ember-persistence/src/snapshot.rs | 46 ++++++++++++++++++++++++ 6 files changed, 85 insertions(+), 1 deletion(-) diff --git a/crates/ember-core/Cargo.toml b/crates/ember-core/Cargo.toml index b61d1d07..c91b5f1c 100644 --- a/crates/ember-core/Cargo.toml +++ b/crates/ember-core/Cargo.toml @@ -11,7 +11,7 @@ readme = "README.md" [features] encryption = ["ember-persistence/encryption"] -protobuf = ["prost-reflect"] +protobuf = ["prost-reflect", "ember-persistence/protobuf"] [lib] name = "ember_core" diff --git a/crates/ember-core/src/memory.rs b/crates/ember-core/src/memory.rs index a6082cd3..77018d9c 100644 --- a/crates/ember-core/src/memory.rs +++ b/crates/ember-core/src/memory.rs @@ -160,6 +160,10 @@ pub fn is_large_value(value: &Value) -> bool { Value::SortedSet(ss) => ss.len() > LAZY_FREE_THRESHOLD, Value::Hash(m) => m.len() > LAZY_FREE_THRESHOLD, Value::Set(s) => s.len() > LAZY_FREE_THRESHOLD, + // Proto values use Bytes (ref-counted, O(1) drop) + a String. + // Neither is expensive to drop. + #[cfg(feature = "protobuf")] + Value::Proto { .. } => false, } } @@ -219,6 +223,9 @@ 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). + #[cfg(feature = "protobuf")] + Value::Proto { type_name, data } => type_name.len() + data.len() + 24, } } diff --git a/crates/ember-core/src/shard.rs b/crates/ember-core/src/shard.rs index afc10fa2..a216bd37 100644 --- a/crates/ember-core/src/shard.rs +++ b/crates/ember-core/src/shard.rs @@ -418,6 +418,8 @@ async fn run_shard( } RecoveredValue::Hash(map) => Value::Hash(map), RecoveredValue::Set(set) => Value::Set(set), + #[cfg(feature = "protobuf")] + RecoveredValue::Proto { type_name, data } => Value::Proto { type_name, data }, }; keyspace.restore(entry.key, value, entry.ttl); } @@ -1064,6 +1066,11 @@ fn write_snapshot( } Value::Hash(map) => SnapValue::Hash(map.clone()), Value::Set(set) => SnapValue::Set(set.clone()), + #[cfg(feature = "protobuf")] + Value::Proto { type_name, data } => SnapValue::Proto { + type_name: type_name.clone(), + data: data.clone(), + }, }; writer.write_entry(&SnapEntry { key: key.to_owned(), diff --git a/crates/ember-core/src/types/mod.rs b/crates/ember-core/src/types/mod.rs index 24b08bba..d9da9de2 100644 --- a/crates/ember-core/src/types/mod.rs +++ b/crates/ember-core/src/types/mod.rs @@ -36,6 +36,12 @@ pub enum Value { /// Unordered set of unique string members. Set(HashSet), + + /// A protobuf message value. Stores the fully-qualified message type + /// name alongside the serialized bytes. Validation happens at the + /// server layer before storage. + #[cfg(feature = "protobuf")] + Proto { type_name: String, data: Bytes }, } impl PartialEq for Value { @@ -51,6 +57,17 @@ impl PartialEq for Value { } (Value::Hash(a), Value::Hash(b)) => a == b, (Value::Set(a), Value::Set(b)) => a == b, + #[cfg(feature = "protobuf")] + ( + Value::Proto { + type_name: t1, + data: d1, + }, + Value::Proto { + type_name: t2, + data: d2, + }, + ) => t1 == t2 && d1 == d2, _ => false, } } @@ -64,6 +81,8 @@ pub fn type_name(value: &Value) -> &'static str { Value::SortedSet(_) => "zset", Value::Hash(_) => "hash", Value::Set(_) => "set", + #[cfg(feature = "protobuf")] + Value::Proto { .. } => "proto", } } diff --git a/crates/ember-persistence/src/recovery.rs b/crates/ember-persistence/src/recovery.rs index c0900764..b1374e14 100644 --- a/crates/ember-persistence/src/recovery.rs +++ b/crates/ember-persistence/src/recovery.rs @@ -37,6 +37,9 @@ pub enum RecoveredValue { Hash(HashMap), /// Unordered set of unique string members. Set(HashSet), + /// A protobuf message: type name + serialized bytes. + #[cfg(feature = "protobuf")] + Proto { type_name: String, data: Bytes }, } impl From for RecoveredValue { @@ -47,6 +50,8 @@ impl From for RecoveredValue { SnapValue::SortedSet(members) => RecoveredValue::SortedSet(members), SnapValue::Hash(map) => RecoveredValue::Hash(map), SnapValue::Set(set) => RecoveredValue::Set(set), + #[cfg(feature = "protobuf")] + SnapValue::Proto { type_name, data } => RecoveredValue::Proto { type_name, data }, } } } diff --git a/crates/ember-persistence/src/snapshot.rs b/crates/ember-persistence/src/snapshot.rs index b7995f5d..85a7026c 100644 --- a/crates/ember-persistence/src/snapshot.rs +++ b/crates/ember-persistence/src/snapshot.rs @@ -36,6 +36,8 @@ const TYPE_LIST: u8 = 1; const TYPE_SORTED_SET: u8 = 2; const TYPE_HASH: u8 = 3; const TYPE_SET: u8 = 4; +#[cfg(feature = "protobuf")] +const TYPE_PROTO: u8 = 5; /// The value stored in a snapshot entry. #[derive(Debug, Clone, PartialEq)] @@ -50,6 +52,9 @@ pub enum SnapValue { Hash(HashMap), /// An unordered set of unique string members. Set(HashSet), + /// A protobuf message: type name + serialized bytes. + #[cfg(feature = "protobuf")] + Proto { type_name: String, data: Bytes }, } /// A single entry in a snapshot file. @@ -190,6 +195,15 @@ impl SnapshotWriter { format::write_bytes(&mut buf, member.as_bytes())?; } } + #[cfg(feature = "protobuf")] + SnapValue::Proto { + type_name, + data, + } => { + format::write_u8(&mut buf, TYPE_PROTO)?; + format::write_bytes(&mut buf, type_name.as_bytes())?; + format::write_bytes(&mut buf, data)?; + } } format::write_i64(&mut buf, entry.expire_ms)?; @@ -419,6 +433,23 @@ impl SnapshotReader { } SnapValue::Set(set) } + #[cfg(feature = "protobuf")] + TYPE_PROTO => { + let type_name_bytes = format::read_bytes(&mut self.reader)?; + format::write_bytes(&mut buf, &type_name_bytes)?; + let type_name = String::from_utf8(type_name_bytes).map_err(|_| { + FormatError::Io(io::Error::new( + io::ErrorKind::InvalidData, + "proto type_name is not valid utf-8", + )) + })?; + let data = format::read_bytes(&mut self.reader)?; + format::write_bytes(&mut buf, &data)?; + SnapValue::Proto { + type_name, + data: Bytes::from(data), + } + } _ => { return Err(FormatError::UnknownTag(type_tag)); } @@ -551,6 +582,21 @@ impl SnapshotReader { } SnapValue::Set(set) } + #[cfg(feature = "protobuf")] + TYPE_PROTO => { + let type_name_bytes = format::read_bytes(&mut cursor)?; + let type_name = String::from_utf8(type_name_bytes).map_err(|_| { + FormatError::Io(io::Error::new( + io::ErrorKind::InvalidData, + "proto type_name is not valid utf-8", + )) + })?; + let data = format::read_bytes(&mut cursor)?; + SnapValue::Proto { + type_name, + data: Bytes::from(data), + } + } _ => return Err(FormatError::UnknownTag(type_tag)), }; let expire_ms = format::read_i64(&mut cursor)?; From 94262c3005b1c87679b4dd104c8c5b8d06d2f936 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Tue, 10 Feb 2026 11:29:18 -0500 Subject: [PATCH 4/8] feat: add proto_set, proto_get, proto_type keyspace methods MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit type-checked read/write operations for protobuf values. no schema validation at this layer — that's the server's job. follows the same memory enforcement and expiry patterns as existing set/get. --- crates/ember-core/src/keyspace.rs | 86 +++++++++++++++++++++++++++++++ 1 file changed, 86 insertions(+) diff --git a/crates/ember-core/src/keyspace.rs b/crates/ember-core/src/keyspace.rs index 810da3a4..6528cd34 100644 --- a/crates/ember-core/src/keyspace.rs +++ b/crates/ember-core/src/keyspace.rs @@ -1842,6 +1842,92 @@ impl Keyspace { expired } + // -- protobuf operations -- + + /// Stores a protobuf value. No schema validation here — that's the + /// server's responsibility. Follows the same pattern as `set()`. + #[cfg(feature = "protobuf")] + pub fn proto_set( + &mut self, + key: String, + type_name: String, + data: Bytes, + expire: Option, + ) -> SetResult { + let has_expiry = expire.is_some(); + let new_value = Value::Proto { type_name, data }; + + let new_size = memory::entry_size(&key, &new_value); + let old_size = self + .entries + .get(&key) + .map(|e| memory::entry_size(&key, &e.value)) + .unwrap_or(0); + let net_increase = new_size.saturating_sub(old_size); + + if !self.enforce_memory_limit(net_increase) { + return SetResult::OutOfMemory; + } + + if let Some(old_entry) = self.entries.get(&key) { + self.memory.replace(&key, &old_entry.value, &new_value); + let had_expiry = old_entry.expires_at_ms != 0; + match (had_expiry, has_expiry) { + (false, true) => self.expiry_count += 1, + (true, false) => self.expiry_count = self.expiry_count.saturating_sub(1), + _ => {} + } + } else { + self.memory.add(&key, &new_value); + if has_expiry { + self.expiry_count += 1; + } + } + + self.entries.insert(key, Entry::new(new_value, expire)); + SetResult::Ok + } + + /// Retrieves a proto value, returning `(type_name, data)` or `None`. + /// + /// Returns `Err(WrongType)` if the key holds a different value type. + #[cfg(feature = "protobuf")] + pub fn proto_get(&mut self, key: &str) -> Result, WrongType> { + if self.remove_if_expired(key) { + return Ok(None); + } + match self.entries.get_mut(key) { + Some(e) => { + if let Value::Proto { type_name, data } = &e.value { + let result = (type_name.clone(), data.clone()); + e.touch(); + Ok(Some(result)) + } else { + Err(WrongType) + } + } + None => Ok(None), + } + } + + /// Returns the protobuf message type name for a key, or `None` if + /// the key doesn't exist. + /// + /// Returns `Err(WrongType)` if the key holds a non-proto value. + #[cfg(feature = "protobuf")] + pub fn proto_type(&mut self, key: &str) -> Result, WrongType> { + if self.remove_if_expired(key) { + return Ok(None); + } + match self.entries.get(key) { + Some(e) => match &e.value { + Value::Proto { type_name, .. } => Ok(Some(type_name.clone())), + _ => Err(WrongType), + }, + None => Ok(None), + } + } + /// Sends a value to the background drop thread if one is configured /// and the value is large enough to justify the overhead. fn defer_drop(&self, value: Value) { From 7c60223c10e8e6ae27294a092a2a080844e27aad Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Tue, 10 Feb 2026 11:31:42 -0500 Subject: [PATCH 5/8] feat: add shard dispatch and AOF persistence for proto values adds ProtoSet/ProtoGet/ProtoType to ShardRequest/ShardResponse with dispatch arms calling the keyspace methods. AOF records ProtoSet and ProtoRegister handle persistence and recovery. also fills in missing read_payload_for_tag arms for tags 14-22 (HSET, HDEL, etc). --- crates/ember-core/src/shard.rs | 76 ++++++++++++++ crates/ember-persistence/src/aof.rs | 122 +++++++++++++++++++++++ crates/ember-persistence/src/recovery.rs | 17 ++++ 3 files changed, 215 insertions(+) diff --git a/crates/ember-core/src/shard.rs b/crates/ember-core/src/shard.rs index a216bd37..975f0455 100644 --- a/crates/ember-core/src/shard.rs +++ b/crates/ember-core/src/shard.rs @@ -260,6 +260,26 @@ pub enum ShardRequest { slot: u16, count: usize, }, + /// Stores a validated protobuf value. + #[cfg(feature = "protobuf")] + ProtoSet { + key: String, + type_name: String, + data: Bytes, + expire: Option, + nx: bool, + xx: bool, + }, + /// Retrieves a protobuf value. + #[cfg(feature = "protobuf")] + ProtoGet { + key: String, + }, + /// Returns the protobuf message type name for a key. + #[cfg(feature = "protobuf")] + ProtoType { + key: String, + }, } /// The shard's response to a request. @@ -316,6 +336,12 @@ pub enum ShardResponse { StringArray(Vec), /// HMGET result: array of optional values. OptionalArray(Vec>), + /// PROTO.GET result: (type_name, data) or None. + #[cfg(feature = "protobuf")] + ProtoValue(Option<(String, Bytes)>), + /// PROTO.TYPE result: message type name or None. + #[cfg(feature = "protobuf")] + ProtoTypeName(Option), } /// A request bundled with its reply channel. @@ -819,6 +845,36 @@ fn dispatch(ks: &mut Keyspace, req: &ShardRequest) -> ShardResponse { ShardRequest::GetKeysInSlot { slot, count } => { ShardResponse::StringArray(ks.get_keys_in_slot(*slot, *count)) } + #[cfg(feature = "protobuf")] + ShardRequest::ProtoSet { + key, + type_name, + data, + expire, + nx, + xx, + } => { + if *nx && ks.exists(key) { + return ShardResponse::Value(None); + } + if *xx && !ks.exists(key) { + return ShardResponse::Value(None); + } + match ks.proto_set(key.clone(), type_name.clone(), data.clone(), *expire) { + SetResult::Ok => ShardResponse::Ok, + SetResult::OutOfMemory => ShardResponse::OutOfMemory, + } + } + #[cfg(feature = "protobuf")] + ShardRequest::ProtoGet { key } => match ks.proto_get(key) { + Ok(val) => ShardResponse::ProtoValue(val), + Err(_) => ShardResponse::WrongType, + }, + #[cfg(feature = "protobuf")] + ShardRequest::ProtoType { key } => match ks.proto_type(key) { + Ok(name) => ShardResponse::ProtoTypeName(name), + Err(_) => ShardResponse::WrongType, + }, // snapshot/rewrite/flush_async are handled in the main loop, not here ShardRequest::Snapshot | ShardRequest::RewriteAof | ShardRequest::FlushDbAsync => { ShardResponse::Ok @@ -961,6 +1017,26 @@ fn to_aof_record(req: &ShardRequest, resp: &ShardResponse) -> Option members: members.clone(), }) } + // Proto commands + #[cfg(feature = "protobuf")] + ( + ShardRequest::ProtoSet { + key, + type_name, + data, + expire, + .. + }, + ShardResponse::Ok, + ) => { + 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, } } diff --git a/crates/ember-persistence/src/aof.rs b/crates/ember-persistence/src/aof.rs index 857ae2a9..0b4a4c6d 100644 --- a/crates/ember-persistence/src/aof.rs +++ b/crates/ember-persistence/src/aof.rs @@ -61,6 +61,10 @@ const TAG_INCRBY: u8 = 19; const TAG_DECRBY: u8 = 20; const TAG_APPEND: u8 = 21; const TAG_RENAME: u8 = 22; +#[cfg(feature = "protobuf")] +const TAG_PROTO_SET: u8 = 23; +#[cfg(feature = "protobuf")] +const TAG_PROTO_REGISTER: u8 = 24; /// A single mutation record stored in the AOF. #[derive(Debug, Clone, PartialEq)] @@ -123,6 +127,17 @@ pub enum AofRecord { Append { key: String, value: Bytes }, /// RENAME key newkey. Rename { key: String, newkey: String }, + /// PROTO.SET key type_name data [expire_ms]. + #[cfg(feature = "protobuf")] + ProtoSet { + key: String, + type_name: String, + data: Bytes, + expire_ms: i64, + }, + /// PROTO.REGISTER name descriptor_bytes (for schema persistence). + #[cfg(feature = "protobuf")] + ProtoRegister { name: String, descriptor: Bytes }, } impl AofRecord { @@ -266,6 +281,25 @@ impl AofRecord { format::write_bytes(&mut buf, key.as_bytes())?; format::write_bytes(&mut buf, newkey.as_bytes())?; } + #[cfg(feature = "protobuf")] + AofRecord::ProtoSet { + key, + type_name, + data, + expire_ms, + } => { + format::write_u8(&mut buf, TAG_PROTO_SET)?; + format::write_bytes(&mut buf, key.as_bytes())?; + format::write_bytes(&mut buf, type_name.as_bytes())?; + format::write_bytes(&mut buf, data)?; + format::write_i64(&mut buf, *expire_ms)?; + } + #[cfg(feature = "protobuf")] + AofRecord::ProtoRegister { name, descriptor } => { + format::write_u8(&mut buf, TAG_PROTO_REGISTER)?; + format::write_bytes(&mut buf, name.as_bytes())?; + format::write_bytes(&mut buf, descriptor)?; + } } Ok(buf) } @@ -416,6 +450,28 @@ impl AofRecord { let newkey = read_string(&mut cursor, "newkey")?; Ok(AofRecord::Rename { key, newkey }) } + #[cfg(feature = "protobuf")] + TAG_PROTO_SET => { + let key = read_string(&mut cursor, "key")?; + let type_name = read_string(&mut cursor, "type_name")?; + let data = format::read_bytes(&mut cursor)?; + let expire_ms = format::read_i64(&mut cursor)?; + Ok(AofRecord::ProtoSet { + key, + type_name, + data: Bytes::from(data), + expire_ms, + }) + } + #[cfg(feature = "protobuf")] + TAG_PROTO_REGISTER => { + let name = read_string(&mut cursor, "name")?; + let descriptor = format::read_bytes(&mut cursor)?; + Ok(AofRecord::ProtoRegister { + name, + descriptor: Bytes::from(descriptor), + }) + } _ => Err(FormatError::UnknownTag(tag)), } } @@ -789,6 +845,72 @@ impl AofReader { let key = format::read_bytes(&mut self.reader)?; format::write_bytes(&mut payload, &key)?; } + TAG_HSET => { + let key = format::read_bytes(&mut self.reader)?; + format::write_bytes(&mut payload, &key)?; + let count = format::read_u32(&mut self.reader)?; + format::write_u32(&mut payload, count)?; + for _ in 0..count { + let field = format::read_bytes(&mut self.reader)?; + format::write_bytes(&mut payload, &field)?; + let value = format::read_bytes(&mut self.reader)?; + format::write_bytes(&mut payload, &value)?; + } + } + TAG_HDEL | TAG_SADD | TAG_SREM => { + let key = format::read_bytes(&mut self.reader)?; + format::write_bytes(&mut payload, &key)?; + let count = format::read_u32(&mut self.reader)?; + format::write_u32(&mut payload, count)?; + for _ in 0..count { + let item = format::read_bytes(&mut self.reader)?; + format::write_bytes(&mut payload, &item)?; + } + } + TAG_HINCRBY => { + let key = format::read_bytes(&mut self.reader)?; + format::write_bytes(&mut payload, &key)?; + let field = format::read_bytes(&mut self.reader)?; + format::write_bytes(&mut payload, &field)?; + let delta = format::read_i64(&mut self.reader)?; + format::write_i64(&mut payload, delta)?; + } + TAG_INCRBY | TAG_DECRBY => { + let key = format::read_bytes(&mut self.reader)?; + format::write_bytes(&mut payload, &key)?; + let delta = format::read_i64(&mut self.reader)?; + format::write_i64(&mut payload, delta)?; + } + TAG_APPEND => { + let key = format::read_bytes(&mut self.reader)?; + format::write_bytes(&mut payload, &key)?; + let value = format::read_bytes(&mut self.reader)?; + format::write_bytes(&mut payload, &value)?; + } + TAG_RENAME => { + let key = format::read_bytes(&mut self.reader)?; + format::write_bytes(&mut payload, &key)?; + let newkey = format::read_bytes(&mut self.reader)?; + format::write_bytes(&mut payload, &newkey)?; + } + #[cfg(feature = "protobuf")] + TAG_PROTO_SET => { + let key = format::read_bytes(&mut self.reader)?; + format::write_bytes(&mut payload, &key)?; + let type_name = format::read_bytes(&mut self.reader)?; + format::write_bytes(&mut payload, &type_name)?; + let data = format::read_bytes(&mut self.reader)?; + format::write_bytes(&mut payload, &data)?; + let expire_ms = format::read_i64(&mut self.reader)?; + format::write_i64(&mut payload, expire_ms)?; + } + #[cfg(feature = "protobuf")] + TAG_PROTO_REGISTER => { + let name = format::read_bytes(&mut self.reader)?; + format::write_bytes(&mut payload, &name)?; + let descriptor = format::read_bytes(&mut self.reader)?; + format::write_bytes(&mut payload, &descriptor)?; + } _ => return Err(FormatError::UnknownTag(tag)), } let stored_crc = format::read_u32(&mut self.reader)?; diff --git a/crates/ember-persistence/src/recovery.rs b/crates/ember-persistence/src/recovery.rs index b1374e14..2745e37f 100644 --- a/crates/ember-persistence/src/recovery.rs +++ b/crates/ember-persistence/src/recovery.rs @@ -423,6 +423,23 @@ fn replay_aof( } } } + #[cfg(feature = "protobuf")] + AofRecord::ProtoSet { + key, + type_name, + data, + expire_ms, + } => { + 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 + } } count += 1; } From 8419f790455320d2f64ba7d8debb8603fa27fe76 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Tue, 10 Feb 2026 11:32:46 -0500 Subject: [PATCH 6/8] feat: wire schema registry through engine config adds optional SharedSchemaRegistry to EngineConfig and Engine, with an accessor method for the server layer. implements Debug for SchemaRegistry manually since DescriptorPool doesn't derive it. --- crates/ember-core/src/engine.rs | 18 +++++++++++++++++- crates/ember-core/src/schema.rs | 11 +++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/crates/ember-core/src/engine.rs b/crates/ember-core/src/engine.rs index b8499514..480e3302 100644 --- a/crates/ember-core/src/engine.rs +++ b/crates/ember-core/src/engine.rs @@ -24,6 +24,10 @@ pub struct EngineConfig { /// Optional persistence configuration. When set, each shard gets /// its own AOF and snapshot files under this directory. pub persistence: Option, + /// Optional schema registry for protobuf value validation. + /// When set, enables PROTO.* commands. + #[cfg(feature = "protobuf")] + pub schema_registry: Option, } /// The sharded engine. Owns handles to all shard tasks and routes @@ -34,6 +38,8 @@ pub struct EngineConfig { #[derive(Debug, Clone)] pub struct Engine { shards: Vec, + #[cfg(feature = "protobuf")] + schema_registry: Option, } impl Engine { @@ -69,7 +75,11 @@ impl Engine { }) .collect(); - Self { shards } + Self { + shards, + #[cfg(feature = "protobuf")] + schema_registry: config.schema_registry, + } } /// Creates an engine with one shard per available CPU core. @@ -88,6 +98,12 @@ impl Engine { Self::with_config(cores, config) } + /// Returns a reference to the schema registry, if protobuf is enabled. + #[cfg(feature = "protobuf")] + pub fn schema_registry(&self) -> Option<&crate::schema::SharedSchemaRegistry> { + self.schema_registry.as_ref() + } + /// Returns the number of shards. pub fn shard_count(&self) -> usize { self.shards.len() diff --git a/crates/ember-core/src/schema.rs b/crates/ember-core/src/schema.rs index 9b817c1b..ab6e58f1 100644 --- a/crates/ember-core/src/schema.rs +++ b/crates/ember-core/src/schema.rs @@ -44,6 +44,9 @@ struct RegisteredSchema { /// /// Each schema is identified by a user-chosen name (e.g. "users/v1") /// and contains one or more message type definitions. +/// +/// Debug is implemented manually because `DescriptorPool` doesn't +/// derive it. pub struct SchemaRegistry { schemas: HashMap, } @@ -173,6 +176,14 @@ impl SchemaRegistry { } } +impl std::fmt::Debug for SchemaRegistry { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("SchemaRegistry") + .field("schema_count", &self.schemas.len()) + .finish() + } +} + impl Default for SchemaRegistry { fn default() -> Self { Self::new() From 4a9c56fd21053789c161a4687c8884d7b2cd668e Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Tue, 10 Feb 2026 11:35:51 -0500 Subject: [PATCH 7/8] fix: add schema_registry field to server config builder --- Cargo.lock | 50 +++++++++++++++++++++++++++++++ crates/ember-server/src/config.rs | 2 ++ 2 files changed, 52 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index 1fa179c4..f537376b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -143,6 +143,12 @@ dependencies = [ "serde", ] +[[package]] +name = "anyhow" +version = "1.0.101" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f0e0fee31ef5ed1ba1316088939cea399010ed7731dba877ed44aeb407a75ea" + [[package]] name = "arrayvec" version = "0.7.6" @@ -730,6 +736,7 @@ dependencies = [ "ember-protocol", "ordered-float", "parking_lot", + "prost-reflect", "rand 0.9.2", "tempfile", "thiserror 2.0.18", @@ -1601,6 +1608,49 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "prost" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2796faa41db3ec313a31f7624d9286acf277b52de526150b7e69f3debf891ee5" +dependencies = [ + "bytes", + "prost-derive", +] + +[[package]] +name = "prost-derive" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" +dependencies = [ + "anyhow", + "itertools", + "proc-macro2", + "quote", + "syn 2.0.114", +] + +[[package]] +name = "prost-reflect" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b5edd582b62f5cde844716e66d92565d7faf7ab1445c8cebce6e00fba83ddb2" +dependencies = [ + "once_cell", + "prost", + "prost-types", +] + +[[package]] +name = "prost-types" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52c2c1bf36ddb1a1c396b3601a3cec27c2462e45f07c386894ec3ccf5332bd16" +dependencies = [ + "prost", +] + [[package]] name = "ptr_meta" version = "0.1.4" diff --git a/crates/ember-server/src/config.rs b/crates/ember-server/src/config.rs index a9516779..3cc6970a 100644 --- a/crates/ember-server/src/config.rs +++ b/crates/ember-server/src/config.rs @@ -95,6 +95,8 @@ pub fn build_engine_config( ..ShardConfig::default() }, persistence, + #[cfg(feature = "protobuf")] + schema_registry: None, } } From 41b36b3a3f00215cfd0c525d5950314a08bfbf53 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Tue, 10 Feb 2026 11:40:42 -0500 Subject: [PATCH 8/8] chore: update prost-reflect to 0.16, run fmt --- Cargo.lock | 17 ++++++++--------- Cargo.toml | 2 +- crates/ember-core/src/schema.rs | 8 ++++++-- crates/ember-persistence/src/recovery.rs | 10 +++++----- crates/ember-persistence/src/snapshot.rs | 5 +---- 5 files changed, 21 insertions(+), 21 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f537376b..8da64a8a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1610,9 +1610,9 @@ dependencies = [ [[package]] name = "prost" -version = "0.13.5" +version = "0.14.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2796faa41db3ec313a31f7624d9286acf277b52de526150b7e69f3debf891ee5" +checksum = "d2ea70524a2f82d518bce41317d0fae74151505651af45faf1ffbd6fd33f0568" dependencies = [ "bytes", "prost-derive", @@ -1620,9 +1620,9 @@ dependencies = [ [[package]] name = "prost-derive" -version = "0.13.5" +version = "0.14.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" +checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b" dependencies = [ "anyhow", "itertools", @@ -1633,20 +1633,19 @@ dependencies = [ [[package]] name = "prost-reflect" -version = "0.14.7" +version = "0.16.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7b5edd582b62f5cde844716e66d92565d7faf7ab1445c8cebce6e00fba83ddb2" +checksum = "b89455ef41ed200cafc47c76c552ee7792370ac420497e551f16123a9135f76e" dependencies = [ - "once_cell", "prost", "prost-types", ] [[package]] name = "prost-types" -version = "0.13.5" +version = "0.14.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52c2c1bf36ddb1a1c396b3601a3cec27c2462e45f07c386894ec3ccf5332bd16" +checksum = "8991c4cbdb8bc5b11f0b074ffe286c30e523de90fee5ba8132f1399f23cb3dd7" dependencies = [ "prost", ] diff --git a/Cargo.toml b/Cargo.toml index 0136fc12..18522977 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -62,7 +62,7 @@ rustls-pemfile = "2" rustls-native-certs = "0.8" # dynamic protobuf messages -prost-reflect = "0.14" +prost-reflect = "0.16" # internal crates (version required for crates.io publishing) emberkv-core = { version = "0.4.2", path = "crates/ember-core" } diff --git a/crates/ember-core/src/schema.rs b/crates/ember-core/src/schema.rs index ab6e58f1..271ecad4 100644 --- a/crates/ember-core/src/schema.rs +++ b/crates/ember-core/src/schema.rs @@ -71,7 +71,11 @@ impl SchemaRegistry { /// /// Returns the list of message type names defined in the schema. /// Fails if the name is already registered or the descriptor is invalid. - pub fn register(&mut self, name: String, descriptor_bytes: Bytes) -> Result, SchemaError> { + pub fn register( + &mut self, + name: String, + descriptor_bytes: Bytes, + ) -> Result, SchemaError> { if self.schemas.contains_key(&name) { return Err(SchemaError::AlreadyExists(name)); } @@ -197,10 +201,10 @@ mod tests { /// Builds a minimal FileDescriptorSet containing a single message type. /// Uses prost-reflect's own encoding rather than shelling out to protoc. fn make_descriptor(package: &str, message_name: &str, field_name: &str) -> Bytes { + use prost_reflect::prost::Message; use prost_reflect::prost_types::{ DescriptorProto, FieldDescriptorProto, FileDescriptorProto, FileDescriptorSet, }; - use prost_reflect::prost::Message; let fds = FileDescriptorSet { file: vec![FileDescriptorProto { diff --git a/crates/ember-persistence/src/recovery.rs b/crates/ember-persistence/src/recovery.rs index 2745e37f..c106e0cf 100644 --- a/crates/ember-persistence/src/recovery.rs +++ b/crates/ember-persistence/src/recovery.rs @@ -39,7 +39,10 @@ pub enum RecoveredValue { Set(HashSet), /// A protobuf message: type name + serialized bytes. #[cfg(feature = "protobuf")] - Proto { type_name: String, data: Bytes }, + Proto { + type_name: String, + data: Bytes, + }, } impl From for RecoveredValue { @@ -430,10 +433,7 @@ fn replay_aof( data, expire_ms, } => { - map.insert( - key, - (RecoveredValue::Proto { type_name, data }, expire_ms), - ); + map.insert(key, (RecoveredValue::Proto { type_name, data }, expire_ms)); } #[cfg(feature = "protobuf")] AofRecord::ProtoRegister { .. } => { diff --git a/crates/ember-persistence/src/snapshot.rs b/crates/ember-persistence/src/snapshot.rs index 85a7026c..b5a594ae 100644 --- a/crates/ember-persistence/src/snapshot.rs +++ b/crates/ember-persistence/src/snapshot.rs @@ -196,10 +196,7 @@ impl SnapshotWriter { } } #[cfg(feature = "protobuf")] - SnapValue::Proto { - type_name, - data, - } => { + SnapValue::Proto { type_name, data } => { format::write_u8(&mut buf, TYPE_PROTO)?; format::write_bytes(&mut buf, type_name.as_bytes())?; format::write_bytes(&mut buf, data)?;