From e385258f7cec71f4ef538d7480cf9a748767397a Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Mon, 9 Feb 2026 20:43:32 -0500 Subject: [PATCH 1/4] feat: add slot-aware key queries to keyspace and shard add count_keys_in_slot and get_keys_in_slot methods to keyspace for cluster slot inspection. add corresponding shard request/response variants and dispatch handlers. these support the CLUSTER COUNTKEYSINSLOT and CLUSTER GETKEYSINSLOT commands. --- crates/ember-core/Cargo.toml | 1 + crates/ember-core/src/keyspace.rs | 81 +++++++++++++++++++++++++++++++ crates/ember-core/src/shard.rs | 10 ++++ 3 files changed, 92 insertions(+) diff --git a/crates/ember-core/Cargo.toml b/crates/ember-core/Cargo.toml index 3c990b39..6abda40a 100644 --- a/crates/ember-core/Cargo.toml +++ b/crates/ember-core/Cargo.toml @@ -13,6 +13,7 @@ readme = "README.md" name = "ember_core" [dependencies] +ember-cluster = { workspace = true } ember-protocol = { workspace = true } ember-persistence = { workspace = true } thiserror = { workspace = true } diff --git a/crates/ember-core/src/keyspace.rs b/crates/ember-core/src/keyspace.rs index 60e03541..810da3a4 100644 --- a/crates/ember-core/src/keyspace.rs +++ b/crates/ember-core/src/keyspace.rs @@ -747,6 +747,30 @@ impl Keyspace { .collect() } + /// Counts live keys in this keyspace that hash to the given cluster slot. + /// + /// O(n) scan over all entries — same cost as KEYS. + pub fn count_keys_in_slot(&self, slot: u16) -> usize { + self.entries + .iter() + .filter(|(_, entry)| !entry.is_expired()) + .filter(|(key, _)| ember_cluster::key_slot(key.as_bytes()) == slot) + .count() + } + + /// Returns up to `count` live keys that hash to the given cluster slot. + /// + /// O(n) scan over all entries — same cost as KEYS. + pub fn get_keys_in_slot(&self, slot: u16, count: usize) -> Vec { + self.entries + .iter() + .filter(|(_, entry)| !entry.is_expired()) + .filter(|(key, _)| ember_cluster::key_slot(key.as_bytes()) == slot) + .take(count) + .map(|(key, _)| key.clone()) + .collect() + } + /// Renames a key to a new name. Returns an error if the source key /// doesn't exist. If the destination key already exists, it is overwritten. pub fn rename(&mut self, key: &str, newkey: &str) -> Result<(), RenameError> { @@ -3957,4 +3981,61 @@ mod tests { assert_eq!(before, after); assert_eq!(ks.stats().key_count, 1); } + + #[test] + fn count_keys_in_slot_empty() { + let ks = Keyspace::new(); + assert_eq!(ks.count_keys_in_slot(0), 0); + } + + #[test] + fn count_keys_in_slot_matches() { + let mut ks = Keyspace::new(); + // insert a few keys and count those in a specific slot + ks.set("a".into(), Bytes::from("1"), None); + ks.set("b".into(), Bytes::from("2"), None); + ks.set("c".into(), Bytes::from("3"), None); + + let slot_a = ember_cluster::key_slot(b"a"); + let count = ks.count_keys_in_slot(slot_a); + // at minimum, "a" should be in its own slot + assert!(count >= 1); + } + + #[test] + fn count_keys_in_slot_skips_expired() { + let mut ks = Keyspace::new(); + let slot = ember_cluster::key_slot(b"temp"); + ks.set( + "temp".into(), + Bytes::from("gone"), + Some(Duration::from_millis(0)), + ); + // key is expired — should not be counted + thread::sleep(Duration::from_millis(5)); + assert_eq!(ks.count_keys_in_slot(slot), 0); + } + + #[test] + fn get_keys_in_slot_returns_matching() { + let mut ks = Keyspace::new(); + ks.set("x".into(), Bytes::from("1"), None); + ks.set("y".into(), Bytes::from("2"), None); + + let slot_x = ember_cluster::key_slot(b"x"); + let keys = ks.get_keys_in_slot(slot_x, 100); + assert!(keys.contains(&"x".to_string())); + } + + #[test] + fn get_keys_in_slot_respects_count_limit() { + let mut ks = Keyspace::new(); + // insert several keys — some might share a slot + for i in 0..100 { + ks.set(format!("key:{i}"), Bytes::from("v"), None); + } + // ask for at most 3 keys from slot 0 + let keys = ks.get_keys_in_slot(0, 3); + assert!(keys.len() <= 3); + } } diff --git a/crates/ember-core/src/shard.rs b/crates/ember-core/src/shard.rs index 05d05531..edc8698c 100644 --- a/crates/ember-core/src/shard.rs +++ b/crates/ember-core/src/shard.rs @@ -247,6 +247,10 @@ pub enum ShardRequest { count: usize, pattern: Option, }, + /// Counts keys in this shard that hash to the given cluster slot. + CountKeysInSlot { slot: u16 }, + /// Returns up to `count` keys that hash to the given cluster slot. + GetKeysInSlot { slot: u16, count: usize }, } /// The shard's response to a request. @@ -783,6 +787,12 @@ fn dispatch(ks: &mut Keyspace, req: &ShardRequest) -> ShardResponse { Ok(count) => ShardResponse::Len(count), Err(_) => ShardResponse::WrongType, }, + ShardRequest::CountKeysInSlot { slot } => { + ShardResponse::KeyCount(ks.count_keys_in_slot(*slot)) + } + ShardRequest::GetKeysInSlot { slot, count } => { + ShardResponse::StringArray(ks.get_keys_in_slot(*slot, *count)) + } // snapshot/rewrite/flush_async are handled in the main loop, not here ShardRequest::Snapshot | ShardRequest::RewriteAof | ShardRequest::FlushDbAsync => { ShardResponse::Ok From f415b763e37dcb1851475fad22758df95a957220 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Mon, 9 Feb 2026 20:44:18 -0500 Subject: [PATCH 2/4] feat: add SETSLOT command handlers to cluster coordinator add MigrationManager to ClusterCoordinator and implement the four SETSLOT subcommands: IMPORTING, MIGRATING, NODE, and STABLE. validates slot ranges, node IDs, and ownership before starting migrations. --- crates/ember-server/src/cluster.rs | 117 ++++++++++++++++++++++++++++- 1 file changed, 116 insertions(+), 1 deletion(-) diff --git a/crates/ember-server/src/cluster.rs b/crates/ember-server/src/cluster.rs index 74b23923..a81ca843 100644 --- a/crates/ember-server/src/cluster.rs +++ b/crates/ember-server/src/cluster.rs @@ -10,7 +10,7 @@ use std::sync::Arc; use bytes::Bytes; use ember_cluster::{ key_slot, ClusterNode, ClusterState, GossipConfig, GossipEngine, GossipEvent, GossipMessage, - NodeId, SLOT_COUNT, + MigrationManager, NodeId, SLOT_COUNT, }; use ember_protocol::Frame; use tokio::net::UdpSocket; @@ -24,6 +24,7 @@ use tracing::{debug, error, info, warn}; pub struct ClusterCoordinator { state: RwLock, gossip: Mutex, + migration: Mutex, local_id: NodeId, /// bound UDP socket for gossip, set after spawn_gossip udp_socket: Mutex>>, @@ -72,6 +73,7 @@ impl ClusterCoordinator { let coordinator = Self { state: RwLock::new(state), gossip: Mutex::new(gossip), + migration: Mutex::new(MigrationManager::new()), local_id, udp_socket: Mutex::new(None), }; @@ -251,6 +253,119 @@ impl ClusterCoordinator { } } + // -- slot migration (SETSLOT) commands -- + + /// CLUSTER SETSLOT IMPORTING + /// + /// Marks a slot as importing from the given source node. The local node + /// becomes the target of the migration. + pub async fn cluster_setslot_importing(&self, slot: u16, node_id_str: &str) -> Frame { + if slot >= SLOT_COUNT { + return Frame::Error(format!("ERR Invalid or out of range slot {slot}")); + } + let source_id = match NodeId::parse(node_id_str) { + Ok(id) => id, + Err(_) => return Frame::Error("ERR Invalid node ID".into()), + }; + if source_id == self.local_id { + return Frame::Error("ERR can't import from myself".into()); + } + + let mut migration = self.migration.lock().await; + match migration.start_import(slot, source_id, self.local_id) { + Ok(_) => Frame::Simple("OK".into()), + Err(e) => Frame::Error(format!("ERR {e}")), + } + } + + /// CLUSTER SETSLOT MIGRATING + /// + /// Marks a slot as migrating to the given target node. The local node + /// must currently own the slot. + pub async fn cluster_setslot_migrating(&self, slot: u16, node_id_str: &str) -> Frame { + if slot >= SLOT_COUNT { + return Frame::Error(format!("ERR Invalid or out of range slot {slot}")); + } + let target_id = match NodeId::parse(node_id_str) { + Ok(id) => id, + Err(_) => return Frame::Error("ERR Invalid node ID".into()), + }; + if target_id == self.local_id { + return Frame::Error("ERR can't migrate to myself".into()); + } + + // verify we own the slot before allowing migration + { + let state = self.state.read().await; + if !state.owns_slot(slot) { + return Frame::Error(format!( + "ERR I'm not the owner of hash slot {slot}" + )); + } + } + + let mut migration = self.migration.lock().await; + match migration.start_migrate(slot, self.local_id, target_id) { + Ok(_) => Frame::Simple("OK".into()), + Err(e) => Frame::Error(format!("ERR {e}")), + } + } + + /// CLUSTER SETSLOT NODE + /// + /// Completes migration by assigning the slot to the given node. + /// Cleans up any in-progress migration state. + pub async fn cluster_setslot_node(&self, slot: u16, node_id_str: &str) -> Frame { + if slot >= SLOT_COUNT { + return Frame::Error(format!("ERR Invalid or out of range slot {slot}")); + } + let node_id = match NodeId::parse(node_id_str) { + Ok(id) => id, + Err(_) => return Frame::Error("ERR Invalid node ID".into()), + }; + + // complete any in-progress migration for this slot + { + let mut migration = self.migration.lock().await; + migration.complete_migration(slot); + } + + // assign the slot to the specified node + let mut state = self.state.write().await; + state.slot_map.assign(slot, node_id); + + // update the node's slot list + let new_slots = state.slot_map.slots_for_node(node_id); + if let Some(node) = state.nodes.get_mut(&node_id) { + node.slots = new_slots; + } + + // also update the local node's slot list if it changed + if node_id != self.local_id { + let local_slots = state.slot_map.slots_for_node(self.local_id); + if let Some(node) = state.nodes.get_mut(&self.local_id) { + node.slots = local_slots; + } + } + + state.update_health(); + Frame::Simple("OK".into()) + } + + /// CLUSTER SETSLOT STABLE + /// + /// Aborts any in-progress migration for the slot, clearing + /// importing/migrating state without changing slot ownership. + pub async fn cluster_setslot_stable(&self, slot: u16) -> Frame { + if slot >= SLOT_COUNT { + return Frame::Error(format!("ERR Invalid or out of range slot {slot}")); + } + + let mut migration = self.migration.lock().await; + migration.abort_migration(slot); + Frame::Simple("OK".into()) + } + // -- slot ownership check -- /// Checks if the local node owns the slot for the given key. From ab9a27b170eb4e6953ecae074c73e1f38a45b32b Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Mon, 9 Feb 2026 20:45:36 -0500 Subject: [PATCH 3/4] feat: wire remaining cluster commands in connection handler MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit replace "not yet implemented" stubs with real dispatch: - SETSLOT IMPORTING/MIGRATING/NODE/STABLE → cluster coordinator - COUNTKEYSINSLOT → broadcast to shards, sum counts - GETKEYSINSLOT → broadcast to shards, aggregate and truncate - REPLICATE/FAILOVER → honest "not yet supported" error - MIGRATE → kept as not implemented (needs remote connections) --- crates/ember-server/src/connection.rs | 79 ++++++++++++++++++++++++--- 1 file changed, 71 insertions(+), 8 deletions(-) diff --git a/crates/ember-server/src/connection.rs b/crates/ember-server/src/connection.rs index 1706a701..73b666ab 100644 --- a/crates/ember-server/src/connection.rs +++ b/crates/ember-server/src/connection.rs @@ -1490,14 +1490,77 @@ async fn execute( None => Frame::Error("ERR This instance has cluster support disabled".into()), }, - Command::ClusterSetSlotImporting { .. } - | Command::ClusterSetSlotMigrating { .. } - | Command::ClusterSetSlotNode { .. } - | Command::ClusterSetSlotStable { .. } - | Command::ClusterReplicate { .. } - | Command::ClusterFailover { .. } - | Command::ClusterCountKeysInSlot { .. } - | Command::ClusterGetKeysInSlot { .. } => Frame::Error("ERR not yet implemented".into()), + Command::ClusterSetSlotImporting { slot, node_id } => match &ctx.cluster { + Some(c) => c.cluster_setslot_importing(slot, &node_id).await, + None => Frame::Error("ERR This instance has cluster support disabled".into()), + }, + + Command::ClusterSetSlotMigrating { slot, node_id } => match &ctx.cluster { + Some(c) => c.cluster_setslot_migrating(slot, &node_id).await, + None => Frame::Error("ERR This instance has cluster support disabled".into()), + }, + + Command::ClusterSetSlotNode { slot, node_id } => match &ctx.cluster { + Some(c) => c.cluster_setslot_node(slot, &node_id).await, + None => Frame::Error("ERR This instance has cluster support disabled".into()), + }, + + Command::ClusterSetSlotStable { slot } => match &ctx.cluster { + Some(c) => c.cluster_setslot_stable(slot).await, + None => Frame::Error("ERR This instance has cluster support disabled".into()), + }, + + Command::ClusterCountKeysInSlot { slot } => { + match engine + .broadcast(|| ShardRequest::CountKeysInSlot { slot }) + .await + { + Ok(responses) => { + let total: usize = responses + .iter() + .map(|r| match r { + ShardResponse::KeyCount(n) => *n, + _ => 0, + }) + .sum(); + Frame::Integer(total as i64) + } + Err(e) => Frame::Error(format!("ERR {e}")), + } + } + + Command::ClusterGetKeysInSlot { slot, count } => { + let count = count as usize; + match engine + .broadcast(|| ShardRequest::GetKeysInSlot { slot, count }) + .await + { + Ok(responses) => { + let mut all_keys = Vec::new(); + for r in responses { + if let ShardResponse::StringArray(keys) = r { + all_keys.extend(keys); + } + } + all_keys.truncate(count); + Frame::Array( + all_keys + .into_iter() + .map(|k| Frame::Bulk(Bytes::from(k))) + .collect(), + ) + } + Err(e) => Frame::Error(format!("ERR {e}")), + } + } + + Command::ClusterReplicate { .. } => { + Frame::Error("ERR REPLICATE not yet supported".into()) + } + + Command::ClusterFailover { .. } => { + Frame::Error("ERR FAILOVER not yet supported".into()) + } Command::Migrate { .. } => Frame::Error("ERR not yet implemented".into()), From 65dd1b6b494f71e257f9daf674c8cbe82f7fc309 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Mon, 9 Feb 2026 20:46:16 -0500 Subject: [PATCH 4/4] test: add unit tests for SETSLOT command handlers covers importing, migrating, node, and stable subcommands including edge cases: invalid slot, self-import/migrate rejection, duplicate migration, non-owner migration, node assignment, migration cleanup, and stable noop. --- crates/ember-server/src/cluster.rs | 172 +++++++++++++++++++++++++++++ 1 file changed, 172 insertions(+) diff --git a/crates/ember-server/src/cluster.rs b/crates/ember-server/src/cluster.rs index a81ca843..e91a4e08 100644 --- a/crates/ember-server/src/cluster.rs +++ b/crates/ember-server/src/cluster.rs @@ -534,3 +534,175 @@ impl ClusterCoordinator { }); } } + +#[cfg(test)] +mod tests { + use super::*; + + /// Creates a test coordinator with a single node that owns no slots. + fn test_coordinator() -> (ClusterCoordinator, mpsc::Receiver) { + let local_id = NodeId::new(); + let addr: SocketAddr = "127.0.0.1:6379".parse().unwrap(); + let config = GossipConfig::default(); + ClusterCoordinator::new(local_id, addr, config, false) + } + + /// Creates a test coordinator bootstrapped with all 16384 slots. + fn test_coordinator_bootstrapped() -> (ClusterCoordinator, mpsc::Receiver) { + let local_id = NodeId::new(); + let addr: SocketAddr = "127.0.0.1:6379".parse().unwrap(); + let config = GossipConfig::default(); + ClusterCoordinator::new(local_id, addr, config, true) + } + + #[tokio::test] + async fn setslot_importing_valid() { + let (coord, _rx) = test_coordinator(); + let source = NodeId::new(); + let resp = coord + .cluster_setslot_importing(100, &source.0.to_string()) + .await; + assert!(matches!(resp, Frame::Simple(_))); + } + + #[tokio::test] + async fn setslot_importing_invalid_slot() { + let (coord, _rx) = test_coordinator(); + let source = NodeId::new(); + let resp = coord + .cluster_setslot_importing(16384, &source.0.to_string()) + .await; + assert!(matches!(resp, Frame::Error(_))); + } + + #[tokio::test] + async fn setslot_importing_self_rejected() { + let (coord, _rx) = test_coordinator(); + let resp = coord + .cluster_setslot_importing(100, &coord.local_id.0.to_string()) + .await; + match resp { + Frame::Error(msg) => assert!(msg.contains("can't import from myself")), + other => panic!("expected error, got {other:?}"), + } + } + + #[tokio::test] + async fn setslot_importing_duplicate_rejected() { + let (coord, _rx) = test_coordinator(); + let source = NodeId::new(); + let id_str = source.0.to_string(); + coord.cluster_setslot_importing(100, &id_str).await; + let resp = coord.cluster_setslot_importing(100, &id_str).await; + assert!(matches!(resp, Frame::Error(_))); + } + + #[tokio::test] + async fn setslot_migrating_valid() { + let (coord, _rx) = test_coordinator_bootstrapped(); + let target = NodeId::new(); + let resp = coord + .cluster_setslot_migrating(0, &target.0.to_string()) + .await; + assert!(matches!(resp, Frame::Simple(_))); + } + + #[tokio::test] + async fn setslot_migrating_not_owner() { + let (coord, _rx) = test_coordinator(); // no slots owned + let target = NodeId::new(); + let resp = coord + .cluster_setslot_migrating(100, &target.0.to_string()) + .await; + match resp { + Frame::Error(msg) => assert!(msg.contains("not the owner")), + other => panic!("expected error, got {other:?}"), + } + } + + #[tokio::test] + async fn setslot_migrating_self_rejected() { + let (coord, _rx) = test_coordinator_bootstrapped(); + let resp = coord + .cluster_setslot_migrating(0, &coord.local_id.0.to_string()) + .await; + match resp { + Frame::Error(msg) => assert!(msg.contains("can't migrate to myself")), + other => panic!("expected error, got {other:?}"), + } + } + + #[tokio::test] + async fn setslot_node_assigns_slot() { + let (coord, _rx) = test_coordinator(); + let target = NodeId::new(); + + // add the target node to cluster state + { + let mut state = coord.state.write().await; + let node = ClusterNode::new_primary(target, "127.0.0.1:6380".parse().unwrap()); + state.add_node(node); + } + + let resp = coord + .cluster_setslot_node(100, &target.0.to_string()) + .await; + assert!(matches!(resp, Frame::Simple(_))); + + // verify the slot is now owned by the target + let state = coord.state.read().await; + assert_eq!(state.slot_map.owner(100), Some(target)); + } + + #[tokio::test] + async fn setslot_node_completes_migration() { + let (coord, _rx) = test_coordinator_bootstrapped(); + let target = NodeId::new(); + + // start a migration + coord + .cluster_setslot_migrating(0, &target.0.to_string()) + .await; + + // add target to state + { + let mut state = coord.state.write().await; + let node = ClusterNode::new_primary(target, "127.0.0.1:6380".parse().unwrap()); + state.add_node(node); + } + + // complete with NODE — should clean up migration state + let resp = coord + .cluster_setslot_node(0, &target.0.to_string()) + .await; + assert!(matches!(resp, Frame::Simple(_))); + + // migration should be cleaned up + let migration = coord.migration.lock().await; + assert!(!migration.is_migrating(0)); + } + + #[tokio::test] + async fn setslot_stable_aborts_migration() { + let (coord, _rx) = test_coordinator(); + let source = NodeId::new(); + coord + .cluster_setslot_importing(100, &source.0.to_string()) + .await; + + let resp = coord.cluster_setslot_stable(100).await; + assert!(matches!(resp, Frame::Simple(_))); + + // migration should be cleaned up + let migration = coord.migration.lock().await; + assert!(!migration.is_importing(100)); + } + + #[tokio::test] + async fn setslot_stable_noop_when_no_migration() { + let (coord, _rx) = test_coordinator(); + // should succeed even with no active migration + let resp = coord.cluster_setslot_stable(100).await; + assert!(matches!(resp, Frame::Simple(_))); + } +}