From 6bc079a08bb2dad8182734b21e772ccbd46c89cb Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Wed, 18 Feb 2026 19:48:05 -0500 Subject: [PATCH 1/3] feat(core): add ReplicationEvent broadcast and SerializeSnapshot to shard - make AofRecord::to_bytes/from_bytes pub for wire protocol use - add write_snapshot_bytes/read_snapshot_from_bytes to snapshot.rs for in-memory snapshot serialization without filesystem I/O - add ReplicationEvent struct: shard_id, offset, record - add ShardRequest::SerializeSnapshot and ShardResponse::SnapshotData - thread optional broadcast::Sender through spawn_shard, run_shard, and process_message; publishes after each mutation - add Engine::subscribe_replication() to create broadcast receivers - thread replication_tx through EngineConfig for wiring in main.rs --- crates/ember-core/src/engine.rs | 22 +- crates/ember-core/src/lib.rs | 2 +- crates/ember-core/src/shard.rs | 93 +++++++- crates/ember-persistence/src/aof.rs | 4 +- crates/ember-persistence/src/snapshot.rs | 274 +++++++++++++++++++++++ 5 files changed, 385 insertions(+), 10 deletions(-) diff --git a/crates/ember-core/src/engine.rs b/crates/ember-core/src/engine.rs index cdadf8ef..ea3bdee4 100644 --- a/crates/ember-core/src/engine.rs +++ b/crates/ember-core/src/engine.rs @@ -4,10 +4,12 @@ //! of the key. Each shard is an independent tokio task — no locks on //! the hot path. +use tokio::sync::broadcast; + use crate::dropper::DropHandle; use crate::error::ShardError; use crate::keyspace::ShardConfig; -use crate::shard::{self, ShardHandle, ShardPersistenceConfig, ShardRequest, ShardResponse}; +use crate::shard::{self, ReplicationEvent, ShardHandle, ShardPersistenceConfig, ShardRequest, ShardResponse}; /// Channel buffer size per shard. 256 is large enough to absorb /// bursts without putting meaningful back-pressure on connections. @@ -21,6 +23,12 @@ pub struct EngineConfig { /// Optional persistence configuration. When set, each shard gets /// its own AOF and snapshot files under this directory. pub persistence: Option, + /// Optional broadcast sender for replication events. + /// + /// When set, every successful mutation is published as a + /// [`ReplicationEvent`] so replication clients can stream it to + /// replicas. + pub replication_tx: Option>, /// Optional schema registry for protobuf value validation. /// When set, enables PROTO.* commands. #[cfg(feature = "protobuf")] @@ -35,6 +43,7 @@ pub struct EngineConfig { #[derive(Debug, Clone)] pub struct Engine { shards: Vec, + replication_tx: Option>, #[cfg(feature = "protobuf")] schema_registry: Option, } @@ -72,6 +81,7 @@ impl Engine { shard_config, config.persistence.clone(), Some(drop_handle.clone()), + config.replication_tx.clone(), #[cfg(feature = "protobuf")] config.schema_registry.clone(), ) @@ -80,6 +90,7 @@ impl Engine { Self { shards, + replication_tx: config.replication_tx, #[cfg(feature = "protobuf")] schema_registry: config.schema_registry, } @@ -112,6 +123,15 @@ impl Engine { self.shards.len() } + /// Creates a new broadcast receiver for replication events. + /// + /// Returns `None` if no replication channel was configured. Each + /// caller gets an independent receiver starting from the current + /// broadcast position — not from the beginning of the stream. + pub fn subscribe_replication(&self) -> Option> { + self.replication_tx.as_ref().map(|tx| tx.subscribe()) + } + /// Sends a request to a specific shard by index. /// /// Used by SCAN to iterate through shards sequentially. diff --git a/crates/ember-core/src/lib.rs b/crates/ember-core/src/lib.rs index 254c234e..5a799c0e 100644 --- a/crates/ember-core/src/lib.rs +++ b/crates/ember-core/src/lib.rs @@ -27,5 +27,5 @@ pub use keyspace::{ }; #[cfg(feature = "vector")] pub use keyspace::{VAddResult, VectorWriteError}; -pub use shard::{ShardPersistenceConfig, ShardRequest, ShardResponse}; +pub use shard::{ReplicationEvent, ShardHandle, ShardPersistenceConfig, ShardRequest, ShardResponse}; pub use types::Value; diff --git a/crates/ember-core/src/shard.rs b/crates/ember-core/src/shard.rs index e6902835..565f915e 100644 --- a/crates/ember-core/src/shard.rs +++ b/crates/ember-core/src/shard.rs @@ -12,7 +12,7 @@ use bytes::Bytes; use ember_persistence::aof::{AofRecord, AofWriter, FsyncPolicy}; use ember_persistence::recovery::{self, RecoveredValue}; use ember_persistence::snapshot::{self, SnapEntry, SnapValue, SnapshotWriter}; -use tokio::sync::{mpsc, oneshot}; +use tokio::sync::{broadcast, mpsc, oneshot}; use tracing::{info, warn}; use crate::dropper::DropHandle; @@ -32,6 +32,21 @@ const EXPIRY_TICK: Duration = Duration::from_millis(100); /// How often to fsync when using the `EverySec` policy. const FSYNC_INTERVAL: Duration = Duration::from_secs(1); +/// A mutation event broadcast to replication subscribers. +/// +/// Published after every successful mutation on the hot path. The +/// `offset` is per-shard and monotonically increasing — replicas use it +/// to detect gaps and trigger re-sync when they fall behind. +#[derive(Debug, Clone)] +pub struct ReplicationEvent { + /// The shard that produced this event. + pub shard_id: u16, + /// Monotonically increasing per-shard offset. + pub offset: u64, + /// The mutation record, ready to replay on a replica. + pub record: AofRecord, +} + /// Optional persistence configuration for a shard. #[derive(Debug, Clone)] pub struct ShardPersistenceConfig { @@ -239,6 +254,11 @@ pub enum ShardRequest { Stats, /// Triggers a snapshot write. Snapshot, + /// Serializes the current shard state to bytes (in-memory snapshot). + /// + /// Used by the replication server to capture a consistent shard + /// snapshot for transmission to a new replica without filesystem I/O. + SerializeSnapshot, /// Triggers an AOF rewrite (snapshot + truncate AOF). RewriteAof, /// Clears all keys from the keyspace. @@ -427,6 +447,8 @@ pub enum ShardResponse { StringArray(Vec), /// Serialized key dump with remaining TTL (for MIGRATE/DUMP). KeyDump { data: Vec, ttl_ms: i64 }, + /// In-memory snapshot of the full shard state (for replication). + SnapshotData { shard_id: u16, data: Vec }, /// HMGET result: array of optional values. OptionalArray(Vec>), /// VADD result: element, vector, and whether it was newly added. @@ -523,6 +545,7 @@ pub fn spawn_shard( config: ShardConfig, persistence: Option, drop_handle: Option, + replication_tx: Option>, #[cfg(feature = "protobuf")] schema_registry: Option, ) -> ShardHandle { let (tx, rx) = mpsc::channel(buffer); @@ -531,6 +554,7 @@ pub fn spawn_shard( config, persistence, drop_handle, + replication_tx, #[cfg(feature = "protobuf")] schema_registry, )); @@ -544,6 +568,7 @@ async fn run_shard( config: ShardConfig, persistence: Option, drop_handle: Option, + replication_tx: Option>, #[cfg(feature = "protobuf")] schema_registry: Option, ) { let shard_id = config.shard_id; @@ -670,6 +695,9 @@ async fn run_shard( .map(|p| p.fsync_policy) .unwrap_or(FsyncPolicy::No); + // monotonically increasing per-shard replication offset + let mut replication_offset: u64 = 0; + // -- tickers -- let mut expiry_tick = tokio::time::interval(EXPIRY_TICK); expiry_tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); @@ -690,6 +718,8 @@ async fn run_shard( &persistence, &drop_handle, shard_id, + &replication_tx, + &mut replication_offset, #[cfg(feature = "protobuf")] &schema_registry, ); @@ -707,6 +737,8 @@ async fn run_shard( &persistence, &drop_handle, shard_id, + &replication_tx, + &mut replication_offset, #[cfg(feature = "protobuf")] &schema_registry, ); @@ -748,6 +780,8 @@ fn process_message( persistence: &Option, drop_handle: &Option, shard_id: u16, + replication_tx: &Option>, + replication_offset: &mut u64, #[cfg(feature = "protobuf")] schema_registry: &Option, ) { let request_kind = describe_request(&msg.request); @@ -758,9 +792,11 @@ fn process_message( schema_registry, ); + // collect mutation records once; used for both AOF and replication + let records = to_aof_records(&msg.request, &response); + // write AOF records for successful mutations if let Some(ref mut writer) = aof_writer { - let records = to_aof_records(&msg.request, &response); for record in &records { if let Err(e) = writer.write_record(record) { warn!(shard_id, "aof write failed: {e}"); @@ -773,6 +809,19 @@ fn process_message( } } + // broadcast mutation events to replication subscribers + if let Some(ref tx) = replication_tx { + for record in records { + *replication_offset += 1; + // ignore send errors — no subscribers or lagged consumers + let _ = tx.send(ReplicationEvent { + shard_id, + offset: *replication_offset, + record, + }); + } + } + // handle special requests that need access to persistence state match request_kind { RequestKind::Snapshot => { @@ -780,6 +829,11 @@ fn process_message( let _ = msg.reply.send(resp); return; } + RequestKind::SerializeSnapshot => { + let resp = handle_serialize_snapshot(keyspace, shard_id); + let _ = msg.reply.send(resp); + return; + } RequestKind::RewriteAof => { let resp = handle_rewrite( keyspace, @@ -810,6 +864,7 @@ fn process_message( /// handling after dispatch without borrowing the request again. enum RequestKind { Snapshot, + SerializeSnapshot, RewriteAof, FlushDbAsync, Other, @@ -818,6 +873,7 @@ enum RequestKind { fn describe_request(req: &ShardRequest) -> RequestKind { match req { ShardRequest::Snapshot => RequestKind::Snapshot, + ShardRequest::SerializeSnapshot => RequestKind::SerializeSnapshot, ShardRequest::RewriteAof => RequestKind::RewriteAof, ShardRequest::FlushDbAsync => RequestKind::FlushDbAsync, _ => RequestKind::Other, @@ -1269,10 +1325,11 @@ fn dispatch( }) }) } - // snapshot/rewrite/flush_async are handled in the main loop, not here - ShardRequest::Snapshot | ShardRequest::RewriteAof | ShardRequest::FlushDbAsync => { - ShardResponse::Ok - } + // these requests are intercepted in process_message, not handled here + ShardRequest::Snapshot + | ShardRequest::SerializeSnapshot + | ShardRequest::RewriteAof + | ShardRequest::FlushDbAsync => ShardResponse::Ok, } } @@ -1611,6 +1668,30 @@ fn handle_snapshot( } } +/// Serializes the current shard state to bytes without filesystem I/O. +/// +/// Used by the replication server to capture a snapshot for transmission +/// to a new replica. The format matches the file-based snapshot and can +/// be loaded via [`ember_persistence::snapshot::read_snapshot_from_bytes`]. +fn handle_serialize_snapshot(keyspace: &Keyspace, shard_id: u16) -> ShardResponse { + let entries: Vec = keyspace + .iter_entries() + .map(|(key, value, expire_ms)| SnapEntry { + key: key.to_owned(), + value: value_to_snap(value), + expire_ms, + }) + .collect(); + + match snapshot::write_snapshot_bytes(shard_id, &entries) { + Ok(data) => ShardResponse::SnapshotData { shard_id, data }, + Err(e) => { + warn!(shard_id, "snapshot serialization failed: {e}"); + ShardResponse::Err(format!("snapshot failed: {e}")) + } + } +} + /// Writes a snapshot and then truncates the AOF. /// /// When protobuf is enabled, re-persists all registered schemas to the diff --git a/crates/ember-persistence/src/aof.rs b/crates/ember-persistence/src/aof.rs index ab8ea063..4db7c5bd 100644 --- a/crates/ember-persistence/src/aof.rs +++ b/crates/ember-persistence/src/aof.rs @@ -333,7 +333,7 @@ impl AofRecord { } /// Serializes this record into a byte vector (tag + payload, no CRC). - fn to_bytes(&self) -> Result, FormatError> { + pub fn to_bytes(&self) -> Result, FormatError> { let mut buf = Vec::with_capacity(self.estimated_size()); format::write_u8(&mut buf, self.tag())?; @@ -489,7 +489,7 @@ impl AofRecord { } /// Deserializes a record from a byte slice (tag + payload, no CRC). - fn from_bytes(data: &[u8]) -> Result { + pub fn from_bytes(data: &[u8]) -> Result { let mut cursor = io::Cursor::new(data); let tag = format::read_u8(&mut cursor)?; match tag { diff --git a/crates/ember-persistence/src/snapshot.rs b/crates/ember-persistence/src/snapshot.rs index 4983ea9a..162af018 100644 --- a/crates/ember-persistence/src/snapshot.rs +++ b/crates/ember-persistence/src/snapshot.rs @@ -885,6 +885,280 @@ impl SnapshotReader { } } +/// Serializes snapshot entries to an in-memory byte buffer. +/// +/// The format is identical to the file-based snapshot so the bytes can be +/// sent over the network and loaded with [`read_snapshot_from_bytes`]. +/// Only unencrypted v2 format is produced — encryption is not used for the +/// in-memory replication path. +pub fn write_snapshot_bytes(shard_id: u16, entries: &[SnapEntry]) -> Result, FormatError> { + use std::io::{Seek, SeekFrom, Write as _}; + + let mut buf = io::Cursor::new(Vec::::new()); + let mut hasher = crc32fast::Hasher::new(); + + format::write_header(&mut buf, format::SNAP_MAGIC)?; + format::write_u16(&mut buf, shard_id)?; + // remember where the entry count lives so we can patch it at the end + let count_pos = buf.position(); + format::write_u32(&mut buf, 0u32)?; + + let mut count = 0u32; + for entry in entries { + let entry_bytes = serialize_entry(entry)?; + hasher.update(&entry_bytes); + buf.write_all(&entry_bytes)?; + count += 1; + } + + // patch entry count in the header + let end_pos = buf.position(); + buf.seek(SeekFrom::Start(count_pos))?; + format::write_u32(&mut buf, count)?; + buf.seek(SeekFrom::Start(end_pos))?; + + // footer CRC + let checksum = hasher.finalize(); + format::write_u32(&mut buf, checksum)?; + + Ok(buf.into_inner()) +} + +/// Deserializes snapshot entries from bytes produced by [`write_snapshot_bytes`]. +/// +/// Returns the shard ID from the header and all live entries. Validates the +/// footer CRC to detect transmission errors. +pub fn read_snapshot_from_bytes(data: &[u8]) -> Result<(u16, Vec), FormatError> { + let mut r = io::Cursor::new(data); + let mut hasher = crc32fast::Hasher::new(); + + let version = format::read_header(&mut r, format::SNAP_MAGIC)?; + if version != format::FORMAT_VERSION { + return Err(FormatError::UnsupportedVersion(version)); + } + let shard_id = format::read_u16(&mut r)?; + let entry_count = format::read_u32(&mut r)?; + + let mut entries = Vec::with_capacity(entry_count.min(65536) as usize); + for _ in 0..entry_count { + let (entry, entry_bytes) = read_entry_with_bytes(&mut r)?; + hasher.update(&entry_bytes); + entries.push(entry); + } + + // verify footer CRC + let expected = hasher.finalize(); + let stored = format::read_u32(&mut r)?; + format::verify_crc32_values(expected, stored)?; + + Ok((shard_id, entries)) +} + +/// Serializes a single snapshot entry to raw bytes (no encryption). +/// +/// Used by [`write_snapshot_bytes`] for in-memory serialization. +fn serialize_entry(entry: &SnapEntry) -> Result, FormatError> { + let mut buf = Vec::with_capacity(entry.estimated_size()); + format::write_bytes(&mut buf, entry.key.as_bytes())?; + match &entry.value { + SnapValue::String(data) => { + format::write_u8(&mut buf, TYPE_STRING)?; + format::write_bytes(&mut buf, data)?; + } + SnapValue::List(deque) => { + format::write_u8(&mut buf, TYPE_LIST)?; + format::write_len(&mut buf, deque.len())?; + for item in deque { + format::write_bytes(&mut buf, item)?; + } + } + SnapValue::SortedSet(members) => { + format::write_u8(&mut buf, TYPE_SORTED_SET)?; + format::write_len(&mut buf, members.len())?; + for (score, member) in members { + format::write_f64(&mut buf, *score)?; + format::write_bytes(&mut buf, member.as_bytes())?; + } + } + SnapValue::Hash(map) => { + format::write_u8(&mut buf, TYPE_HASH)?; + format::write_len(&mut buf, map.len())?; + for (field, value) in map { + format::write_bytes(&mut buf, field.as_bytes())?; + format::write_bytes(&mut buf, value)?; + } + } + SnapValue::Set(set) => { + format::write_u8(&mut buf, TYPE_SET)?; + format::write_len(&mut buf, set.len())?; + for member in set { + format::write_bytes(&mut buf, member.as_bytes())?; + } + } + #[cfg(feature = "vector")] + SnapValue::Vector { + metric, + quantization, + connectivity, + expansion_add, + dim, + elements, + } => { + format::write_u8(&mut buf, TYPE_VECTOR)?; + format::write_u8(&mut buf, *metric)?; + format::write_u8(&mut buf, *quantization)?; + format::write_u32(&mut buf, *connectivity)?; + format::write_u32(&mut buf, *expansion_add)?; + format::write_u32(&mut buf, *dim)?; + format::write_len(&mut buf, elements.len())?; + for (name, vector) in elements { + format::write_bytes(&mut buf, name.as_bytes())?; + for &v in vector { + format::write_f32(&mut buf, v)?; + } + } + } + #[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)?; + Ok(buf) +} + +/// Reads a single entry from a cursor and also returns the raw bytes +/// used for CRC computation. +fn read_entry_with_bytes( + r: &mut io::Cursor<&[u8]>, +) -> Result<(SnapEntry, Vec), FormatError> { + let mut entry_bytes = Vec::new(); + + let key_bytes = format::read_bytes(r)?; + format::write_bytes(&mut entry_bytes, &key_bytes)?; + let key = parse_utf8(key_bytes, "key")?; + + let type_tag = format::read_u8(r)?; + format::write_u8(&mut entry_bytes, type_tag)?; + + let value = match type_tag { + TYPE_STRING => { + let v = format::read_bytes(r)?; + format::write_bytes(&mut entry_bytes, &v)?; + SnapValue::String(Bytes::from(v)) + } + TYPE_LIST => { + let count = format::read_u32(r)?; + format::validate_collection_count(count, "list")?; + format::write_u32(&mut entry_bytes, count)?; + let mut deque = VecDeque::with_capacity(format::capped_capacity(count)); + for _ in 0..count { + let item = format::read_bytes(r)?; + format::write_bytes(&mut entry_bytes, &item)?; + deque.push_back(Bytes::from(item)); + } + SnapValue::List(deque) + } + TYPE_SORTED_SET => { + let count = format::read_u32(r)?; + format::validate_collection_count(count, "sorted set")?; + format::write_u32(&mut entry_bytes, count)?; + let mut members = Vec::with_capacity(format::capped_capacity(count)); + for _ in 0..count { + let score = format::read_f64(r)?; + format::write_f64(&mut entry_bytes, score)?; + let mb = format::read_bytes(r)?; + format::write_bytes(&mut entry_bytes, &mb)?; + members.push((score, parse_utf8(mb, "member")?)); + } + SnapValue::SortedSet(members) + } + TYPE_HASH => { + let count = format::read_u32(r)?; + format::validate_collection_count(count, "hash")?; + format::write_u32(&mut entry_bytes, count)?; + let mut map = HashMap::with_capacity(format::capped_capacity(count)); + for _ in 0..count { + let fb = format::read_bytes(r)?; + format::write_bytes(&mut entry_bytes, &fb)?; + let field = parse_utf8(fb, "hash field")?; + let vb = format::read_bytes(r)?; + format::write_bytes(&mut entry_bytes, &vb)?; + map.insert(field, Bytes::from(vb)); + } + SnapValue::Hash(map) + } + TYPE_SET => { + let count = format::read_u32(r)?; + format::validate_collection_count(count, "set")?; + format::write_u32(&mut entry_bytes, count)?; + let mut set = HashSet::with_capacity(format::capped_capacity(count)); + for _ in 0..count { + let mb = format::read_bytes(r)?; + format::write_bytes(&mut entry_bytes, &mb)?; + set.insert(parse_utf8(mb, "set member")?); + } + SnapValue::Set(set) + } + #[cfg(feature = "vector")] + TYPE_VECTOR => { + let metric = format::read_u8(r)?; + format::write_u8(&mut entry_bytes, metric)?; + let quantization = format::read_u8(r)?; + format::write_u8(&mut entry_bytes, quantization)?; + let connectivity = format::read_u32(r)?; + format::write_u32(&mut entry_bytes, connectivity)?; + let expansion_add = format::read_u32(r)?; + format::write_u32(&mut entry_bytes, expansion_add)?; + let dim = format::read_u32(r)?; + format::write_u32(&mut entry_bytes, dim)?; + let count = format::read_u32(r)?; + format::write_u32(&mut entry_bytes, count)?; + let mut elements = Vec::with_capacity(format::capped_capacity(count)); + for _ in 0..count { + let nb = format::read_bytes(r)?; + format::write_bytes(&mut entry_bytes, &nb)?; + let name = parse_utf8(nb, "vector element name")?; + let mut vector = Vec::with_capacity(dim as usize); + for _ in 0..dim { + let v = format::read_f32(r)?; + format::write_f32(&mut entry_bytes, v)?; + vector.push(v); + } + elements.push((name, vector)); + } + SnapValue::Vector { + metric, + quantization, + connectivity, + expansion_add, + dim, + elements, + } + } + #[cfg(feature = "protobuf")] + TYPE_PROTO => { + let tn_bytes = format::read_bytes(r)?; + format::write_bytes(&mut entry_bytes, &tn_bytes)?; + let type_name = parse_utf8(tn_bytes, "proto type_name")?; + let data = format::read_bytes(r)?; + format::write_bytes(&mut entry_bytes, &data)?; + SnapValue::Proto { + type_name, + data: Bytes::from(data), + } + } + _ => return Err(FormatError::UnknownTag(type_tag)), + }; + + let expire_ms = format::read_i64(r)?; + format::write_i64(&mut entry_bytes, expire_ms)?; + + Ok((SnapEntry { key, value, expire_ms }, entry_bytes)) +} + /// Returns the snapshot file path for a given shard in a data directory. pub fn snapshot_path(data_dir: &Path, shard_id: u16) -> PathBuf { data_dir.join(format!("shard-{shard_id}.snap")) From 0939d06aa61149af6fc06698b4368919d40eeda4 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Wed, 18 Feb 2026 20:23:20 -0500 Subject: [PATCH 2/3] fix(core): update shard test calls for replication_tx param --- crates/ember-core/src/shard.rs | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/crates/ember-core/src/shard.rs b/crates/ember-core/src/shard.rs index 565f915e..7700a6d1 100644 --- a/crates/ember-core/src/shard.rs +++ b/crates/ember-core/src/shard.rs @@ -1973,6 +1973,7 @@ mod tests { ShardConfig::default(), None, None, + None, #[cfg(feature = "protobuf")] None, ); @@ -2010,6 +2011,7 @@ mod tests { ShardConfig::default(), None, None, + None, #[cfg(feature = "protobuf")] None, ); @@ -2041,6 +2043,7 @@ mod tests { ShardConfig::default(), None, None, + None, #[cfg(feature = "protobuf")] None, ); @@ -2114,6 +2117,7 @@ mod tests { config.clone(), Some(pcfg.clone()), None, + None, #[cfg(feature = "protobuf")] None, ); @@ -2162,6 +2166,7 @@ mod tests { config, Some(pcfg), None, + None, #[cfg(feature = "protobuf")] None, ); From 936b2a0f4d71c567c9ac9451ae2c5045e5ef256c Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Wed, 18 Feb 2026 20:23:30 -0500 Subject: [PATCH 3/3] =?UTF-8?q?feat(server):=20implement=20replication=20s?= =?UTF-8?q?tream=20(primary=20=E2=86=92=20replica=20sync)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit adds a dedicated TCP replication channel between primary and replica nodes: - replication.rs: new module with ReplicationServer and ReplicationClient - primary side: per-shard snapshot handshake + incremental AofRecord stream - replica side: snapshot load + record replay with exponential backoff reconnect - wire protocol: little-endian framing with version, shard sync, and resync msgs - aof_record_to_shard_request: maps all 20+ record types to ShardRequest - cluster.rs: set_engine, start_replication_server, start_replication_client - replication port = data_port + gossip_offset + 2 (default: 16381) - cluster_replicate now starts the client after updating topology - replication_info for INFO replication output - server.rs: wire engine into coordinator + start replication server at startup - main.rs: create broadcast channel, inject into EngineConfig when cluster enabled - config.rs: include replication_tx in EngineConfig initializer - connection.rs: INFO replication section (role, replicas, master addr) --- crates/ember-server/src/cluster.rs | 128 +++++ crates/ember-server/src/config.rs | 1 + crates/ember-server/src/connection.rs | 28 + crates/ember-server/src/main.rs | 11 +- crates/ember-server/src/replication.rs | 744 +++++++++++++++++++++++++ crates/ember-server/src/server.rs | 7 + 6 files changed, 918 insertions(+), 1 deletion(-) create mode 100644 crates/ember-server/src/replication.rs diff --git a/crates/ember-server/src/cluster.rs b/crates/ember-server/src/cluster.rs index f798fe1d..c2907d10 100644 --- a/crates/ember-server/src/cluster.rs +++ b/crates/ember-server/src/cluster.rs @@ -18,6 +18,7 @@ use ember_cluster::{ ClusterStateData, ConfigParseError, GossipConfig, GossipEngine, GossipEvent, GossipMessage, MigrationManager, NodeId, NodeRole, RaftNode, RaftProposalError, SlotRange, SLOT_COUNT, }; +use ember_core::Engine; use ember_protocol::Frame; use tokio::net::UdpSocket; use tokio::sync::{mpsc, watch, Mutex, RwLock}; @@ -35,12 +36,16 @@ pub struct ClusterCoordinator { migration: Mutex, local_id: NodeId, gossip_port_offset: u16, + /// local data-plane bind address (used to compute replication port) + bind_addr: SocketAddr, /// bound UDP socket for gossip, set after spawn_gossip udp_socket: Mutex>>, /// directory for nodes.conf persistence (None disables saving) data_dir: Option, /// raft node for linearizable topology mutations; set once after startup raft_node: std::sync::OnceLock>, + /// engine handle for replication; set once during startup + engine: std::sync::OnceLock>, } impl std::fmt::Debug for ClusterCoordinator { @@ -99,9 +104,11 @@ impl ClusterCoordinator { migration: Mutex::new(MigrationManager::new()), local_id, gossip_port_offset: port_offset, + bind_addr, udp_socket: Mutex::new(None), data_dir, raft_node: std::sync::OnceLock::new(), + engine: std::sync::OnceLock::new(), }; Ok((coordinator, event_rx)) @@ -154,9 +161,11 @@ impl ClusterCoordinator { migration: Mutex::new(MigrationManager::new()), local_id, gossip_port_offset: port_offset, + bind_addr, udp_socket: Mutex::new(None), data_dir: Some(data_dir), raft_node: std::sync::OnceLock::new(), + engine: std::sync::OnceLock::new(), }; Ok((coordinator, event_rx)) @@ -783,6 +792,10 @@ impl ClusterCoordinator { let _ = incarnation; // used only to hold the lock briefly self.save_config().await; + + // connect to the primary's replication port and start streaming + self.start_replication_client(primary_id).await; + Frame::Simple("OK".into()) } @@ -811,6 +824,111 @@ impl ClusterCoordinator { } } + /// Attaches the engine so replication can start on demand. + /// + /// Must be called once after the engine is built, before any `CLUSTER REPLICATE` + /// commands are processed. + pub fn set_engine(&self, engine: Arc) { + let _ = self.engine.set(engine); + } + + /// Starts the replication server for this node. + /// + /// Binds a TCP listener on `bind_addr.port() + gossip_port_offset + 2` + /// and accepts replica connections indefinitely. This is a no-op if + /// no engine has been attached via `set_engine`. + pub async fn start_replication_server(self: &Arc) { + let Some(engine) = self.engine.get() else { + warn!("start_replication_server called before set_engine; skipping"); + return; + }; + + let repl_port = match self + .bind_addr + .port() + .checked_add(self.gossip_port_offset) + .and_then(|p| p.checked_add(2)) + { + Some(p) => p, + None => { + error!("replication port overflows u16; not starting replication server"); + return; + } + }; + + let local_id = self.local_id.to_string(); + if let Err(e) = + crate::replication::ReplicationServer::start(Arc::clone(engine), local_id, repl_port) + .await + { + error!("failed to start replication server on port {repl_port}: {e}"); + } + } + + /// Starts the replication client, connecting to the primary's replication port. + /// + /// The replication port of the primary is derived from its data-plane address + /// by adding `gossip_port_offset + 2`. + async fn start_replication_client(&self, primary_id: NodeId) { + let Some(engine) = self.engine.get() else { + warn!("start_replication_client called before set_engine; skipping"); + return; + }; + + let primary_addr = { + let state = self.state.read().await; + state.nodes.get(&primary_id).map(|n| n.addr) + }; + + let Some(addr) = primary_addr else { + warn!(%primary_id, "cannot start replication client: primary not found in state"); + return; + }; + + let repl_port = match addr + .port() + .checked_add(self.gossip_port_offset) + .and_then(|p| p.checked_add(2)) + { + Some(p) => p, + None => { + error!(%primary_id, "primary replication port overflows u16; not connecting"); + return; + } + }; + + let repl_addr = std::net::SocketAddr::new(addr.ip(), repl_port); + info!(%primary_id, %repl_addr, "starting replication client"); + crate::replication::ReplicationClient::start(Arc::clone(engine), repl_addr); + } + + /// Returns replication status for the `INFO replication` section. + /// + /// Returns `(role, Option)`. + pub async fn replication_info(&self) -> ReplicationInfo { + let state = self.state.read().await; + let local = state.nodes.get(&self.local_id); + let role = local.map(|n| n.role).unwrap_or(NodeRole::Primary); + let primary_addr = if role == NodeRole::Replica { + local + .and_then(|n| n.replicates) + .and_then(|id| state.nodes.get(&id)) + .map(|n| n.addr) + } else { + None + }; + let replica_count = if role == NodeRole::Primary { + local.map(|n| n.replicas.len()).unwrap_or(0) + } else { + 0 + }; + ReplicationInfo { + role, + primary_addr, + replica_count, + } + } + /// Pushes the local node's current slot ownership into the gossip engine /// so it propagates to the rest of the cluster. async fn broadcast_local_slots(&self, slots: Vec) { @@ -1135,6 +1253,16 @@ impl ClusterCoordinator { } } +/// Snapshot of replication status for the `INFO replication` command. +#[derive(Debug)] +pub struct ReplicationInfo { + pub role: NodeRole, + /// Address of the primary this node replicates from (replica only). + pub primary_addr: Option, + /// Number of connected replicas (primary only). + pub replica_count: usize, +} + /// Compacts a flat list of slot numbers into contiguous `SlotRange` values. fn compact_slots(slots: &[u16]) -> Vec { let mut sorted: Vec = slots.to_vec(); diff --git a/crates/ember-server/src/config.rs b/crates/ember-server/src/config.rs index 3cc6970a..afd89718 100644 --- a/crates/ember-server/src/config.rs +++ b/crates/ember-server/src/config.rs @@ -95,6 +95,7 @@ pub fn build_engine_config( ..ShardConfig::default() }, persistence, + replication_tx: None, #[cfg(feature = "protobuf")] schema_registry: None, } diff --git a/crates/ember-server/src/connection.rs b/crates/ember-server/src/connection.rs index 31104379..c75311c2 100644 --- a/crates/ember-server/src/connection.rs +++ b/crates/ember-server/src/connection.rs @@ -3485,6 +3485,34 @@ async fn render_info(engine: &Engine, ctx: &Arc, section: Option< } } + if want("REPLICATION") { + out.push_str("# Replication\r\n"); + if let Some(ref cluster) = ctx.cluster { + let info = cluster.replication_info().await; + use ember_cluster::NodeRole; + match info.role { + NodeRole::Primary => { + out.push_str("role:primary\r\n"); + out.push_str(&format!("connected_replicas:{}\r\n", info.replica_count)); + } + NodeRole::Replica => { + out.push_str("role:replica\r\n"); + if let Some(addr) = info.primary_addr { + out.push_str(&format!("master_host:{}\r\n", addr.ip())); + out.push_str(&format!("master_port:{}\r\n", addr.port())); + out.push_str("master_link_status:up\r\n"); + } else { + out.push_str("master_link_status:down\r\n"); + } + } + } + } else { + out.push_str("role:primary\r\n"); + out.push_str("connected_replicas:0\r\n"); + } + out.push_str("\r\n"); + } + // trim trailing blank line if out.ends_with("\r\n\r\n") { out.truncate(out.len() - 2); diff --git a/crates/ember-server/src/main.rs b/crates/ember-server/src/main.rs index 0d7c02ac..6133e4b7 100644 --- a/crates/ember-server/src/main.rs +++ b/crates/ember-server/src/main.rs @@ -13,6 +13,7 @@ mod connection_common; mod grpc; mod metrics; mod pubsub; +mod replication; mod server; mod slowlog; mod tls; @@ -23,7 +24,7 @@ use std::sync::Arc; use clap::Parser; use ember_cluster::{raft_id_from_node_id, GossipConfig, NodeId, RaftNode, RaftStorage}; -use ember_core::ShardPersistenceConfig; +use ember_core::{ReplicationEvent, ShardPersistenceConfig}; use tracing::info; #[cfg(feature = "protobuf")] use tracing::warn; @@ -414,6 +415,14 @@ async fn main() { exit_err("error: --cluster-bootstrap requires --cluster-enabled"); } + // set up the replication broadcast channel when cluster is enabled. + // the sender goes into engine_config so shards can publish mutations; + // replicas subscribe to it when they connect for full sync + streaming. + if args.cluster_enabled { + let (repl_tx, _) = tokio::sync::broadcast::channel::(65536); + engine_config.replication_tx = Some(repl_tx); + } + // build cluster coordinator if cluster mode is enabled let cluster: Option> = if args.cluster_enabled { if args.port.checked_add(args.cluster_port_offset).is_none() { diff --git a/crates/ember-server/src/replication.rs b/crates/ember-server/src/replication.rs new file mode 100644 index 00000000..54b1113e --- /dev/null +++ b/crates/ember-server/src/replication.rs @@ -0,0 +1,744 @@ +//! Replication stream: primary → replica data sync. +//! +//! The primary side (`ReplicationServer`) accepts TCP connections from +//! replicas and streams all mutations as `AofRecord` frames after an +//! initial full-sync snapshot. The replica side (`ReplicationClient`) +//! connects, loads the snapshot, and applies incremental records. +//! +//! # Wire protocol +//! +//! All integers are little-endian. +//! +//! ```text +//! // Replica → primary (handshake request): +//! [version: 1B][num_shards: 2B] +//! +//! // Primary → replica (handshake response): +//! [version: 1B][num_shards: 2B][primary_id_len: 1B][primary_id: N bytes] +//! [status: 1B] (0 = ok, 1 = shard count mismatch) +//! +//! // For each shard (if status = 0): +//! [MSG_SHARD_SYNC: 1B][shard_id: 2B][snapshot_len: 4B][snapshot_bytes] +//! [MSG_SHARD_OFFSET: 1B][shard_id: 2B][offset: 8B] +//! +//! // Incremental records (unbounded stream): +//! [MSG_RECORD: 1B][shard_id: 2B][offset: 8B][record_len: 4B][record_bytes] +//! +//! // When replica falls behind (broadcast lag): +//! [MSG_RESYNC: 1B] primary closes the connection; replica reconnects +//! ``` + +use std::net::SocketAddr; +use std::sync::Arc; +use std::time::Duration; + +use bytes::Bytes; +use ember_core::{Engine, ShardRequest, ShardResponse}; +use ember_persistence::aof::AofRecord; +use ember_persistence::snapshot::{self, SnapValue}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; +use tokio::sync::broadcast; +use tracing::{debug, error, info, warn}; + +// -- protocol constants -- + +const REPL_VERSION: u8 = 1; +const STATUS_OK: u8 = 0; +const STATUS_SHARD_MISMATCH: u8 = 1; + +const MSG_SHARD_SYNC: u8 = 2; +const MSG_SHARD_OFFSET: u8 = 3; +const MSG_RECORD: u8 = 4; +const MSG_RESYNC: u8 = 5; + +// -- helpers -- + +async fn write_u8(w: &mut TcpStream, val: u8) -> std::io::Result<()> { + w.write_all(&[val]).await +} + +async fn write_u16_le(w: &mut TcpStream, val: u16) -> std::io::Result<()> { + w.write_all(&val.to_le_bytes()).await +} + +async fn write_u32_le(w: &mut TcpStream, val: u32) -> std::io::Result<()> { + w.write_all(&val.to_le_bytes()).await +} + +async fn write_u64_le(w: &mut TcpStream, val: u64) -> std::io::Result<()> { + w.write_all(&val.to_le_bytes()).await +} + +async fn read_u8(r: &mut TcpStream) -> std::io::Result { + let mut buf = [0u8; 1]; + r.read_exact(&mut buf).await?; + Ok(buf[0]) +} + +async fn read_u16_le(r: &mut TcpStream) -> std::io::Result { + let mut buf = [0u8; 2]; + r.read_exact(&mut buf).await?; + Ok(u16::from_le_bytes(buf)) +} + +async fn read_u32_le(r: &mut TcpStream) -> std::io::Result { + let mut buf = [0u8; 4]; + r.read_exact(&mut buf).await?; + Ok(u32::from_le_bytes(buf)) +} + +async fn read_u64_le(r: &mut TcpStream) -> std::io::Result { + let mut buf = [0u8; 8]; + r.read_exact(&mut buf).await?; + Ok(u64::from_le_bytes(buf)) +} + +// -- primary-side server -- + +/// Listens for incoming replica connections and drives replication. +/// +/// Each accepted connection performs a full-sync snapshot handshake +/// followed by an incremental stream from the broadcast channel. +pub struct ReplicationServer { + engine: Arc, + primary_id: String, +} + +impl ReplicationServer { + /// Binds the TCP listener and starts accepting replica connections. + /// + /// Runs indefinitely in the background; returns immediately after + /// spawning the accept loop task. + pub async fn start( + engine: Arc, + primary_id: String, + port: u16, + ) -> std::io::Result<()> { + let bind_addr = format!("0.0.0.0:{port}"); + let listener = TcpListener::bind(&bind_addr).await?; + info!(port, "replication server listening"); + + let server = Arc::new(Self { + engine, + primary_id, + }); + + tokio::spawn(async move { + loop { + match listener.accept().await { + Ok((stream, peer)) => { + debug!(%peer, "replica connected"); + let server = Arc::clone(&server); + tokio::spawn(async move { + if let Err(e) = server.handle_replica(stream).await { + debug!(%peer, "replication connection closed: {e}"); + } + }); + } + Err(e) => { + error!("replication accept error: {e}"); + tokio::time::sleep(Duration::from_millis(100)).await; + } + } + } + }); + + Ok(()) + } + + /// Handles a single replica connection: handshake + full sync + stream. + async fn handle_replica(&self, mut stream: TcpStream) -> std::io::Result<()> { + // read replica handshake + let replica_version = read_u8(&mut stream).await?; + if replica_version != REPL_VERSION { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("unsupported replication version: {replica_version}"), + )); + } + let replica_shards = read_u16_le(&mut stream).await?; + let our_shards = self.engine.shard_count() as u16; + + // send primary handshake response + write_u8(&mut stream, REPL_VERSION).await?; + write_u16_le(&mut stream, our_shards).await?; + let id_bytes = self.primary_id.as_bytes(); + write_u8(&mut stream, id_bytes.len() as u8).await?; + stream.write_all(id_bytes).await?; + + if replica_shards != our_shards { + write_u8(&mut stream, STATUS_SHARD_MISMATCH).await?; + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!( + "shard count mismatch: replica={replica_shards} primary={our_shards}" + ), + )); + } + write_u8(&mut stream, STATUS_OK).await?; + + // subscribe to the broadcast channel before snapshotting so we don't + // miss events that happen between snapshot and stream start + let mut rx = match self.engine.subscribe_replication() { + Some(rx) => rx, + None => { + return Err(std::io::Error::other("replication channel not configured on this engine")); + } + }; + + // full sync: snapshot each shard and send + for shard_idx in 0..self.engine.shard_count() { + let resp = self + .engine + .send_to_shard(shard_idx, ShardRequest::SerializeSnapshot) + .await + .map_err(|e| { + std::io::Error::other(format!("shard {shard_idx} serialize failed: {e:?}")) + })?; + + let (shard_id, data) = match resp { + ShardResponse::SnapshotData { shard_id, data } => (shard_id, data), + other => { + return Err(std::io::Error::other(format!("unexpected shard response: {other:?}"))); + } + }; + + let data_len = u32::try_from(data.len()).map_err(|_| { + std::io::Error::new(std::io::ErrorKind::InvalidData, "snapshot too large") + })?; + + write_u8(&mut stream, MSG_SHARD_SYNC).await?; + write_u16_le(&mut stream, shard_id).await?; + write_u32_le(&mut stream, data_len).await?; + stream.write_all(&data).await?; + + // send the current replication offset for this shard (0 until we + // track per-shard offsets; good enough for gap detection) + write_u8(&mut stream, MSG_SHARD_OFFSET).await?; + write_u16_le(&mut stream, shard_id).await?; + write_u64_le(&mut stream, 0u64).await?; + } + + stream.flush().await?; + info!("full sync complete, starting incremental stream"); + + // incremental stream: relay events from the broadcast channel + loop { + match rx.recv().await { + Ok(event) => { + let record_bytes = event.record.to_bytes().map_err(|e| { + std::io::Error::other(format!("record serialization failed: {e}")) + })?; + let record_len = u32::try_from(record_bytes.len()).map_err(|_| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + "record too large", + ) + })?; + + write_u8(&mut stream, MSG_RECORD).await?; + write_u16_le(&mut stream, event.shard_id).await?; + write_u64_le(&mut stream, event.offset).await?; + write_u32_le(&mut stream, record_len).await?; + stream.write_all(&record_bytes).await?; + } + Err(broadcast::error::RecvError::Lagged(count)) => { + warn!("replication stream lagged by {count} events; triggering resync"); + let _ = write_u8(&mut stream, MSG_RESYNC).await; + let _ = stream.flush().await; + return Ok(()); + } + Err(broadcast::error::RecvError::Closed) => { + info!("replication broadcast channel closed; disconnecting replica"); + return Ok(()); + } + } + } + } +} + +// -- replica-side client -- + +/// Connects to a primary's replication port and applies the incoming +/// snapshot and incremental record stream to the local engine. +/// +/// Reconnects automatically with exponential backoff on failure. +pub struct ReplicationClient { + engine: Arc, + primary_addr: SocketAddr, +} + +impl ReplicationClient { + /// Starts the replication client in a background task. + /// + /// Connects to `primary_addr` and applies the stream indefinitely, + /// reconnecting with backoff on any error. + pub fn start(engine: Arc, primary_addr: SocketAddr) { + let client = Arc::new(Self { + engine, + primary_addr, + }); + tokio::spawn(async move { + client.run().await; + }); + } + + async fn run(&self) { + let mut backoff = Duration::from_millis(500); + const MAX_BACKOFF: Duration = Duration::from_secs(30); + + loop { + info!(primary = %self.primary_addr, "connecting to primary for replication"); + match TcpStream::connect(self.primary_addr).await { + Ok(stream) => { + match self.sync(stream).await { + Ok(()) => { + info!("replication connection ended cleanly"); + } + Err(e) => { + warn!("replication error: {e}"); + } + } + } + Err(e) => { + warn!(primary = %self.primary_addr, "failed to connect to primary: {e}"); + } + } + + // back off before reconnecting + tokio::time::sleep(backoff).await; + backoff = (backoff * 2).min(MAX_BACKOFF); + } + } + + /// Performs full sync + incremental stream for one connection session. + async fn sync(&self, mut stream: TcpStream) -> std::io::Result<()> { + let our_shards = self.engine.shard_count() as u16; + + // send handshake + write_u8(&mut stream, REPL_VERSION).await?; + write_u16_le(&mut stream, our_shards).await?; + + // read primary response + let primary_version = read_u8(&mut stream).await?; + if primary_version != REPL_VERSION { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("unsupported primary replication version: {primary_version}"), + )); + } + let primary_shards = read_u16_le(&mut stream).await?; + let id_len = read_u8(&mut stream).await? as usize; + let mut id_buf = vec![0u8; id_len]; + stream.read_exact(&mut id_buf).await?; + let primary_id = String::from_utf8_lossy(&id_buf).into_owned(); + + let status = read_u8(&mut stream).await?; + if status == STATUS_SHARD_MISMATCH { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!( + "shard count mismatch with primary {primary_id}: \ + ours={our_shards} primary={primary_shards}" + ), + )); + } + + info!(primary_id = %primary_id, "handshake ok, loading full sync"); + + // receive per-shard snapshots + for _ in 0..primary_shards { + let msg = read_u8(&mut stream).await?; + if msg != MSG_SHARD_SYNC { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("expected MSG_SHARD_SYNC, got {msg}"), + )); + } + let shard_id = read_u16_le(&mut stream).await?; + let snap_len = read_u32_le(&mut stream).await? as usize; + + let mut snap_bytes = vec![0u8; snap_len]; + stream.read_exact(&mut snap_bytes).await?; + + // apply snapshot to the engine + self.apply_snapshot(shard_id, &snap_bytes).await?; + + // read (and ignore for now) the shard offset tag + let offset_msg = read_u8(&mut stream).await?; + if offset_msg != MSG_SHARD_OFFSET { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("expected MSG_SHARD_OFFSET, got {offset_msg}"), + )); + } + let _recv_shard_id = read_u16_le(&mut stream).await?; + let _recv_offset = read_u64_le(&mut stream).await?; + } + + info!("full sync applied, starting incremental replay"); + + // incremental stream + loop { + let msg = read_u8(&mut stream).await?; + match msg { + MSG_RECORD => { + let _shard_id = read_u16_le(&mut stream).await?; + let _offset = read_u64_le(&mut stream).await?; + let record_len = read_u32_le(&mut stream).await? as usize; + let mut record_bytes = vec![0u8; record_len]; + stream.read_exact(&mut record_bytes).await?; + + let record = AofRecord::from_bytes(&record_bytes).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("invalid AOF record: {e}"), + ) + })?; + + if let Some(request) = aof_record_to_shard_request(&record) { + let key = primary_key_for_request(&request).map(|k| k.to_owned()); + if let Some(key) = key { + if let Err(e) = self.engine.route(&key, request).await { + warn!("replication apply failed: {e:?}"); + } + } + } + } + MSG_RESYNC => { + info!("primary requested resync; reconnecting"); + return Ok(()); + } + other => { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("unknown replication message type: {other}"), + )); + } + } + } + } + + /// Applies a snapshot blob to the local engine for the given shard. + async fn apply_snapshot(&self, _shard_id: u16, data: &[u8]) -> std::io::Result<()> { + let (_, entries) = snapshot::read_snapshot_from_bytes(data).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("snapshot parse error: {e}"), + ) + })?; + + // flush the existing shard state before loading the snapshot + // we route by key so route() picks the right shard automatically + for entry in entries { + let value: Bytes = match &entry.value { + SnapValue::String(data) => data.clone(), + // for non-string types, reconstruct via the appropriate request + _ => { + self.apply_snap_entry(entry).await; + continue; + } + }; + + let expire = if entry.expire_ms > 0 { + Some(std::time::Duration::from_millis(entry.expire_ms as u64)) + } else { + None + }; + + let request = ShardRequest::Set { + key: entry.key.clone(), + value, + expire, + nx: false, + xx: false, + }; + if let Err(e) = self.engine.route(&entry.key, request).await { + warn!(key = %entry.key, "snapshot restore failed: {e:?}"); + } + } + + Ok(()) + } + + /// Restores a non-string snapshot entry by converting it to the + /// appropriate write request(s). + async fn apply_snap_entry(&self, entry: ember_persistence::snapshot::SnapEntry) { + use ember_persistence::snapshot::SnapValue; + + let key = entry.key.clone(); + let expire = if entry.expire_ms > 0 { + Some(std::time::Duration::from_millis(entry.expire_ms as u64)) + } else { + None + }; + + match entry.value { + SnapValue::List(deque) => { + let values: Vec = deque.into_iter().collect(); + let req = ShardRequest::RPush { + key: key.clone(), + values, + }; + if let Err(e) = self.engine.route(&key, req).await { + warn!(%key, "list restore failed: {e:?}"); + } + } + SnapValue::SortedSet(members) => { + let req = ShardRequest::ZAdd { + key: key.clone(), + members, + nx: false, + xx: false, + gt: false, + lt: false, + ch: false, + }; + if let Err(e) = self.engine.route(&key, req).await { + warn!(%key, "sorted set restore failed: {e:?}"); + } + } + SnapValue::Hash(map) => { + let fields: Vec<(String, Bytes)> = map.into_iter().collect(); + let req = ShardRequest::HSet { + key: key.clone(), + fields, + }; + if let Err(e) = self.engine.route(&key, req).await { + warn!(%key, "hash restore failed: {e:?}"); + } + } + SnapValue::Set(set) => { + let members: Vec = set.into_iter().collect(); + let req = ShardRequest::SAdd { + key: key.clone(), + members, + }; + if let Err(e) = self.engine.route(&key, req).await { + warn!(%key, "set restore failed: {e:?}"); + } + } + // strings are handled by the caller + SnapValue::String(_) => {} + #[cfg(feature = "vector")] + SnapValue::Vector { .. } => { + // vector restoration is complex; skip for now + warn!(%key, "vector snapshot restore not yet supported in replication"); + } + #[cfg(feature = "protobuf")] + SnapValue::Proto { type_name: _, data } => { + // proto values serialize to Set (raw bytes) + let req = ShardRequest::Set { + key: key.clone(), + value: data, + expire, + nx: false, + xx: false, + }; + if let Err(e) = self.engine.route(&key, req).await { + warn!(%key, "proto snapshot restore failed: {e:?}"); + } + } + } + + // apply TTL if any + if let Some(expire_duration) = expire { + let ms = expire_duration.as_millis() as u64; + let req = ShardRequest::Pexpire { + key: key.clone(), + milliseconds: ms, + }; + if let Err(e) = self.engine.route(&key, req).await { + warn!(%key, "pexpire after restore failed: {e:?}"); + } + } + } +} + +// -- AofRecord → ShardRequest conversion -- + +/// Converts an `AofRecord` into the equivalent `ShardRequest` for replay. +/// +/// Returns `None` for record types that have no meaningful replay action +/// (e.g. schema registration, which is handled at startup). +pub fn aof_record_to_shard_request(record: &AofRecord) -> Option { + match record { + AofRecord::Set { + key, + value, + expire_ms, + } => { + let expire = if *expire_ms > 0 { + Some(Duration::from_millis(*expire_ms as u64)) + } else { + None + }; + Some(ShardRequest::Set { + key: key.clone(), + value: value.clone(), + expire, + nx: false, + xx: false, + }) + } + AofRecord::Del { key } => Some(ShardRequest::Del { key: key.clone() }), + AofRecord::Expire { key, seconds } => Some(ShardRequest::Expire { + key: key.clone(), + seconds: *seconds, + }), + AofRecord::LPush { key, values } => Some(ShardRequest::LPush { + key: key.clone(), + values: values.clone(), + }), + AofRecord::RPush { key, values } => Some(ShardRequest::RPush { + key: key.clone(), + values: values.clone(), + }), + AofRecord::LPop { key } => Some(ShardRequest::LPop { key: key.clone() }), + AofRecord::RPop { key } => Some(ShardRequest::RPop { key: key.clone() }), + AofRecord::ZAdd { key, members } => Some(ShardRequest::ZAdd { + key: key.clone(), + members: members.clone(), + nx: false, + xx: false, + gt: false, + lt: false, + ch: false, + }), + AofRecord::ZRem { key, members } => Some(ShardRequest::ZRem { + key: key.clone(), + members: members.clone(), + }), + AofRecord::Persist { key } => Some(ShardRequest::Persist { key: key.clone() }), + AofRecord::Pexpire { key, milliseconds } => Some(ShardRequest::Pexpire { + key: key.clone(), + milliseconds: *milliseconds, + }), + AofRecord::Incr { key } => Some(ShardRequest::Incr { key: key.clone() }), + AofRecord::Decr { key } => Some(ShardRequest::Decr { key: key.clone() }), + AofRecord::HSet { key, fields } => Some(ShardRequest::HSet { + key: key.clone(), + fields: fields.clone(), + }), + AofRecord::HDel { key, fields } => Some(ShardRequest::HDel { + key: key.clone(), + fields: fields.clone(), + }), + AofRecord::HIncrBy { key, field, delta } => Some(ShardRequest::HIncrBy { + key: key.clone(), + field: field.clone(), + delta: *delta, + }), + AofRecord::SAdd { key, members } => Some(ShardRequest::SAdd { + key: key.clone(), + members: members.clone(), + }), + AofRecord::SRem { key, members } => Some(ShardRequest::SRem { + key: key.clone(), + members: members.clone(), + }), + AofRecord::IncrBy { key, delta } => Some(ShardRequest::IncrBy { + key: key.clone(), + delta: *delta, + }), + AofRecord::DecrBy { key, delta } => Some(ShardRequest::DecrBy { + key: key.clone(), + delta: *delta, + }), + AofRecord::Append { key, value } => Some(ShardRequest::Append { + key: key.clone(), + value: value.clone(), + }), + AofRecord::Rename { key, newkey } => Some(ShardRequest::Rename { + key: key.clone(), + newkey: newkey.clone(), + }), + #[cfg(feature = "vector")] + AofRecord::VAdd { .. } | AofRecord::VRem { .. } => { + // vector replication not yet supported + None + } + #[cfg(feature = "protobuf")] + AofRecord::ProtoSet { .. } + | AofRecord::ProtoRegister { .. } => { + // protobuf replication not yet supported + None + } + } +} + +/// Returns a reference to the primary key of a `ShardRequest` for routing. +fn primary_key_for_request(req: &ShardRequest) -> Option<&str> { + match req { + ShardRequest::Set { key, .. } + | ShardRequest::Del { key } + | ShardRequest::Unlink { key } + | ShardRequest::Expire { key, .. } + | ShardRequest::Persist { key } + | ShardRequest::Pexpire { key, .. } + | ShardRequest::Incr { key } + | ShardRequest::Decr { key } + | ShardRequest::IncrBy { key, .. } + | ShardRequest::DecrBy { key, .. } + | ShardRequest::IncrByFloat { key, .. } + | ShardRequest::Append { key, .. } + | ShardRequest::LPush { key, .. } + | ShardRequest::RPush { key, .. } + | ShardRequest::LPop { key } + | ShardRequest::RPop { key } + | ShardRequest::ZAdd { key, .. } + | ShardRequest::ZRem { key, .. } + | ShardRequest::HSet { key, .. } + | ShardRequest::HDel { key, .. } + | ShardRequest::HIncrBy { key, .. } + | ShardRequest::SAdd { key, .. } + | ShardRequest::SRem { key, .. } + | ShardRequest::Rename { key, .. } => Some(key), + _ => None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use bytes::Bytes; + + #[test] + fn aof_set_roundtrip() { + let record = AofRecord::Set { + key: "foo".into(), + value: Bytes::from("bar"), + expire_ms: 5000, + }; + let req = aof_record_to_shard_request(&record).expect("Set should map to ShardRequest"); + match req { + ShardRequest::Set { key, value, expire, nx, xx } => { + assert_eq!(key, "foo"); + assert_eq!(value, Bytes::from("bar")); + assert_eq!(expire, Some(Duration::from_millis(5000))); + assert!(!nx && !xx); + } + other => panic!("expected ShardRequest::Set, got {other:?}"), + } + } + + #[test] + fn aof_del_roundtrip() { + let record = AofRecord::Del { key: "gone".into() }; + let req = aof_record_to_shard_request(&record).unwrap(); + assert!(matches!(req, ShardRequest::Del { key } if key == "gone")); + } + + #[test] + fn primary_key_set() { + let req = ShardRequest::Set { + key: "mykey".into(), + value: Bytes::new(), + expire: None, + nx: false, + xx: false, + }; + assert_eq!(primary_key_for_request(&req), Some("mykey")); + } +} diff --git a/crates/ember-server/src/server.rs b/crates/ember-server/src/server.rs index c7cd1971..6bf74280 100644 --- a/crates/ember-server/src/server.rs +++ b/crates/ember-server/src/server.rs @@ -90,6 +90,13 @@ pub async fn run( let engine = Engine::with_config(shard_count, config); + // wire replication: give the cluster coordinator access to the engine + // and start the replication server so replicas can connect + if let Some(ref coordinator) = cluster { + coordinator.set_engine(Arc::new(engine.clone())); + coordinator.start_replication_server().await; + } + if metrics_enabled { crate::metrics::spawn_stats_poller(engine.clone()); }