diff --git a/Cargo.lock b/Cargo.lock index 1fa179c4..8da64a8a 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,48 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "prost" +version = "0.14.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2ea70524a2f82d518bce41317d0fae74151505651af45faf1ffbd6fd33f0568" +dependencies = [ + "bytes", + "prost-derive", +] + +[[package]] +name = "prost-derive" +version = "0.14.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b" +dependencies = [ + "anyhow", + "itertools", + "proc-macro2", + "quote", + "syn 2.0.114", +] + +[[package]] +name = "prost-reflect" +version = "0.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b89455ef41ed200cafc47c76c552ee7792370ac420497e551f16123a9135f76e" +dependencies = [ + "prost", + "prost-types", +] + +[[package]] +name = "prost-types" +version = "0.14.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8991c4cbdb8bc5b11f0b074ffe286c30e523de90fee5ba8132f1399f23cb3dd7" +dependencies = [ + "prost", +] + [[package]] name = "ptr_meta" version = "0.1.4" diff --git a/Cargo.toml b/Cargo.toml index d2eac414..18522977 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.16" + # 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..c91b5f1c 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", "ember-persistence/protobuf"] [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-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/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) { 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/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/schema.rs b/crates/ember-core/src/schema.rs new file mode 100644 index 00000000..271ecad4 --- /dev/null +++ b/crates/ember-core/src/schema.rs @@ -0,0 +1,337 @@ +//! 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. +/// +/// Debug is implemented manually because `DescriptorPool` doesn't +/// derive it. +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 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() + } +} + +#[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::Message; + use prost_reflect::prost_types::{ + DescriptorProto, FieldDescriptorProto, FileDescriptorProto, FileDescriptorSet, + }; + + 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"]); + } +} diff --git a/crates/ember-core/src/shard.rs b/crates/ember-core/src/shard.rs index afc10fa2..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. @@ -418,6 +444,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); } @@ -817,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 @@ -959,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, } } @@ -1064,6 +1142,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/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-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 c0900764..c106e0cf 100644 --- a/crates/ember-persistence/src/recovery.rs +++ b/crates/ember-persistence/src/recovery.rs @@ -37,6 +37,12 @@ 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 +53,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 }, } } } @@ -418,6 +426,20 @@ 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; } diff --git a/crates/ember-persistence/src/snapshot.rs b/crates/ember-persistence/src/snapshot.rs index b7995f5d..b5a594ae 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,12 @@ 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 +430,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 +579,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)?; 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 } 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, } }