diff --git a/crates/ember-cluster/src/topology.rs b/crates/ember-cluster/src/topology.rs index b5d1f4e6..d63d799b 100644 --- a/crates/ember-cluster/src/topology.rs +++ b/crates/ember-cluster/src/topology.rs @@ -10,7 +10,7 @@ use std::time::Instant; use serde::{Deserialize, Serialize}; use uuid::Uuid; -use crate::slots::{SlotMap, SlotRange}; +use crate::slots::{SlotMap, SlotRange, SLOT_COUNT}; use crate::ClusterError; /// Unique identifier for a cluster node. @@ -408,8 +408,8 @@ impl ClusterState { /// Generates the response for CLUSTER INFO command. pub fn cluster_info(&self) -> String { - let primaries: Vec<_> = self.primaries().collect(); - let assigned_slots: u16 = primaries.iter().map(|n| n.slot_count()).sum(); + let assigned_slots = (SLOT_COUNT as usize - self.slot_map.unassigned_count()) as u16; + let primaries_count = self.primaries().count(); format!( "cluster_state:{}\r\n\ @@ -429,7 +429,7 @@ impl ClusterState { 0 }, self.nodes.len(), - primaries.len(), + primaries_count, self.config_epoch, self.local_node().map(|n| n.config_epoch).unwrap_or(0), ) @@ -566,7 +566,7 @@ mod tests { let info = state.cluster_info(); assert!(info.contains("cluster_state:ok")); - assert!(info.contains("cluster_slots_assigned:0")); // slots in node.slots, not slot_map + assert!(info.contains("cluster_slots_assigned:16384")); assert!(info.contains("cluster_known_nodes:1")); } diff --git a/crates/ember-server/src/cluster.rs b/crates/ember-server/src/cluster.rs new file mode 100644 index 00000000..38182540 --- /dev/null +++ b/crates/ember-server/src/cluster.rs @@ -0,0 +1,436 @@ +//! Cluster coordination layer for the ember server. +//! +//! Wraps the ember-cluster crate's types into a server-integrated +//! coordinator that handles gossip networking, cluster commands, +//! and slot ownership validation. + +use std::net::SocketAddr; +use std::sync::Arc; + +use bytes::Bytes; +use ember_cluster::{ + key_slot, ClusterNode, ClusterState, GossipConfig, GossipEngine, GossipEvent, GossipMessage, + NodeId, SLOT_COUNT, +}; +use ember_protocol::Frame; +use tokio::net::UdpSocket; +use tokio::sync::{mpsc, Mutex, RwLock}; +use tracing::{debug, error, info, warn}; + +/// Integration struct wrapping cluster crate types for the running server. +/// +/// Thread-safe via interior mutability: `RwLock` for state (many readers, +/// rare writers) and `Mutex` for gossip (single writer during ticks). +pub struct ClusterCoordinator { + state: RwLock, + gossip: Mutex, + local_id: NodeId, + /// bound UDP socket for gossip, set after spawn_gossip + udp_socket: Mutex>>, +} + +impl std::fmt::Debug for ClusterCoordinator { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ClusterCoordinator") + .field("local_id", &self.local_id) + .finish_non_exhaustive() + } +} + +impl ClusterCoordinator { + /// Creates a new cluster coordinator. + /// + /// Returns the coordinator and a receiver for gossip events that + /// should be consumed by a background task. + pub fn new( + local_id: NodeId, + bind_addr: SocketAddr, + gossip_config: GossipConfig, + bootstrap: bool, + ) -> (Self, mpsc::Receiver) { + let (event_tx, event_rx) = mpsc::channel(256); + + let gossip_addr = SocketAddr::new( + bind_addr.ip(), + bind_addr.port() + gossip_config.gossip_port_offset, + ); + + let gossip = GossipEngine::new(local_id, gossip_addr, gossip_config, event_tx); + + let state = if bootstrap { + let mut node = ClusterNode::new_primary(local_id, bind_addr); + node.set_myself(); + ClusterState::single_node(node) + } else { + let mut cs = ClusterState::new(local_id); + let mut node = ClusterNode::new_primary(local_id, bind_addr); + node.set_myself(); + cs.add_node(node); + cs + }; + + let coordinator = Self { + state: RwLock::new(state), + gossip: Mutex::new(gossip), + local_id, + udp_socket: Mutex::new(None), + }; + + (coordinator, event_rx) + } + + // -- cluster command handlers -- + + /// CLUSTER INFO + pub async fn cluster_info(&self) -> Frame { + let state = self.state.read().await; + Frame::Bulk(Bytes::from(state.cluster_info())) + } + + /// CLUSTER NODES + pub async fn cluster_nodes(&self) -> Frame { + let state = self.state.read().await; + Frame::Bulk(Bytes::from(state.cluster_nodes())) + } + + /// CLUSTER MYID + pub fn cluster_myid(&self) -> Frame { + Frame::Bulk(Bytes::from(self.local_id.0.to_string())) + } + + /// CLUSTER SLOTS — returns slot ranges in the Redis array format. + pub async fn cluster_slots(&self) -> Frame { + let state = self.state.read().await; + + let mut result = Vec::new(); + for node in state.primaries() { + let ranges = state.slot_map.slots_for_node(node.id); + for range in ranges { + let mut entry = vec![ + Frame::Integer(range.start as i64), + Frame::Integer(range.end as i64), + // node info: [ip, port, id] + Frame::Array(vec![ + Frame::Bulk(Bytes::from(node.addr.ip().to_string())), + Frame::Integer(node.addr.port() as i64), + Frame::Bulk(Bytes::from(node.id.0.to_string())), + ]), + ]; + + // add replicas + for replica in state.replicas_of(node.id) { + entry.push(Frame::Array(vec![ + Frame::Bulk(Bytes::from(replica.addr.ip().to_string())), + Frame::Integer(replica.addr.port() as i64), + Frame::Bulk(Bytes::from(replica.id.0.to_string())), + ])); + } + + result.push(Frame::Array(entry)); + } + } + + Frame::Array(result) + } + + /// CLUSTER MEET ip port + pub async fn cluster_meet(&self, ip: &str, port: u16) -> Frame { + let addr: SocketAddr = match format!("{ip}:{port}").parse() { + Ok(a) => a, + Err(e) => return Frame::Error(format!("ERR invalid address: {e}")), + }; + + let mut gossip = self.gossip.lock().await; + let new_id = NodeId::new(); + let gossip_port_offset = 10000u16; // default + let gossip_addr = + SocketAddr::new(addr.ip(), addr.port().saturating_add(gossip_port_offset)); + + gossip.add_seed(new_id, gossip_addr); + + // send join message via UDP + let join_msg = gossip.create_join_message(); + let encoded = join_msg.encode(); + + let socket = self.udp_socket.lock().await; + if let Some(ref sock) = *socket { + if let Err(e) = sock.send_to(&encoded, gossip_addr).await { + warn!("failed to send join to {gossip_addr}: {e}"); + return Frame::Error(format!("ERR failed to send join: {e}")); + } + } else { + return Frame::Error("ERR gossip socket not ready".into()); + } + + // add to cluster state as well + let mut state = self.state.write().await; + let node = ClusterNode::new_primary(new_id, addr); + state.add_node(node); + + Frame::Simple("OK".into()) + } + + /// CLUSTER ADDSLOTS slot [slot ...] + pub async fn cluster_addslots(&self, slots: &[u16]) -> Frame { + let mut state = self.state.write().await; + + // validate: all slots must be unassigned + for &slot in slots { + if slot >= SLOT_COUNT { + return Frame::Error(format!( + "ERR Invalid or out of range slot {slot}" + )); + } + if state.slot_map.owner(slot).is_some() { + return Frame::Error(format!( + "ERR Slot {slot} is already busy" + )); + } + } + + // assign all slots to local node + for &slot in slots { + state.slot_map.assign(slot, self.local_id); + } + + // update node.slots from slot_map + let new_slots = state.slot_map.slots_for_node(self.local_id); + if let Some(node) = state.nodes.get_mut(&self.local_id) { + node.slots = new_slots; + } + + state.update_health(); + Frame::Simple("OK".into()) + } + + /// CLUSTER DELSLOTS slot [slot ...] + pub async fn cluster_delslots(&self, slots: &[u16]) -> Frame { + let mut state = self.state.write().await; + + // validate: all slots must be owned by us + for &slot in slots { + if slot >= SLOT_COUNT { + return Frame::Error(format!( + "ERR Invalid or out of range slot {slot}" + )); + } + match state.slot_map.owner(slot) { + Some(owner) if owner != self.local_id => { + return Frame::Error(format!( + "ERR Slot {slot} is not owned by this node" + )); + } + None => { + return Frame::Error(format!( + "ERR Slot {slot} is already unassigned" + )); + } + _ => {} + } + } + + for &slot in slots { + state.slot_map.unassign(slot); + } + + // update node.slots from slot_map + let new_slots = state.slot_map.slots_for_node(self.local_id); + if let Some(node) = state.nodes.get_mut(&self.local_id) { + node.slots = new_slots; + } + + state.update_health(); + Frame::Simple("OK".into()) + } + + /// CLUSTER FORGET node-id + pub async fn cluster_forget(&self, node_id_str: &str) -> Frame { + let node_id = match NodeId::parse(node_id_str) { + Ok(id) => id, + Err(_) => return Frame::Error("ERR Invalid node ID".into()), + }; + + if node_id == self.local_id { + return Frame::Error("ERR I tried hard but I can't forget myself...".into()); + } + + let mut state = self.state.write().await; + match state.remove_node(node_id) { + Some(_) => Frame::Simple("OK".into()), + None => Frame::Error("ERR Unknown node ID".into()), + } + } + + // -- slot ownership check -- + + /// Checks if the local node owns the slot for the given key. + /// + /// Returns `None` if local node owns the slot (proceed normally). + /// Returns `Some(Frame::Error("MOVED ..."))` if another node owns it. + /// Returns `Some(Frame::Error("CLUSTERDOWN ..."))` if slot is unassigned. + pub async fn check_slot(&self, key: &[u8]) -> Option { + let slot = key_slot(key); + let state = self.state.read().await; + + if state.owns_slot(slot) { + return None; + } + + match state.slot_owner(slot) { + Some(owner) => Some(Frame::Error(format!( + "MOVED {} {}", + slot, owner.addr + ))), + None => Some(Frame::Error( + "CLUSTERDOWN Hash slot not served".into(), + )), + } + } + + /// Checks that all keys hash to the same slot. + /// + /// Returns `Ok(())` if all keys are in the same slot. + /// Returns `Err(Frame)` with a CROSSSLOT error if they span multiple slots. + pub fn check_crossslot(&self, keys: &[String]) -> Result<(), Frame> { + if keys.len() <= 1 { + return Ok(()); + } + let first_slot = key_slot(keys[0].as_bytes()); + for key in &keys[1..] { + if key_slot(key.as_bytes()) != first_slot { + return Err(Frame::Error( + "CROSSSLOT Keys in request don't hash to the same slot".into(), + )); + } + } + Ok(()) + } + + // -- gossip networking -- + + /// Spawns the gossip network tasks: UDP send/receive and event consumer. + pub async fn spawn_gossip( + self: &Arc, + bind_addr: SocketAddr, + mut event_rx: mpsc::Receiver, + ) { + let gossip_addr = { + let gossip = self.gossip.lock().await; + // gossip engine was initialized with the gossip address + // we need to bind to the same port + drop(gossip); + + // compute gossip address from bind_addr + offset + // the offset was baked into the GossipEngine, but we need + // to derive it for the UDP bind + SocketAddr::new(bind_addr.ip(), bind_addr.port() + 10000) + }; + + let socket = match UdpSocket::bind(gossip_addr).await { + Ok(s) => Arc::new(s), + Err(e) => { + error!("failed to bind gossip UDP socket on {gossip_addr}: {e}"); + return; + } + }; + + info!("gossip listening on {gossip_addr}"); + + // store socket for cluster_meet + { + let mut guard = self.udp_socket.lock().await; + *guard = Some(Arc::clone(&socket)); + } + + // task 1: gossip tick + UDP recv/send loop + let coordinator = Arc::clone(self); + let sock = Arc::clone(&socket); + tokio::spawn(async move { + let mut recv_buf = vec![0u8; 65535]; + let mut tick_interval = tokio::time::interval(std::time::Duration::from_secs(1)); + + loop { + tokio::select! { + _ = tick_interval.tick() => { + let mut gossip = coordinator.gossip.lock().await; + if let Some((target_addr, msg)) = gossip.tick() { + let encoded = msg.encode(); + if let Err(e) = sock.send_to(&encoded, target_addr).await { + debug!("gossip send error to {target_addr}: {e}"); + } + } + } + + result = sock.recv_from(&mut recv_buf) => { + match result { + Ok((len, from)) => { + match GossipMessage::decode(&recv_buf[..len]) { + Ok(msg) => { + let mut gossip = coordinator.gossip.lock().await; + if let Some(reply) = gossip.handle_message(msg, from).await { + let encoded = reply.encode(); + if let Err(e) = sock.send_to(&encoded, from).await { + debug!("gossip reply error to {from}: {e}"); + } + } + } + Err(e) => { + debug!("gossip decode error from {from}: {e}"); + } + } + } + Err(e) => { + warn!("gossip recv error: {e}"); + } + } + } + } + } + }); + + // task 2: gossip event consumer — updates cluster state + let coordinator = Arc::clone(self); + tokio::spawn(async move { + while let Some(event) = event_rx.recv().await { + let mut state = coordinator.state.write().await; + match event { + GossipEvent::MemberJoined(id, addr) => { + info!("cluster: node {} joined at {}", id, addr); + if !state.nodes.contains_key(&id) { + let node = ClusterNode::new_primary(id, addr); + state.add_node(node); + } + } + GossipEvent::MemberSuspected(id) => { + info!("cluster: node {} suspected", id); + if let Some(node) = state.nodes.get_mut(&id) { + node.flags.pfail = true; + } + state.update_health(); + } + GossipEvent::MemberFailed(id) => { + warn!("cluster: node {} confirmed failed", id); + if let Some(node) = state.nodes.get_mut(&id) { + node.flags.fail = true; + node.flags.pfail = false; + } + state.update_health(); + } + GossipEvent::MemberLeft(id) => { + info!("cluster: node {} left", id); + state.remove_node(id); + state.update_health(); + } + GossipEvent::MemberAlive(id) => { + debug!("cluster: node {} alive", id); + if let Some(node) = state.nodes.get_mut(&id) { + node.flags.pfail = false; + node.flags.fail = false; + } + state.update_health(); + } + } + } + }); + } +} diff --git a/crates/ember-server/src/connection.rs b/crates/ember-server/src/connection.rs index 29843451..95ba884d 100644 --- a/crates/ember-server/src/connection.rs +++ b/crates/ember-server/src/connection.rs @@ -503,6 +503,23 @@ async fn process( } } +/// Checks if the cluster owns the slot for the given key. +/// +/// Returns `None` if we should proceed (either not in cluster mode, or +/// we own the slot). Returns `Some(Frame)` with a MOVED/CLUSTERDOWN error +/// if the command should not be executed locally. +async fn check_cluster_slot(ctx: &ServerContext, key: &str) -> Option { + ctx.cluster.as_ref()?.check_slot(key.as_bytes()).await +} + +/// Checks that all keys hash to the same slot when in cluster mode. +/// +/// Returns `None` if ok. Returns `Some(Frame)` with CROSSSLOT error if not. +fn check_crossslot(ctx: &ServerContext, keys: &[String]) -> Option { + let cluster = ctx.cluster.as_ref()?; + cluster.check_crossslot(keys).err() +} + /// Executes a parsed command and returns the response frame. /// /// Ping and Echo are handled inline (no shard routing needed). @@ -523,6 +540,9 @@ async fn execute( // -- single-key commands -- Command::Get { key } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::Get { key: key.clone() }; match engine.route(&key, req).await { Ok(ShardResponse::Value(Some(Value::String(data)))) => Frame::Bulk(data), @@ -540,6 +560,9 @@ async fn execute( nx, xx, } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let duration = expire.map(|e| match e { SetExpire::Ex(secs) => Duration::from_secs(secs), SetExpire::Px(millis) => Duration::from_millis(millis), @@ -561,6 +584,9 @@ async fn execute( } Command::Expire { key, seconds } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::Expire { key: key.clone(), seconds, @@ -573,6 +599,9 @@ async fn execute( } Command::Ttl { key } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::Ttl { key: key.clone() }; match engine.route(&key, req).await { Ok(ShardResponse::Ttl(TtlResult::Seconds(s))) => Frame::Integer(s as i64), @@ -584,6 +613,9 @@ async fn execute( } Command::Incr { key } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::Incr { key: key.clone() }; match engine.route(&key, req).await { Ok(ShardResponse::Integer(n)) => Frame::Integer(n), @@ -596,6 +628,9 @@ async fn execute( } Command::Decr { key } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::Decr { key: key.clone() }; match engine.route(&key, req).await { Ok(ShardResponse::Integer(n)) => Frame::Integer(n), @@ -608,6 +643,9 @@ async fn execute( } Command::IncrBy { key, delta } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::IncrBy { key: key.clone(), delta, @@ -623,6 +661,9 @@ async fn execute( } Command::DecrBy { key, delta } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::DecrBy { key: key.clone(), delta, @@ -638,6 +679,9 @@ async fn execute( } Command::Append { key, value } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::Append { key: key.clone(), value, @@ -652,6 +696,9 @@ async fn execute( } Command::Strlen { key } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::Strlen { key: key.clone() }; match engine.route(&key, req).await { Ok(ShardResponse::Len(n)) => Frame::Integer(n as i64), @@ -662,6 +709,9 @@ async fn execute( } Command::IncrByFloat { key, delta } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::IncrByFloat { key: key.clone(), delta, @@ -677,6 +727,9 @@ async fn execute( } Command::Persist { key } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::Persist { key: key.clone() }; match engine.route(&key, req).await { Ok(ShardResponse::Bool(b)) => Frame::Integer(i64::from(b)), @@ -686,6 +739,9 @@ async fn execute( } Command::Pttl { key } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::Pttl { key: key.clone() }; match engine.route(&key, req).await { Ok(ShardResponse::Ttl(TtlResult::Milliseconds(ms))) => Frame::Integer(ms as i64), @@ -697,6 +753,9 @@ async fn execute( } Command::Pexpire { key, milliseconds } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::Pexpire { key: key.clone(), milliseconds, @@ -710,18 +769,50 @@ async fn execute( // -- multi-key fan-out -- Command::Del { keys } => { + if let Some(err) = check_crossslot(ctx, &keys) { + return err; + } + if let Some(first) = keys.first() { + if let Some(redirect) = check_cluster_slot(ctx, first).await { + return redirect; + } + } multi_key_bool(engine, &keys, |k| ShardRequest::Del { key: k }).await } Command::Unlink { keys } => { + if let Some(err) = check_crossslot(ctx, &keys) { + return err; + } + if let Some(first) = keys.first() { + if let Some(redirect) = check_cluster_slot(ctx, first).await { + return redirect; + } + } multi_key_bool(engine, &keys, |k| ShardRequest::Unlink { key: k }).await } Command::Exists { keys } => { + if let Some(err) = check_crossslot(ctx, &keys) { + return err; + } + if let Some(first) = keys.first() { + if let Some(redirect) = check_cluster_slot(ctx, first).await { + return redirect; + } + } multi_key_bool(engine, &keys, |k| ShardRequest::Exists { key: k }).await } Command::MGet { keys } => { + if let Some(err) = check_crossslot(ctx, &keys) { + return err; + } + if let Some(first) = keys.first() { + if let Some(redirect) = check_cluster_slot(ctx, first).await { + return redirect; + } + } match engine .route_multi(&keys, |k| ShardRequest::Get { key: k }) .await @@ -744,6 +835,18 @@ async fn execute( } Command::MSet { pairs } => { + // crossslot check for cluster mode + if ctx.cluster.is_some() { + let mset_keys: Vec = pairs.iter().map(|(k, _)| k.clone()).collect(); + if let Some(err) = check_crossslot(ctx, &mset_keys) { + return err; + } + if let Some(first) = mset_keys.first() { + if let Some(redirect) = check_cluster_slot(ctx, first).await { + return redirect; + } + } + } // Fan out individual SET requests — MSET always succeeds (or OOMs). // We build a HashMap for O(1) value lookups during routing. If there // are duplicate keys in pairs, the HashMap keeps the last value, which @@ -843,10 +946,14 @@ async fn execute( } Command::Rename { key, newkey } => { - // Route to the source key's shard. Both keys must hash to the - // same shard for a correct rename (same as Redis Cluster's - // cross-slot restriction). If they don't, the rename operates - // on the source shard and the newkey will live there. + if ctx.cluster.is_some() { + if let Some(err) = check_crossslot(ctx, &[key.clone(), newkey.clone()]) { + return err; + } + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } + } let req = ShardRequest::Rename { key: key.clone(), newkey, @@ -945,6 +1052,9 @@ async fn execute( // -- list commands -- Command::LPush { key, values } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::LPush { key: key.clone(), values, @@ -959,6 +1069,9 @@ async fn execute( } Command::RPush { key, values } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::RPush { key: key.clone(), values, @@ -973,6 +1086,9 @@ async fn execute( } Command::LPop { key } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::LPop { key: key.clone() }; match engine.route(&key, req).await { Ok(ShardResponse::Value(Some(Value::String(data)))) => Frame::Bulk(data), @@ -984,6 +1100,9 @@ async fn execute( } Command::RPop { key } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::RPop { key: key.clone() }; match engine.route(&key, req).await { Ok(ShardResponse::Value(Some(Value::String(data)))) => Frame::Bulk(data), @@ -995,6 +1114,9 @@ async fn execute( } Command::LRange { key, start, stop } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::LRange { key: key.clone(), start, @@ -1012,6 +1134,9 @@ async fn execute( } Command::LLen { key } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::LLen { key: key.clone() }; match engine.route(&key, req).await { Ok(ShardResponse::Len(n)) => Frame::Integer(n as i64), @@ -1022,6 +1147,9 @@ async fn execute( } Command::Type { key } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::Type { key: key.clone() }; match engine.route(&key, req).await { Ok(ShardResponse::TypeName(name)) => Frame::Simple(name.into()), @@ -1036,6 +1164,9 @@ async fn execute( flags, members, } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::ZAdd { key: key.clone(), members, @@ -1055,6 +1186,9 @@ async fn execute( } Command::ZRem { key, members } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::ZRem { key: key.clone(), members, @@ -1068,6 +1202,9 @@ async fn execute( } Command::ZScore { key, member } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::ZScore { key: key.clone(), member, @@ -1082,6 +1219,9 @@ async fn execute( } Command::ZRank { key, member } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::ZRank { key: key.clone(), member, @@ -1101,6 +1241,9 @@ async fn execute( stop, with_scores, } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::ZRange { key: key.clone(), start, @@ -1125,6 +1268,9 @@ async fn execute( } Command::ZCard { key } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::ZCard { key: key.clone() }; match engine.route(&key, req).await { Ok(ShardResponse::Len(n)) => Frame::Integer(n as i64), @@ -1136,6 +1282,9 @@ async fn execute( // --- hash commands --- Command::HSet { key, fields } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::HSet { key: key.clone(), fields, @@ -1150,6 +1299,9 @@ async fn execute( } Command::HGet { key, field } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::HGet { key: key.clone(), field, @@ -1164,6 +1316,9 @@ async fn execute( } Command::HGetAll { key } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::HGetAll { key: key.clone() }; match engine.route(&key, req).await { Ok(ShardResponse::HashFields(fields)) => { @@ -1181,6 +1336,9 @@ async fn execute( } Command::HDel { key, fields } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::HDel { key: key.clone(), fields, @@ -1194,6 +1352,9 @@ async fn execute( } Command::HExists { key, field } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::HExists { key: key.clone(), field, @@ -1207,6 +1368,9 @@ async fn execute( } Command::HLen { key } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::HLen { key: key.clone() }; match engine.route(&key, req).await { Ok(ShardResponse::Len(n)) => Frame::Integer(n as i64), @@ -1217,6 +1381,9 @@ async fn execute( } Command::HIncrBy { key, field, delta } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::HIncrBy { key: key.clone(), field, @@ -1233,6 +1400,9 @@ async fn execute( } Command::HKeys { key } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::HKeys { key: key.clone() }; match engine.route(&key, req).await { Ok(ShardResponse::StringArray(keys)) => Frame::Array( @@ -1247,6 +1417,9 @@ async fn execute( } Command::HVals { key } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::HVals { key: key.clone() }; match engine.route(&key, req).await { Ok(ShardResponse::Array(vals)) => { @@ -1259,6 +1432,9 @@ async fn execute( } Command::HMGet { key, fields } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::HMGet { key: key.clone(), fields, @@ -1280,6 +1456,9 @@ async fn execute( // --- set commands --- Command::SAdd { key, members } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::SAdd { key: key.clone(), members, @@ -1294,6 +1473,9 @@ async fn execute( } Command::SRem { key, members } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::SRem { key: key.clone(), members, @@ -1307,6 +1489,9 @@ async fn execute( } Command::SMembers { key } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::SMembers { key: key.clone() }; match engine.route(&key, req).await { Ok(ShardResponse::StringArray(members)) => Frame::Array( @@ -1322,6 +1507,9 @@ async fn execute( } Command::SIsMember { key, member } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::SIsMember { key: key.clone(), member, @@ -1335,6 +1523,9 @@ async fn execute( } Command::SCard { key } => { + if let Some(redirect) = check_cluster_slot(ctx, &key).await { + return redirect; + } let req = ShardRequest::SCard { key: key.clone() }; match engine.route(&key, req).await { Ok(ShardResponse::Len(n)) => Frame::Integer(n as i64), @@ -1345,56 +1536,66 @@ async fn execute( } // --- cluster commands --- - // Note: Full cluster support requires integration with ember-cluster crate. - // For now, CLUSTER KEYSLOT works, and other commands return stub responses. Command::ClusterKeySlot { key } => { let slot = ember_cluster::key_slot(key.as_bytes()); Frame::Integer(slot as i64) } - Command::ClusterInfo => { - // Return minimal info indicating cluster mode is disabled - let info = "cluster_enabled:0\r\n"; - Frame::Bulk(Bytes::from(info)) - } + Command::ClusterInfo => match &ctx.cluster { + Some(c) => c.cluster_info().await, + None => Frame::Bulk(Bytes::from("cluster_enabled:0\r\n")), + }, - Command::ClusterNodes => { - // In non-cluster mode, return empty string - Frame::Bulk(Bytes::from("")) - } + Command::ClusterNodes => match &ctx.cluster { + Some(c) => c.cluster_nodes().await, + None => Frame::Bulk(Bytes::from("")), + }, - Command::ClusterSlots => { - // In non-cluster mode, return empty array - Frame::Array(vec![]) - } + Command::ClusterSlots => match &ctx.cluster { + Some(c) => c.cluster_slots().await, + None => Frame::Array(vec![]), + }, - Command::ClusterMyId => { - // In non-cluster mode, return an error - Frame::Error("ERR This instance has cluster support disabled".into()) - } + Command::ClusterMyId => match &ctx.cluster { + Some(c) => c.cluster_myid(), + None => Frame::Error("ERR This instance has cluster support disabled".into()), + }, - Command::Asking => { - // ASKING is a no-op in non-cluster mode, just return OK - Frame::Simple("OK".into()) - } + Command::Asking => Frame::Simple("OK".into()), + + Command::ClusterMeet { ip, port } => match &ctx.cluster { + Some(c) => c.cluster_meet(&ip, port).await, + None => Frame::Error("ERR This instance has cluster support disabled".into()), + }, + + Command::ClusterAddSlots { slots } => match &ctx.cluster { + Some(c) => c.cluster_addslots(&slots).await, + None => Frame::Error("ERR This instance has cluster support disabled".into()), + }, + + Command::ClusterDelSlots { slots } => match &ctx.cluster { + Some(c) => c.cluster_delslots(&slots).await, + None => Frame::Error("ERR This instance has cluster support disabled".into()), + }, + + Command::ClusterForget { node_id } => match &ctx.cluster { + Some(c) => c.cluster_forget(&node_id).await, + None => Frame::Error("ERR This instance has cluster support disabled".into()), + }, Command::ClusterSetSlotImporting { .. } | Command::ClusterSetSlotMigrating { .. } | Command::ClusterSetSlotNode { .. } | 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()) + Frame::Error("ERR not yet implemented".into()) } Command::Migrate { .. } => { - Frame::Error("ERR This instance has cluster support disabled".into()) + Frame::Error("ERR not yet implemented".into()) } // -- slow log commands -- diff --git a/crates/ember-server/src/main.rs b/crates/ember-server/src/main.rs index 31b7822a..10351fd2 100644 --- a/crates/ember-server/src/main.rs +++ b/crates/ember-server/src/main.rs @@ -4,6 +4,7 @@ #[global_allocator] static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; +mod cluster; mod concurrent_handler; mod config; mod connection; @@ -16,11 +17,14 @@ mod tls; use std::net::SocketAddr; use std::path::PathBuf; +use std::sync::Arc; use clap::Parser; +use ember_cluster::{GossipConfig, NodeId}; use ember_core::ShardPersistenceConfig; use tracing::info; +use crate::cluster::ClusterCoordinator; use crate::config::{ build_engine_config, parse_byte_size, parse_eviction_policy, parse_fsync_policy, }; @@ -110,6 +114,23 @@ struct Args { /// accepts: yes, no. default: no #[arg(long, default_value = "no", env = "EMBER_TLS_AUTH_CLIENTS")] tls_auth_clients: String, + + // -- cluster options -- + /// enable cluster mode with gossip-based discovery and slot routing + #[arg(long, env = "EMBER_CLUSTER_ENABLED")] + cluster_enabled: bool, + + /// bootstrap a new cluster as a single node owning all 16384 slots + #[arg(long, env = "EMBER_CLUSTER_BOOTSTRAP")] + cluster_bootstrap: bool, + + /// port offset for the cluster gossip bus (data_port + offset) + #[arg(long, default_value_t = 10000, env = "EMBER_CLUSTER_PORT_OFFSET")] + cluster_port_offset: u16, + + /// node timeout in milliseconds for failure detection + #[arg(long, default_value_t = 5000, env = "EMBER_CLUSTER_NODE_TIMEOUT")] + cluster_node_timeout: u64, } #[tokio::main] @@ -288,6 +309,44 @@ async fn main() { None }; + // validate cluster mode + if args.cluster_enabled && args.concurrent { + eprintln!("error: --cluster-enabled and --concurrent are mutually exclusive"); + std::process::exit(1); + } + + if args.cluster_bootstrap && !args.cluster_enabled { + eprintln!("error: --cluster-bootstrap requires --cluster-enabled"); + std::process::exit(1); + } + + // build cluster coordinator if cluster mode is enabled + let cluster: Option> = if args.cluster_enabled { + let local_id = NodeId::new(); + let gossip_config = GossipConfig { + gossip_port_offset: args.cluster_port_offset, + probe_timeout: std::time::Duration::from_millis(args.cluster_node_timeout / 2), + ..GossipConfig::default() + }; + + let (coordinator, event_rx) = + ClusterCoordinator::new(local_id, addr, gossip_config, args.cluster_bootstrap); + let coordinator = Arc::new(coordinator); + + // spawn gossip networking tasks + coordinator.spawn_gossip(addr, event_rx).await; + + if args.cluster_bootstrap { + info!("cluster mode: bootstrapped with all 16384 slots"); + } else { + info!("cluster mode: waiting for CLUSTER MEET"); + } + + Some(coordinator) + } else { + None + }; + let result = if args.concurrent { server::run_concurrent( addr, @@ -312,6 +371,7 @@ async fn main() { slowlog_config, args.requirepass, tls_config, + cluster, ) .await }; diff --git a/crates/ember-server/src/server.rs b/crates/ember-server/src/server.rs index 0ae758c4..5d559c56 100644 --- a/crates/ember-server/src/server.rs +++ b/crates/ember-server/src/server.rs @@ -15,6 +15,7 @@ use tokio::sync::Semaphore; use tokio_rustls::TlsAcceptor; use tracing::{error, info, warn}; +use crate::cluster::ClusterCoordinator; use crate::connection; use crate::pubsub::PubSubManager; use crate::slowlog::{SlowLog, SlowLogConfig}; @@ -43,6 +44,8 @@ pub struct ServerContext { pub requirepass: Option, /// The address the server is bound to (for protected mode checks). pub bind_addr: SocketAddr, + /// Cluster coordinator, present when --cluster-enabled is set. + pub cluster: Option>, } /// Binds to `addr` and runs the accept loop. @@ -67,6 +70,7 @@ pub async fn run( slowlog_config: SlowLogConfig, requirepass: Option, tls: Option<(SocketAddr, TlsConfig)>, + cluster: Option>, ) -> Result<(), Box> { // ensure data directory exists if persistence is configured if let Some(ref pcfg) = config.persistence { @@ -117,6 +121,7 @@ pub async fn run( commands_processed: AtomicU64::new(0), requirepass, bind_addr: addr, + cluster, }); let slow_log = Arc::new(SlowLog::new(slowlog_config)); @@ -345,6 +350,7 @@ pub async fn run_concurrent( commands_processed: AtomicU64::new(0), requirepass, bind_addr: addr, + cluster: None, }); let slow_log = Arc::new(SlowLog::new(slowlog_config));