From 4e6664e8b390d7ab24bf48ada4628ad0b9cc09ca Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Wed, 18 Feb 2026 21:14:53 -0500 Subject: [PATCH] refactor(cluster): polish ember-cluster crate - message.rs: add ADDR_IPV4/ADDR_IPV6 constants, extract encode_slice helper to de-dup the truncate+write-count+iterate pattern across five encoding sites, add byte-size comments to the decode path, and replace bare .unwrap() in tests with .expect("roundtrip should succeed") - gossip.rs: collapse check_probe_timeouts to a single retain pass with two collection vecs; fix double-clone in SlotsChanged arm and Welcome branch to clone once and reuse the owned value - election.rs: collapse nested if/else in record_vote to a single assign-then-return - topology.rs: lift ping_ms/pong_ms out of the format! call in to_cluster_nodes_line; clarify promote_replica mutation comments - lib.rs: use rust,no_run instead of rust,ignore in crate-level example --- crates/ember-cluster/src/election.rs | 8 +-- crates/ember-cluster/src/gossip.rs | 46 ++++++------ crates/ember-cluster/src/lib.rs | 2 +- crates/ember-cluster/src/message.rs | 101 +++++++++++++++------------ crates/ember-cluster/src/topology.rs | 25 ++++--- 5 files changed, 96 insertions(+), 86 deletions(-) diff --git a/crates/ember-cluster/src/election.rs b/crates/ember-cluster/src/election.rs index 3fe6598c..a6d9aa84 100644 --- a/crates/ember-cluster/src/election.rs +++ b/crates/ember-cluster/src/election.rs @@ -54,12 +54,8 @@ impl Election { return false; } self.votes.insert(from); - if self.votes.len() >= Self::quorum(total_primaries) { - self.promoted = true; - true - } else { - false - } + self.promoted = self.votes.len() >= Self::quorum(total_primaries); + self.promoted } /// Minimum votes required for a majority. diff --git a/crates/ember-cluster/src/gossip.rs b/crates/ember-cluster/src/gossip.rs index ffcd94fe..a9d7a5f8 100644 --- a/crates/ember-cluster/src/gossip.rs +++ b/crates/ember-cluster/src/gossip.rs @@ -478,10 +478,10 @@ impl GossipEngine { if member.id == self.local_id { continue; } - let slots = member.slots.clone(); if let std::collections::hash_map::Entry::Vacant(e) = self.members.entry(member.id) { + let slots = member.slots.clone(); e.insert(MemberState { id: member.id, addr: member.addr, @@ -759,9 +759,11 @@ impl GossipEngine { if let Some(member) = self.members.get_mut(node) { // only accept if incarnation is at least as recent if *incarnation >= member.incarnation { - member.slots = slots.clone(); - self.emit(GossipEvent::SlotsChanged(*node, slots.clone())) - .await; + let owned = slots.clone(); + member.slots = owned.clone(); + // release the mutable borrow before the async call + drop(member); + self.emit(GossipEvent::SlotsChanged(*node, owned)).await; } } } @@ -849,17 +851,23 @@ impl GossipEngine { let now = Instant::now(); let mut outgoing = Vec::new(); - // Phase 2: indirect probe timeouts → mark Suspect - let indirect_timed_out: Vec<_> = self - .pending_probes - .iter() - .filter(|(_, probe)| probe.indirect && now.duration_since(probe.sent_at) > timeout) - .map(|(seq, probe)| (*seq, probe.target)) - .collect(); - - for (seq, target) in indirect_timed_out { - self.pending_probes.remove(&seq); + // Single pass: collect timed-out entries by kind and remove them. + let mut timed_out_indirect: Vec<(u64, NodeId)> = Vec::new(); + let mut timed_out_direct: Vec<(u64, NodeId)> = Vec::new(); + self.pending_probes.retain(|seq, probe| { + if now.duration_since(probe.sent_at) <= timeout { + return true; + } + if probe.indirect { + timed_out_indirect.push((*seq, probe.target)); + } else { + timed_out_direct.push((*seq, probe.target)); + } + false + }); + // Phase 2: indirect probe timeouts → mark Suspect + for (_seq, target) in timed_out_indirect { let incarnation = self .members .get(&target) @@ -880,15 +888,7 @@ impl GossipEngine { } // Phase 1: direct ping timeouts → send PingReq - let direct_timed_out: Vec<_> = self - .pending_probes - .iter() - .filter(|(_, probe)| !probe.indirect && now.duration_since(probe.sent_at) > timeout) - .map(|(seq, probe)| (*seq, probe.target)) - .collect(); - - for (seq, target) in direct_timed_out { - self.pending_probes.remove(&seq); + for (_seq, target) in timed_out_direct { let target_addr = match self.members.get(&target) { Some(m) if m.state == MemberStatus::Alive => m.addr, diff --git a/crates/ember-cluster/src/lib.rs b/crates/ember-cluster/src/lib.rs index 87234508..8e2f9bcb 100644 --- a/crates/ember-cluster/src/lib.rs +++ b/crates/ember-cluster/src/lib.rs @@ -16,7 +16,7 @@ //! //! # Quick Start //! -//! ```rust,ignore +//! ```rust,no_run //! use ember_cluster::{ClusterState, ClusterNode, NodeId, key_slot}; //! //! // Create a single-node cluster diff --git a/crates/ember-cluster/src/message.rs b/crates/ember-cluster/src/message.rs index e9bac705..843f0513 100644 --- a/crates/ember-cluster/src/message.rs +++ b/crates/ember-cluster/src/message.rs @@ -14,6 +14,11 @@ use crate::{NodeId, SlotRange}; /// Prevents allocation bombs from crafted messages. const MAX_COLLECTION_COUNT: usize = 1024; +/// Address family discriminant for IPv4 (standard AF_INET byte count: 4). +const ADDR_IPV4: u8 = 4; +/// Address family discriminant for IPv6 (standard AF_INET6 byte count: 6 hex groups, 16 bytes). +const ADDR_IPV6: u8 = 6; + // Safe read helpers that return io::Error instead of panicking on truncated input. fn safe_get_u8(buf: &mut &[u8]) -> io::Result { @@ -163,6 +168,22 @@ const UPDATE_ROLE_CHANGED: u8 = 6; const UPDATE_VOTE_REQUEST: u8 = 7; const UPDATE_VOTE_GRANTED: u8 = 8; +/// Writes a count-capped slice into `buf` using a per-item encoder. +/// +/// Writes `min(items.len(), MAX_COLLECTION_COUNT)` as a u16, then calls +/// `encode_one` for each item in the truncated slice. Used everywhere a +/// repeated list is encoded to avoid duplicating the truncation logic. +fn encode_slice(buf: &mut BytesMut, items: &[T], encode_one: F) +where + F: Fn(&mut BytesMut, &T), +{ + let count = items.len().min(MAX_COLLECTION_COUNT); + buf.put_u16_le(count as u16); + for item in &items[..count] { + encode_one(buf, item); + } +} + impl GossipMessage { /// Serializes the message to bytes. pub fn encode(&self) -> Bytes { @@ -217,11 +238,7 @@ impl GossipMessage { GossipMessage::Welcome { sender, members } => { buf.put_u8(MSG_WELCOME); encode_node_id(buf, sender); - let count = members.len().min(MAX_COLLECTION_COUNT); - buf.put_u16_le(count as u16); - for member in &members[..count] { - encode_member_info(buf, member); - } + encode_slice(buf, members, encode_member_info); } GossipMessage::SlotsAnnounce { sender, @@ -231,12 +248,10 @@ impl GossipMessage { buf.put_u8(MSG_SLOTS_ANNOUNCE); encode_node_id(buf, sender); buf.put_u64_le(*incarnation); - let count = slots.len().min(MAX_COLLECTION_COUNT); - buf.put_u16_le(count as u16); - for slot in &slots[..count] { - buf.put_u16_le(slot.start); - buf.put_u16_le(slot.end); - } + encode_slice(buf, slots, |b, slot| { + b.put_u16_le(slot.start); + b.put_u16_le(slot.end); + }); } } } @@ -356,12 +371,12 @@ fn decode_node_id(buf: &mut &[u8]) -> io::Result { fn encode_socket_addr(buf: &mut BytesMut, addr: &SocketAddr) { match addr { SocketAddr::V4(v4) => { - buf.put_u8(4); + buf.put_u8(ADDR_IPV4); buf.put_slice(&v4.ip().octets()); buf.put_u16_le(v4.port()); } SocketAddr::V6(v6) => { - buf.put_u8(6); + buf.put_u8(ADDR_IPV6); buf.put_slice(&v6.ip().octets()); buf.put_u16_le(v6.port()); } @@ -377,7 +392,8 @@ fn decode_socket_addr(buf: &mut &[u8]) -> io::Result { } let addr_type = buf.get_u8(); match addr_type { - 4 => { + ADDR_IPV4 => { + // 4 octets + 2-byte port = 6 bytes if buf.len() < 6 { return Err(io::Error::new( io::ErrorKind::UnexpectedEof, @@ -389,7 +405,8 @@ fn decode_socket_addr(buf: &mut &[u8]) -> io::Result { let port = buf.get_u16_le(); Ok(SocketAddr::from((octets, port))) } - 6 => { + ADDR_IPV6 => { + // 16 octets + 2-byte port = 18 bytes if buf.len() < 18 { return Err(io::Error::new( io::ErrorKind::UnexpectedEof, @@ -409,11 +426,7 @@ fn decode_socket_addr(buf: &mut &[u8]) -> io::Result { } fn encode_updates(buf: &mut BytesMut, updates: &[NodeUpdate]) { - let count = updates.len().min(MAX_COLLECTION_COUNT); - buf.put_u16_le(count as u16); - for update in &updates[..count] { - encode_update(buf, update); - } + encode_slice(buf, updates, encode_update); } fn encode_update(buf: &mut BytesMut, update: &NodeUpdate) { @@ -450,12 +463,10 @@ fn encode_update(buf: &mut BytesMut, update: &NodeUpdate) { buf.put_u8(UPDATE_SLOTS_CHANGED); encode_node_id(buf, node); buf.put_u64_le(*incarnation); - let count = slots.len().min(MAX_COLLECTION_COUNT); - buf.put_u16_le(count as u16); - for slot in &slots[..count] { - buf.put_u16_le(slot.start); - buf.put_u16_le(slot.end); - } + encode_slice(buf, slots, |b, slot| { + b.put_u16_le(slot.start); + b.put_u16_le(slot.end); + }); } NodeUpdate::RoleChanged { node, @@ -613,12 +624,10 @@ fn encode_member_info(buf: &mut BytesMut, member: &MemberInfo) { encode_socket_addr(buf, &member.addr); buf.put_u64_le(member.incarnation); buf.put_u8(if member.is_primary { 1 } else { 0 }); - let slot_count = member.slots.len().min(MAX_COLLECTION_COUNT); - buf.put_u16_le(slot_count as u16); - for slot in &member.slots[..slot_count] { - buf.put_u16_le(slot.start); - buf.put_u16_le(slot.end); - } + encode_slice(buf, &member.slots, |b, slot| { + b.put_u16_le(slot.start); + b.put_u16_le(slot.end); + }); } fn decode_member_info(buf: &mut &[u8]) -> io::Result { @@ -669,7 +678,7 @@ mod tests { updates: vec![], }; let encoded = msg.encode(); - let decoded = GossipMessage::decode(&encoded).unwrap(); + let decoded = GossipMessage::decode(&encoded).expect("roundtrip should succeed"); assert_eq!(msg, decoded); } @@ -693,7 +702,7 @@ mod tests { ], }; let encoded = msg.encode(); - let decoded = GossipMessage::decode(&encoded).unwrap(); + let decoded = GossipMessage::decode(&encoded).expect("roundtrip should succeed"); assert_eq!(msg, decoded); } @@ -706,7 +715,7 @@ mod tests { target_addr: test_addr(), }; let encoded = msg.encode(); - let decoded = GossipMessage::decode(&encoded).unwrap(); + let decoded = GossipMessage::decode(&encoded).expect("roundtrip should succeed"); assert_eq!(msg, decoded); } @@ -721,7 +730,7 @@ mod tests { }], }; let encoded = msg.encode(); - let decoded = GossipMessage::decode(&encoded).unwrap(); + let decoded = GossipMessage::decode(&encoded).expect("roundtrip should succeed"); assert_eq!(msg, decoded); } @@ -732,7 +741,7 @@ mod tests { sender_addr: test_addr(), }; let encoded = msg.encode(); - let decoded = GossipMessage::decode(&encoded).unwrap(); + let decoded = GossipMessage::decode(&encoded).expect("roundtrip should succeed"); assert_eq!(msg, decoded); } @@ -758,7 +767,7 @@ mod tests { ], }; let encoded = msg.encode(); - let decoded = GossipMessage::decode(&encoded).unwrap(); + let decoded = GossipMessage::decode(&encoded).expect("roundtrip should succeed"); assert_eq!(msg, decoded); } @@ -769,7 +778,7 @@ mod tests { sender_addr: test_addr_v6(), }; let encoded = msg.encode(); - let decoded = GossipMessage::decode(&encoded).unwrap(); + let decoded = GossipMessage::decode(&encoded).expect("roundtrip should succeed"); assert_eq!(msg, decoded); } @@ -790,7 +799,7 @@ mod tests { }], }; let encoded = msg.encode(); - let decoded = GossipMessage::decode(&encoded).unwrap(); + let decoded = GossipMessage::decode(&encoded).expect("roundtrip should succeed"); assert_eq!(msg, decoded); // primary variant (no replicates field) @@ -805,7 +814,7 @@ mod tests { }], }; let encoded2 = msg2.encode(); - let decoded2 = GossipMessage::decode(&encoded2).unwrap(); + let decoded2 = GossipMessage::decode(&encoded2).expect("roundtrip should succeed"); assert_eq!(msg2, decoded2); } @@ -839,7 +848,7 @@ mod tests { updates, }; let encoded = msg.encode(); - let decoded = GossipMessage::decode(&encoded).unwrap(); + let decoded = GossipMessage::decode(&encoded).expect("roundtrip should succeed"); assert_eq!(msg, decoded); } @@ -856,7 +865,7 @@ mod tests { }], }; let encoded = msg.encode(); - let decoded = GossipMessage::decode(&encoded).unwrap(); + let decoded = GossipMessage::decode(&encoded).expect("roundtrip should succeed"); assert_eq!(msg, decoded); } @@ -877,7 +886,7 @@ mod tests { }], }; let encoded = msg.encode(); - let decoded = GossipMessage::decode(&encoded).unwrap(); + let decoded = GossipMessage::decode(&encoded).expect("roundtrip should succeed"); assert_eq!(msg, decoded); } @@ -925,7 +934,7 @@ mod tests { }], }; let encoded = msg.encode(); - let decoded = GossipMessage::decode(&encoded).unwrap(); + let decoded = GossipMessage::decode(&encoded).expect("roundtrip should succeed"); assert_eq!(msg, decoded); } @@ -943,7 +952,7 @@ mod tests { }], }; let encoded = msg.encode(); - let decoded = GossipMessage::decode(&encoded).unwrap(); + let decoded = GossipMessage::decode(&encoded).expect("roundtrip should succeed"); assert_eq!(msg, decoded); } diff --git a/crates/ember-cluster/src/topology.rs b/crates/ember-cluster/src/topology.rs index 6e540667..0897a494 100644 --- a/crates/ember-cluster/src/topology.rs +++ b/crates/ember-cluster/src/topology.rs @@ -250,6 +250,15 @@ impl ClusterNode { .map(|id| id.0.to_string()) .unwrap_or_else(|| "-".to_string()); + let ping_ms = self + .last_ping_sent + .map(|t| t.elapsed().as_millis() as u64) + .unwrap_or(0); + let pong_ms = self + .last_pong_received + .map(|t| t.elapsed().as_millis() as u64) + .unwrap_or(0); + // Format: @ format!( "{} {}@{} {} {} {} {} {} connected {}", @@ -258,12 +267,8 @@ impl ClusterNode { self.cluster_bus_addr.port(), self.format_flags(), replicates_str, - self.last_ping_sent - .map(|t| t.elapsed().as_millis() as u64) - .unwrap_or(0), - self.last_pong_received - .map(|t| t.elapsed().as_millis() as u64) - .unwrap_or(0), + ping_ms, + pong_ms, self.config_epoch, slots_str ) @@ -470,7 +475,7 @@ impl ClusterState { .ok_or_else(|| format!("replica {replica_id} has no primary configured"))? }; - // transfer every slot currently owned by the old primary + // transfer every slot the old primary owned to the promoted replica for slot in 0..SLOT_COUNT { if self.slot_map.owner(slot) == Some(primary_id) { self.slot_map.assign(slot, replica_id); @@ -478,11 +483,11 @@ impl ClusterState { } let new_primary_slots = self.slot_map.slots_for_node(replica_id); - // bump epoch before touching node state + // bump epoch before touching node state so both nodes land on the same epoch self.config_epoch += 1; let new_epoch = self.config_epoch; - // demote old primary → it now replicates the new primary + // demote the old primary: clear its slots and make it a replica of the winner if let Some(old_primary) = self.nodes.get_mut(&primary_id) { old_primary.role = NodeRole::Replica; old_primary.replicates = Some(replica_id); @@ -491,7 +496,7 @@ impl ClusterState { old_primary.config_epoch = new_epoch; } - // promote the replica → it becomes the new primary + // promote the winner: take ownership of all slots and clear the replicates pointer if let Some(new_primary) = self.nodes.get_mut(&replica_id) { new_primary.role = NodeRole::Primary; new_primary.replicates = None;