From f35959372d2f1c0e5c9cd433f2b346debc271985 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Fri, 6 Feb 2026 08:22:38 -0500 Subject: [PATCH 1/3] feat: add migration module for live slot resharding implements the core migration types and state management: - Migration struct with source/target nodes and state tracking - MigrationManager for tracking active incoming/outgoing migrations - MigrationBatch for streaming keys between nodes - MigrationRedirect for MOVED/ASK responses - MigrationConfig for tuning batch sizes and timeouts --- crates/ember-cluster/src/lib.rs | 5 + crates/ember-cluster/src/migration.rs | 679 ++++++++++++++++++++++++++ 2 files changed, 684 insertions(+) create mode 100644 crates/ember-cluster/src/migration.rs diff --git a/crates/ember-cluster/src/lib.rs b/crates/ember-cluster/src/lib.rs index 7605d858..8bccb18c 100644 --- a/crates/ember-cluster/src/lib.rs +++ b/crates/ember-cluster/src/lib.rs @@ -32,6 +32,7 @@ mod error; mod gossip; mod message; +mod migration; mod raft; mod slots; mod topology; @@ -39,6 +40,10 @@ mod topology; pub use error::ClusterError; pub use gossip::{GossipConfig, GossipEngine, GossipEvent, MemberState, MemberStatus}; pub use message::{GossipMessage, MemberInfo, NodeUpdate}; +pub use migration::{ + Migration, MigrationBatch, MigrationConfig, MigrationEntry, MigrationError, MigrationId, + MigrationManager, MigrationRedirect, MigrationState, +}; pub use raft::{ ClusterCommand, ClusterResponse, ClusterSnapshot, ClusterStateData, Storage as RaftStorage, TypeConfig, diff --git a/crates/ember-cluster/src/migration.rs b/crates/ember-cluster/src/migration.rs new file mode 100644 index 00000000..dc221668 --- /dev/null +++ b/crates/ember-cluster/src/migration.rs @@ -0,0 +1,679 @@ +//! Live slot migration for cluster resharding. +//! +//! This module implements Redis-compatible slot migration that allows moving +//! slots between nodes without downtime. The protocol follows these steps: +//! +//! 1. Target marks slot as IMPORTING from source +//! 2. Source marks slot as MIGRATING to target +//! 3. Keys are streamed in batches from source to target +//! 4. During migration, reads go to source, writes get ASK redirect to target +//! 5. Final batch completes, ownership transfers via Raft +//! +//! # Example +//! +//! ```ignore +//! // On target node: +//! CLUSTER SETSLOT 100 IMPORTING +//! +//! // On source node: +//! CLUSTER SETSLOT 100 MIGRATING +//! +//! // Migrate keys (repeated for each key in slot): +//! MIGRATE 0 5000 +//! +//! // Finalize on both nodes: +//! CLUSTER SETSLOT 100 NODE +//! ``` + +use std::collections::{HashMap, HashSet}; +use std::net::SocketAddr; +use std::time::{Duration, Instant}; + +use serde::{Deserialize, Serialize}; + +use crate::NodeId; + +/// Unique identifier for a migration operation. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub struct MigrationId(pub u64); + +impl MigrationId { + /// Generate a new migration ID from timestamp and slot. + pub fn new(slot: u16) -> Self { + use std::time::{SystemTime, UNIX_EPOCH}; + let ts = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_nanos() as u64; + // Combine timestamp with slot for uniqueness + Self(ts ^ (slot as u64)) + } +} + +/// Current state of a slot migration. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum MigrationState { + /// Target has marked slot as importing, waiting for source. + Importing, + /// Source has marked slot as migrating, ready to stream keys. + Migrating, + /// Keys are being transferred in batches. + Streaming, + /// Final batch being sent, brief pause on writes. + Finalizing, + /// Migration complete, ownership transferred. + Complete, +} + +impl std::fmt::Display for MigrationState { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Importing => write!(f, "importing"), + Self::Migrating => write!(f, "migrating"), + Self::Streaming => write!(f, "streaming"), + Self::Finalizing => write!(f, "finalizing"), + Self::Complete => write!(f, "complete"), + } + } +} + +/// A single slot migration operation. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Migration { + /// Unique migration identifier. + pub id: MigrationId, + /// The slot being migrated. + pub slot: u16, + /// Node currently owning the slot. + pub source: NodeId, + /// Node receiving the slot. + pub target: NodeId, + /// Current migration state. + pub state: MigrationState, + /// When migration started. + #[serde(skip)] + pub started_at: Option, + /// Number of keys migrated so far. + pub keys_migrated: u64, + /// Total keys to migrate (if known). + pub keys_total: Option, +} + +impl Migration { + /// Create a new migration in the importing state (target perspective). + pub fn new_importing(slot: u16, source: NodeId, target: NodeId) -> Self { + Self { + id: MigrationId::new(slot), + slot, + source, + target, + state: MigrationState::Importing, + started_at: Some(Instant::now()), + keys_migrated: 0, + keys_total: None, + } + } + + /// Create a new migration in the migrating state (source perspective). + pub fn new_migrating(slot: u16, source: NodeId, target: NodeId) -> Self { + Self { + id: MigrationId::new(slot), + slot, + source, + target, + state: MigrationState::Migrating, + started_at: Some(Instant::now()), + keys_migrated: 0, + keys_total: None, + } + } + + /// Check if this migration involves a specific node. + pub fn involves(&self, node: &NodeId) -> bool { + self.source == *node || self.target == *node + } + + /// Get progress as a percentage (0-100). + pub fn progress(&self) -> Option { + self.keys_total.map(|total| { + if total == 0 { + 100 + } else { + ((self.keys_migrated * 100) / total).min(100) as u8 + } + }) + } + + /// Transition to streaming state. + pub fn start_streaming(&mut self, total_keys: u64) { + self.state = MigrationState::Streaming; + self.keys_total = Some(total_keys); + } + + /// Record migrated keys. + pub fn record_migrated(&mut self, count: u64) { + self.keys_migrated += count; + } + + /// Transition to finalizing state. + pub fn start_finalizing(&mut self) { + self.state = MigrationState::Finalizing; + } + + /// Mark migration as complete. + pub fn complete(&mut self) { + self.state = MigrationState::Complete; + } +} + +/// Tracks all active migrations for a node. +#[derive(Debug, Default)] +pub struct MigrationManager { + /// Migrations where this node is the source (slot is migrating out). + outgoing: HashMap, + /// Migrations where this node is the target (slot is importing in). + incoming: HashMap, + /// Keys that have been migrated but not yet confirmed. + pending_keys: HashMap>>, +} + +impl MigrationManager { + /// Create a new migration manager. + pub fn new() -> Self { + Self::default() + } + + /// Check if a slot is currently migrating out. + pub fn is_migrating(&self, slot: u16) -> bool { + self.outgoing.contains_key(&slot) + } + + /// Check if a slot is currently being imported. + pub fn is_importing(&self, slot: u16) -> bool { + self.incoming.contains_key(&slot) + } + + /// Get migration info for a slot that's migrating out. + pub fn get_outgoing(&self, slot: u16) -> Option<&Migration> { + self.outgoing.get(&slot) + } + + /// Get migration info for a slot that's being imported. + pub fn get_incoming(&self, slot: u16) -> Option<&Migration> { + self.incoming.get(&slot) + } + + /// Start importing a slot from another node. + /// + /// Returns error if slot is already involved in a migration. + pub fn start_import( + &mut self, + slot: u16, + source: NodeId, + local_id: NodeId, + ) -> Result<&Migration, MigrationError> { + if self.outgoing.contains_key(&slot) { + return Err(MigrationError::SlotAlreadyMigrating { slot }); + } + if self.incoming.contains_key(&slot) { + return Err(MigrationError::SlotAlreadyImporting { slot }); + } + + let migration = Migration::new_importing(slot, source, local_id); + self.incoming.insert(slot, migration); + self.pending_keys.insert(slot, HashSet::new()); + Ok(self.incoming.get(&slot).unwrap()) + } + + /// Start migrating a slot to another node. + /// + /// Returns error if slot is already involved in a migration. + pub fn start_migrate( + &mut self, + slot: u16, + local_id: NodeId, + target: NodeId, + ) -> Result<&Migration, MigrationError> { + if self.outgoing.contains_key(&slot) { + return Err(MigrationError::SlotAlreadyMigrating { slot }); + } + if self.incoming.contains_key(&slot) { + return Err(MigrationError::SlotAlreadyImporting { slot }); + } + + let migration = Migration::new_migrating(slot, local_id, target); + self.outgoing.insert(slot, migration); + self.pending_keys.insert(slot, HashSet::new()); + Ok(self.outgoing.get(&slot).unwrap()) + } + + /// Record that a key has been migrated. + pub fn key_migrated(&mut self, slot: u16, key: Vec) { + if let Some(keys) = self.pending_keys.get_mut(&slot) { + keys.insert(key); + } + if let Some(migration) = self.outgoing.get_mut(&slot) { + migration.record_migrated(1); + } + } + + /// Check if a specific key has been migrated. + pub fn is_key_migrated(&self, slot: u16, key: &[u8]) -> bool { + self.pending_keys + .get(&slot) + .is_some_and(|keys| keys.contains(key)) + } + + /// Complete a migration and clean up state. + pub fn complete_migration(&mut self, slot: u16) -> Option { + self.pending_keys.remove(&slot); + // Try outgoing first, then incoming + self.outgoing + .remove(&slot) + .or_else(|| self.incoming.remove(&slot)) + .map(|mut m| { + m.complete(); + m + }) + } + + /// Abort a migration and clean up state. + pub fn abort_migration(&mut self, slot: u16) -> Option { + self.pending_keys.remove(&slot); + self.outgoing + .remove(&slot) + .or_else(|| self.incoming.remove(&slot)) + } + + /// Get all active outgoing migrations. + pub fn outgoing_migrations(&self) -> impl Iterator { + self.outgoing.values() + } + + /// Get all active incoming migrations. + pub fn incoming_migrations(&self) -> impl Iterator { + self.incoming.values() + } + + /// Get total number of active migrations. + pub fn active_count(&self) -> usize { + self.outgoing.len() + self.incoming.len() + } +} + +/// Represents a batch of keys to migrate. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MigrationBatch { + /// The slot being migrated. + pub slot: u16, + /// Keys and their values in this batch. + pub entries: Vec, + /// Whether this is the final batch. + pub is_final: bool, + /// Sequence number for ordering. + pub sequence: u64, +} + +/// A single key-value entry being migrated. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MigrationEntry { + /// The key being migrated. + pub key: Vec, + /// Serialized value data. + pub value: Vec, + /// TTL remaining in milliseconds (0 = no expiry). + pub ttl_ms: u64, +} + +impl MigrationBatch { + /// Create a new migration batch. + pub fn new(slot: u16, sequence: u64) -> Self { + Self { + slot, + entries: Vec::new(), + is_final: false, + sequence, + } + } + + /// Add an entry to the batch. + pub fn add_entry(&mut self, key: Vec, value: Vec, ttl_ms: u64) { + self.entries.push(MigrationEntry { key, value, ttl_ms }); + } + + /// Mark this as the final batch. + pub fn mark_final(&mut self) { + self.is_final = true; + } + + /// Get the number of entries in this batch. + pub fn len(&self) -> usize { + self.entries.len() + } + + /// Check if the batch is empty. + pub fn is_empty(&self) -> bool { + self.entries.is_empty() + } + + /// Estimate the size of this batch in bytes. + pub fn size_bytes(&self) -> usize { + self.entries + .iter() + .map(|e| e.key.len() + e.value.len() + 8) + .sum() + } +} + +/// Errors that can occur during migration. +#[derive(Debug, Clone, thiserror::Error)] +pub enum MigrationError { + /// Slot is already being migrated out. + #[error("slot {slot} is already migrating")] + SlotAlreadyMigrating { slot: u16 }, + + /// Slot is already being imported. + #[error("slot {slot} is already importing")] + SlotAlreadyImporting { slot: u16 }, + + /// No migration in progress for this slot. + #[error("no migration in progress for slot {slot}")] + NoMigrationInProgress { slot: u16 }, + + /// Migration target is unreachable. + #[error("cannot reach migration target {addr}: {reason}")] + TargetUnreachable { addr: SocketAddr, reason: String }, + + /// Migration was aborted. + #[error("migration for slot {slot} was aborted")] + Aborted { slot: u16 }, + + /// Invalid migration state transition. + #[error("invalid state transition from {from} to {to}")] + InvalidStateTransition { + from: MigrationState, + to: MigrationState, + }, + + /// Timeout during migration. + #[error("migration timeout after {elapsed:?}")] + Timeout { elapsed: Duration }, +} + +/// Result of checking whether a command should be redirected during migration. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum MigrationRedirect { + /// No redirect needed, handle locally. + None, + /// Send MOVED redirect (slot permanently moved). + Moved { slot: u16, addr: SocketAddr }, + /// Send ASK redirect (slot temporarily at another node). + Ask { slot: u16, addr: SocketAddr }, +} + +impl MigrationRedirect { + /// Format as RESP error string. + pub fn to_error_string(&self) -> Option { + match self { + Self::None => None, + Self::Moved { slot, addr } => Some(format!("MOVED {} {}", slot, addr)), + Self::Ask { slot, addr } => Some(format!("ASK {} {}", slot, addr)), + } + } +} + +/// Configuration for migration behavior. +#[derive(Debug, Clone)] +pub struct MigrationConfig { + /// Maximum keys per batch. + pub batch_size: usize, + /// Maximum batch size in bytes. + pub batch_bytes: usize, + /// Timeout for individual key migration. + pub key_timeout: Duration, + /// Timeout for entire migration. + pub migration_timeout: Duration, + /// Delay between batches to avoid overwhelming the network. + pub batch_delay: Duration, +} + +impl Default for MigrationConfig { + fn default() -> Self { + Self { + batch_size: 100, + batch_bytes: 1024 * 1024, // 1MB + key_timeout: Duration::from_secs(5), + migration_timeout: Duration::from_secs(3600), // 1 hour + batch_delay: Duration::from_millis(10), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn node_id() -> NodeId { + NodeId::new() + } + + #[test] + fn migration_new_importing() { + let source = node_id(); + let target = node_id(); + let m = Migration::new_importing(100, source, target); + + assert_eq!(m.slot, 100); + assert_eq!(m.source, source); + assert_eq!(m.target, target); + assert_eq!(m.state, MigrationState::Importing); + assert_eq!(m.keys_migrated, 0); + } + + #[test] + fn migration_new_migrating() { + let source = node_id(); + let target = node_id(); + let m = Migration::new_migrating(100, source, target); + + assert_eq!(m.state, MigrationState::Migrating); + } + + #[test] + fn migration_involves() { + let source = node_id(); + let target = node_id(); + let other = node_id(); + let m = Migration::new_importing(100, source, target); + + assert!(m.involves(&source)); + assert!(m.involves(&target)); + assert!(!m.involves(&other)); + } + + #[test] + fn migration_progress() { + let mut m = Migration::new_migrating(100, node_id(), node_id()); + + // No total set + assert_eq!(m.progress(), None); + + // Set total and migrate some + m.start_streaming(100); + assert_eq!(m.progress(), Some(0)); + + m.record_migrated(50); + assert_eq!(m.progress(), Some(50)); + + m.record_migrated(50); + assert_eq!(m.progress(), Some(100)); + } + + #[test] + fn migration_state_transitions() { + let mut m = Migration::new_migrating(100, node_id(), node_id()); + + assert_eq!(m.state, MigrationState::Migrating); + + m.start_streaming(50); + assert_eq!(m.state, MigrationState::Streaming); + + m.start_finalizing(); + assert_eq!(m.state, MigrationState::Finalizing); + + m.complete(); + assert_eq!(m.state, MigrationState::Complete); + } + + #[test] + fn manager_start_import() { + let mut manager = MigrationManager::new(); + let source = node_id(); + let local = node_id(); + + let result = manager.start_import(100, source, local); + assert!(result.is_ok()); + assert!(manager.is_importing(100)); + assert!(!manager.is_migrating(100)); + } + + #[test] + fn manager_start_migrate() { + let mut manager = MigrationManager::new(); + let local = node_id(); + let target = node_id(); + + let result = manager.start_migrate(100, local, target); + assert!(result.is_ok()); + assert!(manager.is_migrating(100)); + assert!(!manager.is_importing(100)); + } + + #[test] + fn manager_double_migration_error() { + let mut manager = MigrationManager::new(); + let local = node_id(); + let target = node_id(); + + manager.start_migrate(100, local, target).unwrap(); + + // Can't migrate same slot again + let result = manager.start_migrate(100, local, node_id()); + assert!(matches!( + result, + Err(MigrationError::SlotAlreadyMigrating { slot: 100 }) + )); + + // Can't import a slot that's migrating + let result = manager.start_import(100, node_id(), local); + assert!(matches!( + result, + Err(MigrationError::SlotAlreadyMigrating { slot: 100 }) + )); + } + + #[test] + fn manager_key_tracking() { + let mut manager = MigrationManager::new(); + let local = node_id(); + let target = node_id(); + + manager.start_migrate(100, local, target).unwrap(); + + assert!(!manager.is_key_migrated(100, b"key1")); + + manager.key_migrated(100, b"key1".to_vec()); + assert!(manager.is_key_migrated(100, b"key1")); + assert!(!manager.is_key_migrated(100, b"key2")); + } + + #[test] + fn manager_complete_migration() { + let mut manager = MigrationManager::new(); + let local = node_id(); + let target = node_id(); + + manager.start_migrate(100, local, target).unwrap(); + manager.key_migrated(100, b"key1".to_vec()); + + let completed = manager.complete_migration(100); + assert!(completed.is_some()); + assert_eq!(completed.unwrap().state, MigrationState::Complete); + + // State should be cleaned up + assert!(!manager.is_migrating(100)); + assert!(!manager.is_key_migrated(100, b"key1")); + } + + #[test] + fn manager_abort_migration() { + let mut manager = MigrationManager::new(); + let local = node_id(); + let target = node_id(); + + manager.start_migrate(100, local, target).unwrap(); + + let aborted = manager.abort_migration(100); + assert!(aborted.is_some()); + + assert!(!manager.is_migrating(100)); + } + + #[test] + fn batch_operations() { + let mut batch = MigrationBatch::new(100, 1); + + assert!(batch.is_empty()); + assert_eq!(batch.len(), 0); + + batch.add_entry(b"key1".to_vec(), b"value1".to_vec(), 0); + batch.add_entry(b"key2".to_vec(), b"value2".to_vec(), 5000); + + assert!(!batch.is_empty()); + assert_eq!(batch.len(), 2); + assert!(!batch.is_final); + + batch.mark_final(); + assert!(batch.is_final); + } + + #[test] + fn redirect_formatting() { + let moved = MigrationRedirect::Moved { + slot: 100, + addr: "127.0.0.1:6379".parse().unwrap(), + }; + assert_eq!( + moved.to_error_string(), + Some("MOVED 100 127.0.0.1:6379".to_string()) + ); + + let ask = MigrationRedirect::Ask { + slot: 200, + addr: "127.0.0.1:6380".parse().unwrap(), + }; + assert_eq!( + ask.to_error_string(), + Some("ASK 200 127.0.0.1:6380".to_string()) + ); + + assert_eq!(MigrationRedirect::None.to_error_string(), None); + } + + #[test] + fn migration_state_display() { + assert_eq!(MigrationState::Importing.to_string(), "importing"); + assert_eq!(MigrationState::Migrating.to_string(), "migrating"); + assert_eq!(MigrationState::Streaming.to_string(), "streaming"); + assert_eq!(MigrationState::Finalizing.to_string(), "finalizing"); + assert_eq!(MigrationState::Complete.to_string(), "complete"); + } + + #[test] + fn config_defaults() { + let config = MigrationConfig::default(); + assert_eq!(config.batch_size, 100); + assert_eq!(config.batch_bytes, 1024 * 1024); + assert_eq!(config.key_timeout, Duration::from_secs(5)); + } +} From 0be14058e91d0eed9e43d29c19c7bd9cfad99756 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Fri, 6 Feb 2026 08:24:52 -0500 Subject: [PATCH 2/3] feat: add CLUSTER SETSLOT and MIGRATE command parsing adds protocol support for slot migration commands: - CLUSTER SETSLOT IMPORTING - CLUSTER SETSLOT MIGRATING - CLUSTER SETSLOT NODE - CLUSTER SETSLOT STABLE - MIGRATE host port key db timeout [COPY] [REPLACE] handlers return cluster-disabled error in non-cluster mode --- crates/ember-protocol/src/command.rs | 239 ++++++++++++++++++++++++++ crates/ember-server/src/connection.rs | 11 ++ 2 files changed, 250 insertions(+) diff --git a/crates/ember-protocol/src/command.rs b/crates/ember-protocol/src/command.rs index c70004f7..0e98f2a3 100644 --- a/crates/ember-protocol/src/command.rs +++ b/crates/ember-protocol/src/command.rs @@ -212,6 +212,30 @@ pub enum Command { /// CLUSTER MYID. Returns the node's ID. ClusterMyId, + /// CLUSTER SETSLOT IMPORTING . Mark slot as importing from node. + ClusterSetSlotImporting { slot: u16, node_id: String }, + + /// CLUSTER SETSLOT MIGRATING . Mark slot as migrating to node. + ClusterSetSlotMigrating { slot: u16, node_id: String }, + + /// CLUSTER SETSLOT NODE . Assign slot to node. + ClusterSetSlotNode { slot: u16, node_id: String }, + + /// CLUSTER SETSLOT STABLE. Clear importing/migrating state. + ClusterSetSlotStable { slot: u16 }, + + /// MIGRATE [COPY] [REPLACE] [KEYS key...]. + /// Migrate a key to another node. + Migrate { + host: String, + port: u16, + key: String, + db: u32, + timeout_ms: u64, + copy: bool, + replace: bool, + }, + /// ASKING. Signals that the next command is for a migrating slot. Asking, @@ -312,6 +336,7 @@ impl Command { "SCARD" => parse_scard(&frames[1..]), "CLUSTER" => parse_cluster(&frames[1..]), "ASKING" => parse_asking(&frames[1..]), + "MIGRATE" => parse_migrate(&frames[1..]), _ => Ok(Command::Unknown(name)), } } @@ -1077,6 +1102,7 @@ fn parse_cluster(args: &[Frame]) -> Result { } Ok(Command::ClusterMyId) } + "SETSLOT" => parse_cluster_setslot(&args[1..]), _ => Err(ProtocolError::InvalidCommandFrame(format!( "unknown CLUSTER subcommand '{subcommand}'" ))), @@ -1090,6 +1116,107 @@ fn parse_asking(args: &[Frame]) -> Result { Ok(Command::Asking) } +fn parse_cluster_setslot(args: &[Frame]) -> Result { + if args.is_empty() { + return Err(ProtocolError::WrongArity("CLUSTER SETSLOT".into())); + } + + let slot_str = extract_string(&args[0])?; + let slot: u16 = slot_str + .parse() + .map_err(|_| ProtocolError::InvalidCommandFrame("invalid slot number".into()))?; + + if args.len() < 2 { + return Err(ProtocolError::WrongArity("CLUSTER SETSLOT".into())); + } + + let action = extract_string(&args[1])?.to_ascii_uppercase(); + match action.as_str() { + "IMPORTING" => { + if args.len() != 3 { + return Err(ProtocolError::WrongArity( + "CLUSTER SETSLOT IMPORTING".into(), + )); + } + let node_id = extract_string(&args[2])?; + Ok(Command::ClusterSetSlotImporting { slot, node_id }) + } + "MIGRATING" => { + if args.len() != 3 { + return Err(ProtocolError::WrongArity( + "CLUSTER SETSLOT MIGRATING".into(), + )); + } + let node_id = extract_string(&args[2])?; + Ok(Command::ClusterSetSlotMigrating { slot, node_id }) + } + "NODE" => { + if args.len() != 3 { + return Err(ProtocolError::WrongArity("CLUSTER SETSLOT NODE".into())); + } + let node_id = extract_string(&args[2])?; + Ok(Command::ClusterSetSlotNode { slot, node_id }) + } + "STABLE" => { + if args.len() != 2 { + return Err(ProtocolError::WrongArity("CLUSTER SETSLOT STABLE".into())); + } + Ok(Command::ClusterSetSlotStable { slot }) + } + _ => Err(ProtocolError::InvalidCommandFrame(format!( + "unknown CLUSTER SETSLOT action '{action}'" + ))), + } +} + +fn parse_migrate(args: &[Frame]) -> Result { + // MIGRATE host port key db timeout [COPY] [REPLACE] + if args.len() < 5 { + return Err(ProtocolError::WrongArity("MIGRATE".into())); + } + + let host = extract_string(&args[0])?; + let port_str = extract_string(&args[1])?; + let port: u16 = port_str + .parse() + .map_err(|_| ProtocolError::InvalidCommandFrame("invalid port number".into()))?; + let key = extract_string(&args[2])?; + let db_str = extract_string(&args[3])?; + let db: u32 = db_str + .parse() + .map_err(|_| ProtocolError::InvalidCommandFrame("invalid db number".into()))?; + let timeout_str = extract_string(&args[4])?; + let timeout_ms: u64 = timeout_str + .parse() + .map_err(|_| ProtocolError::InvalidCommandFrame("invalid timeout".into()))?; + + let mut copy = false; + let mut replace = false; + + for arg in &args[5..] { + let opt = extract_string(arg)?.to_ascii_uppercase(); + match opt.as_str() { + "COPY" => copy = true, + "REPLACE" => replace = true, + _ => { + return Err(ProtocolError::InvalidCommandFrame(format!( + "unknown MIGRATE option '{opt}'" + ))) + } + } + } + + Ok(Command::Migrate { + host, + port, + key, + db, + timeout_ms, + copy, + replace, + }) +} + #[cfg(test)] mod tests { use super::*; @@ -2668,4 +2795,116 @@ mod tests { let err = Command::from_frame(cmd(&["ASKING", "extra"])).unwrap_err(); assert!(matches!(err, ProtocolError::WrongArity(_))); } + + #[test] + fn cluster_setslot_importing() { + assert_eq!( + Command::from_frame(cmd(&["CLUSTER", "SETSLOT", "100", "IMPORTING", "node123"])) + .unwrap(), + Command::ClusterSetSlotImporting { + slot: 100, + node_id: "node123".into() + }, + ); + } + + #[test] + fn cluster_setslot_migrating() { + assert_eq!( + Command::from_frame(cmd(&["CLUSTER", "SETSLOT", "200", "MIGRATING", "node456"])) + .unwrap(), + Command::ClusterSetSlotMigrating { + slot: 200, + node_id: "node456".into() + }, + ); + } + + #[test] + fn cluster_setslot_node() { + assert_eq!( + Command::from_frame(cmd(&["CLUSTER", "SETSLOT", "300", "NODE", "node789"])).unwrap(), + Command::ClusterSetSlotNode { + slot: 300, + node_id: "node789".into() + }, + ); + } + + #[test] + fn cluster_setslot_stable() { + assert_eq!( + Command::from_frame(cmd(&["CLUSTER", "SETSLOT", "400", "STABLE"])).unwrap(), + Command::ClusterSetSlotStable { slot: 400 }, + ); + } + + #[test] + fn cluster_setslot_invalid_slot() { + let err = + Command::from_frame(cmd(&["CLUSTER", "SETSLOT", "notanumber", "STABLE"])).unwrap_err(); + assert!(matches!(err, ProtocolError::InvalidCommandFrame(_))); + } + + #[test] + fn cluster_setslot_wrong_arity() { + let err = Command::from_frame(cmd(&["CLUSTER", "SETSLOT", "100"])).unwrap_err(); + assert!(matches!(err, ProtocolError::WrongArity(_))); + } + + #[test] + fn migrate_basic() { + assert_eq!( + Command::from_frame(cmd(&["MIGRATE", "127.0.0.1", "6379", "mykey", "0", "5000"])) + .unwrap(), + Command::Migrate { + host: "127.0.0.1".into(), + port: 6379, + key: "mykey".into(), + db: 0, + timeout_ms: 5000, + copy: false, + replace: false, + }, + ); + } + + #[test] + fn migrate_with_options() { + assert_eq!( + Command::from_frame(cmd(&[ + "MIGRATE", + "192.168.1.1", + "6380", + "testkey", + "1", + "10000", + "COPY", + "REPLACE" + ])) + .unwrap(), + Command::Migrate { + host: "192.168.1.1".into(), + port: 6380, + key: "testkey".into(), + db: 1, + timeout_ms: 10000, + copy: true, + replace: true, + }, + ); + } + + #[test] + fn migrate_wrong_arity() { + let err = Command::from_frame(cmd(&["MIGRATE", "host", "port", "key"])).unwrap_err(); + assert!(matches!(err, ProtocolError::WrongArity(_))); + } + + #[test] + fn migrate_invalid_port() { + let err = Command::from_frame(cmd(&["MIGRATE", "host", "notaport", "key", "0", "1000"])) + .unwrap_err(); + assert!(matches!(err, ProtocolError::InvalidCommandFrame(_))); + } } diff --git a/crates/ember-server/src/connection.rs b/crates/ember-server/src/connection.rs index cd8b7361..24f53c37 100644 --- a/crates/ember-server/src/connection.rs +++ b/crates/ember-server/src/connection.rs @@ -872,6 +872,17 @@ async fn execute(cmd: Command, engine: &Engine) -> Frame { Frame::Simple("OK".into()) } + Command::ClusterSetSlotImporting { .. } + | Command::ClusterSetSlotMigrating { .. } + | Command::ClusterSetSlotNode { .. } + | Command::ClusterSetSlotStable { .. } => { + Frame::Error("ERR This instance has cluster support disabled".into()) + } + + Command::Migrate { .. } => { + Frame::Error("ERR This instance has cluster support disabled".into()) + } + Command::Unknown(name) => Frame::Error(format!("ERR unknown command '{name}'")), } } From 55fdab0540f799d02ab01a7040ecd13320000b04 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Fri, 6 Feb 2026 08:26:35 -0500 Subject: [PATCH 3/3] feat: add remaining CLUSTER command variants adds protocol support for cluster management: - CLUSTER MEET - CLUSTER ADDSLOTS [slot...] - CLUSTER DELSLOTS [slot...] - CLUSTER FORGET - CLUSTER REPLICATE - CLUSTER FAILOVER [FORCE|TAKEOVER] - CLUSTER COUNTKEYSINSLOT - CLUSTER GETKEYSINSLOT all handlers return cluster-disabled error in non-cluster mode --- crates/ember-protocol/src/command.rs | 219 ++++++++++++++++++++++++++ crates/ember-server/src/connection.rs | 10 +- 2 files changed, 228 insertions(+), 1 deletion(-) diff --git a/crates/ember-protocol/src/command.rs b/crates/ember-protocol/src/command.rs index 0e98f2a3..d665b1f3 100644 --- a/crates/ember-protocol/src/command.rs +++ b/crates/ember-protocol/src/command.rs @@ -224,6 +224,30 @@ pub enum Command { /// CLUSTER SETSLOT STABLE. Clear importing/migrating state. ClusterSetSlotStable { slot: u16 }, + /// CLUSTER MEET . Add a node to the cluster. + ClusterMeet { ip: String, port: u16 }, + + /// CLUSTER ADDSLOTS [slot...]. Assign slots to the local node. + ClusterAddSlots { slots: Vec }, + + /// CLUSTER DELSLOTS [slot...]. Remove slots from the local node. + ClusterDelSlots { slots: Vec }, + + /// CLUSTER FORGET . Remove a node from the cluster. + ClusterForget { node_id: String }, + + /// CLUSTER REPLICATE . Make this node a replica of another. + ClusterReplicate { node_id: String }, + + /// CLUSTER FAILOVER [FORCE|TAKEOVER]. Trigger a manual failover. + ClusterFailover { force: bool, takeover: bool }, + + /// CLUSTER COUNTKEYSINSLOT . Return the number of keys in a slot. + ClusterCountKeysInSlot { slot: u16 }, + + /// CLUSTER GETKEYSINSLOT . Return keys in a slot. + ClusterGetKeysInSlot { slot: u16, count: u32 }, + /// MIGRATE [COPY] [REPLACE] [KEYS key...]. /// Migrate a key to another node. Migrate { @@ -1103,6 +1127,86 @@ fn parse_cluster(args: &[Frame]) -> Result { Ok(Command::ClusterMyId) } "SETSLOT" => parse_cluster_setslot(&args[1..]), + "MEET" => { + if args.len() != 3 { + return Err(ProtocolError::WrongArity("CLUSTER MEET".into())); + } + let ip = extract_string(&args[1])?; + let port_str = extract_string(&args[2])?; + let port: u16 = port_str + .parse() + .map_err(|_| ProtocolError::InvalidCommandFrame("invalid port number".into()))?; + Ok(Command::ClusterMeet { ip, port }) + } + "ADDSLOTS" => { + if args.len() < 2 { + return Err(ProtocolError::WrongArity("CLUSTER ADDSLOTS".into())); + } + let slots = parse_slot_list(&args[1..])?; + Ok(Command::ClusterAddSlots { slots }) + } + "DELSLOTS" => { + if args.len() < 2 { + return Err(ProtocolError::WrongArity("CLUSTER DELSLOTS".into())); + } + let slots = parse_slot_list(&args[1..])?; + Ok(Command::ClusterDelSlots { slots }) + } + "FORGET" => { + if args.len() != 2 { + return Err(ProtocolError::WrongArity("CLUSTER FORGET".into())); + } + let node_id = extract_string(&args[1])?; + Ok(Command::ClusterForget { node_id }) + } + "REPLICATE" => { + if args.len() != 2 { + return Err(ProtocolError::WrongArity("CLUSTER REPLICATE".into())); + } + let node_id = extract_string(&args[1])?; + Ok(Command::ClusterReplicate { node_id }) + } + "FAILOVER" => { + let mut force = false; + let mut takeover = false; + for arg in &args[1..] { + let opt = extract_string(arg)?.to_ascii_uppercase(); + match opt.as_str() { + "FORCE" => force = true, + "TAKEOVER" => takeover = true, + _ => { + return Err(ProtocolError::InvalidCommandFrame(format!( + "unknown CLUSTER FAILOVER option '{opt}'" + ))) + } + } + } + Ok(Command::ClusterFailover { force, takeover }) + } + "COUNTKEYSINSLOT" => { + if args.len() != 2 { + return Err(ProtocolError::WrongArity("CLUSTER COUNTKEYSINSLOT".into())); + } + let slot_str = extract_string(&args[1])?; + let slot: u16 = slot_str + .parse() + .map_err(|_| ProtocolError::InvalidCommandFrame("invalid slot number".into()))?; + Ok(Command::ClusterCountKeysInSlot { slot }) + } + "GETKEYSINSLOT" => { + if args.len() != 3 { + return Err(ProtocolError::WrongArity("CLUSTER GETKEYSINSLOT".into())); + } + let slot_str = extract_string(&args[1])?; + let slot: u16 = slot_str + .parse() + .map_err(|_| ProtocolError::InvalidCommandFrame("invalid slot number".into()))?; + let count_str = extract_string(&args[2])?; + let count: u32 = count_str + .parse() + .map_err(|_| ProtocolError::InvalidCommandFrame("invalid count".into()))?; + Ok(Command::ClusterGetKeysInSlot { slot, count }) + } _ => Err(ProtocolError::InvalidCommandFrame(format!( "unknown CLUSTER subcommand '{subcommand}'" ))), @@ -1116,6 +1220,18 @@ fn parse_asking(args: &[Frame]) -> Result { Ok(Command::Asking) } +fn parse_slot_list(args: &[Frame]) -> Result, ProtocolError> { + let mut slots = Vec::with_capacity(args.len()); + for arg in args { + let slot_str = extract_string(arg)?; + let slot: u16 = slot_str + .parse() + .map_err(|_| ProtocolError::InvalidCommandFrame("invalid slot number".into()))?; + slots.push(slot); + } + Ok(slots) +} + fn parse_cluster_setslot(args: &[Frame]) -> Result { if args.is_empty() { return Err(ProtocolError::WrongArity("CLUSTER SETSLOT".into())); @@ -2907,4 +3023,107 @@ mod tests { .unwrap_err(); assert!(matches!(err, ProtocolError::InvalidCommandFrame(_))); } + + #[test] + fn cluster_meet_basic() { + assert_eq!( + Command::from_frame(cmd(&["CLUSTER", "MEET", "192.168.1.1", "6379"])).unwrap(), + Command::ClusterMeet { + ip: "192.168.1.1".into(), + port: 6379 + }, + ); + } + + #[test] + fn cluster_addslots_basic() { + assert_eq!( + Command::from_frame(cmd(&["CLUSTER", "ADDSLOTS", "0", "1", "2"])).unwrap(), + Command::ClusterAddSlots { + slots: vec![0, 1, 2] + }, + ); + } + + #[test] + fn cluster_delslots_basic() { + assert_eq!( + Command::from_frame(cmd(&["CLUSTER", "DELSLOTS", "100", "101"])).unwrap(), + Command::ClusterDelSlots { + slots: vec![100, 101] + }, + ); + } + + #[test] + fn cluster_forget_basic() { + assert_eq!( + Command::from_frame(cmd(&["CLUSTER", "FORGET", "abc123"])).unwrap(), + Command::ClusterForget { + node_id: "abc123".into() + }, + ); + } + + #[test] + fn cluster_replicate_basic() { + assert_eq!( + Command::from_frame(cmd(&["CLUSTER", "REPLICATE", "master-id"])).unwrap(), + Command::ClusterReplicate { + node_id: "master-id".into() + }, + ); + } + + #[test] + fn cluster_failover_basic() { + assert_eq!( + Command::from_frame(cmd(&["CLUSTER", "FAILOVER"])).unwrap(), + Command::ClusterFailover { + force: false, + takeover: false + }, + ); + } + + #[test] + fn cluster_failover_force() { + assert_eq!( + Command::from_frame(cmd(&["CLUSTER", "FAILOVER", "FORCE"])).unwrap(), + Command::ClusterFailover { + force: true, + takeover: false + }, + ); + } + + #[test] + fn cluster_failover_takeover() { + assert_eq!( + Command::from_frame(cmd(&["CLUSTER", "FAILOVER", "TAKEOVER"])).unwrap(), + Command::ClusterFailover { + force: false, + takeover: true + }, + ); + } + + #[test] + fn cluster_countkeysinslot_basic() { + assert_eq!( + Command::from_frame(cmd(&["CLUSTER", "COUNTKEYSINSLOT", "100"])).unwrap(), + Command::ClusterCountKeysInSlot { slot: 100 }, + ); + } + + #[test] + fn cluster_getkeysinslot_basic() { + assert_eq!( + Command::from_frame(cmd(&["CLUSTER", "GETKEYSINSLOT", "200", "10"])).unwrap(), + Command::ClusterGetKeysInSlot { + slot: 200, + count: 10 + }, + ); + } } diff --git a/crates/ember-server/src/connection.rs b/crates/ember-server/src/connection.rs index 24f53c37..2949650a 100644 --- a/crates/ember-server/src/connection.rs +++ b/crates/ember-server/src/connection.rs @@ -875,7 +875,15 @@ async fn execute(cmd: Command, engine: &Engine) -> Frame { Command::ClusterSetSlotImporting { .. } | Command::ClusterSetSlotMigrating { .. } | Command::ClusterSetSlotNode { .. } - | Command::ClusterSetSlotStable { .. } => { + | Command::ClusterSetSlotStable { .. } + | Command::ClusterMeet { .. } + | Command::ClusterAddSlots { .. } + | Command::ClusterDelSlots { .. } + | Command::ClusterForget { .. } + | Command::ClusterReplicate { .. } + | Command::ClusterFailover { .. } + | Command::ClusterCountKeysInSlot { .. } + | Command::ClusterGetKeysInSlot { .. } => { Frame::Error("ERR This instance has cluster support disabled".into()) }