From bcb1853ca75e42b076763a0c6b10a8c9b327b6a4 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Fri, 17 Jul 2026 17:35:01 +0400 Subject: [PATCH 1/9] outline --- src/control_plane/consensus/multi_raft.rs | 3 +- src/control_plane/consensus/raft/mod.rs | 1 + .../consensus/raft/states/mod.rs | 3 + src/control_plane/consensus/seal_recovery.rs | 2 +- src/control_plane/metadata/mod.rs | 1 - src/control_plane/metadata/state_machine.rs | 1560 ----------------- 6 files changed, 7 insertions(+), 1563 deletions(-) create mode 100644 src/control_plane/consensus/raft/states/mod.rs delete mode 100644 src/control_plane/metadata/state_machine.rs diff --git a/src/control_plane/consensus/multi_raft.rs b/src/control_plane/consensus/multi_raft.rs index c4676325..b837697f 100644 --- a/src/control_plane/consensus/multi_raft.rs +++ b/src/control_plane/consensus/multi_raft.rs @@ -4,7 +4,8 @@ use crate::control_plane::consensus::messages::{ MultiRaftActorCommand, ProposeSegmentRoll, RaftEvent, RaftProtocolMessage, RaftTimeoutCallback, }; use crate::control_plane::consensus::raft::errors::ProposalError; -use crate::control_plane::consensus::raft::state::{LeaderlessSegments, Raft, TimerSeqs}; +use crate::control_plane::consensus::raft::state::{Raft, TimerSeqs}; +use crate::control_plane::consensus::raft::states::transient_state::LeaderlessSegments; use crate::control_plane::consensus::raft::storage::RaftStorage; use crate::control_plane::consensus::raft::{compute_replacement_replica_set, now_ms}; use crate::control_plane::consensus::seal_recovery::{SealEndRecovery, SealEndStep}; diff --git a/src/control_plane/consensus/raft/mod.rs b/src/control_plane/consensus/raft/mod.rs index 81677f48..43053821 100644 --- a/src/control_plane/consensus/raft/mod.rs +++ b/src/control_plane/consensus/raft/mod.rs @@ -5,6 +5,7 @@ pub(crate) mod command; pub(crate) mod errors; pub(crate) mod log; pub(crate) mod state; +pub(crate) mod states; pub(crate) mod storage; pub(crate) fn now_ms() -> u64 { diff --git a/src/control_plane/consensus/raft/states/mod.rs b/src/control_plane/consensus/raft/states/mod.rs new file mode 100644 index 00000000..e8f1c366 --- /dev/null +++ b/src/control_plane/consensus/raft/states/mod.rs @@ -0,0 +1,3 @@ +pub(crate) mod log_state; +pub(crate) mod metadata_state; +pub(crate) mod transient_state; diff --git a/src/control_plane/consensus/seal_recovery.rs b/src/control_plane/consensus/seal_recovery.rs index d6009e33..0728629d 100644 --- a/src/control_plane/consensus/seal_recovery.rs +++ b/src/control_plane/consensus/seal_recovery.rs @@ -11,7 +11,7 @@ use std::collections::{HashMap, HashSet}; use crate::control_plane::NodeId; -use crate::control_plane::consensus::raft::state::LeaderlessSegments; +use crate::control_plane::consensus::raft::states::transient_state::LeaderlessSegments; use crate::control_plane::membership::ShardGroupId; use crate::control_plane::metadata::EntryId; use crate::data_plane::SegmentKey; diff --git a/src/control_plane/metadata/mod.rs b/src/control_plane/metadata/mod.rs index c5effe84..6255f737 100644 --- a/src/control_plane/metadata/mod.rs +++ b/src/control_plane/metadata/mod.rs @@ -6,7 +6,6 @@ pub mod error; pub(crate) mod event; pub(crate) mod range; -pub(crate) mod state_machine; pub mod strategy; pub(crate) mod topic; diff --git a/src/control_plane/metadata/state_machine.rs b/src/control_plane/metadata/state_machine.rs deleted file mode 100644 index b261f345..00000000 --- a/src/control_plane/metadata/state_machine.rs +++ /dev/null @@ -1,1560 +0,0 @@ -use super::command::*; -use super::event::*; - -use super::topic::{TopicMeta, TopicState, TopicStats}; -use crate::control_plane::NodeId; -use crate::control_plane::Replicas; -use crate::control_plane::membership::ShardGroupId; -use crate::control_plane::metadata::ConsumerGroupAssignment; -use crate::control_plane::metadata::{EntryId, RangeId, SegmentId, TopicId, error::MetadataError}; -use crate::data_plane::SegmentKey; -#[cfg(any(test, debug_assertions))] -use crate::test_traits::TAssertInvariant; -use MetadataError::*; -use std::collections::HashMap; -use uuid::Uuid; - -pub struct MetadataStateMachine { - pub(crate) topics: HashMap, - topic_name_index: HashMap, - next_topic_id: u64, - pending_proposals: Vec, -} - -impl MetadataStateMachine { - pub(crate) fn new(shard_group_id: ShardGroupId) -> Self { - MetadataStateMachine { - topics: HashMap::new(), - topic_name_index: HashMap::new(), - next_topic_id: shard_group_id.0 << 32, - pending_proposals: Vec::new(), - } - } - - pub(crate) fn get_topic(&self, id: &TopicId) -> Option<&TopicMeta> { - self.topics.get(id) - } - - pub(crate) fn get_topic_by_name(&self, name: &str) -> Option<&TopicMeta> { - self.topic_name_index - .get(name) - .and_then(|id| self.topics.get(id)) - } - - pub(crate) fn get_consumer_group_assignment( - &self, - topic_name: &str, - group_id: &str, - member_id: Uuid, - ) -> Option { - Some( - self.get_topic_by_name(topic_name)? - .consumer_groups - .get(group_id)? - .assignment_for(member_id), - ) - } - - pub(crate) fn topic_names(&self) -> Box<[String]> { - self.topic_name_index.keys().cloned().collect() - } - - #[cfg(test)] - pub fn topic_count(&self) -> usize { - self.topics.len() - } - - pub(crate) fn topic_stats(&self) -> Box<[TopicStats]> { - self.topics.values().map(|t| t.stats()).collect() - } - - pub(crate) fn take_pending_proposals(&mut self) -> Box<[MetadataCommand]> { - std::mem::take(&mut self.pending_proposals).into_boxed_slice() - } - - pub(crate) fn active_segments_for_node( - &self, - node_id: &NodeId, - ) -> Box<[(SegmentKey, Replicas)]> { - self.topics - .values() - .flat_map(|t| t.active_segments_for_node(node_id)) - .collect() - } - - /// Every active segment across all topics with its replica set and start - /// offset, for the leader's periodic assignment re-drive. - pub(crate) fn active_segment_assignments(&self) -> Box<[(SegmentKey, Replicas, EntryId)]> { - self.topics - .values() - .flat_map(|t| t.active_segment_assignments()) - .collect() - } - - pub(crate) fn apply(&mut self, command: MetadataCommand) -> Result { - use MetadataCommand::*; - let result = match command { - CreateTopic(cmd) => self.create_topic(cmd)?.into(), - RollSegment(cmd) => self.roll_segment(cmd)?, - SplitRange(cmd) => self.split_range(cmd)?.into(), - MergeRange(cmd) => self.merge_range(cmd)?.into(), - DeleteTopic(cmd) => self.delete_topic(cmd)?.into(), - ReassignSegment(cmd) => self.reassign_segment(cmd)?, - DeleteSegments(cmd) => self.delete_segments(cmd)?, - SyncConsumerGroup(cmd) => self - .sync_consumer_group(cmd)? - .map(ApplyResult::ConsumerGroupChanged) - .unwrap_or(ApplyResult::Noop), - }; - #[cfg(any(test, debug_assertions))] - self.assert_invariants(); - Ok(result) - } - - fn create_topic(&mut self, cmd: CreateTopic) -> Result { - if self.topic_name_index.contains_key(&cmd.name) { - return Err(TopicNameAlreadyExists(cmd.name)); - } - let topic_id = TopicId(self.next_topic_id); - let replica_set = cmd.replica_set.clone(); - let topic = TopicMeta::new( - cmd.name, - topic_id, - cmd.replica_set, - cmd.created_at, - cmd.storage_policy, - ); - self.topic_name_index.insert(topic.name.clone(), topic_id); - self.topics.insert(topic.id, topic); - self.next_topic_id += 1; - Ok(TopicCreated { - segment_key: SegmentKey::new(topic_id, RangeId(0), SegmentId(0)), - replica_set, - }) - } - - fn roll_segment(&mut self, cmd: RollSegment) -> Result { - let topic = self.get_active_topic_mut(cmd.segment_key.topic_id)?; - let can_split = topic.can_split(); - let range = topic.get_range_mut(&cmd.segment_key.range_id)?; - - let is_active = range.active_segment == Some(cmd.segment_key.segment_id); - - // If Inactive, correction path - if !is_active { - let Some(end_entry_id) = cmd.end_entry_id else { - return Ok(ApplyResult::Noop); - }; - let Some(replica_set) = - range.correct_end_offset(cmd.segment_key.segment_id, end_entry_id) - else { - return Ok(ApplyResult::Noop); - }; - - return Ok(ApplyResult::SegmentSealCorrected(SegmentSealCorrected { - segment_key: cmd.segment_key, - replica_set, - committed_entry_id: cmd.end_entry_id, - })); - } - - // If Active, Roll - let new_segment_id = range.roll_segment(cmd.clone())?; - let split_proposal = (range.should_split(cmd.sealed_at) && can_split) - .then(|| range.build_split_proposal(&cmd)); - - // For segment roll, unless data nodes got changed, consumer - // TODO For now, it loops over EVERY consumger groups and take epoch snapshot for every topic. - // TODO To reduce the load, it should specifically target groups that are affected by the possible range split - let consumer_group_epochs = topic - .consumer_groups - .keys() - .filter_map(|group_id| topic.consumer_group_epoch(group_id)) - .collect(); - tracing::debug!("Consumer groups: {:?}", consumer_group_epochs); - - if let Some(proposal) = split_proposal { - match proposal { - Ok(proposal) => self - .pending_proposals - .push(MetadataCommand::SplitRange(proposal)), - Err(error) => tracing::debug!( - "Split proposal skipped for range {:?}: {:?}", - cmd.segment_key.range_id, - error - ), - } - } - Ok(SegmentRolled { - new_segment_key: cmd.segment_key.with_segment_id(new_segment_id), - new_replica_set: cmd.new_replica_set, - end_entry_id: cmd.end_entry_id, - consumer_group_epochs, - } - .into()) - } - - /// Re-points a sealed segment's replica set. - /// `Noop` when the set is unchanged, otherwise `SegmentReassigned`. - /// The segment must be sealed; an active/deleting/unknown one is rejected. - fn reassign_segment(&mut self, cmd: ReassignSegment) -> Result { - let segment = self - .topics - .get_mut(&cmd.segment_key.topic_id) - .ok_or(TopicNotFound(cmd.segment_key.topic_id))? - .get_mut(cmd.segment_key)?; - - // The dispatch announces the desired replica set; receivers reconcile, so - // we only carry the sealed bounds (the catch-up target) alongside it. - if segment.reassign(cmd.replica_set.clone())? { - Ok(SegmentReassigned { - segment_key: cmd.segment_key, - start_entry_id: segment.start_entry_id, - sealed_end: segment.end_entry_id, - new_replica_set: cmd.replica_set, - } - .into()) - } else { - Ok(ApplyResult::Noop) - } - } - - /// Retention: mark an oldest-first prefix of a range's sealed segments `Deleting`. - /// `Noop` when nothing transitions (all named ids already `Deleting`/absent), else - /// `SegmentsDeleted` carrying the deleted segments grouped by `replica_set` for - /// batched dispatch. Uses plain `get_mut` (not `validate_active`) so a command - /// applied after the topic is being deleted is a harmless no-op (already `Deleting`). - fn delete_segments(&mut self, cmd: DeleteSegments) -> Result { - let range = self - .topics - .get_mut(&cmd.topic_id) - .ok_or(TopicNotFound(cmd.topic_id))? - .get_range_mut(&cmd.range_id)?; - - let deleted_ids = range.delete_segments(&cmd.segment_ids); - if deleted_ids.is_empty() { - return Ok(ApplyResult::Noop); - } - // Group the deleted segments by replica_set here. - let mut groups: Vec<(Replicas, Vec)> = Vec::new(); - for sid in &deleted_ids { - let Some(seg) = range.segments.get(sid) else { - continue; - }; - let key = SegmentKey::new(cmd.topic_id, cmd.range_id, *sid); - match groups.iter_mut().find(|(rs, _)| rs == &seg.replica_set) { - Some((_, keys)) => keys.push(key), - None => groups.push((seg.replica_set.clone(), vec![key])), - } - } - Ok(SegmentsDeleted { groups }.into()) - } - - fn get_active_topic_mut(&mut self, id: TopicId) -> Result<&mut TopicMeta, MetadataError> { - let topic = self.topics.get_mut(&id).ok_or(TopicNotFound(id))?; - topic.validate_active()?; - Ok(topic) - } - - fn split_range(&mut self, cmd: SplitRange) -> Result { - let topic = self - .topics - .get_mut(&cmd.topic_id) - .ok_or(TopicNotFound(cmd.topic_id))?; - topic.validate_active()?; - if !topic.can_split() { - return Err(SplitNotAllowed(cmd.topic_id)); - } - - let parent_range = topic - .ranges - .get(&cmd.range_id) - .ok_or(MetadataError::RangeNotFound)?; - - let parent_active_segment = parent_range.active_segment.and_then(|seg_id| { - let seg = parent_range.segments.get(&seg_id)?; - Some(( - SegmentKey::new(cmd.topic_id, cmd.range_id, seg_id), - seg.replica_set.clone(), - )) - }); - - let (left_id, right_id) = topic.execute_split(cmd.clone())?; - - let consumer_group_epochs = topic.rebalance_consumer_groups(); - - Ok(RangeSplit { - topic_id: cmd.topic_id, - children: [ - (left_id, SegmentId(0), cmd.left_replica_set), - (right_id, SegmentId(0), cmd.right_replica_set), - ], - parent_active_segment, - consumer_group_epochs, - }) - } - - fn merge_range(&mut self, cmd: MergeRange) -> Result { - let topic_id = cmd.topic_id; - let replica_set = cmd.merged_replica_set.clone(); - let topic = self - .topics - .get_mut(&topic_id) - .ok_or(TopicNotFound(topic_id))?; - topic.validate_active()?; - let merged_id = topic.execute_merge(cmd)?; - let consumer_group_epochs = topic.rebalance_consumer_groups(); - - Ok(RangeMerged { - segment_key: SegmentKey::new(topic_id, merged_id, SegmentId(0)), - replica_set, - consumer_group_epochs, - }) - } - - fn sync_consumer_group( - &mut self, - cmd: SyncConsumerGroup, - ) -> Result, MetadataError> { - let group_id = cmd.group_id.clone(); - - let topic = self - .topic_name_index - .get(&cmd.topic_name) - .and_then(|topic_id| self.topics.get_mut(topic_id)) - .ok_or_else(|| TopicNameNotFound(cmd.topic_name.clone()))?; - topic.validate_active()?; - if !topic.sync_consumer_group(cmd) { - return Ok(None); - } - Ok(topic.consumer_group_epoch(&group_id)) - } - - // ! SAFETY: When the loop evaluates the [1, 2] pair and decides it is mergeable, - // ! it generates the event and immediately stops looking at the rest of that topic's ranges - // ! - // ! Caution on race condition. Take the following example: - // ! 1. the following method proposes merging (1, 2). - // ! 2. The MergeRange command goes into a queue (or a Raft log) to be processed. - // ! 3. Before the command is executed, range 2 receives a massive burst of traffic and splits into 2A and 2B. - // ! 4. The MergeRange(1, 2) command is finally executed by your merge function. - // ! By the time the command executes, range 2 might be split, already sealed, or completely deleted - // * This is already safe as the 'stale' proposal fails gracefully at apply time, following "stale proposals are safe" invariant - pub(crate) fn evaluate_merges(&self, now: u64) -> Vec { - self.topics - .values() - .filter_map(|topic| topic.find_mergeable_pair(now)) - .collect() - } - - fn delete_topic(&mut self, cmd: DeleteTopic) -> Result { - let topic_id = self - .topic_name_index - .get(&cmd.name) - .copied() - .ok_or(MetadataError::TopicNameNotFound(cmd.name.clone()))?; - - // Safety: topic_name_index and topics are always in sync — - // see invariant below. - let topic = self.topics.get_mut(&topic_id).unwrap(); - topic.delete(); - - let consumer_group_epochs = topic.rebalance_consumer_groups(); - - self.topic_name_index.remove(&cmd.name); - - Ok(TopicDeleted { - consumer_group_epochs, - }) - } -} - -#[cfg(any(test, debug_assertions))] -impl crate::test_traits::TAssertInvariant for MetadataStateMachine { - fn assert_invariants(&self) { - assert_eq!( - self.topic_name_index.len(), - self.topics - .values() - .filter(|t| t.state != TopicState::Deleted) - .count(), - "topic_name_index out of sync with non-deleted topics" - ); - for (name, id) in &self.topic_name_index { - let topic = self - .topics - .get(id) - .expect("name index points to missing topic"); - assert_eq!(&topic.name, name); - } - - for id in self.topics.keys() { - assert!(id.0 < self.next_topic_id, "topic ID >= next_topic_id"); - } - for topic in self.topics.values() { - topic.assert_invariants(); - } - } -} - -#[cfg(test)] -mod tests { - use super::super::constants::*; - use super::super::range::*; - use super::super::segment::*; - use super::*; - use crate::connections::protocol::ConsumerGroupSyncAction; - use crate::control_plane::membership::ShardGroupId; - use crate::control_plane::{ - NodeId, - metadata::{ - SegmentId, - strategy::{PartitionStrategy, StoragePolicy}, - }, - }; - use std::collections::VecDeque; - - fn default_policy() -> StoragePolicy { - StoragePolicy { - retention_ms: Some(3_600_000), - replication_factor: 3, - partition_strategy: PartitionStrategy::AutoSplit, - } - } - - fn fixed_policy() -> StoragePolicy { - StoragePolicy { - retention_ms: Some(3_600_000), - replication_factor: 3, - partition_strategy: PartitionStrategy::Fixed, - } - } - - fn replica_set() -> Replicas { - Replicas::new(vec![ - NodeId::new("node-1"), - NodeId::new("node-2"), - NodeId::new("node-3"), - ]) - } - - fn create_topic(sm: &mut MetadataStateMachine, name: &str) -> TopicId { - let result = sm.apply(MetadataCommand::CreateTopic(CreateTopic { - name: name.to_string(), - storage_policy: default_policy(), - replica_set: replica_set(), - created_at: 1000, - })); - match result.unwrap() { - ApplyResult::TopicCreated(tc) => tc.segment_key.topic_id, - other => panic!("expected TopicCreated, got {:?}", other), - } - } - - #[test] - fn consumer_group_generation_is_applied_through_metadata_log_command() { - let mut sm = MetadataStateMachine::new(ShardGroupId(1)); - create_topic(&mut sm, "orders"); - let member = uuid::Uuid::new_v4(); - let command = SyncConsumerGroup { - req: SyncConsumerGroupRequest { - topic_name: "orders".into(), - group_id: "workers".into(), - member_id: member, - action: ConsumerGroupSyncAction::Heartbeat, - }, - observed_at: 100, - session_timeout_ms: 10_000, - }; - - let ApplyResult::ConsumerGroupChanged(epoch) = sm.apply(command.clone().into()).unwrap() - else { - panic!("first member must create a committed generation"); - }; - assert_eq!(*epoch.generation, 1); - assert_eq!(epoch.ranges.len(), 1); - let assignment = sm - .get_consumer_group_assignment("orders", "workers", member) - .unwrap(); - assert_eq!(*assignment.generation, 1); - assert_eq!(assignment.ranges.as_ref(), &[RangeId(0)]); - - let mut heartbeat = command; - heartbeat.observed_at = 200; - assert!(matches!( - sm.apply(heartbeat.into()).unwrap(), - ApplyResult::Noop - )); - assert_eq!( - *sm.get_consumer_group_assignment("orders", "workers", member) - .unwrap() - .generation, - 1 - ); - } - - fn roll_segment( - sm: &mut MetadataStateMachine, - topic_id: TopicId, - range_id: RangeId, - segment_id: SegmentId, - sealed_at: u64, - ) { - let result = sm.apply(MetadataCommand::RollSegment(RollSegment { - segment_key: SegmentKey::new(topic_id, range_id, segment_id), - sealed_at, - new_replica_set: replica_set(), - end_entry_id: None, - })); - assert!(matches!(result.unwrap(), ApplyResult::SegmentRolled(_))); - } - - // ── D7 retention ─────────────────────────────────────────────────────── - - /// Roll the active segment, sealing it at `end_entry_id` / `sealed_at` so the - /// resulting sealed segment has a known end and seal time (unlike the death-roll - /// `roll_segment` helper above which seals with `None`). - fn roll_with_end( - sm: &mut MetadataStateMachine, - topic_id: TopicId, - segment_id: SegmentId, - end_entry_id: u64, - sealed_at: u64, - ) { - let result = sm.apply(MetadataCommand::RollSegment(RollSegment { - segment_key: SegmentKey::new(topic_id, RangeId(0), segment_id), - sealed_at, - new_replica_set: replica_set(), - end_entry_id: Some(EntryId(end_entry_id)), - })); - assert!(matches!(result.unwrap(), ApplyResult::SegmentRolled(_))); - } - - fn seg_state( - sm: &MetadataStateMachine, - topic_id: TopicId, - segment_id: SegmentId, - ) -> SegmentMetaState { - sm.get_topic(&topic_id).unwrap().ranges[&RangeId(0)].segments[&segment_id] - .state - .clone() - } - - /// Build a topic with sealed segments 0,1,2 (ends 9/19/29, sealed at 100/200/300) - /// and an active head 3. - fn topic_with_three_sealed(sm: &mut MetadataStateMachine) -> TopicId { - let t = create_topic(sm, "t"); - roll_with_end(sm, t, SegmentId(0), 9, 100); - roll_with_end(sm, t, SegmentId(1), 19, 200); - roll_with_end(sm, t, SegmentId(2), 29, 300); - t - } - - fn delete_segments( - sm: &mut MetadataStateMachine, - topic_id: TopicId, - ids: &[u64], - ) -> Result { - sm.apply(MetadataCommand::DeleteSegments(DeleteSegments { - topic_id, - range_id: RangeId(0), - segment_ids: ids.iter().map(|&i| SegmentId(i)).collect(), - })) - } - - #[test] - fn delete_segments_marks_oldest_prefix_deleting() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let t = topic_with_three_sealed(&mut sm); - - let result = delete_segments(&mut sm, t, &[0, 1]).unwrap(); - let ApplyResult::SegmentsDeleted(d) = result else { - panic!("expected SegmentsDeleted, got {result:?}"); - }; - // All three sealed segments share one replica_set → a single group of 2 keys. - assert_eq!(d.groups.len(), 1); - assert_eq!(d.groups[0].1.len(), 2); - assert_eq!(seg_state(&sm, t, SegmentId(0)), SegmentMetaState::Deleting); - assert_eq!(seg_state(&sm, t, SegmentId(1)), SegmentMetaState::Deleting); - assert_eq!(seg_state(&sm, t, SegmentId(2)), SegmentMetaState::Sealed); - assert_eq!(seg_state(&sm, t, SegmentId(3)), SegmentMetaState::Active); - } - - #[test] - fn delete_segments_skips_the_active_head() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let t = topic_with_three_sealed(&mut sm); - // Naming the active head (seg 3) alongside the sealed prefix: only the sealed - // ones transition; the write head is skipped, never deleted. - let ApplyResult::SegmentsDeleted(d) = delete_segments(&mut sm, t, &[0, 1, 2, 3]).unwrap() - else { - panic!("expected SegmentsDeleted"); - }; - let total_keys: usize = d.groups.iter().map(|(_, keys)| keys.len()).sum(); - assert_eq!(total_keys, 3); - assert_eq!(seg_state(&sm, t, SegmentId(3)), SegmentMetaState::Active); - } - - /// The no-hole property is a structural invariant, not a hot-path check: a - /// non-prefix deletion (seg 1 while seg 0 survives) trips `assert_retention_prefix`. - #[test] - #[should_panic(expected = "not an oldest-first prefix")] - fn delete_segments_non_prefix_trips_invariant() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let t = topic_with_three_sealed(&mut sm); - let _ = delete_segments(&mut sm, t, &[1]); - } - - #[test] - fn delete_segments_is_idempotent() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let t = topic_with_three_sealed(&mut sm); - assert!(matches!( - delete_segments(&mut sm, t, &[0]).unwrap(), - ApplyResult::SegmentsDeleted(_) - )); - // Re-applying for an already-Deleting segment is a no-op. - assert!(matches!( - delete_segments(&mut sm, t, &[0]).unwrap(), - ApplyResult::Noop - )); - } - - #[test] - fn expired_prefix_selects_by_age_oldest_first() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let t = topic_with_three_sealed(&mut sm); // sealed_at 100/200/300, retention 3_600_000 - let topic = sm.get_topic(&t).unwrap(); - - // now such that segs 0,1 are past the window but seg 2 isn't. - let now = 200 + 3_600_000 + 1; - let prefixes = topic.expired_segment_prefixes(now); - assert_eq!(prefixes.len(), 1); - let (range_id, ids) = &prefixes[0]; - assert_eq!(*range_id, RangeId(0)); - assert_eq!(ids.as_ref(), &[SegmentId(0), SegmentId(1)]); - } - - #[test] - fn expired_prefix_empty_without_retention() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let t = sm - .apply(MetadataCommand::CreateTopic(CreateTopic { - name: "no-retention".into(), - storage_policy: StoragePolicy { - retention_ms: None, - replication_factor: 3, - partition_strategy: PartitionStrategy::AutoSplit, - }, - replica_set: replica_set(), - created_at: 1000, - })) - .map(|r| match r { - ApplyResult::TopicCreated(tc) => tc.segment_key.topic_id, - other => panic!("{other:?}"), - }) - .unwrap(); - roll_with_end(&mut sm, t, SegmentId(0), 9, 100); - // Far past any window, but no policy → nothing expires. - assert!( - sm.get_topic(&t) - .unwrap() - .expired_segment_prefixes(u64::MAX) - .is_empty() - ); - } - - fn split_range( - sm: &mut MetadataStateMachine, - topic_id: TopicId, - range_id: RangeId, - split_point: Vec, - created_at: u64, - ) -> (RangeId, RangeId) { - let result = sm.apply(MetadataCommand::SplitRange(SplitRange { - topic_id, - range_id, - split_point, - created_at, - left_replica_set: replica_set(), - right_replica_set: replica_set(), - })); - match result.unwrap() { - ApplyResult::RangeSplit(rs) => (rs.children[0].0, rs.children[1].0), - other => panic!("expected RangeSplit, got {:?}", other), - } - } - - fn merge_range( - sm: &mut MetadataStateMachine, - topic_id: TopicId, - range_id_1: RangeId, - range_id_2: RangeId, - created_at: u64, - ) -> RangeId { - let result = sm.apply(MetadataCommand::MergeRange(MergeRange { - topic_id, - range_id_1, - range_id_2, - created_at, - merged_replica_set: replica_set(), - })); - match result.unwrap() { - ApplyResult::RangeMerged(rm) => rm.segment_key.range_id, - other => panic!("expected RangeMerged, got {:?}", other), - } - } - - /// A surviving subset plus a fresh replacement — what the coordinator picks - /// when a replica of a sealed segment dies (node-3 → node-4 here). - fn replacement_set() -> Replicas { - Replicas::new(vec![ - NodeId::new("node-1"), - NodeId::new("node-2"), - NodeId::new("node-4"), - ]) - } - - // --- ReassignSegment --- - - #[test] - fn reassign_swaps_a_sealed_segments_replica_set() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let topic_id = create_topic(&mut sm, "blue"); - // Roll so SegmentId(0) becomes Sealed (SegmentId(1) is the new write head). - roll_segment(&mut sm, topic_id, RangeId(0), SegmentId(0), 2000); - let sealed = SegmentKey::new(topic_id, RangeId(0), SegmentId(0)); - - let result = sm - .apply(MetadataCommand::ReassignSegment(ReassignSegment { - segment_key: sealed, - replica_set: replacement_set(), - })) - .unwrap(); - - match result { - ApplyResult::SegmentReassigned(r) => { - assert_eq!(r.segment_key, sealed); - assert_eq!(r.new_replica_set, replacement_set()); - } - other => panic!("expected SegmentReassigned, got {other:?}"), - } - - let seg = &sm.get_topic(&topic_id).unwrap().ranges[&RangeId(0)].segments[&SegmentId(0)]; - assert_eq!(seg.replica_set, replacement_set()); - assert_eq!(seg.state, SegmentMetaState::Sealed); // stays sealed - } - - #[test] - fn reassign_same_set_is_a_noop() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let topic_id = create_topic(&mut sm, "blue"); - roll_segment(&mut sm, topic_id, RangeId(0), SegmentId(0), 2000); - let sealed = SegmentKey::new(topic_id, RangeId(0), SegmentId(0)); - - sm.apply(MetadataCommand::ReassignSegment(ReassignSegment { - segment_key: sealed, - replica_set: replacement_set(), - })) - .unwrap(); - - // Re-applying the identical set (duplicate death detection / re-proposal) - // changes nothing. - let again = sm - .apply(MetadataCommand::ReassignSegment(ReassignSegment { - segment_key: sealed, - replica_set: replacement_set(), - })) - .unwrap(); - assert_eq!(again, ApplyResult::Noop); - } - - #[test] - fn reassign_rejects_an_active_segment() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let topic_id = create_topic(&mut sm, "blue"); - // SegmentId(0) is the active write head — no roll yet. - let active = SegmentKey::new(topic_id, RangeId(0), SegmentId(0)); - - let result = sm.apply(MetadataCommand::ReassignSegment(ReassignSegment { - segment_key: active, - replica_set: replacement_set(), - })); - assert!(matches!(result, Err(MetadataError::SegmentNotSealed))); - } - - #[test] - fn reassign_rejects_an_unknown_segment() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let topic_id = create_topic(&mut sm, "blue"); - let unknown = SegmentKey::new(topic_id, RangeId(0), SegmentId(99)); - - let result = sm.apply(MetadataCommand::ReassignSegment(ReassignSegment { - segment_key: unknown, - replica_set: replacement_set(), - })); - assert!(matches!(result, Err(MetadataError::SegmentNotFound))); - } - - // --- CreateTopic --- - - #[test] - fn create_topic_basic() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let id = create_topic(&mut sm, "blue"); - - let topic = sm.get_topic(&id).unwrap(); - assert_eq!(topic.name, "blue"); - assert_eq!(topic.state, TopicState::Active); - assert_eq!(topic.active_ranges.len(), 1); - assert_eq!(topic.ranges.len(), 1); - - let range = &topic.ranges[&RangeId(0)]; - assert_eq!(range.state, RangeState::Active); - assert!(range.active_segment.is_some()); - assert_eq!(range.segments.len(), 1); - } - - #[test] - fn create_topic_duplicate_name_rejected() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - create_topic(&mut sm, "blue"); - - let result = sm.apply(MetadataCommand::CreateTopic(CreateTopic { - name: "blue".to_string(), - storage_policy: default_policy(), - replica_set: replica_set(), - created_at: 2000, - })); - assert_eq!(result, Err(TopicNameAlreadyExists("blue".to_string()))); - } - - #[test] - fn create_topic_increments_id() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let id1 = create_topic(&mut sm, "alpha"); - let id2 = create_topic(&mut sm, "beta"); - - assert_eq!(id1, TopicId(0)); - assert_eq!(id2, TopicId(1)); - assert_eq!(sm.topic_count(), 2); - } - - #[test] - fn create_topic_initial_offsets() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let id = create_topic(&mut sm, "blue"); - - let topic = sm.get_topic(&id).unwrap(); - let range = &topic.ranges[&RangeId(0)]; - assert_eq!(range.next_offset, EntryId(0)); - - let seg = &range.segments[&SegmentId(0)]; - assert_eq!(seg.start_entry_id, EntryId(0)); - assert_eq!(seg.end_entry_id, None); - } - - #[test] - fn create_topic_name_index() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let id = create_topic(&mut sm, "blue"); - - let found = sm.get_topic_by_name("blue").unwrap(); - assert_eq!(found.id, id); - assert!(sm.get_topic_by_name("red").is_none()); - } - - // --- RollSegment --- - - #[test] - fn roll_segment_creates_next() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - - roll_segment(&mut sm, tid, RangeId(0), SegmentId(0), 2000); - - let range = &sm.get_topic(&tid).unwrap().ranges[&RangeId(0)]; - assert_eq!(range.active_segment, Some(SegmentId(1))); - assert_eq!(range.segments.len(), 2); - } - - #[test] - fn roll_segment_increments_segment_id() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - - roll_segment(&mut sm, tid, RangeId(0), SegmentId(0), 2000); - roll_segment(&mut sm, tid, RangeId(0), SegmentId(1), 3000); - - let range = &sm.get_topic(&tid).unwrap().ranges[&RangeId(0)]; - assert_eq!(range.active_segment, Some(SegmentId(2))); - assert_eq!(range.segments.len(), 3); - assert_eq!(range.next_segment_id, 3); - } - - #[test] - fn roll_segment_bad_topic() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let result = sm.apply(MetadataCommand::RollSegment(RollSegment { - segment_key: SegmentKey::new(TopicId(99), RangeId(0), SegmentId(0)), - sealed_at: 2000, - new_replica_set: replica_set(), - end_entry_id: None, - })); - assert_eq!(result, Err(TopicNotFound(TopicId(99)))); - } - - #[test] - fn roll_segment_bad_range() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - - let result = sm.apply(MetadataCommand::RollSegment(RollSegment { - segment_key: SegmentKey::new(tid, RangeId(99), SegmentId(0)), - sealed_at: 2000, - new_replica_set: replica_set(), - end_entry_id: None, - })); - assert_eq!(result, Err(RangeNotFound)); - } - - #[test] - fn roll_segment_stale_is_rejected() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - - let result = sm.apply(MetadataCommand::RollSegment(RollSegment { - segment_key: SegmentKey::new(tid, RangeId(0), SegmentId(99)), - sealed_at: 2000, - new_replica_set: replica_set(), - end_entry_id: None, - })); - assert_eq!(result, Ok(ApplyResult::Noop)); - - let range = &sm.get_topic(&tid).unwrap().ranges[&RangeId(0)]; - assert_eq!(range.active_segment, Some(SegmentId(0))); - assert_eq!(range.segments.len(), 1); - } - - // --- SplitRange --- - - #[test] - fn split_range_basic() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - - let (c1, c2) = split_range(&mut sm, tid, RangeId(0), vec![0x80], 2000); - - let topic = sm.get_topic(&tid).unwrap(); - assert_eq!(topic.ranges.len(), 3); - - let child1 = &topic.ranges[&c1]; - assert_eq!(child1.keyspace_start, KEYSPACE_MIN); - assert_eq!(child1.keyspace_end, vec![0x80]); - assert_eq!(child1.state, RangeState::Active); - - let child2 = &topic.ranges[&c2]; - assert_eq!(child2.keyspace_start, vec![0x80]); - assert_eq!(child2.keyspace_end, KEYSPACE_MAX); - assert_eq!(child2.state, RangeState::Active); - } - - #[test] - fn split_range_updates_active_ranges() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - - let (c1, c2) = split_range(&mut sm, tid, RangeId(0), vec![0x80], 2000); - - let topic = sm.get_topic(&tid).unwrap(); - assert_eq!(topic.active_ranges, vec![c1, c2]); - assert!(!topic.active_ranges.contains(&RangeId(0))); - } - - #[test] - fn split_range_lineage() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - - let (c1, c2) = split_range(&mut sm, tid, RangeId(0), vec![0x80], 2000); - - let topic = sm.get_topic(&tid).unwrap(); - assert_eq!(topic.ranges[&RangeId(0)].split_into, Some([c1, c2])); - } - - #[test] - fn split_range_fixed_rejected() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let result = sm.apply(MetadataCommand::CreateTopic(CreateTopic { - name: "ordered".to_string(), - storage_policy: fixed_policy(), - replica_set: replica_set(), - created_at: 1000, - })); - let tid = match result.unwrap() { - ApplyResult::TopicCreated(tc) => tc.segment_key.topic_id, - other => panic!("expected TopicCreated, got {:?}", other), - }; - - let split_result = sm.apply(MetadataCommand::SplitRange(SplitRange { - topic_id: tid, - range_id: RangeId(0), - split_point: vec![0x80], - created_at: 2000, - left_replica_set: replica_set(), - right_replica_set: replica_set(), - })); - assert_eq!(split_result, Err(SplitNotAllowed(tid))); - } - - #[test] - fn split_range_invalid_split_point() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - - let upper_bound = sm.apply(MetadataCommand::SplitRange(SplitRange { - topic_id: tid, - range_id: RangeId(0), - split_point: vec![0xFF], - created_at: 2000, - left_replica_set: replica_set(), - right_replica_set: replica_set(), - })); - assert_eq!(upper_bound, Err(InvalidSplitPoint)); - - let lower_bound = sm.apply(MetadataCommand::SplitRange(SplitRange { - topic_id: tid, - range_id: RangeId(0), - split_point: vec![], - created_at: 2000, - left_replica_set: replica_set(), - right_replica_set: replica_set(), - })); - assert_eq!(lower_bound, Err(InvalidSplitPoint)); - } - - #[test] - fn split_range_keyspace_coverage() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - - let (c1, c2) = split_range(&mut sm, tid, RangeId(0), vec![0x80], 2000); - - let topic = sm.get_topic(&tid).unwrap(); - let r1 = &topic.ranges[&c1]; - let r2 = &topic.ranges[&c2]; - - assert_eq!(r1.keyspace_end, r2.keyspace_start); - assert_eq!(r1.keyspace_start, KEYSPACE_MIN); - assert_eq!(r2.keyspace_end, KEYSPACE_MAX); - } - - // --- MergeRange --- - - #[test] - fn merge_range_basic() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - let (c1, c2) = split_range(&mut sm, tid, RangeId(0), vec![0x80], 2000); - - let merged_id = merge_range(&mut sm, tid, c1, c2, 3000); - - let topic = sm.get_topic(&tid).unwrap(); - let merged = &topic.ranges[&merged_id]; - assert_eq!(merged.keyspace_start, KEYSPACE_MIN); - assert_eq!(merged.keyspace_end, KEYSPACE_MAX); - assert_eq!(merged.state, RangeState::Active); - } - - #[test] - fn merge_range_active_ranges_updated() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - let (c1, c2) = split_range(&mut sm, tid, RangeId(0), vec![0x80], 2000); - - let merged_id = merge_range(&mut sm, tid, c1, c2, 3000); - - let topic = sm.get_topic(&tid).unwrap(); - assert_eq!(topic.active_ranges, vec![merged_id]); - } - - #[test] - fn merge_range_lineage() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - let (c1, c2) = split_range(&mut sm, tid, RangeId(0), vec![0x80], 2000); - - let merged_id = merge_range(&mut sm, tid, c1, c2, 3000); - - let topic = sm.get_topic(&tid).unwrap(); - assert_eq!(topic.ranges[&c1].merged_into, Some(merged_id)); - assert_eq!(topic.ranges[&c2].merged_into, Some(merged_id)); - assert_eq!(topic.ranges[&merged_id].merged_from, Some([c1, c2])); - } - - #[test] - fn merge_range_non_adjacent_rejected() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - let (c1, c2) = split_range(&mut sm, tid, RangeId(0), vec![0x80], 2000); - let (c1a, _c1b) = split_range(&mut sm, tid, c1, vec![0x40], 3000); - - let result = sm.apply(MetadataCommand::MergeRange(MergeRange { - topic_id: tid, - range_id_1: c1a, - range_id_2: c2, - created_at: 4000, - merged_replica_set: replica_set(), - })); - assert_eq!(result, Err(RangesNotAdjacent)); - } - - // --- DeleteTopic --- - - #[test] - fn delete_topic_cascades() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - - sm.apply(MetadataCommand::DeleteTopic(DeleteTopic { - name: "blue".into(), - })) - .unwrap(); - - let topic = sm.get_topic(&tid).unwrap(); - assert_eq!(topic.state, TopicState::Deleted); - assert!(topic.active_ranges.is_empty()); - - for range in topic.ranges.values() { - assert_eq!(range.state, RangeState::Deleting); - for seg in range.segments.values() { - assert_eq!(seg.state, SegmentMetaState::Deleting); - } - } - } - - #[test] - fn delete_topic_removes_name_index() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let _tid = create_topic(&mut sm, "blue"); - - sm.apply(MetadataCommand::DeleteTopic(DeleteTopic { - name: "blue".into(), - })) - .unwrap(); - - assert!(sm.get_topic_by_name("blue").is_none()); - } - - #[test] - fn delete_topic_nonexistent() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let result = sm.apply(MetadataCommand::DeleteTopic(DeleteTopic { - name: "nope".into(), - })); - assert_eq!(result, Err(MetadataError::TopicNameNotFound("nope".into()))); - } - - // --- Integration --- - - #[test] - fn create_split_seal() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - let (c1, _c2) = split_range(&mut sm, tid, RangeId(0), vec![0x80], 2000); - - roll_segment(&mut sm, tid, c1, SegmentId(0), 3000); - - let child = &sm.get_topic(&tid).unwrap().ranges[&c1]; - assert_eq!(child.active_segment, Some(SegmentId(1))); - assert_eq!(child.segments.len(), 2); - } - - #[test] - fn split_merge_roundtrip() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - let (c1, c2) = split_range(&mut sm, tid, RangeId(0), vec![0x80], 2000); - - let merged = merge_range(&mut sm, tid, c1, c2, 3000); - - let topic = sm.get_topic(&tid).unwrap(); - let range = &topic.ranges[&merged]; - assert_eq!(range.keyspace_start, KEYSPACE_MIN); - assert_eq!(range.keyspace_end, KEYSPACE_MAX); - assert_eq!(topic.active_ranges, vec![merged]); - } - - #[test] - fn multiple_topics_independent() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let t1 = create_topic(&mut sm, "alpha"); - let t2 = create_topic(&mut sm, "beta"); - - split_range(&mut sm, t1, RangeId(0), vec![0x80], 2000); - - let alpha = sm.get_topic(&t1).unwrap(); - assert_eq!(alpha.active_ranges.len(), 2); - - let beta = sm.get_topic(&t2).unwrap(); - assert_eq!(beta.active_ranges.len(), 1); - assert_eq!(beta.ranges.len(), 1); - } - - // --- Invariant: Segment immutability after seal --- - - #[test] - fn roll_already_rolled_segment_is_stale() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - - roll_segment(&mut sm, tid, RangeId(0), SegmentId(0), 2000); - - let result = sm.apply(MetadataCommand::RollSegment(RollSegment { - segment_key: SegmentKey::new(tid, RangeId(0), SegmentId(0)), - sealed_at: 3000, - new_replica_set: replica_set(), - end_entry_id: None, - })); - assert_eq!(result, Ok(ApplyResult::Noop)); - - let range = &sm.get_topic(&tid).unwrap().ranges[&RangeId(0)]; - assert_eq!(range.active_segment, Some(SegmentId(1))); - assert_eq!(range.segments.len(), 2); - } - - // --- Hot Range Detection --- - - #[test] - fn seal_history_records_timestamps() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - - roll_segment(&mut sm, tid, RangeId(0), SegmentId(0), 1000); - roll_segment(&mut sm, tid, RangeId(0), SegmentId(1), 2000); - roll_segment(&mut sm, tid, RangeId(0), SegmentId(2), 3000); - - let range = &sm.get_topic(&tid).unwrap().ranges[&RangeId(0)]; - assert_eq!(range.seal_history.seal_count(), 3); - } - - #[test] - fn seal_history_prunes_old_entries() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - - roll_segment(&mut sm, tid, RangeId(0), SegmentId(0), 1000); - roll_segment(&mut sm, tid, RangeId(0), SegmentId(1), 2000); - // Jump far beyond the window — both old entries pruned - let far_future = 2000 + MEASUREMENT_WINDOW_MS + 1; - roll_segment(&mut sm, tid, RangeId(0), SegmentId(2), far_future); - - let range = &sm.get_topic(&tid).unwrap().ranges[&RangeId(0)]; - assert_eq!(range.seal_history.seal_count(), 1); - } - - #[test] - fn seal_history_orders_delayed_proposal_timestamps() { - let mut history = RangeSealHistory::default(); - history.record_seal(2000); - history.record_seal(1000); - history.record_seal(1500); - - assert_eq!(history.seal_timestamps, VecDeque::from([1000, 1500, 2000])); - } - - #[test] - fn delayed_old_seal_is_pruned_against_newest_timestamp() { - let mut history = RangeSealHistory::default(); - let newest = MEASUREMENT_WINDOW_MS + 2000; - history.record_seal(newest); - history.record_seal(1000); - - assert_eq!(history.seal_timestamps, VecDeque::from([newest])); - } - - #[test] - fn should_split_threshold_met() { - let mut history = RangeSealHistory::default(); - history.record_seal(1000); - history.record_seal(2000); - history.record_seal(3000); - - assert!(history.should_split(3000)); - } - - #[test] - fn should_split_below_threshold() { - let mut history = RangeSealHistory::default(); - history.record_seal(1000); - history.record_seal(2000); - - assert!(!history.should_split(2000)); - } - - #[test] - fn should_split_cooldown_blocks() { - let mut history = RangeSealHistory { - seal_timestamps: VecDeque::new(), - created_by_split_at: Some(1000), - }; - history.record_seal(1100); - history.record_seal(1200); - history.record_seal(1300); - - // Within cooldown — blocked - assert!(!history.should_split(1300)); - - // After cooldown — allowed - assert!(history.should_split(1000 + SPLIT_COOLDOWN_MS)); - } - - #[test] - fn auto_proposal_on_hot_range() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - - for i in 0..SPLIT_SEAL_THRESHOLD { - roll_segment( - &mut sm, - tid, - RangeId(0), - SegmentId(i as u64), - 1000 * (i as u64 + 1), - ); - } - - let proposals = sm.take_pending_proposals(); - assert_eq!(proposals.len(), 1); - assert!(matches!(proposals[0], MetadataCommand::SplitRange(_))); - } - - #[test] - fn no_auto_proposal_below_threshold() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - - for i in 0..(SPLIT_SEAL_THRESHOLD - 1) { - roll_segment( - &mut sm, - tid, - RangeId(0), - SegmentId(i as u64), - 1000 * (i as u64 + 1), - ); - } - - let proposals = sm.take_pending_proposals(); - assert!(proposals.is_empty()); - } - - #[test] - fn no_auto_proposal_for_fixed_strategy() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let result = sm.apply(MetadataCommand::CreateTopic(CreateTopic { - name: "ordered".to_string(), - storage_policy: fixed_policy(), - replica_set: replica_set(), - created_at: 1000, - })); - let tid = match result.unwrap() { - ApplyResult::TopicCreated(tc) => tc.segment_key.topic_id, - other => panic!("expected TopicCreated, got {:?}", other), - }; - - for i in 0..SPLIT_SEAL_THRESHOLD { - roll_segment( - &mut sm, - tid, - RangeId(0), - SegmentId(i as u64), - 2000 * (i as u64 + 1), - ); - } - - let proposals = sm.take_pending_proposals(); - assert!(proposals.is_empty()); - } - - #[test] - fn evaluate_merges_cold_adjacent() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - split_range(&mut sm, tid, RangeId(0), vec![0x80], 2000); - - // Both children are cold (no seals) - let proposals = sm.evaluate_merges(2000 + SPLIT_COOLDOWN_MS + 1); - assert_eq!(proposals.len(), 1); - assert!(matches!(proposals[0], MetadataCommand::MergeRange(_))); - } - - #[test] - fn evaluate_merges_one_hot() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - let (c1, _c2) = split_range(&mut sm, tid, RangeId(0), vec![0x80], 2000); - sm.take_pending_proposals(); // discard any split proposals - - // Seal one child — makes it hot - roll_segment(&mut sm, tid, c1, SegmentId(0), 3000); - - let proposals = sm.evaluate_merges(3000); - assert!(proposals.is_empty()); - } - - #[test] - fn split_clears_seal_history() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - - roll_segment(&mut sm, tid, RangeId(0), SegmentId(0), 2000); - let (c1, c2) = split_range(&mut sm, tid, RangeId(0), vec![0x80], 3000); - - let topic = sm.get_topic(&tid).unwrap(); - assert!(topic.ranges[&c1].seal_history.seal_timestamps.is_empty()); - assert!(topic.ranges[&c2].seal_history.seal_timestamps.is_empty()); - } - - #[test] - fn split_sets_cooldown() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - let (c1, c2) = split_range(&mut sm, tid, RangeId(0), vec![0x80], 2000); - - let topic = sm.get_topic(&tid).unwrap(); - assert_eq!( - topic.ranges[&c1].seal_history.created_by_split_at, - Some(2000) - ); - assert_eq!( - topic.ranges[&c2].seal_history.created_by_split_at, - Some(2000) - ); - } - - // --- D3: active_segments_for_node --- - - #[test] - fn active_segments_for_node_returns_matching() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - - let segments = sm.active_segments_for_node(&NodeId::new("node-1")); - assert_eq!(segments.len(), 1); - let (key, rs) = &segments[0]; - assert_eq!(key.topic_id, tid); - assert_eq!(key.range_id, RangeId(0)); - assert_eq!(key.segment_id, SegmentId(0)); - assert!(rs.contains(&NodeId::new("node-1"))); - } - - #[test] - fn active_segments_for_node_excludes_non_member() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - create_topic(&mut sm, "blue"); - - let segments = sm.active_segments_for_node(&NodeId::new("node-99")); - assert!(segments.is_empty()); - } - - #[test] - fn active_segments_for_node_excludes_sealed() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - roll_segment(&mut sm, tid, RangeId(0), SegmentId(0), 2000); - - let segments = sm.active_segments_for_node(&NodeId::new("node-1")); - assert_eq!(segments.len(), 1); - assert_eq!(segments[0].0.segment_id, SegmentId(1)); - } - - // --- D3: end_entry_id in RollSegment --- - - #[test] - fn roll_segment_uses_end_entry_id() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - - let result = sm.apply(MetadataCommand::RollSegment(RollSegment { - segment_key: SegmentKey::new(tid, RangeId(0), SegmentId(0)), - sealed_at: 2000, - new_replica_set: replica_set(), - end_entry_id: Some(EntryId(42000)), - })); - - let result = result.unwrap(); - assert!(matches!( - result, - ApplyResult::SegmentRolled(SegmentRolled { - end_entry_id: Some(EntryId(42000)), - .. - }) - )); - - let range = &sm.get_topic(&tid).unwrap().ranges[&RangeId(0)]; - let sealed = &range.segments[&SegmentId(0)]; - assert_eq!(sealed.end_entry_id, Some(EntryId(42000))); - let new_seg = &range.segments[&SegmentId(1)]; - assert_eq!(new_seg.start_entry_id, EntryId(42001)); - } - - // --- D3: end-offset correction --- - - #[test] - fn end_offset_correction_updates_placeholder() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - - // Death-triggered roll with end_entry_id=0 (placeholder) - let _ = sm.apply(MetadataCommand::RollSegment(RollSegment { - segment_key: SegmentKey::new(tid, RangeId(0), SegmentId(0)), - sealed_at: 2000, - new_replica_set: replica_set(), - end_entry_id: None, - })); - - assert_eq!( - sm.get_topic(&tid).unwrap().ranges[&RangeId(0)].segments[&SegmentId(0)].end_entry_id, - None - ); - - // Segment leader's RollSegment arrives with correct end_entry_id - let result = sm.apply(MetadataCommand::RollSegment(RollSegment { - segment_key: SegmentKey::new(tid, RangeId(0), SegmentId(0)), - sealed_at: 2500, - new_replica_set: replica_set(), - end_entry_id: Some(EntryId(42000)), - })); - assert!(result.is_ok()); - - let range = &sm.get_topic(&tid).unwrap().ranges[&RangeId(0)]; - assert_eq!( - range.segments[&SegmentId(0)].end_entry_id, - Some(EntryId(42000)) - ); - assert_eq!(range.segments[&SegmentId(1)].start_entry_id, EntryId(42001)); - } - - #[test] - fn end_offset_correction_rejected_when_already_set() { - let mut sm = MetadataStateMachine::new(ShardGroupId(0)); - let tid = create_topic(&mut sm, "blue"); - - // Normal roll with actual end_entry_id - let _ = sm.apply(MetadataCommand::RollSegment(RollSegment { - segment_key: SegmentKey::new(tid, RangeId(0), SegmentId(0)), - sealed_at: 2000, - new_replica_set: replica_set(), - end_entry_id: Some(EntryId(1000)), - })); - - // Duplicate roll is rejected (end_offset already set) - let result = sm.apply(MetadataCommand::RollSegment(RollSegment { - segment_key: SegmentKey::new(tid, RangeId(0), SegmentId(0)), - sealed_at: 2500, - new_replica_set: replica_set(), - end_entry_id: Some(EntryId(42000)), - })); - assert_eq!(result, Ok(ApplyResult::Noop)); - } -} From 5e713a96ddde7553c138041d0ae9af6bda443555 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Fri, 17 Jul 2026 21:27:40 +0400 Subject: [PATCH 2/9] docs --- .../consensus/raft/states/mod.rs | 41 +++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/src/control_plane/consensus/raft/states/mod.rs b/src/control_plane/consensus/raft/states/mod.rs index e8f1c366..a88e1e9a 100644 --- a/src/control_plane/consensus/raft/states/mod.rs +++ b/src/control_plane/consensus/raft/states/mod.rs @@ -1,3 +1,44 @@ +/*! +Ownership boundaries for one Raft group's state. + +[`LogState`] owns the in-memory Raft logs and its persistence bookkeeping: +[`current_term`], [`voted_for`], the in-memory [`LogState::log`], +and the local durability watermark [`stabled_index`]. +Its [`last_index`] is derived from the final in-memory Raft log entry; it is not a data-plane WAL position. + +[`TransientState`] owns knowledge that may be reconstructed after restart: +the current role, leader identity, replication progress, and [`commit_index`]. +A vote is not transient: forgetting [`voted_for`] after a crash could let one replica vote for two candidates in the same term. + +[`MetadataState`] owns the committed application state and [`last_applied_index`]. + +Numeric progress watermarks are ordered as: +```text +last_applied_index <= commit_index <= stabled_index <= last_index +``` + +Higher indices mean further progress. Expressed as entry sets, the same +relationship is `applied ⊆ committed ⊆ stable ⊆ in-memory log`. + +- [`last_index`]: newest entry currently present in the in-memory Raft log. +- [`stabled_index`]: newest Raft log entry successfully persisted to RocksDB. +- [`commit_index`]: newest durable entry known to have quorum agreement. +- [`last_applied_index`]: newest committed entry applied to metadata. + +The data-plane WAL is separate: it stores segment records, while these +boundaries describe the control-plane Raft metadata log. + +[`LogState`]: self::log_state::LogState +[`current_term`]: self::log_state::LogState::current_term +[`voted_for`]: self::log_state::LogState::voted_for +[`LogState::log`]: self::log_state::LogState::log +[`stabled_index`]: self::log_state::LogState::stabled_index +[`last_index`]: self::log_state::LogState::last_index +[`TransientState`]: self::transient_state::TransientState +[`commit_index`]: self::transient_state::TransientState::commit_index +[`MetadataState`]: self::metadata_state::MetadataState +[`last_applied_index`]: self::metadata_state::MetadataState::last_applied_index +*/ pub(crate) mod log_state; pub(crate) mod metadata_state; pub(crate) mod transient_state; From 6cad5318f3489e72e5b8266d1df62eabe601a66e Mon Sep 17 00:00:00 2001 From: Migorithm Date: Fri, 17 Jul 2026 21:38:36 +0400 Subject: [PATCH 3/9] rn --- src/control_plane/consensus/raft/states/mod.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/control_plane/consensus/raft/states/mod.rs b/src/control_plane/consensus/raft/states/mod.rs index a88e1e9a..c90a9a9b 100644 --- a/src/control_plane/consensus/raft/states/mod.rs +++ b/src/control_plane/consensus/raft/states/mod.rs @@ -2,10 +2,11 @@ Ownership boundaries for one Raft group's state. [`LogState`] owns the in-memory Raft logs and its persistence bookkeeping: -[`current_term`], [`voted_for`], the in-memory [`LogState::log`], +[`current_term`], [`voted_for`], the in-memory [`LogState::entries`], and the local durability watermark [`stabled_index`]. Its [`last_index`] is derived from the final in-memory Raft log entry; it is not a data-plane WAL position. + [`TransientState`] owns knowledge that may be reconstructed after restart: the current role, leader identity, replication progress, and [`commit_index`]. A vote is not transient: forgetting [`voted_for`] after a crash could let one replica vote for two candidates in the same term. @@ -31,7 +32,8 @@ boundaries describe the control-plane Raft metadata log. [`LogState`]: self::log_state::LogState [`current_term`]: self::log_state::LogState::current_term [`voted_for`]: self::log_state::LogState::voted_for -[`LogState::log`]: self::log_state::LogState::log +[`LogState::entries`]: self::log_state::LogState::entries +[`LogState::unflushed_mutations`]: self::log_state::LogState::unflushed_mutations [`stabled_index`]: self::log_state::LogState::stabled_index [`last_index`]: self::log_state::LogState::last_index [`TransientState`]: self::transient_state::TransientState From 1248e9723b0e208c8d29cbf6484585ed504cbc5b Mon Sep 17 00:00:00 2001 From: Migorithm Date: Fri, 17 Jul 2026 21:41:14 +0400 Subject: [PATCH 4/9] log state --- .../consensus/raft/states/log_state.rs | 187 ++++++++++++++++++ 1 file changed, 187 insertions(+) create mode 100644 src/control_plane/consensus/raft/states/log_state.rs diff --git a/src/control_plane/consensus/raft/states/log_state.rs b/src/control_plane/consensus/raft/states/log_state.rs new file mode 100644 index 00000000..c53a522a --- /dev/null +++ b/src/control_plane/consensus/raft/states/log_state.rs @@ -0,0 +1,187 @@ +use crate::control_plane::NodeId; +use crate::control_plane::consensus::messages::LogMutation; +use crate::control_plane::consensus::raft::log::LogEntry; +use crate::control_plane::consensus::raft::storage::RaftPersistentState; + +pub(crate) struct LogState { + current_term: u64, + voted_for: Option, + entries: Vec, + stabled_index: u64, + unflushed_mutations: Vec, +} + +impl LogState { + pub(crate) fn from_persistent(persistent: RaftPersistentState) -> Self { + Self { + stabled_index: persistent.stabled_index(), + current_term: persistent.term, + voted_for: persistent.voted_for, + entries: persistent.log, + unflushed_mutations: Vec::new(), + } + } + + pub(crate) fn current_term(&self) -> u64 { + self.current_term + } + + pub(crate) fn voted_for(&self) -> Option<&NodeId> { + self.voted_for.as_ref() + } + + pub(crate) fn vote_available_for(&self, candidate_id: &NodeId) -> bool { + self.voted_for + .as_ref() + .is_none_or(|voted_for| voted_for == candidate_id) + } + + pub(crate) fn begin_election(&mut self, node_id: &NodeId) -> u64 { + self.current_term += 1; + self.voted_for = Some(node_id.clone()); + self.buffer_hard_state(); + self.current_term + } + + pub(crate) fn grant_vote(&mut self, candidate_id: NodeId) { + self.voted_for = Some(candidate_id); + self.buffer_hard_state(); + } + + pub(crate) fn advance_term(&mut self, new_term: u64) -> bool { + if new_term <= self.current_term { + return false; + } + self.current_term = new_term; + self.voted_for = None; + self.buffer_hard_state(); + true + } + + pub(crate) fn stabled_index(&self) -> u64 { + self.stabled_index + } + + pub(crate) fn advance_stabled_index(&mut self, index: u64) { + self.stabled_index = self.stabled_index.max(index); + } + + pub(crate) fn entries(&self) -> &[LogEntry] { + &self.entries + } + + pub(crate) fn last_index(&self) -> u64 { + self.entries.last().map_or(0, |entry| entry.index) + } + + pub(crate) fn last_term(&self) -> u64 { + self.entries.last().map_or(0, |entry| entry.term) + } + + pub(crate) fn term_at(&self, index: u64) -> u64 { + if index == 0 { + return 0; + } + self.get(index).map_or(0, |entry| entry.term) + } + + pub(crate) fn get(&self, index: u64) -> Option<&LogEntry> { + if index == 0 { + return None; + } + self.entries.get((index - 1) as usize) + } + + pub(crate) fn entries_from(&self, start_index: u64) -> Box<[LogEntry]> { + if start_index == 0 || start_index > self.last_index() { + return Box::new([]); + } + self.entries[(start_index - 1) as usize..].into() + } + + pub(crate) fn append(&mut self, entry: LogEntry) { + debug_assert_eq!( + entry.index, + self.last_index() + 1, + "log entry index must be contiguous" + ); + self.unflushed_mutations + .push(LogMutation::Append(entry.clone())); + self.entries.push(entry); + } + + pub(crate) fn truncate_from(&mut self, from_index: u64) { + if from_index == 0 || from_index > self.last_index() + 1 { + return; + } + self.unflushed_mutations + .push(LogMutation::TruncateFrom(from_index)); + self.entries.truncate((from_index - 1) as usize); + } + + fn buffer_hard_state(&mut self) { + self.unflushed_mutations.push(LogMutation::HardState { + term: self.current_term, + voted_for: self.voted_for.clone(), + }); + } + + pub(crate) fn take_mutations(&mut self) -> Vec { + std::mem::take(&mut self.unflushed_mutations) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::control_plane::consensus::raft::command::RaftCommand; + + #[test] + fn append_and_truncate_emit_persistence_mutations() { + let mut state = LogState::from_persistent(RaftPersistentState::default()); + let entry = LogEntry { + term: 1, + index: 1, + command: RaftCommand::Noop, + }; + + state.append(entry.clone()); + state.truncate_from(1); + + assert_eq!(state.last_index(), 0); + let mutations = state.take_mutations(); + assert!(matches!( + &mutations[..], + [ + LogMutation::Append(appended), + LogMutation::TruncateFrom(1) + ] if appended == &entry + )); + } + + #[test] + fn election_transitions_buffer_hard_state() { + let node = NodeId::new("node-1"); + let candidate = NodeId::new("node-2"); + let mut state = LogState::from_persistent(RaftPersistentState::default()); + + assert_eq!(state.begin_election(&node), 1); + assert_eq!(state.voted_for(), Some(&node)); + assert!(!state.vote_available_for(&candidate)); + + assert!(state.advance_term(2)); + assert!(state.voted_for().is_none()); + state.grant_vote(candidate.clone()); + assert_eq!(state.voted_for(), Some(&candidate)); + + let mutations = state.take_mutations(); + assert_eq!(mutations.len(), 3); + assert!(matches!( + mutations.last(), + Some(LogMutation::HardState { + term: 2, + voted_for: Some(voted_node), + }) if voted_node == &candidate + )); + } +} From 08a2476c70d63f4cb5c0a00be6e34e8941d82eb0 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Fri, 17 Jul 2026 22:14:15 +0400 Subject: [PATCH 5/9] refactor: state delegation --- src/control_plane/consensus/raft/state.rs | 941 ++++++++-------------- 1 file changed, 339 insertions(+), 602 deletions(-) diff --git a/src/control_plane/consensus/raft/state.rs b/src/control_plane/consensus/raft/state.rs index 8e74df43..c197644a 100644 --- a/src/control_plane/consensus/raft/state.rs +++ b/src/control_plane/consensus/raft/state.rs @@ -1,19 +1,22 @@ #![allow(dead_code)] use crate::control_plane::consensus::messages::*; -use crate::control_plane::consensus::raft::catch_up::CatchUpRepairs; use crate::control_plane::consensus::raft::command::RaftCommand; use crate::control_plane::consensus::raft::errors::{EvictionError, ProposalError}; use crate::control_plane::consensus::raft::log::LogEntry; +use crate::control_plane::consensus::raft::states::log_state::LogState; +use crate::control_plane::consensus::raft::states::metadata_state::MetadataState; +use crate::control_plane::consensus::raft::states::transient_state::{ + PeerState, Role, TransientState, +}; use crate::control_plane::consensus::raft::storage::RaftPersistentState; use crate::control_plane::consensus::raft::{compute_replacement_replica_set, now_ms}; use crate::control_plane::membership::{ShardGroupId, TopologyReader}; use crate::control_plane::metadata::command::DeleteSegments; use crate::control_plane::metadata::event::ApplyResult; -use crate::control_plane::metadata::state_machine::MetadataStateMachine; use crate::control_plane::metadata::{ - ConsumerGroupAssignment, ConsumerMemberId, EntryId, MetadataCommand, RangeId, ReassignSegment, - RollSegment, SegmentId, TopicId, TopicMeta, TopicStats, + ConsumerGroupAssignment, ConsumerMemberId, MetadataCommand, ReassignSegment, RollSegment, + TopicId, TopicMeta, TopicStats, }; use crate::control_plane::{NodeId, Replicas}; use crate::data_plane::SegmentKey; @@ -22,10 +25,7 @@ use crate::data_plane::transport::command::DataTransportCommand; use crate::schedulers::ticker_message::TimerCommand; #[cfg(any(test, debug_assertions))] use crate::test_traits::TAssertInvariant; -use std::collections::{BTreeSet, HashMap, HashSet}; -use std::hash::{Hash, Hasher}; - -const ELECTION_JITTER_RANGE: u32 = 20; +use std::collections::{BTreeSet, HashSet}; /// #135: The number of consecutive, identical ring observations needed /// before we can safely evict a live ex-owner. @@ -37,35 +37,6 @@ const ELECTION_JITTER_RANGE: u32 = 20; /// during a rebalance. pub(crate) const RING_STABLE_OBSERVATIONS: u32 = 3; -/// Segments whose write leader crashed (sole death), -pub(crate) type LeaderlessSegments = Vec<(SegmentKey, Vec)>; - -struct ElectionJitter { - seed: u64, - counter: u64, -} - -impl ElectionJitter { - fn new(seed: u64) -> Self { - Self { seed, counter: 0 } - } - - fn next(&mut self) -> u32 { - let mut hasher = std::collections::hash_map::DefaultHasher::new(); - self.seed.hash(&mut hasher); - self.counter.hash(&mut hasher); - self.counter += 1; - (hasher.finish() % ELECTION_JITTER_RANGE as u64) as u32 - } -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum Role { - Follower, - Candidate { votes_received: u32 }, - Leader, -} - // Peer tracking (LEADER-ONLY) // - next_index: Index of the next log entry to send to this peer. // - match_index: Highest log index known to be replicated on this peer. @@ -73,12 +44,6 @@ pub enum Role { // - When a node becomes a leader: it sets next_index = last_log_index +1 and match_index = 0, so initially next_index could be 10 while match_index being 0. // - After rejection: when a follower rejects AppendEntries, the leader decrements next_index to retry with earlier entries - but match_index stays the same until the peer confirms. // So, `next_index` is the leader's guess and `match_index` is confirmed truth -#[derive(Debug, Clone)] -struct PeerState { - next_index: u64, - match_index: u64, -} - // --------------------------------------------------------------------------- // Raft state machine — pure sync, no I/O // --------------------------------------------------------------------------- @@ -89,63 +54,15 @@ struct PeerState { /// All outbound packets and timer commands are buffered and drained by the /// caller (the actor layer). pub struct Raft { - // Identity pub node_id: NodeId, pub shard_group_id: ShardGroupId, - // -- LOG STATE - current_term: u64, - voted_for: Option, - log: Vec, - stabled_index: u64, // flushed to disk - - // -- APPLICATION STATE - state_machine: MetadataStateMachine, - last_applied_index: u64, // applied to state machine - - // Transient State (Volatile Consensus & Coordinator Metadata) - commit_index: u64, // majority voted - role: Role, - /// Tracks who the current leader is — set when this node becomes leader - /// or when a valid `AppendEntries` is received from a leader. - current_leader: Option, - peers: HashSet, - // LEADER-ONLY volatile state - peer_states: HashMap, - /// LEADER-ONLY: non-voting members the leader is catching up before promotion. - /// Replicated to (like peers) but **excluded from the commit quorum** — a node - /// that can't yet participate (e.g. a freshly ring-assigned host without a local - /// instance) is staged here and never counted, so it can't freeze the group. It is - /// promoted to a voting peer via a committed `AddPeer` only once caught up - /// (`maybe_promote_learner`). Re-derived on takeover (leader-volatile). Disjoint - /// from `peers`; empty on followers. - learner_states: HashMap, - /// Election-timer generation. Bumped whenever the election timer is - /// (re)armed or cancelled; a fired `ElectionTimeout` carrying an older - /// epoch raced its own cancellation in flight and must be ignored - election_epoch: u64, - election_jitter: ElectionJitter, - /// Segments whose data-leader has acked its - /// `PlaceSegment`, mapped to the acking node. The heartbeat sweep skips - /// re-driving a segment whose confirmed node still matches `replica_set[0]` - confirmed_placement: HashMap, - /// In-flight sealed-segment repairs. Seeded at `ReassignSegment` apply; the - /// heartbeat sweep re-drives until acked. Leader-volatile — cleared on step-down. - catch_up: CatchUpRepairs, - pending_proposals: Vec, - leaderless_segments: LeaderlessSegments, - /// Tracks the ring membership seen during the last check, along with a count - /// of how many times we've seen this exact membership in a row (#135). - /// - /// This state is "leader-volatile," meaning it gets wiped clean whenever a - /// new leader takes over. This guarantees that a new leader must observe a - /// stable ring over time to "re-earn confidence" before it is allowed to - /// evict anyone. - ring_observation_streak: Option<(BTreeSet, u32)>, - - pending_log_mutations: Vec, // must persist - events: Vec, // volatile side effects + l_stat: LogState, + t_stat: TransientState, + m_stat: MetadataState, + peers: HashSet, + events: Vec, timer_seqs: TimerSeqs, } @@ -169,46 +86,26 @@ impl Raft { node_id, shard_group_id, peers, - stabled_index: persistent.stabled_index(), - current_term: persistent.term, - voted_for: persistent.voted_for, - log: persistent.log, - pending_log_mutations: Vec::new(), + l_stat: LogState::from_persistent(persistent), + m_stat: MetadataState::new(shard_group_id), + t_stat: TransientState::new(election_jitter_seed), events: Vec::new(), - commit_index: 0, - last_applied_index: 0, - role: Role::Follower, - current_leader: None, - state_machine: MetadataStateMachine::new(shard_group_id), - pending_proposals: Vec::new(), - leaderless_segments: Vec::new(), - peer_states: HashMap::new(), - learner_states: HashMap::new(), - election_jitter: ElectionJitter::new(election_jitter_seed), timer_seqs, - election_epoch: 0, - confirmed_placement: HashMap::new(), - catch_up: CatchUpRepairs::default(), - ring_observation_streak: None, }; raft.reset_election_timer(); raft } - pub(crate) fn heartbeat_seq(&self) -> u64 { - self.timer_seqs.rpc - } - pub(crate) fn topic_names(&self) -> Box<[String]> { - self.state_machine.topic_names() + self.m_stat.topic_names() } pub(crate) fn topic_stats(&self) -> Box<[TopicStats]> { - self.state_machine.topic_stats() + self.m_stat.topic_stats() } pub(crate) fn get_topic_by_name(&self, name: &str) -> Option<&TopicMeta> { - self.state_machine.get_topic_by_name(name) + self.m_stat.get_topic_by_name(name) } pub(crate) fn get_consumer_group_assignment( @@ -217,7 +114,7 @@ impl Raft { group_id: &str, member_id: ConsumerMemberId, ) -> Option { - self.state_machine + self.m_stat .get_consumer_group_assignment(topic_name, group_id, member_id) } @@ -225,14 +122,7 @@ impl Raft { &self, node_id: &NodeId, ) -> Box<[(SegmentKey, Replicas)]> { - self.state_machine.active_segments_for_node(node_id) - } - - /// Every active segment's assignment tuple `(key, replica_set, start_offset)`. - /// The leader's confirmation-gated assignment sweep (`MultiRaft::build_redrive_cmds`) - /// turns these into `PlaceSegment` re-drives for unconfirmed segments. - pub(crate) fn active_segment_assignments(&self) -> Box<[(SegmentKey, Replicas, EntryId)]> { - self.state_machine.active_segment_assignments() + self.m_stat.active_segments_for_node(node_id) } /// Full reconciliation against the current topology: assert the ring-assigned @@ -345,7 +235,7 @@ impl Raft { /// for this group: live ex-owners present, or more voters than the ring /// names. Catching the ratchet in the wild precedes curing it. pub(crate) fn log_ring_drift(&self, topology_reader: &TopologyReader, live: &HashSet) { - if self.role != Role::Leader { + if !self.is_leader() { return; } let Some(ring_members) = topology_reader.group_ring_members(self.shard_group_id) else { @@ -375,9 +265,9 @@ impl Raft { /// is proposed during a check, any `RemovePeer` proposed later in that /// same check is delayed until the add successfully commits. fn has_uncommitted_membership_change(&self) -> bool { - (self.commit_index + 1..=self.log_last_index()).any(|i| { + (self.t_stat.commit_index + 1..=self.log_last_index()).any(|i| { matches!( - self.log_get(i).map(|e| &e.command), + self.l_stat.get(i).map(|e| &e.command), Some(RaftCommand::AddPeer(_) | RaftCommand::RemovePeer(_)) ) }) @@ -409,12 +299,12 @@ impl Raft { topology_reader: &TopologyReader, live: &HashSet, ) -> Result<(), EvictionError> { - if self.role != Role::Leader { + if !self.is_leader() { return Err(EvictionError::NotLeader); } let Some(ring_members) = topology_reader.group_ring_members(self.shard_group_id) else { - self.ring_observation_streak = None; + self.t_stat.ring_observation_streak = None; return Err(EvictionError::GroupNotFound); }; let ring: BTreeSet = ring_members.iter().cloned().collect(); @@ -457,10 +347,7 @@ impl Raft { } // It checks if a specific node has a complete, up-to-date copy of all permanently saved data. - let in_sync = self - .peer_states - .get(member) - .is_some_and(|ps| ps.match_index >= self.commit_index); + let in_sync = self.t_stat.is_peer_caught_up(member); if !in_sync { return Err(EvictionError::FollowersLagging); } @@ -487,12 +374,7 @@ impl Raft { #[inline(always)] fn record_ring_observation(&mut self, ring: &BTreeSet) -> u32 { - let observations = match self.ring_observation_streak.take() { - Some((prev, n)) if prev == *ring => n.saturating_add(1), - _ => 1, - }; - self.ring_observation_streak = Some((ring.clone(), observations)); - observations + self.t_stat.record_ring_observation(ring) } /// Repair this group's segments whose replica set still names a dead node: @@ -503,21 +385,25 @@ impl Raft { /// Runs per-death (`handle_node_death`) and as the takeover backfill sweep. pub(crate) fn reconcile_segments(&mut self, live_set: &HashSet) -> bool { let mut to_roll: Vec<(SegmentKey, Replicas)> = Vec::new(); - for (key, rs) in self.active_segments_with_dead_members(live_set).into_vec() { + for (key, rs) in self + .m_stat + .active_segments_with_dead_members(live_set) + .into_vec() + { if Self::only_replica_dead(&rs, live_set) { let survivors = rs .iter() .filter(|n| live_set.contains(*n)) .cloned() .collect(); - self.leaderless_segments.push((key, survivors)); + self.t_stat.leaderless_segments.push((key, survivors)); } else { to_roll.push((key, rs)); } } let mut leaderless_segments = vec![]; - for (key, rs) in self.boundary_unknown_segments() { + for (key, rs) in self.m_stat.boundary_unknown_segments() { let survivors = rs .iter() .filter(|n| live_set.contains(*n)) @@ -525,9 +411,9 @@ impl Raft { .collect(); leaderless_segments.push((key, survivors)); } - self.leaderless_segments.extend(leaderless_segments); + self.t_stat.leaderless_segments.extend(leaderless_segments); - let sealed = self.sealed_segments_with_dead_members(live_set); + let sealed = self.m_stat.sealed_segments_with_dead_members(live_set); let mut changed = false; // Roll the others: seal + reopen with a healthy set. `end_entry_id = None` @@ -593,12 +479,9 @@ impl Raft { /// Re-fill known-end sealed segments left under-replicated by an earlier death fn refill_under_replicated_segments(&mut self, topology: &TopologyReader) -> bool { - let targets: Box<[(SegmentKey, Replicas)]> = self - .state_machine - .topics - .values() - .flat_map(|t| t.under_replicated_sealed_segments(topology.replication_factor())) - .collect(); + let targets = self + .m_stat + .under_replicated_sealed_segments(topology.replication_factor()); let mut changed = false; for (segment_key, mut replica_set) in targets { @@ -637,17 +520,7 @@ impl Raft { /// `DeleteSegments`. Topics with no retention policy contribute nothing. fn reconcile_retention_deletes(&mut self) -> bool { let now = now_ms(); - let targets: Vec<(TopicId, RangeId, Box<[SegmentId]>)> = self - .state_machine - .topics - .values() - .flat_map(|t| { - let topic_id = t.id; - t.expired_segment_prefixes(now) - .into_iter() - .map(move |(range_id, ids)| (topic_id, range_id, ids)) - }) - .collect(); + let targets = self.m_stat.expired_segment_prefixes(now); let mut changed = false; for (topic_id, range_id, segment_ids) in targets { @@ -707,41 +580,6 @@ impl Raft { self.peers_iter().filter(|p| !live.contains(*p)) } - /// Active segments across all topics whose `replica_set` contains at - /// least one non-live member. - pub(crate) fn active_segments_with_dead_members( - &self, - live: &HashSet, - ) -> Box<[(SegmentKey, Replicas)]> { - self.state_machine - .topics - .values() - .flat_map(|t| t.active_segments_with_dead_members(live)) - .collect() - } - - /// Sealed (known-end) segments across all topics whose `replica_set` names a - /// non-live member — D5 sealed-segment repair candidates. - pub(crate) fn sealed_segments_with_dead_members( - &self, - live: &HashSet, - ) -> Box<[(SegmentKey, Replicas)]> { - self.state_machine - .topics - .values() - .flat_map(|t| t.sealed_segments_with_dead_members(live)) - .collect() - } - - pub(crate) fn boundary_unknown_segments( - &self, - ) -> impl Iterator)> { - self.state_machine - .topics - .values() - .flat_map(|t| t.boundary_unknown_segments()) - } - /// A single missing replica is recoverable: the remaining replicas can /// report the all-ack committed minimum before the successor is opened. fn only_replica_dead(replica_set: &[NodeId], live: &HashSet) -> bool { @@ -749,11 +587,11 @@ impl Raft { } pub(crate) fn has_topic(&self, topic_id: &TopicId) -> bool { - self.state_machine.get_topic(topic_id).is_some() + self.m_stat.get_topic(topic_id).is_some() } pub(crate) fn get_replica_set(&self, key: &SegmentKey) -> Option { - let topic = self.state_machine.get_topic(&key.topic_id)?; + let topic = self.m_stat.get_topic(&key.topic_id)?; let range = topic.ranges.get(&key.range_id)?; let seg = range.segments.get(&key.segment_id)?; Some(seg.replica_set.clone()) @@ -763,79 +601,30 @@ impl Raft { std::mem::take(&mut self.events) } + fn raise(&mut self, event: impl Into) { + self.events.push(event.into()); + } + pub fn take_log_mutations(&mut self) -> Vec { - std::mem::take(&mut self.pending_log_mutations) + self.l_stat.take_mutations() } pub(crate) fn take_pending_proposals(&mut self) -> Vec { - std::mem::take(&mut self.pending_proposals) + std::mem::take(&mut self.t_stat.pending_proposals) } /// Drain the leaderless segments found by `reconcile_segments` — the actor /// drives seal-end recovery (poll survivors, seal at the recovered end). pub(crate) fn take_leaderless_segments(&mut self) -> Vec<(SegmentKey, Vec)> { - std::mem::take(&mut self.leaderless_segments) + std::mem::take(&mut self.t_stat.leaderless_segments) } pub(crate) fn last_applied_index(&self) -> u64 { - self.last_applied_index + self.m_stat.last_applied_index } pub(crate) fn log_last_index(&self) -> u64 { - self.log.last().map_or(0, |e| e.index) - } - - fn log_last_term(&self) -> u64 { - self.log.last().map_or(0, |e| e.term) - } - - fn log_term_at(&self, index: u64) -> u64 { - if index == 0 { - return 0; - } - self.log.get((index - 1) as usize).map_or(0, |e| e.term) - } - - fn log_get(&self, index: u64) -> Option<&LogEntry> { - if index == 0 { - return None; - } - self.log.get((index - 1) as usize) - } - - fn log_entries_from(&self, start_index: u64) -> Box<[LogEntry]> { - let last = self.log_last_index(); - if start_index == 0 || start_index > last { - return Box::new([]); - } - self.log[(start_index - 1) as usize..].into() - } - - fn log_append(&mut self, entry: LogEntry) { - debug_assert_eq!( - entry.index, - self.log_last_index() + 1, - "log entry index must be contiguous" - ); - self.pending_log_mutations - .push(LogMutation::Append(entry.clone())); - self.log.push(entry); - } - - fn log_truncate_from(&mut self, from_index: u64) { - if from_index == 0 || from_index > self.log_last_index() + 1 { - return; - } - self.pending_log_mutations - .push(LogMutation::TruncateFrom(from_index)); - self.log.truncate((from_index - 1) as usize); - } - - fn push_hard_state(&mut self) { - self.pending_log_mutations.push(LogMutation::HardState { - term: self.current_term, - voted_for: self.voted_for.clone(), - }); + self.l_stat.last_index() } pub fn peers_count(&self) -> usize { @@ -847,11 +636,11 @@ impl Raft { } pub fn current_leader(&self) -> Option<&NodeId> { - self.current_leader.as_ref() + self.t_stat.current_leader.as_ref() } pub fn is_leader(&self) -> bool { - self.role == Role::Leader + self.t_stat.role == Role::Leader } pub fn has_peer(&self, node_id: &NodeId) -> bool { @@ -861,12 +650,12 @@ impl Raft { /// A node the leader is catching up as a non-voting learner (not yet a voter). #[cfg(test)] pub(crate) fn is_learner(&self, node_id: &NodeId) -> bool { - self.learner_states.contains_key(node_id) + self.t_stat.learner_states.contains_key(node_id) } #[cfg(test)] pub(crate) fn learner_count(&self) -> usize { - self.learner_states.len() + self.t_stat.learner_states.len() } /// Minimum number of nodes needed for a majority (strict majority). @@ -877,7 +666,7 @@ impl Raft { } pub(crate) fn stabled_index(&self) -> u64 { - self.stabled_index + self.l_stat.stabled_index() } // ------------------------------------------------------------------- @@ -891,7 +680,7 @@ impl Raft { // u64::MAX bypasses the staleness check for direct // invocations in tests; real timers carry the epoch they // were armed with. - if epoch == self.election_epoch || epoch == u64::MAX { + if epoch == self.t_stat.election_epoch || epoch == u64::MAX { self.start_election(); return; } @@ -899,7 +688,7 @@ impl Raft { node = %self.node_id, group = self.shard_group_id.0, stale_epoch = epoch, - epoch = self.election_epoch, + epoch = self.t_stat.election_epoch, "election: dropped stale election timeout" ); } @@ -935,12 +724,10 @@ impl Raft { fn start_election(&mut self) { // Leaders don't run election timers — they send heartbeats instead. - if self.role == Role::Leader { + if self.t_stat.role == Role::Leader { return; } - self.current_term += 1; - self.voted_for = Some(self.node_id.clone()); - self.push_hard_state(); + let term = self.l_stat.begin_election(&self.node_id); if self.peers.is_empty() { // Single-node cluster: elect self immediately. @@ -948,20 +735,20 @@ impl Raft { return; } - self.role = Role::Candidate { votes_received: 1 }; // vote for self + self.t_stat.begin_campaign(); self.reset_election_timer(); tracing::trace!( node = %self.node_id, group = self.shard_group_id.0, - term = self.current_term, + term, "election: became candidate, broadcasting RequestVote" ); let req = RequestVote { - term: self.current_term, + term, candidate_id: self.node_id.clone(), last_log_index: self.log_last_index(), - last_log_term: self.log_last_term(), + last_log_term: self.l_stat.last_term(), }; for peer_id in self.peers.iter() { self.events.push( @@ -974,11 +761,11 @@ impl Raft { // ! Followers grant or deny - All roles must respond fn handle_request_vote(&mut self, from: NodeId, req: RequestVote) { // If the request term is newer, step down. - if req.term > self.current_term { + if req.term > self.l_stat.current_term() { self.step_down(req.term); } - let term_ok = req.term == self.current_term; + let term_ok = req.term == self.l_stat.current_term(); let vote_ok = self.vote_available_for(&req.candidate_id); let log_ok = self.log_is_up_to_date(req.last_log_index, req.last_log_term); let vote_granted = term_ok && vote_ok && log_ok; @@ -987,7 +774,7 @@ impl Raft { group = self.shard_group_id.0, from = %req.candidate_id, req_term = req.term, - term = self.current_term, + term = self.l_stat.current_term(), granted = vote_granted, term_ok, vote_ok, @@ -996,23 +783,19 @@ impl Raft { ); if vote_granted { - self.voted_for = Some(req.candidate_id); - self.push_hard_state(); + self.l_stat.grant_vote(req.candidate_id); self.reset_election_timer(); } - self.events.push( - OutboundRaftPacket::new( - self.shard_group_id, - from, - RequestVoteResponse { - term: self.current_term, - node_id: self.node_id.clone(), - vote_granted, - }, - ) - .into(), - ); + self.raise(OutboundRaftPacket::new( + self.shard_group_id, + from, + RequestVoteResponse { + term: self.l_stat.current_term(), + node_id: self.node_id.clone(), + vote_granted, + }, + )); } fn handle_request_vote_response(&mut self, resp: RequestVoteResponse) { @@ -1021,11 +804,11 @@ impl Raft { group = self.shard_group_id.0, from = %resp.node_id, resp_term = resp.term, - term = self.current_term, + term = self.l_stat.current_term(), granted = resp.vote_granted, "election: RequestVoteResponse received" ); - if resp.term > self.current_term { + if resp.term > self.l_stat.current_term() { self.step_down(resp.term); return; } @@ -1033,29 +816,22 @@ impl Raft { } fn count_vote_if_eligible(&mut self, resp: RequestVoteResponse) { - let Role::Candidate { votes_received } = &mut self.role else { - return; - }; - if resp.term != self.current_term || !resp.vote_granted { + if resp.term != self.l_stat.current_term() || !resp.vote_granted { return; } - *votes_received += 1; - if *votes_received >= self.quorum() { + if self.t_stat.record_vote(self.quorum()) { self.become_leader(); } } fn vote_available_for(&self, candidate_id: &NodeId) -> bool { - match &self.voted_for { - None => true, - Some(id) => id == candidate_id, - } + self.l_stat.vote_available_for(candidate_id) } /// §5.4.1: A candidate's log is "at least as up-to-date" if its last /// entry has a higher term, or the same term with a >= index. fn log_is_up_to_date(&self, last_log_index: u64, last_log_term: u64) -> bool { - let my_last_term = self.log_last_term(); + let my_last_term = self.l_stat.last_term(); let my_last_index = self.log_last_index(); if last_log_term != my_last_term { @@ -1069,23 +845,21 @@ impl Raft { // ------------------------------------------------------------------- fn become_leader(&mut self) { - self.role = Role::Leader; - self.current_leader = Some(self.node_id.clone()); + let next = self.log_last_index() + 1; + self.t_stat + .initialize_leader(&self.node_id, &self.peers, next); tracing::debug!( node = %self.node_id, group = self.shard_group_id.0, - term = self.current_term, + term = self.l_stat.current_term(), "election: became leader" ); - self.events.push( - LeaderChange { - shard_group_id: self.shard_group_id, - leader_node_id: self.node_id.clone(), - term: self.current_term, - } - .into(), - ); + self.raise(LeaderChange { + shard_group_id: self.shard_group_id, + leader_node_id: self.node_id.clone(), + term: self.l_stat.current_term(), + }); // Cancel election timer, start heartbeat + merge/ring check timers. self.cancel_all_timers(); @@ -1093,26 +867,6 @@ impl Raft { self.schedule_merge_check_timer(); self.schedule_ring_check_timer(); - // next_index is set *before* the noop is appended, so it points - // at the noop's index — causing the first AppendEntries to carry it. - let next = self.log_last_index() + 1; - - // Peer state tracker needs to be re-initialized on every leadership transition - self.peer_states.clear(); - // Learners are leader-volatile catch-up state; reconcile re-stages them. - self.learner_states.clear(); - // #135: ring confidence is leader-volatile — re-earn it. - self.ring_observation_streak = None; - for peer_id in self.peers.iter() { - self.peer_states.insert( - peer_id.clone(), - PeerState { - next_index: next, - match_index: 0, - }, - ); - } - // Append a Noop entry at the new term so that all // preceding entries from earlier terms can be committed. // Without this, old-term entries remain in limbo until a real @@ -1128,10 +882,10 @@ impl Raft { fn step_down(&mut self, new_term: u64) { debug_assert!( - new_term >= self.current_term, + new_term >= self.l_stat.current_term(), "step_down must never regress the term" ); - let was_leader = self.role == Role::Leader; + let was_leader = self.t_stat.role == Role::Leader; tracing::debug!( node = %self.node_id, @@ -1147,15 +901,9 @@ impl Raft { // rule `recognize_leader`'s demotion branch follows. All production // callers pass strictly newer terms (guarded `>` at every call // site); tests use equal-term step_down to depose a leader in place. - if new_term > self.current_term { - self.current_term = new_term; - self.voted_for = None; - self.push_hard_state(); - } + self.l_stat.advance_term(new_term); self.cancel_leader_timers(); - // #135: a deposed leader's ring observations die with its term. - self.ring_observation_streak = None; if was_leader { // The election timer resets only on a vote grant or // on AppendEntries from the current leader — never on a mere @@ -1165,12 +913,7 @@ impl Raft { // runs no election timer, arms a fresh one. self.reset_election_timer(); } - self.role = Role::Follower; - self.current_leader = None; - self.peer_states.clear(); - self.learner_states.clear(); - self.confirmed_placement.clear(); - self.catch_up.clear(); + self.t_stat.reset_for_follower(); } // ------------------------------------------------------------------- @@ -1178,7 +921,7 @@ impl Raft { // ------------------------------------------------------------------- fn send_heartbeats(&mut self) { - if self.role != Role::Leader { + if !self.is_leader() { return; } @@ -1195,9 +938,10 @@ impl Raft { fn maybe_redrive_segment_assignments(&mut self) { // rederive Metadata <> Datanode segment assignment - let active = self.active_segment_assignments(); + let active = self.m_stat.active_segment_assignments(); let active_keys: HashSet = active.iter().map(|(k, _, _)| *k).collect(); - self.confirmed_placement + self.t_stat + .confirmed_data_leaders .retain(|k, _| active_keys.contains(k)); let mut redrives = Vec::new(); @@ -1207,7 +951,7 @@ impl Raft { }; // * If data leader acks assignment, it would have been added to confirmed_placement through Raft::handle_segment_placed - if self.confirmed_placement.get(&segment_key) == Some(target) { + if self.t_stat.confirmed_data_leaders.get(&segment_key) == Some(target) { continue; } @@ -1222,27 +966,29 @@ impl Raft { )); } if !redrives.is_empty() { - self.events.push(RaftEvent::RedriveAssignments(redrives)); + self.raise(RaftEvent::RedriveAssignments(redrives)); } } pub(crate) fn handle_segment_placed(&mut self, ack: SegmentPlaced) { - self.confirmed_placement.insert(ack.segment_key, ack.from); + self.t_stat + .confirmed_data_leaders + .insert(ack.segment_key, ack.from); } /// Re-drive the catch-up sweep — the sealed-segment analogue of /// `maybe_redrive_segment_assignments`. See `.claude/rules/raft-actor.md` #9. fn maybe_redrive_catch_ups(&mut self) { - let redrives = self.catch_up.redrives(self.shard_group_id); + let redrives = self.t_stat.catch_up.redrives(self.shard_group_id); if !redrives.is_empty() { - self.events.push(RaftEvent::RedriveAssignments(redrives)); + self.raise(RaftEvent::RedriveAssignments(redrives)); } } /// A member confirmed it holds a reassigned sealed segment; routed here from /// `MultiRaft`. pub(crate) fn handle_catch_up_ack(&mut self, ack: SegmentCaughtUp) { - self.catch_up.confirm(ack.segment_key, ack.from); + self.t_stat.catch_up.confirm(ack.segment_key, ack.from); } /// Takeover backstop: re-seed catch-up for every known-end sealed segment this @@ -1251,14 +997,10 @@ impl Raft { /// heartbeat sweep re-drives the re-seeded set (already-complete members /// full-match-ack cheaply). See `.claude/rules/raft-actor.md` #9. pub(crate) fn reseed_catch_up(&mut self) { - let sealed: Vec<_> = self - .state_machine - .topics - .values() - .flat_map(|t| t.known_end_sealed_segments()) - .collect(); + let sealed = self.m_stat.known_end_sealed_segments(); for (segment_key, start, end, replica_set) in sealed { - self.catch_up + self.t_stat + .catch_up .track_sealed(segment_key, start, end, replica_set); } } @@ -1271,16 +1013,17 @@ impl Raft { RaftEvent::ShardLeaderRefresh(LeaderChange { shard_group_id: self.shard_group_id, leader_node_id: self.node_id.clone(), - term: self.current_term, + term: self.l_stat.current_term(), }) }) } /// Everyone the leader replicates to: voting peers plus catching-up learners. fn replication_targets(&self) -> Vec { - self.peer_states + self.t_stat + .peer_states .keys() - .chain(self.learner_states.keys()) + .chain(self.t_stat.learner_states.keys()) .cloned() .collect() } @@ -1291,14 +1034,14 @@ impl Raft { /// that never participates (e.g. a freshly ring-assigned host with no local group /// instance yet) can never freeze the group. Leader-only. pub(crate) fn stage_learner(&mut self, node: NodeId) -> bool { - if self.role != Role::Leader + if !self.is_leader() || node == self.node_id || self.peers.contains(&node) - || self.learner_states.contains_key(&node) + || self.t_stat.learner_states.contains_key(&node) { return false; } - self.learner_states.insert( + self.t_stat.learner_states.insert( node.clone(), PeerState { next_index: self.log_last_index() + 1, @@ -1317,43 +1060,38 @@ impl Raft { if self.has_uncommitted_membership_change() { return; } - let Some(ls) = self.learner_states.get(node) else { - return; - }; - if self.commit_index > 0 && ls.match_index >= self.commit_index { + if self.t_stat.is_learner_ready_for_promotion(node) { let _ = self.propose(RaftCommand::AddPeer(node.clone())); } } fn send_append_entries(&mut self, peer_id: NodeId) { let peer_state = match self + .t_stat .peer_states .get(&peer_id) - .or_else(|| self.learner_states.get(&peer_id)) + .or_else(|| self.t_stat.learner_states.get(&peer_id)) { Some(ps) => ps, None => return, }; let prev_log_index = peer_state.next_index.saturating_sub(1); - let prev_log_term = self.log_term_at(prev_log_index); - let entries = self.log_entries_from(peer_state.next_index); + let prev_log_term = self.l_stat.term_at(prev_log_index); + let entries = self.l_stat.entries_from(peer_state.next_index); - self.events.push( - OutboundRaftPacket::new( - self.shard_group_id, - peer_id, - AppendEntries { - term: self.current_term, - leader_id: self.node_id.clone(), - prev_log_index, - prev_log_term, - entries, - leader_commit: self.commit_index, - }, - ) - .into(), - ); + self.raise(OutboundRaftPacket::new( + self.shard_group_id, + peer_id, + AppendEntries { + term: self.l_stat.current_term(), + leader_id: self.node_id.clone(), + prev_log_index, + prev_log_term, + entries, + leader_commit: self.t_stat.commit_index, + }, + )); } // ! SAFETY: @@ -1361,7 +1099,7 @@ impl Raft { // ! - Candidates step down on same term : because receiving entries while being a candidate means another node already won. // ! - Followers process normally. All roles must respond. fn handle_append_entries(&mut self, from: NodeId, req: AppendEntries) { - if req.term < self.current_term { + if req.term < self.l_stat.current_term() { self.reject_append_entries(from); return; } @@ -1380,19 +1118,19 @@ impl Raft { } fn recognize_leader(&mut self, req: &AppendEntries) { - if req.term > self.current_term { + if req.term > self.l_stat.current_term() { self.step_down(req.term); - } else if self.role != Role::Follower { - self.role = Role::Follower; - self.peer_states.clear(); - self.learner_states.clear(); + } else if self.t_stat.role != Role::Follower { + self.t_stat.role = Role::Follower; + self.t_stat.peer_states.clear(); + self.t_stat.learner_states.clear(); } - self.current_leader = Some(req.leader_id.clone()); + self.t_stat.current_leader = Some(req.leader_id.clone()); self.reset_election_timer(); tracing::debug!( node = %self.node_id, group = self.shard_group_id.0, - term = self.current_term, + term = self.l_stat.current_term(), leader = %req.leader_id, "election: leader recognized" ); @@ -1402,25 +1140,25 @@ impl Raft { if prev_log_index == 0 { return true; } - let local_term = self.log_term_at(prev_log_index); + let local_term = self.l_stat.term_at(prev_log_index); local_term != 0 && local_term == prev_log_term } fn replicate_entries(&mut self, entries: Box<[LogEntry]>) { for entry in entries { - let existing_term = self.log_term_at(entry.index); + let existing_term = self.l_stat.term_at(entry.index); if existing_term != 0 && existing_term != entry.term { - self.log_truncate_from(entry.index); + self.l_stat.truncate_from(entry.index); } if entry.index > self.log_last_index() { - self.log_append(entry); + self.l_stat.append(entry); } } } fn advance_follower_commit(&mut self, leader_commit: u64) { - if leader_commit > self.commit_index { - self.commit_index = leader_commit.min(self.log_last_index()); + if leader_commit > self.t_stat.commit_index { + self.t_stat.commit_index = leader_commit.min(self.log_last_index()); self.apply_committed_entries(); } } @@ -1429,22 +1167,22 @@ impl Raft { // ! - term guard // ! - only leaders track peer state fn handle_append_entries_response(&mut self, resp: AppendEntriesResponse) { - if resp.term > self.current_term { + if resp.term > self.l_stat.current_term() { self.step_down(resp.term); return; } - if self.role != Role::Leader { + if !self.is_leader() { return; } // The responder may be a voting peer or a catching-up learner. let node_id = resp.node_id.clone(); - let is_voter = self.peer_states.contains_key(&node_id); + let is_voter = self.t_stat.peer_states.contains_key(&node_id); let peer_state = if is_voter { - self.peer_states.get_mut(&node_id) + self.t_stat.peer_states.get_mut(&node_id) } else { - self.learner_states.get_mut(&node_id) + self.t_stat.learner_states.get_mut(&node_id) }; if let Some(peer_state) = peer_state { if resp.success { @@ -1472,19 +1210,16 @@ impl Raft { } fn send_append_entries_response(&mut self, target: NodeId, success: bool) { - self.events.push( - OutboundRaftPacket::new( - self.shard_group_id, - target, - AppendEntriesResponse { - term: self.current_term, - node_id: self.node_id.clone(), - success, - last_log_index: self.log_last_index(), - }, - ) - .into(), - ); + self.raise(OutboundRaftPacket::new( + self.shard_group_id, + target, + AppendEntriesResponse { + term: self.l_stat.current_term(), + node_id: self.node_id.clone(), + success, + last_log_index: self.log_last_index(), + }, + )); } // A log entry is committed once it is replicated on a @@ -1509,19 +1244,14 @@ impl Raft { // Scan top-down: the highest current-term entry with quorum // implicitly commits everything below it (log matching property). - for n in (self.commit_index + 1..=last).rev() { - if self.log_term_at(n) != self.current_term { + for n in (self.t_stat.commit_index + 1..=last).rev() { + if self.l_stat.term_at(n) != self.l_stat.current_term() { continue; } - let replication_count = self - .peer_states - .values() - .filter(|ps| ps.match_index >= n) - .count() as u32 - + 1; // +1 for self + let replication_count = self.t_stat.replicated_voter_count(n); if replication_count >= quorum { - self.commit_index = n; + self.t_stat.commit_index = n; self.apply_committed_entries(); return; } @@ -1529,7 +1259,7 @@ impl Raft { } pub(crate) fn advance_stabled_index(&mut self, value: u64) { - self.stabled_index = self.stabled_index.max(value); + self.l_stat.advance_stabled_index(value); self.apply_committed_entries(); #[cfg(any(test, debug_assertions))] @@ -1538,17 +1268,19 @@ impl Raft { #[tracing::instrument(level = "debug", skip_all, fields( group = self.shard_group_id.0, - from = self.last_applied_index + 1, - to = self.commit_index.min(self.stabled_index), + from = self.m_stat.last_applied_index + 1, + to = self.t_stat.commit_index.min(self.l_stat.stabled_index()), ))] fn apply_committed_entries(&mut self) { - while self.last_applied_index < self.commit_index.min(self.stabled_index) { - self.last_applied_index += 1; - let Some(entry) = self.log_get(self.last_applied_index).cloned() else { + while self.m_stat.last_applied_index + < self.t_stat.commit_index.min(self.l_stat.stabled_index()) + { + self.m_stat.last_applied_index += 1; + let Some(entry) = self.l_stat.get(self.m_stat.last_applied_index).cloned() else { tracing::error!( "[{}] committed entry at index {} missing from log", self.node_id, - self.last_applied_index + self.m_stat.last_applied_index ); break; }; @@ -1564,7 +1296,7 @@ impl Raft { } fn apply_metadata_entry(&mut self, cmd: MetadataCommand, index: u64) { - match self.state_machine.apply(cmd) { + match self.m_stat.apply(cmd) { Ok(result) => { tracing::debug!( "[{}] Applied metadata at index {}: {:?}", @@ -1576,17 +1308,14 @@ impl Raft { if self.is_leader() && let ApplyResult::SegmentReassigned(r) = &result { - self.catch_up.track(r); + self.t_stat.catch_up.track(r); } - self.events.push( - MetadataCommitted { - shard_group_id: self.shard_group_id, - result, - log_index: index, - seal_context: None, - } - .into(), - ); + self.raise(MetadataCommitted { + shard_group_id: self.shard_group_id, + result, + log_index: index, + seal_context: None, + }); } Err(e) => tracing::error!( "[{}] Metadata apply error at index {}: {:?}", @@ -1595,9 +1324,10 @@ impl Raft { e ), } - if self.role == Role::Leader { - self.pending_proposals - .extend(self.state_machine.take_pending_proposals()); + if self.t_stat.role == Role::Leader { + self.t_stat + .pending_proposals + .extend(self.m_stat.take_pending_proposals()); } } @@ -1610,13 +1340,13 @@ impl Raft { } // Promotion: a learner graduating to a voter carries its catch-up progress, so // the new voter isn't reset to match_index 0 (which would stall commits anew). - let carried = self.learner_states.remove(&node_id); - if self.peers.insert(node_id.clone()) && self.role == Role::Leader { + let carried = self.t_stat.learner_states.remove(&node_id); + if self.peers.insert(node_id.clone()) && self.t_stat.role == Role::Leader { let state = carried.unwrap_or(PeerState { next_index: self.log_last_index() + 1, match_index: 0, }); - self.peer_states.insert(node_id, state); + self.t_stat.peer_states.insert(node_id, state); } } @@ -1624,20 +1354,20 @@ impl Raft { /// `RemovePeer` log entry commits. Never call directly — the peer set is part /// of the replicated state machine and must only mutate through the log. fn apply_remove_peer(&mut self, node_id: NodeId) { - self.learner_states.remove(&node_id); + self.t_stat.learner_states.remove(&node_id); if self.peers.remove(&node_id) { - self.peer_states.remove(&node_id); - self.events.push(RaftEvent::DisconnectPeer(node_id)); + self.t_stat.peer_states.remove(&node_id); + self.raise(RaftEvent::DisconnectPeer(node_id)); } } fn add_new_entry(&mut self, command: RaftCommand) { let entry = LogEntry { - term: self.current_term, + term: self.l_stat.current_term(), index: self.log_last_index() + 1, command, }; - self.log_append(entry); + self.l_stat.append(entry); } /// Propose a command to the Raft log. Only the leader can accept proposals. @@ -1648,12 +1378,12 @@ impl Raft { // -> Appended to leader's log // -> Replicated to shard #45's followers // -> Majority ack -> committed - // -> Applied to MetadataStateMachine → topic blue exists + // -> Applied to MetadataState → topic blue exists /// Returns the log index at which the command was appended on success. #[tracing::instrument(level = "trace", skip_all, fields(group = self.shard_group_id.0, command = ?command))] pub fn propose(&mut self, command: RaftCommand) -> Result { - if self.role != Role::Leader { - return Err(ProposalError::NotLeader(self.current_leader.clone())); + if !self.is_leader() { + return Err(ProposalError::NotLeader(self.t_stat.current_leader.clone())); } self.add_new_entry(command); @@ -1717,48 +1447,45 @@ impl Raft { } fn reset_election_timer(&mut self) { - self.election_epoch = self.election_epoch.wrapping_add(1); - self.events - .push(RaftEvent::Timer(TimerCommand::CancelSchedule { - seq: self.timer_seqs.election, - })); - let jitter = self.election_jitter.next(); + let election_epoch = self.t_stat.advance_election_epoch(); + + let jitter = self.t_stat.election_jitter.next(); tracing::trace!( node = %self.node_id, group = self.shard_group_id.0, - epoch = self.election_epoch, + epoch = election_epoch, jitter_ticks = jitter, "election: timer armed" ); - self.events - .push(RaftEvent::Timer(TimerCommand::SetSchedule { - seq: self.timer_seqs.election, - timer: RaftTimer::election(jitter, self.shard_group_id, self.election_epoch), - })); + + self.raise(RaftEvent::Timer(TimerCommand::CancelSchedule { + seq: self.timer_seqs.election, + })); + self.raise(RaftEvent::Timer(TimerCommand::SetSchedule { + seq: self.timer_seqs.election, + timer: RaftTimer::election(jitter, self.shard_group_id, election_epoch), + })); } fn schedule_rpc_timer(&mut self) { - self.events - .push(RaftEvent::Timer(TimerCommand::SetSchedule { - seq: self.timer_seqs.rpc, - timer: RaftTimer::rpc(self.shard_group_id), - })); + self.raise(RaftEvent::Timer(TimerCommand::SetSchedule { + seq: self.timer_seqs.rpc, + timer: RaftTimer::rpc(self.shard_group_id), + })); } fn schedule_merge_check_timer(&mut self) { - self.events - .push(RaftEvent::Timer(TimerCommand::SetSchedule { - seq: self.timer_seqs.merge_check, - timer: RaftTimer::merge_check(self.shard_group_id), - })); + self.raise(RaftEvent::Timer(TimerCommand::SetSchedule { + seq: self.timer_seqs.merge_check, + timer: RaftTimer::merge_check(self.shard_group_id), + })); } fn schedule_ring_check_timer(&mut self) { - self.events - .push(RaftEvent::Timer(TimerCommand::SetSchedule { - seq: self.timer_seqs.ring_check, - timer: RaftTimer::ring_check(self.shard_group_id), - })); + self.raise(RaftEvent::Timer(TimerCommand::SetSchedule { + seq: self.timer_seqs.ring_check, + timer: RaftTimer::ring_check(self.shard_group_id), + })); } /// Re-arm the periodic ring check (#135 trigger gap). The ring diff @@ -1767,20 +1494,20 @@ impl Raft { /// the timer lifecycle, leader-gated like `evaluate_merges`: a deposed /// leader lets the timer die (it is re-armed on the next become_leader). pub(crate) fn reschedule_ring_check(&mut self) { - if self.role != Role::Leader { + if !self.is_leader() { return; } self.schedule_ring_check_timer(); } pub(crate) fn evaluate_merges(&mut self, now: u64) { - if self.role != Role::Leader { + if !self.is_leader() { return; } - let merge_proposals = self.state_machine.evaluate_merges(now); + let merge_proposals = self.m_stat.evaluate_merges(now); for cmd in merge_proposals { - self.pending_proposals.push(cmd); + self.t_stat.pending_proposals.push(cmd); } self.schedule_merge_check_timer(); @@ -1791,38 +1518,31 @@ impl Raft { /// follower or candidate keeps its currently armed election deadline /// (#133). fn cancel_leader_timers(&mut self) { - self.events - .push(RaftEvent::Timer(TimerCommand::CancelSchedule { - seq: self.timer_seqs.rpc, - })); - self.events - .push(RaftEvent::Timer(TimerCommand::CancelSchedule { - seq: self.timer_seqs.merge_check, - })); - self.events - .push(RaftEvent::Timer(TimerCommand::CancelSchedule { - seq: self.timer_seqs.ring_check, - })); + self.raise(RaftEvent::Timer(TimerCommand::CancelSchedule { + seq: self.timer_seqs.rpc, + })); + self.raise(RaftEvent::Timer(TimerCommand::CancelSchedule { + seq: self.timer_seqs.merge_check, + })); + self.raise(RaftEvent::Timer(TimerCommand::CancelSchedule { + seq: self.timer_seqs.ring_check, + })); } pub(crate) fn cancel_all_timers(&mut self) { - self.election_epoch = self.election_epoch.wrapping_add(1); - self.events - .push(RaftEvent::Timer(TimerCommand::CancelSchedule { - seq: self.timer_seqs.election, - })); - self.events - .push(RaftEvent::Timer(TimerCommand::CancelSchedule { - seq: self.timer_seqs.rpc, - })); - self.events - .push(RaftEvent::Timer(TimerCommand::CancelSchedule { - seq: self.timer_seqs.merge_check, - })); - self.events - .push(RaftEvent::Timer(TimerCommand::CancelSchedule { - seq: self.timer_seqs.ring_check, - })); + self.t_stat.advance_election_epoch(); + self.raise(RaftEvent::Timer(TimerCommand::CancelSchedule { + seq: self.timer_seqs.election, + })); + self.raise(RaftEvent::Timer(TimerCommand::CancelSchedule { + seq: self.timer_seqs.rpc, + })); + self.raise(RaftEvent::Timer(TimerCommand::CancelSchedule { + seq: self.timer_seqs.merge_check, + })); + self.raise(RaftEvent::Timer(TimerCommand::CancelSchedule { + seq: self.timer_seqs.ring_check, + })); } } @@ -1830,26 +1550,26 @@ impl Raft { impl crate::test_traits::TAssertInvariant for Raft { fn assert_invariants(&self) { assert!( - self.last_applied_index <= self.commit_index, + self.m_stat.last_applied_index <= self.t_stat.commit_index, "last_applied ({}) > commit_index ({})", - self.last_applied_index, - self.commit_index, + self.m_stat.last_applied_index, + self.t_stat.commit_index, ); assert!( - self.last_applied_index <= self.stabled_index, + self.m_stat.last_applied_index <= self.l_stat.stabled_index(), "last_applied ({}) > stabled_index ({}) — applied a non-durable entry", - self.last_applied_index, - self.stabled_index, + self.m_stat.last_applied_index, + self.l_stat.stabled_index(), ); assert!( - self.commit_index <= self.log_last_index(), + self.t_stat.commit_index <= self.log_last_index(), "commit_index ({}) > log_last_index ({})", - self.commit_index, + self.t_stat.commit_index, self.log_last_index(), ); // Log indices are contiguous and 1-based - for (i, entry) in self.log.iter().enumerate() { + for (i, entry) in self.l_stat.entries().iter().enumerate() { assert_eq!( entry.index, (i + 1) as u64, @@ -1857,36 +1577,36 @@ impl crate::test_traits::TAssertInvariant for Raft { entry.index, ); assert!( - entry.term <= self.current_term, + entry.term <= self.l_stat.current_term(), "log entry at index {} has term {} > current_term {}", entry.index, entry.term, - self.current_term, + self.l_stat.current_term(), ); } // Invariant: peer_states exists only on the leader (and matches the peer // set when leader). Followers/candidates carry an empty peer_states. - match self.role { + match self.t_stat.role { Role::Leader => { for peer in &self.peers { assert!( - self.peer_states.contains_key(peer), + self.t_stat.peer_states.contains_key(peer), "leader missing peer_state for {:?}", peer, ); } assert_eq!( - self.peer_states.len(), + self.t_stat.peer_states.len(), self.peers.len(), "leader peer_states size ({}) != peers size ({})", - self.peer_states.len(), + self.t_stat.peer_states.len(), self.peers.len(), ); // Invariant: learners are non-voting and disjoint from voters — a node // is never both — and self is never a learner. Learners are replicated // to but excluded from the commit quorum until promoted via `AddPeer`. - for learner in self.learner_states.keys() { + for learner in self.t_stat.learner_states.keys() { assert!( !self.peers.contains(learner), "node {:?} is both a voter and a learner", @@ -1898,23 +1618,23 @@ impl crate::test_traits::TAssertInvariant for Raft { // the local fragment of this: a leader must have voted for itself this // term (and is therefore the only node that could have won this term). assert_eq!( - self.voted_for.as_ref(), + self.l_stat.voted_for(), Some(&self.node_id), "leader has voted_for {:?}, expected self ({:?})", - self.voted_for, + self.l_stat.voted_for(), self.node_id, ); } Role::Follower | Role::Candidate { .. } => { assert!( - self.peer_states.is_empty(), + self.t_stat.peer_states.is_empty(), "non-leader carries peer_states ({} entries)", - self.peer_states.len(), + self.t_stat.peer_states.len(), ); assert!( - self.learner_states.is_empty(), + self.t_stat.learner_states.is_empty(), "non-leader carries learner_states ({} entries)", - self.learner_states.len(), + self.t_stat.learner_states.len(), ); } } @@ -1947,11 +1667,11 @@ mod tests { } pub(crate) fn current_term(&self) -> u64 { - self.current_term + self.l_stat.current_term() } pub(crate) fn voted_for(&self) -> Option { - self.voted_for.clone() + self.l_stat.voted_for().cloned() } pub(crate) fn simulate_flush_and_apply(&mut self) { @@ -1959,8 +1679,8 @@ mod tests { self.apply_committed_entries(); } - pub(crate) fn state_machine(&self) -> &MetadataStateMachine { - &self.state_machine + pub(crate) fn state_machine(&self) -> &MetadataState { + &self.m_stat } } fn node(id: &str) -> NodeId { @@ -2033,16 +1753,19 @@ mod tests { #[test] fn single_node_elects_self_on_timeout() { let mut raft = single_node_raft(); - assert_eq!(raft.role, Role::Follower); + assert_eq!(raft.t_stat.role, Role::Follower); raft.handle_timeout(RaftTimeoutCallback::ElectionTimeout { shard_group_id: TEST_SHARD, epoch: u64::MAX, }); - assert_eq!(raft.role, Role::Leader); - assert_eq!(raft.current_term, 1); - assert_eq!(raft.voted_for, Some(NodeId::new("node-1"))); + assert_eq!(raft.t_stat.role, Role::Leader); + assert_eq!(raft.l_stat.current_term(), 1); + assert_eq!( + raft.l_stat.voted_for().cloned(), + Some(NodeId::new("node-1")) + ); } #[test] @@ -2052,7 +1775,7 @@ mod tests { shard_group_id: TEST_SHARD, epoch: u64::MAX, }); - assert_eq!(raft.current_term, 1); + assert_eq!(raft.l_stat.current_term(), 1); // Step down and trigger another election raft.step_down(1); @@ -2060,7 +1783,7 @@ mod tests { shard_group_id: TEST_SHARD, epoch: u64::MAX, }); - assert_eq!(raft.current_term, 2); + assert_eq!(raft.l_stat.current_term(), 2); } // ------------------------------------------------------------------- @@ -2076,8 +1799,11 @@ mod tests { epoch: u64::MAX, }); - assert!(matches!(raft.role, Role::Candidate { votes_received: 1 })); - assert_eq!(raft.current_term, 1); + assert!(matches!( + raft.t_stat.role, + Role::Candidate { votes_received: 1 } + )); + assert_eq!(raft.l_stat.current_term(), 1); let out = packets(&mut raft); assert_eq!(out.len(), 2); // one per peer @@ -2107,7 +1833,10 @@ mod tests { } _ => panic!("expected RequestVoteResponse"), } - assert_eq!(raft.voted_for, Some(NodeId::new("node-1"))); + assert_eq!( + raft.l_stat.voted_for().cloned(), + Some(NodeId::new("node-1")) + ); } #[test] @@ -2159,7 +1888,7 @@ mod tests { }; raft.handle_rpc(node("node-2"), resp); - assert_eq!(raft.role, Role::Leader); + assert_eq!(raft.t_stat.role, Role::Leader); } #[test] @@ -2178,8 +1907,8 @@ mod tests { }; raft.handle_rpc(node("node-2"), resp); - assert_eq!(raft.role, Role::Follower); - assert_eq!(raft.current_term, 5); + assert_eq!(raft.t_stat.role, Role::Follower); + assert_eq!(raft.l_stat.current_term(), 5); } // ------------------------------------------------------------------- @@ -2190,7 +1919,7 @@ mod tests { fn rejects_vote_if_candidate_log_is_stale() { let mut raft = three_node_raft("node-2"); // Give node-2 a log entry at term 2 - raft.log_append(LogEntry { + raft.l_stat.append(LogEntry { term: 2, index: 1, command: RaftCommand::Noop, @@ -2309,7 +2038,7 @@ mod tests { #[test] fn follower_rejects_append_entries_with_stale_term() { let mut raft = three_node_raft("node-2"); - raft.current_term = 5; + raft.l_stat.advance_term(5); let ae = AppendEntries { term: 3, @@ -2383,7 +2112,7 @@ mod tests { raft.propose_noop().unwrap(); drain(&mut raft); assert_eq!(raft.log_last_index(), 2); - assert_eq!(raft.commit_index, 0); + assert_eq!(raft.t_stat.commit_index, 0); // node-2 acknowledges both entries (noop + proposal) let resp = AppendEntriesResponse { @@ -2395,7 +2124,7 @@ mod tests { raft.handle_rpc(node("node-2"), resp); // Majority achieved (self + node-2 = 2 out of 3) - assert_eq!(raft.commit_index, 2); + assert_eq!(raft.t_stat.commit_index, 2); } #[test] @@ -2454,7 +2183,7 @@ mod tests { raft.handle_rpc(node("node-1"), ae); drain(&mut raft); - assert_eq!(raft.commit_index, 1); + assert_eq!(raft.t_stat.commit_index, 1); } // ------------------------------------------------------------------- @@ -2523,7 +2252,7 @@ mod tests { }), ); drain(&mut raft); - assert_eq!(raft.role, Role::Leader); + assert_eq!(raft.t_stat.role, Role::Leader); // Receive AppendEntries from a leader with higher term let ae = AppendEntries { @@ -2536,8 +2265,8 @@ mod tests { }; raft.handle_rpc(node("node-3"), ae); - assert_eq!(raft.role, Role::Follower); - assert_eq!(raft.current_term, 3); + assert_eq!(raft.t_stat.role, Role::Follower); + assert_eq!(raft.l_stat.current_term(), 3); } // ------------------------------------------------------------------- @@ -2563,8 +2292,8 @@ mod tests { }), ); drain(&mut raft); - assert_eq!(raft.role, Role::Leader); - assert_eq!(raft.current_term, 1); + assert_eq!(raft.t_stat.role, Role::Leader); + assert_eq!(raft.l_stat.current_term(), 1); // Stale election timeout arrives — should be ignored raft.handle_timeout(RaftTimeoutCallback::ElectionTimeout { @@ -2573,17 +2302,17 @@ mod tests { }); assert_eq!( - raft.role, + raft.t_stat.role, Role::Leader, "leader must not start a new election" ); - assert_eq!(raft.current_term, 1, "term must not increment"); + assert_eq!(raft.l_stat.current_term(), 1, "term must not increment"); } #[test] fn follower_ignores_rpc_timeout() { let mut raft = three_node_raft("node-1"); - assert_eq!(raft.role, Role::Follower); + assert_eq!(raft.t_stat.role, Role::Follower); raft.handle_timeout(RaftTimeoutCallback::RpcTimeout { shard_group_id: TEST_SHARD, @@ -2601,7 +2330,7 @@ mod tests { epoch: u64::MAX, }); drain(&mut raft); - assert!(matches!(raft.role, Role::Candidate { .. })); + assert!(matches!(raft.t_stat.role, Role::Candidate { .. })); raft.handle_timeout(RaftTimeoutCallback::RpcTimeout { shard_group_id: TEST_SHARD, @@ -2640,7 +2369,7 @@ mod tests { ); drain(&mut raft); - assert_eq!(raft.role, Role::Leader); + assert_eq!(raft.t_stat.role, Role::Leader); assert_eq!(raft.current_leader(), Some(&node("node-1"))); } @@ -2930,7 +2659,7 @@ mod tests { let result = raft.propose(cmd.into()); assert!(result.is_ok()); raft.simulate_flush(); - assert!(raft.last_applied_index > 0); + assert!(raft.m_stat.last_applied_index > 0); } // ------------------------------------------------------------------- @@ -3039,7 +2768,7 @@ mod tests { } // ------------------------------------------------------------------- - // Phase 3 — MetadataStateMachine apply + // Phase 3 — MetadataState apply // ------------------------------------------------------------------- use crate::control_plane::metadata::command::{CreateTopic, MetadataCommand}; @@ -3158,7 +2887,7 @@ mod tests { drain(&mut raft); raft.simulate_flush(); assert_eq!(raft.log_last_index(), 1); - assert_eq!(raft.commit_index, 0); + assert_eq!(raft.t_stat.commit_index, 0); // node-1 wins election at term 2 raft.handle_timeout(RaftTimeoutCallback::ElectionTimeout { @@ -3175,12 +2904,12 @@ mod tests { }), ); drain(&mut raft); - assert_eq!(raft.role, Role::Leader); - assert_eq!(raft.current_term, 2); + assert_eq!(raft.t_stat.role, Role::Leader); + assert_eq!(raft.l_stat.current_term(), 2); // become_leader appends a Noop at term 2 (index 2) assert_eq!(raft.log_last_index(), 2); - assert_eq!(raft.log_term_at(1), 1); - assert_eq!(raft.log_term_at(2), 2); + assert_eq!(raft.l_stat.term_at(1), 1); + assert_eq!(raft.l_stat.term_at(2), 2); // node-3 acks only the old term-1 entry (index 1) but not the term-2 entry raft.handle_rpc( @@ -3194,7 +2923,7 @@ mod tests { ); // commit_index must NOT advance — the replicated entry is term 1, not current term assert_eq!( - raft.commit_index, 0, + raft.t_stat.commit_index, 0, "term-1 entry must not be directly committed even with majority" ); @@ -3210,7 +2939,7 @@ mod tests { ); // Now both entries committed (term-2 entry at index 2 has quorum, // implicitly committing the term-1 entry at index 1) - assert_eq!(raft.commit_index, 2); + assert_eq!(raft.t_stat.commit_index, 2); } #[test] @@ -3369,7 +3098,7 @@ mod tests { .next() .expect("three_node_raft must have peers") .clone(); - let term = raft.current_term; + let term = raft.l_stat.current_term(); raft.handle_rpc( peer.clone(), RaftRpc::RequestVoteResponse(RequestVoteResponse { @@ -3379,7 +3108,11 @@ mod tests { }), ); drain(&mut raft); - assert_eq!(raft.role, Role::Leader, "must be leader after election"); + assert_eq!( + raft.t_stat.role, + Role::Leader, + "must be leader after election" + ); raft } @@ -3393,14 +3126,14 @@ mod tests { epoch: u64::MAX, }); drain(&mut raft); - assert_eq!(raft.role, Role::Leader); + assert_eq!(raft.t_stat.role, Role::Leader); raft } /// Proposals in the log after the become_leader noop (index 1). fn proposals_after_become_leader(raft: &Raft) -> Vec { (2..=raft.log_last_index()) - .filter_map(|i| raft.log_get(i).map(|e| e.command.clone())) + .filter_map(|i| raft.l_stat.get(i).map(|e| e.command.clone())) .collect() } @@ -3552,7 +3285,11 @@ mod tests { ); } drain(&mut raft); - assert_eq!(raft.role, Role::Leader, "must be leader after election"); + assert_eq!( + raft.t_stat.role, + Role::Leader, + "must be leader after election" + ); let members = group.replicas.clone(); (raft, reader, members) @@ -3952,7 +3689,7 @@ mod tests { epoch: u64::MAX, }); drain(&mut raft); - assert_eq!(raft.role, Role::Leader); + assert_eq!(raft.t_stat.role, Role::Leader); raft } @@ -4167,8 +3904,8 @@ mod tests { fn catch_up_repairs_dropped_on_step_down() { let (mut raft, _) = raft_with_seeded_catch_up(vec![node("node-1"), node("y")]); // A higher term deposes the leader → leader-volatile tracker is cleared. - raft.step_down(raft.current_term + 1); - assert!(raft.catch_up.is_empty()); + raft.step_down(raft.l_stat.current_term() + 1); + assert!(raft.t_stat.catch_up.is_empty()); assert!(drain_catch_up_redrives(&mut raft).is_empty()); } @@ -4179,7 +3916,7 @@ mod tests { // Simulate the takeover gap: the leader-volatile tracker is empty, but the // sealed segment is still under-replicated in the state machine. - raft.catch_up.clear(); + raft.t_stat.catch_up.clear(); assert!(drain_catch_up_redrives(&mut raft).is_empty()); raft.reseed_catch_up(); From e2ce2ab5f9eace349a6f0e24c0f989c9543b6f45 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Fri, 17 Jul 2026 22:16:41 +0400 Subject: [PATCH 6/9] rename --- src/control_plane/consensus/raft/state.rs | 2 +- src/control_plane/metadata/topic.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/control_plane/consensus/raft/state.rs b/src/control_plane/consensus/raft/state.rs index c197644a..f5bbd14d 100644 --- a/src/control_plane/consensus/raft/state.rs +++ b/src/control_plane/consensus/raft/state.rs @@ -520,7 +520,7 @@ impl Raft { /// `DeleteSegments`. Topics with no retention policy contribute nothing. fn reconcile_retention_deletes(&mut self) -> bool { let now = now_ms(); - let targets = self.m_stat.expired_segment_prefixes(now); + let targets = self.m_stat.expipred_segments(now); let mut changed = false; for (topic_id, range_id, segment_ids) in targets { diff --git a/src/control_plane/metadata/topic.rs b/src/control_plane/metadata/topic.rs index 83e36403..a5f66446 100644 --- a/src/control_plane/metadata/topic.rs +++ b/src/control_plane/metadata/topic.rs @@ -323,7 +323,7 @@ impl TopicMeta { /// Retention (D7): per-range oldest-first prefixes of sealed segments expired /// under this topic's policy as of `now`. Empty when the topic has no retention /// set (`retention_ms = None`, the default — keep everything). - pub(crate) fn expired_segment_prefixes(&self, now: u64) -> Vec<(RangeId, Box<[SegmentId]>)> { + pub(crate) fn expired_segments(&self, now: u64) -> Vec<(RangeId, Box<[SegmentId]>)> { let Some(retention_ms) = self.storage_policy.retention_ms else { return Vec::new(); }; From 0035721706c750c79883ef04177bcf63930fc317 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Fri, 17 Jul 2026 22:20:27 +0400 Subject: [PATCH 7/9] metadata/transient stats --- .../consensus/raft/states/metadata_state.rs | 1619 +++++++++++++++++ .../consensus/raft/states/transient_state.rs | 186 ++ src/control_plane/metadata/topic.rs | 2 +- 3 files changed, 1806 insertions(+), 1 deletion(-) create mode 100644 src/control_plane/consensus/raft/states/metadata_state.rs create mode 100644 src/control_plane/consensus/raft/states/transient_state.rs diff --git a/src/control_plane/consensus/raft/states/metadata_state.rs b/src/control_plane/consensus/raft/states/metadata_state.rs new file mode 100644 index 00000000..0b14dc2c --- /dev/null +++ b/src/control_plane/consensus/raft/states/metadata_state.rs @@ -0,0 +1,1619 @@ +use crate::control_plane::metadata::command::*; +use crate::control_plane::metadata::event::*; + +use crate::control_plane::NodeId; +use crate::control_plane::Replicas; +use crate::control_plane::membership::ShardGroupId; +use crate::control_plane::metadata::ConsumerGroupAssignment; +use crate::control_plane::metadata::topic::{TopicMeta, TopicState, TopicStats}; +use crate::control_plane::metadata::{EntryId, RangeId, SegmentId, TopicId, error::MetadataError}; +use crate::data_plane::SegmentKey; +#[cfg(any(test, debug_assertions))] +use crate::test_traits::TAssertInvariant; +use MetadataError::*; +use std::collections::HashMap; +use uuid::Uuid; + +pub struct MetadataState { + pub(crate) topics: HashMap, + pub(crate) last_applied_index: u64, + topic_name_index: HashMap, + next_topic_id: u64, + pending_proposals: Vec, +} + +impl MetadataState { + pub(crate) fn new(shard_group_id: ShardGroupId) -> Self { + MetadataState { + topics: HashMap::new(), + last_applied_index: 0, + topic_name_index: HashMap::new(), + next_topic_id: shard_group_id.0 << 32, + pending_proposals: Vec::new(), + } + } + + pub(crate) fn get_topic(&self, id: &TopicId) -> Option<&TopicMeta> { + self.topics.get(id) + } + + pub(crate) fn get_topic_by_name(&self, name: &str) -> Option<&TopicMeta> { + self.topic_name_index + .get(name) + .and_then(|id| self.topics.get(id)) + } + + pub(crate) fn get_consumer_group_assignment( + &self, + topic_name: &str, + group_id: &str, + member_id: Uuid, + ) -> Option { + Some( + self.get_topic_by_name(topic_name)? + .consumer_groups + .get(group_id)? + .assignment_for(member_id), + ) + } + + pub(crate) fn topic_names(&self) -> Box<[String]> { + self.topic_name_index.keys().cloned().collect() + } + + pub(crate) fn topic_stats(&self) -> Box<[TopicStats]> { + self.topics.values().map(|t| t.stats()).collect() + } + + pub(crate) fn take_pending_proposals(&mut self) -> Box<[MetadataCommand]> { + std::mem::take(&mut self.pending_proposals).into_boxed_slice() + } + + pub(crate) fn active_segments_for_node( + &self, + node_id: &NodeId, + ) -> Box<[(SegmentKey, Replicas)]> { + self.topics + .values() + .flat_map(|t| t.active_segments_for_node(node_id)) + .collect() + } + + /// Every active segment across all topics with its replica set and start + /// offset, for the leader's periodic assignment re-drive. + pub(crate) fn active_segment_assignments(&self) -> Box<[(SegmentKey, Replicas, EntryId)]> { + self.topics + .values() + .flat_map(|t| t.active_segment_assignments()) + .collect() + } + + /// Active segments across all topics whose `replica_set` contains at + /// least one non-live member. + pub(crate) fn active_segments_with_dead_members( + &self, + live: &std::collections::HashSet, + ) -> Box<[(SegmentKey, Replicas)]> { + self.topics + .values() + .flat_map(|topic| topic.active_segments_with_dead_members(live)) + .collect() + } + + pub(crate) fn sealed_segments_with_dead_members( + &self, + live: &std::collections::HashSet, + ) -> Box<[(SegmentKey, Replicas)]> { + self.topics + .values() + .flat_map(|topic| topic.sealed_segments_with_dead_members(live)) + .collect() + } + + pub(crate) fn boundary_unknown_segments(&self) -> Box<[(SegmentKey, Vec)]> { + self.topics + .values() + .flat_map(TopicMeta::boundary_unknown_segments) + .collect() + } + + pub(crate) fn under_replicated_sealed_segments( + &self, + replication_factor: usize, + ) -> Box<[(SegmentKey, Replicas)]> { + self.topics + .values() + .flat_map(|topic| topic.under_replicated_sealed_segments(replication_factor)) + .collect() + } + + pub(crate) fn known_end_sealed_segments( + &self, + ) -> Vec<(SegmentKey, EntryId, EntryId, Replicas)> { + self.topics + .values() + .flat_map(TopicMeta::known_end_sealed_segments) + .collect() + } + + pub(crate) fn expipred_segments(&self, now: u64) -> Vec<(TopicId, RangeId, Box<[SegmentId]>)> { + self.topics + .values() + .flat_map(|topic| { + let topic_id = topic.id; + topic + .expired_segments(now) + .into_iter() + .map(move |(range_id, ids)| (topic_id, range_id, ids)) + }) + .collect() + } + + pub(crate) fn apply(&mut self, command: MetadataCommand) -> Result { + use MetadataCommand::*; + let result = match command { + CreateTopic(cmd) => self.create_topic(cmd)?.into(), + RollSegment(cmd) => self.roll_segment(cmd)?, + SplitRange(cmd) => self.split_range(cmd)?.into(), + MergeRange(cmd) => self.merge_range(cmd)?.into(), + DeleteTopic(cmd) => self.delete_topic(cmd)?.into(), + ReassignSegment(cmd) => self.reassign_segment(cmd)?, + DeleteSegments(cmd) => self.delete_segments(cmd)?, + SyncConsumerGroup(cmd) => self + .sync_consumer_group(cmd)? + .map(ApplyResult::ConsumerGroupChanged) + .unwrap_or(ApplyResult::Noop), + }; + #[cfg(any(test, debug_assertions))] + self.assert_invariants(); + Ok(result) + } + + fn create_topic(&mut self, cmd: CreateTopic) -> Result { + if self.topic_name_index.contains_key(&cmd.name) { + return Err(TopicNameAlreadyExists(cmd.name)); + } + let topic_id = TopicId(self.next_topic_id); + let replica_set = cmd.replica_set.clone(); + let topic = TopicMeta::new( + cmd.name, + topic_id, + cmd.replica_set, + cmd.created_at, + cmd.storage_policy, + ); + self.topic_name_index.insert(topic.name.clone(), topic_id); + self.topics.insert(topic.id, topic); + self.next_topic_id += 1; + Ok(TopicCreated { + segment_key: SegmentKey::new(topic_id, RangeId(0), SegmentId(0)), + replica_set, + }) + } + + fn roll_segment(&mut self, cmd: RollSegment) -> Result { + let topic = self.get_active_topic_mut(cmd.segment_key.topic_id)?; + let can_split = topic.can_split(); + let range = topic.get_range_mut(&cmd.segment_key.range_id)?; + + let is_active = range.active_segment == Some(cmd.segment_key.segment_id); + + // If Inactive, correction path + if !is_active { + let Some(end_entry_id) = cmd.end_entry_id else { + return Ok(ApplyResult::Noop); + }; + let Some(replica_set) = + range.correct_end_offset(cmd.segment_key.segment_id, end_entry_id) + else { + return Ok(ApplyResult::Noop); + }; + + return Ok(ApplyResult::SegmentSealCorrected(SegmentSealCorrected { + segment_key: cmd.segment_key, + replica_set, + committed_entry_id: cmd.end_entry_id, + })); + } + + // If Active, Roll + let new_segment_id = range.roll_segment(cmd.clone())?; + let split_proposal = (range.should_split(cmd.sealed_at) && can_split) + .then(|| range.build_split_proposal(&cmd)); + + // For segment roll, unless data nodes got changed, consumer + // TODO For now, it loops over EVERY consumger groups and take epoch snapshot for every topic. + // TODO To reduce the load, it should specifically target groups that are affected by the possible range split + let consumer_group_epochs = topic + .consumer_groups + .keys() + .filter_map(|group_id| topic.consumer_group_epoch(group_id)) + .collect(); + tracing::debug!("Consumer groups: {:?}", consumer_group_epochs); + + if let Some(proposal) = split_proposal { + match proposal { + Ok(proposal) => self + .pending_proposals + .push(MetadataCommand::SplitRange(proposal)), + Err(error) => tracing::debug!( + "Split proposal skipped for range {:?}: {:?}", + cmd.segment_key.range_id, + error + ), + } + } + Ok(SegmentRolled { + new_segment_key: cmd.segment_key.with_segment_id(new_segment_id), + new_replica_set: cmd.new_replica_set, + end_entry_id: cmd.end_entry_id, + consumer_group_epochs, + } + .into()) + } + + /// Re-points a sealed segment's replica set. + /// `Noop` when the set is unchanged, otherwise `SegmentReassigned`. + /// The segment must be sealed; an active/deleting/unknown one is rejected. + fn reassign_segment(&mut self, cmd: ReassignSegment) -> Result { + let segment = self + .topics + .get_mut(&cmd.segment_key.topic_id) + .ok_or(TopicNotFound(cmd.segment_key.topic_id))? + .get_mut(cmd.segment_key)?; + + // The dispatch announces the desired replica set; receivers reconcile, so + // we only carry the sealed bounds (the catch-up target) alongside it. + if segment.reassign(cmd.replica_set.clone())? { + Ok(SegmentReassigned { + segment_key: cmd.segment_key, + start_entry_id: segment.start_entry_id, + sealed_end: segment.end_entry_id, + new_replica_set: cmd.replica_set, + } + .into()) + } else { + Ok(ApplyResult::Noop) + } + } + + /// Retention: mark an oldest-first prefix of a range's sealed segments `Deleting`. + /// `Noop` when nothing transitions (all named ids already `Deleting`/absent), else + /// `SegmentsDeleted` carrying the deleted segments grouped by `replica_set` for + /// batched dispatch. Uses plain `get_mut` (not `validate_active`) so a command + /// applied after the topic is being deleted is a harmless no-op (already `Deleting`). + fn delete_segments(&mut self, cmd: DeleteSegments) -> Result { + let range = self + .topics + .get_mut(&cmd.topic_id) + .ok_or(TopicNotFound(cmd.topic_id))? + .get_range_mut(&cmd.range_id)?; + + let deleted_ids = range.delete_segments(&cmd.segment_ids); + if deleted_ids.is_empty() { + return Ok(ApplyResult::Noop); + } + // Group the deleted segments by replica_set here. + let mut groups: Vec<(Replicas, Vec)> = Vec::new(); + for sid in &deleted_ids { + let Some(seg) = range.segments.get(sid) else { + continue; + }; + let key = SegmentKey::new(cmd.topic_id, cmd.range_id, *sid); + match groups.iter_mut().find(|(rs, _)| rs == &seg.replica_set) { + Some((_, keys)) => keys.push(key), + None => groups.push((seg.replica_set.clone(), vec![key])), + } + } + Ok(SegmentsDeleted { groups }.into()) + } + + fn get_active_topic_mut(&mut self, id: TopicId) -> Result<&mut TopicMeta, MetadataError> { + let topic = self.topics.get_mut(&id).ok_or(TopicNotFound(id))?; + topic.validate_active()?; + Ok(topic) + } + + fn split_range(&mut self, cmd: SplitRange) -> Result { + let topic = self + .topics + .get_mut(&cmd.topic_id) + .ok_or(TopicNotFound(cmd.topic_id))?; + topic.validate_active()?; + if !topic.can_split() { + return Err(SplitNotAllowed(cmd.topic_id)); + } + + let parent_range = topic + .ranges + .get(&cmd.range_id) + .ok_or(MetadataError::RangeNotFound)?; + + let parent_active_segment = parent_range.active_segment.and_then(|seg_id| { + let seg = parent_range.segments.get(&seg_id)?; + Some(( + SegmentKey::new(cmd.topic_id, cmd.range_id, seg_id), + seg.replica_set.clone(), + )) + }); + + let (left_id, right_id) = topic.execute_split(cmd.clone())?; + + let consumer_group_epochs = topic.rebalance_consumer_groups(); + + Ok(RangeSplit { + topic_id: cmd.topic_id, + children: [ + (left_id, SegmentId(0), cmd.left_replica_set), + (right_id, SegmentId(0), cmd.right_replica_set), + ], + parent_active_segment, + consumer_group_epochs, + }) + } + + fn merge_range(&mut self, cmd: MergeRange) -> Result { + let topic_id = cmd.topic_id; + let replica_set = cmd.merged_replica_set.clone(); + let topic = self + .topics + .get_mut(&topic_id) + .ok_or(TopicNotFound(topic_id))?; + topic.validate_active()?; + let merged_id = topic.execute_merge(cmd)?; + let consumer_group_epochs = topic.rebalance_consumer_groups(); + + Ok(RangeMerged { + segment_key: SegmentKey::new(topic_id, merged_id, SegmentId(0)), + replica_set, + consumer_group_epochs, + }) + } + + fn sync_consumer_group( + &mut self, + cmd: SyncConsumerGroup, + ) -> Result, MetadataError> { + let group_id = cmd.group_id.clone(); + + let topic = self + .topic_name_index + .get(&cmd.topic_name) + .and_then(|topic_id| self.topics.get_mut(topic_id)) + .ok_or_else(|| TopicNameNotFound(cmd.topic_name.clone()))?; + topic.validate_active()?; + if !topic.sync_consumer_group(cmd) { + return Ok(None); + } + Ok(topic.consumer_group_epoch(&group_id)) + } + + // ! SAFETY: When the loop evaluates the [1, 2] pair and decides it is mergeable, + // ! it generates the event and immediately stops looking at the rest of that topic's ranges + // ! + // ! Caution on race condition. Take the following example: + // ! 1. the following method proposes merging (1, 2). + // ! 2. The MergeRange command goes into a queue (or a Raft log) to be processed. + // ! 3. Before the command is executed, range 2 receives a massive burst of traffic and splits into 2A and 2B. + // ! 4. The MergeRange(1, 2) command is finally executed by your merge function. + // ! By the time the command executes, range 2 might be split, already sealed, or completely deleted + // * This is already safe as the 'stale' proposal fails gracefully at apply time, following "stale proposals are safe" invariant + pub(crate) fn evaluate_merges(&self, now: u64) -> Vec { + self.topics + .values() + .filter_map(|topic| topic.find_mergeable_range_pair(now)) + .collect() + } + + fn delete_topic(&mut self, cmd: DeleteTopic) -> Result { + let topic_id = self + .topic_name_index + .get(&cmd.name) + .copied() + .ok_or(MetadataError::TopicNameNotFound(cmd.name.clone()))?; + + // Safety: topic_name_index and topics are always in sync — + // see invariant below. + let topic = self.topics.get_mut(&topic_id).unwrap(); + topic.delete(); + + let consumer_group_epochs = topic.rebalance_consumer_groups(); + + self.topic_name_index.remove(&cmd.name); + + Ok(TopicDeleted { + consumer_group_epochs, + }) + } + + #[cfg(test)] + pub fn topic_count(&self) -> usize { + self.topics.len() + } +} + +#[cfg(any(test, debug_assertions))] +impl crate::test_traits::TAssertInvariant for MetadataState { + fn assert_invariants(&self) { + assert_eq!( + self.topic_name_index.len(), + self.topics + .values() + .filter(|t| t.state != TopicState::Deleted) + .count(), + "topic_name_index out of sync with non-deleted topics" + ); + for (name, id) in &self.topic_name_index { + let topic = self + .topics + .get(id) + .expect("name index points to missing topic"); + assert_eq!(&topic.name, name); + } + + for id in self.topics.keys() { + assert!(id.0 < self.next_topic_id, "topic ID >= next_topic_id"); + } + for topic in self.topics.values() { + topic.assert_invariants(); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::connections::protocol::ConsumerGroupSyncAction; + use crate::control_plane::membership::ShardGroupId; + use crate::control_plane::metadata::constants::*; + use crate::control_plane::metadata::range::*; + use crate::control_plane::metadata::segment::*; + use crate::control_plane::{ + NodeId, + metadata::{ + SegmentId, + strategy::{PartitionStrategy, StoragePolicy}, + }, + }; + use std::collections::VecDeque; + + fn default_policy() -> StoragePolicy { + StoragePolicy { + retention_ms: Some(3_600_000), + replication_factor: 3, + partition_strategy: PartitionStrategy::AutoSplit, + } + } + + fn fixed_policy() -> StoragePolicy { + StoragePolicy { + retention_ms: Some(3_600_000), + replication_factor: 3, + partition_strategy: PartitionStrategy::Fixed, + } + } + + fn replica_set() -> Replicas { + Replicas::new(vec![ + NodeId::new("node-1"), + NodeId::new("node-2"), + NodeId::new("node-3"), + ]) + } + + fn create_topic(sm: &mut MetadataState, name: &str) -> TopicId { + let result = sm.apply(MetadataCommand::CreateTopic(CreateTopic { + name: name.to_string(), + storage_policy: default_policy(), + replica_set: replica_set(), + created_at: 1000, + })); + match result.unwrap() { + ApplyResult::TopicCreated(tc) => tc.segment_key.topic_id, + other => panic!("expected TopicCreated, got {:?}", other), + } + } + + #[test] + fn consumer_group_generation_is_applied_through_metadata_log_command() { + let mut sm = MetadataState::new(ShardGroupId(1)); + create_topic(&mut sm, "orders"); + let member = uuid::Uuid::new_v4(); + let command = SyncConsumerGroup { + req: SyncConsumerGroupRequest { + topic_name: "orders".into(), + group_id: "workers".into(), + member_id: member, + action: ConsumerGroupSyncAction::Heartbeat, + }, + observed_at: 100, + session_timeout_ms: 10_000, + }; + + let ApplyResult::ConsumerGroupChanged(epoch) = sm.apply(command.clone().into()).unwrap() + else { + panic!("first member must create a committed generation"); + }; + assert_eq!(*epoch.generation, 1); + assert_eq!(epoch.ranges.len(), 1); + let assignment = sm + .get_consumer_group_assignment("orders", "workers", member) + .unwrap(); + assert_eq!(*assignment.generation, 1); + assert_eq!(assignment.ranges.as_ref(), &[RangeId(0)]); + + let mut heartbeat = command; + heartbeat.observed_at = 200; + assert!(matches!( + sm.apply(heartbeat.into()).unwrap(), + ApplyResult::Noop + )); + assert_eq!( + *sm.get_consumer_group_assignment("orders", "workers", member) + .unwrap() + .generation, + 1 + ); + } + + fn roll_segment( + sm: &mut MetadataState, + topic_id: TopicId, + range_id: RangeId, + segment_id: SegmentId, + sealed_at: u64, + ) { + let result = sm.apply(MetadataCommand::RollSegment(RollSegment { + segment_key: SegmentKey::new(topic_id, range_id, segment_id), + sealed_at, + new_replica_set: replica_set(), + end_entry_id: None, + })); + assert!(matches!(result.unwrap(), ApplyResult::SegmentRolled(_))); + } + + // ── D7 retention ─────────────────────────────────────────────────────── + + /// Roll the active segment, sealing it at `end_entry_id` / `sealed_at` so the + /// resulting sealed segment has a known end and seal time (unlike the death-roll + /// `roll_segment` helper above which seals with `None`). + fn roll_with_end( + sm: &mut MetadataState, + topic_id: TopicId, + segment_id: SegmentId, + end_entry_id: u64, + sealed_at: u64, + ) { + let result = sm.apply(MetadataCommand::RollSegment(RollSegment { + segment_key: SegmentKey::new(topic_id, RangeId(0), segment_id), + sealed_at, + new_replica_set: replica_set(), + end_entry_id: Some(EntryId(end_entry_id)), + })); + assert!(matches!(result.unwrap(), ApplyResult::SegmentRolled(_))); + } + + fn seg_state(sm: &MetadataState, topic_id: TopicId, segment_id: SegmentId) -> SegmentMetaState { + sm.get_topic(&topic_id).unwrap().ranges[&RangeId(0)].segments[&segment_id] + .state + .clone() + } + + /// Build a topic with sealed segments 0,1,2 (ends 9/19/29, sealed at 100/200/300) + /// and an active head 3. + fn topic_with_three_sealed(sm: &mut MetadataState) -> TopicId { + let t = create_topic(sm, "t"); + roll_with_end(sm, t, SegmentId(0), 9, 100); + roll_with_end(sm, t, SegmentId(1), 19, 200); + roll_with_end(sm, t, SegmentId(2), 29, 300); + t + } + + fn delete_segments( + sm: &mut MetadataState, + topic_id: TopicId, + ids: &[u64], + ) -> Result { + sm.apply(MetadataCommand::DeleteSegments(DeleteSegments { + topic_id, + range_id: RangeId(0), + segment_ids: ids.iter().map(|&i| SegmentId(i)).collect(), + })) + } + + #[test] + fn delete_segments_marks_oldest_prefix_deleting() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let t = topic_with_three_sealed(&mut sm); + + let result = delete_segments(&mut sm, t, &[0, 1]).unwrap(); + let ApplyResult::SegmentsDeleted(d) = result else { + panic!("expected SegmentsDeleted, got {result:?}"); + }; + // All three sealed segments share one replica_set → a single group of 2 keys. + assert_eq!(d.groups.len(), 1); + assert_eq!(d.groups[0].1.len(), 2); + assert_eq!(seg_state(&sm, t, SegmentId(0)), SegmentMetaState::Deleting); + assert_eq!(seg_state(&sm, t, SegmentId(1)), SegmentMetaState::Deleting); + assert_eq!(seg_state(&sm, t, SegmentId(2)), SegmentMetaState::Sealed); + assert_eq!(seg_state(&sm, t, SegmentId(3)), SegmentMetaState::Active); + } + + #[test] + fn delete_segments_skips_the_active_head() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let t = topic_with_three_sealed(&mut sm); + // Naming the active head (seg 3) alongside the sealed prefix: only the sealed + // ones transition; the write head is skipped, never deleted. + let ApplyResult::SegmentsDeleted(d) = delete_segments(&mut sm, t, &[0, 1, 2, 3]).unwrap() + else { + panic!("expected SegmentsDeleted"); + }; + let total_keys: usize = d.groups.iter().map(|(_, keys)| keys.len()).sum(); + assert_eq!(total_keys, 3); + assert_eq!(seg_state(&sm, t, SegmentId(3)), SegmentMetaState::Active); + } + + /// The no-hole property is a structural invariant, not a hot-path check: a + /// non-prefix deletion (seg 1 while seg 0 survives) trips `assert_retention_prefix`. + #[test] + #[should_panic(expected = "not an oldest-first prefix")] + fn delete_segments_non_prefix_trips_invariant() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let t = topic_with_three_sealed(&mut sm); + let _ = delete_segments(&mut sm, t, &[1]); + } + + #[test] + fn delete_segments_is_idempotent() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let t = topic_with_three_sealed(&mut sm); + assert!(matches!( + delete_segments(&mut sm, t, &[0]).unwrap(), + ApplyResult::SegmentsDeleted(_) + )); + // Re-applying for an already-Deleting segment is a no-op. + assert!(matches!( + delete_segments(&mut sm, t, &[0]).unwrap(), + ApplyResult::Noop + )); + } + + #[test] + fn expired_prefix_selects_by_age_oldest_first() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let t = topic_with_three_sealed(&mut sm); // sealed_at 100/200/300, retention 3_600_000 + let topic = sm.get_topic(&t).unwrap(); + + // now such that segs 0,1 are past the window but seg 2 isn't. + let now = 200 + 3_600_000 + 1; + let prefixes = topic.expired_segments(now); + assert_eq!(prefixes.len(), 1); + let (range_id, ids) = &prefixes[0]; + assert_eq!(*range_id, RangeId(0)); + assert_eq!(ids.as_ref(), &[SegmentId(0), SegmentId(1)]); + } + + #[test] + fn expired_prefix_empty_without_retention() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let t = sm + .apply(MetadataCommand::CreateTopic(CreateTopic { + name: "no-retention".into(), + storage_policy: StoragePolicy { + retention_ms: None, + replication_factor: 3, + partition_strategy: PartitionStrategy::AutoSplit, + }, + replica_set: replica_set(), + created_at: 1000, + })) + .map(|r| match r { + ApplyResult::TopicCreated(tc) => tc.segment_key.topic_id, + other => panic!("{other:?}"), + }) + .unwrap(); + roll_with_end(&mut sm, t, SegmentId(0), 9, 100); + // Far past any window, but no policy → nothing expires. + assert!( + sm.get_topic(&t) + .unwrap() + .expired_segments(u64::MAX) + .is_empty() + ); + } + + fn split_range( + sm: &mut MetadataState, + topic_id: TopicId, + range_id: RangeId, + split_point: Vec, + created_at: u64, + ) -> (RangeId, RangeId) { + let result = sm.apply(MetadataCommand::SplitRange(SplitRange { + topic_id, + range_id, + split_point, + created_at, + left_replica_set: replica_set(), + right_replica_set: replica_set(), + })); + match result.unwrap() { + ApplyResult::RangeSplit(rs) => (rs.children[0].0, rs.children[1].0), + other => panic!("expected RangeSplit, got {:?}", other), + } + } + + fn merge_range( + sm: &mut MetadataState, + topic_id: TopicId, + range_id_1: RangeId, + range_id_2: RangeId, + created_at: u64, + ) -> RangeId { + let result = sm.apply(MetadataCommand::MergeRange(MergeRange { + topic_id, + range_id_1, + range_id_2, + created_at, + merged_replica_set: replica_set(), + })); + match result.unwrap() { + ApplyResult::RangeMerged(rm) => rm.segment_key.range_id, + other => panic!("expected RangeMerged, got {:?}", other), + } + } + + /// A surviving subset plus a fresh replacement — what the coordinator picks + /// when a replica of a sealed segment dies (node-3 → node-4 here). + fn replacement_set() -> Replicas { + Replicas::new(vec![ + NodeId::new("node-1"), + NodeId::new("node-2"), + NodeId::new("node-4"), + ]) + } + + // --- ReassignSegment --- + + #[test] + fn reassign_swaps_a_sealed_segments_replica_set() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let topic_id = create_topic(&mut sm, "blue"); + // Roll so SegmentId(0) becomes Sealed (SegmentId(1) is the new write head). + roll_segment(&mut sm, topic_id, RangeId(0), SegmentId(0), 2000); + let sealed = SegmentKey::new(topic_id, RangeId(0), SegmentId(0)); + + let result = sm + .apply(MetadataCommand::ReassignSegment(ReassignSegment { + segment_key: sealed, + replica_set: replacement_set(), + })) + .unwrap(); + + match result { + ApplyResult::SegmentReassigned(r) => { + assert_eq!(r.segment_key, sealed); + assert_eq!(r.new_replica_set, replacement_set()); + } + other => panic!("expected SegmentReassigned, got {other:?}"), + } + + let seg = &sm.get_topic(&topic_id).unwrap().ranges[&RangeId(0)].segments[&SegmentId(0)]; + assert_eq!(seg.replica_set, replacement_set()); + assert_eq!(seg.state, SegmentMetaState::Sealed); // stays sealed + } + + #[test] + fn reassign_same_set_is_a_noop() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let topic_id = create_topic(&mut sm, "blue"); + roll_segment(&mut sm, topic_id, RangeId(0), SegmentId(0), 2000); + let sealed = SegmentKey::new(topic_id, RangeId(0), SegmentId(0)); + + sm.apply(MetadataCommand::ReassignSegment(ReassignSegment { + segment_key: sealed, + replica_set: replacement_set(), + })) + .unwrap(); + + // Re-applying the identical set (duplicate death detection / re-proposal) + // changes nothing. + let again = sm + .apply(MetadataCommand::ReassignSegment(ReassignSegment { + segment_key: sealed, + replica_set: replacement_set(), + })) + .unwrap(); + assert_eq!(again, ApplyResult::Noop); + } + + #[test] + fn reassign_rejects_an_active_segment() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let topic_id = create_topic(&mut sm, "blue"); + // SegmentId(0) is the active write head — no roll yet. + let active = SegmentKey::new(topic_id, RangeId(0), SegmentId(0)); + + let result = sm.apply(MetadataCommand::ReassignSegment(ReassignSegment { + segment_key: active, + replica_set: replacement_set(), + })); + assert!(matches!(result, Err(MetadataError::SegmentNotSealed))); + } + + #[test] + fn reassign_rejects_an_unknown_segment() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let topic_id = create_topic(&mut sm, "blue"); + let unknown = SegmentKey::new(topic_id, RangeId(0), SegmentId(99)); + + let result = sm.apply(MetadataCommand::ReassignSegment(ReassignSegment { + segment_key: unknown, + replica_set: replacement_set(), + })); + assert!(matches!(result, Err(MetadataError::SegmentNotFound))); + } + + // --- CreateTopic --- + + #[test] + fn create_topic_basic() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let id = create_topic(&mut sm, "blue"); + + let topic = sm.get_topic(&id).unwrap(); + assert_eq!(topic.name, "blue"); + assert_eq!(topic.state, TopicState::Active); + assert_eq!(topic.active_ranges.len(), 1); + assert_eq!(topic.ranges.len(), 1); + + let range = &topic.ranges[&RangeId(0)]; + assert_eq!(range.state, RangeState::Active); + assert!(range.active_segment.is_some()); + assert_eq!(range.segments.len(), 1); + } + + #[test] + fn create_topic_duplicate_name_rejected() { + let mut sm = MetadataState::new(ShardGroupId(0)); + create_topic(&mut sm, "blue"); + + let result = sm.apply(MetadataCommand::CreateTopic(CreateTopic { + name: "blue".to_string(), + storage_policy: default_policy(), + replica_set: replica_set(), + created_at: 2000, + })); + assert_eq!(result, Err(TopicNameAlreadyExists("blue".to_string()))); + } + + #[test] + fn create_topic_increments_id() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let id1 = create_topic(&mut sm, "alpha"); + let id2 = create_topic(&mut sm, "beta"); + + assert_eq!(id1, TopicId(0)); + assert_eq!(id2, TopicId(1)); + assert_eq!(sm.topic_count(), 2); + } + + #[test] + fn create_topic_initial_offsets() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let id = create_topic(&mut sm, "blue"); + + let topic = sm.get_topic(&id).unwrap(); + let range = &topic.ranges[&RangeId(0)]; + assert_eq!(range.next_offset, EntryId(0)); + + let seg = &range.segments[&SegmentId(0)]; + assert_eq!(seg.start_entry_id, EntryId(0)); + assert_eq!(seg.end_entry_id, None); + } + + #[test] + fn create_topic_name_index() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let id = create_topic(&mut sm, "blue"); + + let found = sm.get_topic_by_name("blue").unwrap(); + assert_eq!(found.id, id); + assert!(sm.get_topic_by_name("red").is_none()); + } + + // --- RollSegment --- + + #[test] + fn roll_segment_creates_next() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + + roll_segment(&mut sm, tid, RangeId(0), SegmentId(0), 2000); + + let range = &sm.get_topic(&tid).unwrap().ranges[&RangeId(0)]; + assert_eq!(range.active_segment, Some(SegmentId(1))); + assert_eq!(range.segments.len(), 2); + } + + #[test] + fn roll_segment_increments_segment_id() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + + roll_segment(&mut sm, tid, RangeId(0), SegmentId(0), 2000); + roll_segment(&mut sm, tid, RangeId(0), SegmentId(1), 3000); + + let range = &sm.get_topic(&tid).unwrap().ranges[&RangeId(0)]; + assert_eq!(range.active_segment, Some(SegmentId(2))); + assert_eq!(range.segments.len(), 3); + assert_eq!(range.next_segment_id, 3); + } + + #[test] + fn roll_segment_bad_topic() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let result = sm.apply(MetadataCommand::RollSegment(RollSegment { + segment_key: SegmentKey::new(TopicId(99), RangeId(0), SegmentId(0)), + sealed_at: 2000, + new_replica_set: replica_set(), + end_entry_id: None, + })); + assert_eq!(result, Err(TopicNotFound(TopicId(99)))); + } + + #[test] + fn roll_segment_bad_range() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + + let result = sm.apply(MetadataCommand::RollSegment(RollSegment { + segment_key: SegmentKey::new(tid, RangeId(99), SegmentId(0)), + sealed_at: 2000, + new_replica_set: replica_set(), + end_entry_id: None, + })); + assert_eq!(result, Err(RangeNotFound)); + } + + #[test] + fn roll_segment_stale_is_rejected() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + + let result = sm.apply(MetadataCommand::RollSegment(RollSegment { + segment_key: SegmentKey::new(tid, RangeId(0), SegmentId(99)), + sealed_at: 2000, + new_replica_set: replica_set(), + end_entry_id: None, + })); + assert_eq!(result, Ok(ApplyResult::Noop)); + + let range = &sm.get_topic(&tid).unwrap().ranges[&RangeId(0)]; + assert_eq!(range.active_segment, Some(SegmentId(0))); + assert_eq!(range.segments.len(), 1); + } + + // --- SplitRange --- + + #[test] + fn split_range_basic() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + + let (c1, c2) = split_range(&mut sm, tid, RangeId(0), vec![0x80], 2000); + + let topic = sm.get_topic(&tid).unwrap(); + assert_eq!(topic.ranges.len(), 3); + + let child1 = &topic.ranges[&c1]; + assert_eq!(child1.keyspace_start, KEYSPACE_MIN); + assert_eq!(child1.keyspace_end, vec![0x80]); + assert_eq!(child1.state, RangeState::Active); + + let child2 = &topic.ranges[&c2]; + assert_eq!(child2.keyspace_start, vec![0x80]); + assert_eq!(child2.keyspace_end, KEYSPACE_MAX); + assert_eq!(child2.state, RangeState::Active); + } + + #[test] + fn split_range_updates_active_ranges() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + + let (c1, c2) = split_range(&mut sm, tid, RangeId(0), vec![0x80], 2000); + + let topic = sm.get_topic(&tid).unwrap(); + assert_eq!(topic.active_ranges, vec![c1, c2]); + assert!(!topic.active_ranges.contains(&RangeId(0))); + } + + #[test] + fn split_range_lineage() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + + let (c1, c2) = split_range(&mut sm, tid, RangeId(0), vec![0x80], 2000); + + let topic = sm.get_topic(&tid).unwrap(); + assert_eq!(topic.ranges[&RangeId(0)].split_into, Some([c1, c2])); + } + + #[test] + fn split_range_fixed_rejected() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let result = sm.apply(MetadataCommand::CreateTopic(CreateTopic { + name: "ordered".to_string(), + storage_policy: fixed_policy(), + replica_set: replica_set(), + created_at: 1000, + })); + let tid = match result.unwrap() { + ApplyResult::TopicCreated(tc) => tc.segment_key.topic_id, + other => panic!("expected TopicCreated, got {:?}", other), + }; + + let split_result = sm.apply(MetadataCommand::SplitRange(SplitRange { + topic_id: tid, + range_id: RangeId(0), + split_point: vec![0x80], + created_at: 2000, + left_replica_set: replica_set(), + right_replica_set: replica_set(), + })); + assert_eq!(split_result, Err(SplitNotAllowed(tid))); + } + + #[test] + fn split_range_invalid_split_point() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + + let upper_bound = sm.apply(MetadataCommand::SplitRange(SplitRange { + topic_id: tid, + range_id: RangeId(0), + split_point: vec![0xFF], + created_at: 2000, + left_replica_set: replica_set(), + right_replica_set: replica_set(), + })); + assert_eq!(upper_bound, Err(InvalidSplitPoint)); + + let lower_bound = sm.apply(MetadataCommand::SplitRange(SplitRange { + topic_id: tid, + range_id: RangeId(0), + split_point: vec![], + created_at: 2000, + left_replica_set: replica_set(), + right_replica_set: replica_set(), + })); + assert_eq!(lower_bound, Err(InvalidSplitPoint)); + } + + #[test] + fn split_range_keyspace_coverage() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + + let (c1, c2) = split_range(&mut sm, tid, RangeId(0), vec![0x80], 2000); + + let topic = sm.get_topic(&tid).unwrap(); + let r1 = &topic.ranges[&c1]; + let r2 = &topic.ranges[&c2]; + + assert_eq!(r1.keyspace_end, r2.keyspace_start); + assert_eq!(r1.keyspace_start, KEYSPACE_MIN); + assert_eq!(r2.keyspace_end, KEYSPACE_MAX); + } + + // --- MergeRange --- + + #[test] + fn merge_range_basic() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + let (c1, c2) = split_range(&mut sm, tid, RangeId(0), vec![0x80], 2000); + + let merged_id = merge_range(&mut sm, tid, c1, c2, 3000); + + let topic = sm.get_topic(&tid).unwrap(); + let merged = &topic.ranges[&merged_id]; + assert_eq!(merged.keyspace_start, KEYSPACE_MIN); + assert_eq!(merged.keyspace_end, KEYSPACE_MAX); + assert_eq!(merged.state, RangeState::Active); + } + + #[test] + fn merge_range_active_ranges_updated() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + let (c1, c2) = split_range(&mut sm, tid, RangeId(0), vec![0x80], 2000); + + let merged_id = merge_range(&mut sm, tid, c1, c2, 3000); + + let topic = sm.get_topic(&tid).unwrap(); + assert_eq!(topic.active_ranges, vec![merged_id]); + } + + #[test] + fn merge_range_lineage() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + let (c1, c2) = split_range(&mut sm, tid, RangeId(0), vec![0x80], 2000); + + let merged_id = merge_range(&mut sm, tid, c1, c2, 3000); + + let topic = sm.get_topic(&tid).unwrap(); + assert_eq!(topic.ranges[&c1].merged_into, Some(merged_id)); + assert_eq!(topic.ranges[&c2].merged_into, Some(merged_id)); + assert_eq!(topic.ranges[&merged_id].merged_from, Some([c1, c2])); + } + + #[test] + fn merge_range_non_adjacent_rejected() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + let (c1, c2) = split_range(&mut sm, tid, RangeId(0), vec![0x80], 2000); + let (c1a, _c1b) = split_range(&mut sm, tid, c1, vec![0x40], 3000); + + let result = sm.apply(MetadataCommand::MergeRange(MergeRange { + topic_id: tid, + range_id_1: c1a, + range_id_2: c2, + created_at: 4000, + merged_replica_set: replica_set(), + })); + assert_eq!(result, Err(RangesNotAdjacent)); + } + + // --- DeleteTopic --- + + #[test] + fn delete_topic_cascades() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + + sm.apply(MetadataCommand::DeleteTopic(DeleteTopic { + name: "blue".into(), + })) + .unwrap(); + + let topic = sm.get_topic(&tid).unwrap(); + assert_eq!(topic.state, TopicState::Deleted); + assert!(topic.active_ranges.is_empty()); + + for range in topic.ranges.values() { + assert_eq!(range.state, RangeState::Deleting); + for seg in range.segments.values() { + assert_eq!(seg.state, SegmentMetaState::Deleting); + } + } + } + + #[test] + fn delete_topic_removes_name_index() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let _tid = create_topic(&mut sm, "blue"); + + sm.apply(MetadataCommand::DeleteTopic(DeleteTopic { + name: "blue".into(), + })) + .unwrap(); + + assert!(sm.get_topic_by_name("blue").is_none()); + } + + #[test] + fn delete_topic_nonexistent() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let result = sm.apply(MetadataCommand::DeleteTopic(DeleteTopic { + name: "nope".into(), + })); + assert_eq!(result, Err(MetadataError::TopicNameNotFound("nope".into()))); + } + + // --- Integration --- + + #[test] + fn create_split_seal() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + let (c1, _c2) = split_range(&mut sm, tid, RangeId(0), vec![0x80], 2000); + + roll_segment(&mut sm, tid, c1, SegmentId(0), 3000); + + let child = &sm.get_topic(&tid).unwrap().ranges[&c1]; + assert_eq!(child.active_segment, Some(SegmentId(1))); + assert_eq!(child.segments.len(), 2); + } + + #[test] + fn split_merge_roundtrip() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + let (c1, c2) = split_range(&mut sm, tid, RangeId(0), vec![0x80], 2000); + + let merged = merge_range(&mut sm, tid, c1, c2, 3000); + + let topic = sm.get_topic(&tid).unwrap(); + let range = &topic.ranges[&merged]; + assert_eq!(range.keyspace_start, KEYSPACE_MIN); + assert_eq!(range.keyspace_end, KEYSPACE_MAX); + assert_eq!(topic.active_ranges, vec![merged]); + } + + #[test] + fn multiple_topics_independent() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let t1 = create_topic(&mut sm, "alpha"); + let t2 = create_topic(&mut sm, "beta"); + + split_range(&mut sm, t1, RangeId(0), vec![0x80], 2000); + + let alpha = sm.get_topic(&t1).unwrap(); + assert_eq!(alpha.active_ranges.len(), 2); + + let beta = sm.get_topic(&t2).unwrap(); + assert_eq!(beta.active_ranges.len(), 1); + assert_eq!(beta.ranges.len(), 1); + } + + // --- Invariant: Segment immutability after seal --- + + #[test] + fn roll_already_rolled_segment_is_stale() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + + roll_segment(&mut sm, tid, RangeId(0), SegmentId(0), 2000); + + let result = sm.apply(MetadataCommand::RollSegment(RollSegment { + segment_key: SegmentKey::new(tid, RangeId(0), SegmentId(0)), + sealed_at: 3000, + new_replica_set: replica_set(), + end_entry_id: None, + })); + assert_eq!(result, Ok(ApplyResult::Noop)); + + let range = &sm.get_topic(&tid).unwrap().ranges[&RangeId(0)]; + assert_eq!(range.active_segment, Some(SegmentId(1))); + assert_eq!(range.segments.len(), 2); + } + + // --- Hot Range Detection --- + + #[test] + fn seal_history_records_timestamps() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + + roll_segment(&mut sm, tid, RangeId(0), SegmentId(0), 1000); + roll_segment(&mut sm, tid, RangeId(0), SegmentId(1), 2000); + roll_segment(&mut sm, tid, RangeId(0), SegmentId(2), 3000); + + let range = &sm.get_topic(&tid).unwrap().ranges[&RangeId(0)]; + assert_eq!(range.seal_history.seal_count(), 3); + } + + #[test] + fn seal_history_prunes_old_entries() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + + roll_segment(&mut sm, tid, RangeId(0), SegmentId(0), 1000); + roll_segment(&mut sm, tid, RangeId(0), SegmentId(1), 2000); + // Jump far beyond the window — both old entries pruned + let far_future = 2000 + MEASUREMENT_WINDOW_MS + 1; + roll_segment(&mut sm, tid, RangeId(0), SegmentId(2), far_future); + + let range = &sm.get_topic(&tid).unwrap().ranges[&RangeId(0)]; + assert_eq!(range.seal_history.seal_count(), 1); + } + + #[test] + fn seal_history_orders_delayed_proposal_timestamps() { + let mut history = RangeSealHistory::default(); + history.record_seal(2000); + history.record_seal(1000); + history.record_seal(1500); + + assert_eq!(history.seal_timestamps, VecDeque::from([1000, 1500, 2000])); + } + + #[test] + fn delayed_old_seal_is_pruned_against_newest_timestamp() { + let mut history = RangeSealHistory::default(); + let newest = MEASUREMENT_WINDOW_MS + 2000; + history.record_seal(newest); + history.record_seal(1000); + + assert_eq!(history.seal_timestamps, VecDeque::from([newest])); + } + + #[test] + fn should_split_threshold_met() { + let mut history = RangeSealHistory::default(); + history.record_seal(1000); + history.record_seal(2000); + history.record_seal(3000); + + assert!(history.should_split(3000)); + } + + #[test] + fn should_split_below_threshold() { + let mut history = RangeSealHistory::default(); + history.record_seal(1000); + history.record_seal(2000); + + assert!(!history.should_split(2000)); + } + + #[test] + fn should_split_cooldown_blocks() { + let mut history = RangeSealHistory { + seal_timestamps: VecDeque::new(), + created_by_split_at: Some(1000), + }; + history.record_seal(1100); + history.record_seal(1200); + history.record_seal(1300); + + // Within cooldown — blocked + assert!(!history.should_split(1300)); + + // After cooldown — allowed + assert!(history.should_split(1000 + SPLIT_COOLDOWN_MS)); + } + + #[test] + fn auto_proposal_on_hot_range() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + + for i in 0..SPLIT_SEAL_THRESHOLD { + roll_segment( + &mut sm, + tid, + RangeId(0), + SegmentId(i as u64), + 1000 * (i as u64 + 1), + ); + } + + let proposals = sm.take_pending_proposals(); + assert_eq!(proposals.len(), 1); + assert!(matches!(proposals[0], MetadataCommand::SplitRange(_))); + } + + #[test] + fn no_auto_proposal_below_threshold() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + + for i in 0..(SPLIT_SEAL_THRESHOLD - 1) { + roll_segment( + &mut sm, + tid, + RangeId(0), + SegmentId(i as u64), + 1000 * (i as u64 + 1), + ); + } + + let proposals = sm.take_pending_proposals(); + assert!(proposals.is_empty()); + } + + #[test] + fn no_auto_proposal_for_fixed_strategy() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let result = sm.apply(MetadataCommand::CreateTopic(CreateTopic { + name: "ordered".to_string(), + storage_policy: fixed_policy(), + replica_set: replica_set(), + created_at: 1000, + })); + let tid = match result.unwrap() { + ApplyResult::TopicCreated(tc) => tc.segment_key.topic_id, + other => panic!("expected TopicCreated, got {:?}", other), + }; + + for i in 0..SPLIT_SEAL_THRESHOLD { + roll_segment( + &mut sm, + tid, + RangeId(0), + SegmentId(i as u64), + 2000 * (i as u64 + 1), + ); + } + + let proposals = sm.take_pending_proposals(); + assert!(proposals.is_empty()); + } + + #[test] + fn evaluate_merges_cold_adjacent() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + split_range(&mut sm, tid, RangeId(0), vec![0x80], 2000); + + // Both children are cold (no seals) + let proposals = sm.evaluate_merges(2000 + SPLIT_COOLDOWN_MS + 1); + assert_eq!(proposals.len(), 1); + assert!(matches!(proposals[0], MetadataCommand::MergeRange(_))); + } + + #[test] + fn evaluate_merges_one_hot() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + let (c1, _c2) = split_range(&mut sm, tid, RangeId(0), vec![0x80], 2000); + sm.take_pending_proposals(); // discard any split proposals + + // Seal one child — makes it hot + roll_segment(&mut sm, tid, c1, SegmentId(0), 3000); + + let proposals = sm.evaluate_merges(3000); + assert!(proposals.is_empty()); + } + + #[test] + fn split_clears_seal_history() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + + roll_segment(&mut sm, tid, RangeId(0), SegmentId(0), 2000); + let (c1, c2) = split_range(&mut sm, tid, RangeId(0), vec![0x80], 3000); + + let topic = sm.get_topic(&tid).unwrap(); + assert!(topic.ranges[&c1].seal_history.seal_timestamps.is_empty()); + assert!(topic.ranges[&c2].seal_history.seal_timestamps.is_empty()); + } + + #[test] + fn split_sets_cooldown() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + let (c1, c2) = split_range(&mut sm, tid, RangeId(0), vec![0x80], 2000); + + let topic = sm.get_topic(&tid).unwrap(); + assert_eq!( + topic.ranges[&c1].seal_history.created_by_split_at, + Some(2000) + ); + assert_eq!( + topic.ranges[&c2].seal_history.created_by_split_at, + Some(2000) + ); + } + + // --- D3: active_segments_for_node --- + + #[test] + fn active_segments_for_node_returns_matching() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + + let segments = sm.active_segments_for_node(&NodeId::new("node-1")); + assert_eq!(segments.len(), 1); + let (key, rs) = &segments[0]; + assert_eq!(key.topic_id, tid); + assert_eq!(key.range_id, RangeId(0)); + assert_eq!(key.segment_id, SegmentId(0)); + assert!(rs.contains(&NodeId::new("node-1"))); + } + + #[test] + fn active_segments_for_node_excludes_non_member() { + let mut sm = MetadataState::new(ShardGroupId(0)); + create_topic(&mut sm, "blue"); + + let segments = sm.active_segments_for_node(&NodeId::new("node-99")); + assert!(segments.is_empty()); + } + + #[test] + fn active_segments_for_node_excludes_sealed() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + roll_segment(&mut sm, tid, RangeId(0), SegmentId(0), 2000); + + let segments = sm.active_segments_for_node(&NodeId::new("node-1")); + assert_eq!(segments.len(), 1); + assert_eq!(segments[0].0.segment_id, SegmentId(1)); + } + + // --- D3: end_entry_id in RollSegment --- + + #[test] + fn roll_segment_uses_end_entry_id() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + + let result = sm.apply(MetadataCommand::RollSegment(RollSegment { + segment_key: SegmentKey::new(tid, RangeId(0), SegmentId(0)), + sealed_at: 2000, + new_replica_set: replica_set(), + end_entry_id: Some(EntryId(42000)), + })); + + let result = result.unwrap(); + assert!(matches!( + result, + ApplyResult::SegmentRolled(SegmentRolled { + end_entry_id: Some(EntryId(42000)), + .. + }) + )); + + let range = &sm.get_topic(&tid).unwrap().ranges[&RangeId(0)]; + let sealed = &range.segments[&SegmentId(0)]; + assert_eq!(sealed.end_entry_id, Some(EntryId(42000))); + let new_seg = &range.segments[&SegmentId(1)]; + assert_eq!(new_seg.start_entry_id, EntryId(42001)); + } + + // --- D3: end-offset correction --- + + #[test] + fn end_offset_correction_updates_placeholder() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + + // Death-triggered roll with end_entry_id=0 (placeholder) + let _ = sm.apply(MetadataCommand::RollSegment(RollSegment { + segment_key: SegmentKey::new(tid, RangeId(0), SegmentId(0)), + sealed_at: 2000, + new_replica_set: replica_set(), + end_entry_id: None, + })); + + assert_eq!( + sm.get_topic(&tid).unwrap().ranges[&RangeId(0)].segments[&SegmentId(0)].end_entry_id, + None + ); + + // Segment leader's RollSegment arrives with correct end_entry_id + let result = sm.apply(MetadataCommand::RollSegment(RollSegment { + segment_key: SegmentKey::new(tid, RangeId(0), SegmentId(0)), + sealed_at: 2500, + new_replica_set: replica_set(), + end_entry_id: Some(EntryId(42000)), + })); + assert!(result.is_ok()); + + let range = &sm.get_topic(&tid).unwrap().ranges[&RangeId(0)]; + assert_eq!( + range.segments[&SegmentId(0)].end_entry_id, + Some(EntryId(42000)) + ); + assert_eq!(range.segments[&SegmentId(1)].start_entry_id, EntryId(42001)); + } + + #[test] + fn end_offset_correction_rejected_when_already_set() { + let mut sm = MetadataState::new(ShardGroupId(0)); + let tid = create_topic(&mut sm, "blue"); + + // Normal roll with actual end_entry_id + let _ = sm.apply(MetadataCommand::RollSegment(RollSegment { + segment_key: SegmentKey::new(tid, RangeId(0), SegmentId(0)), + sealed_at: 2000, + new_replica_set: replica_set(), + end_entry_id: Some(EntryId(1000)), + })); + + // Duplicate roll is rejected (end_offset already set) + let result = sm.apply(MetadataCommand::RollSegment(RollSegment { + segment_key: SegmentKey::new(tid, RangeId(0), SegmentId(0)), + sealed_at: 2500, + new_replica_set: replica_set(), + end_entry_id: Some(EntryId(42000)), + })); + assert_eq!(result, Ok(ApplyResult::Noop)); + } +} diff --git a/src/control_plane/consensus/raft/states/transient_state.rs b/src/control_plane/consensus/raft/states/transient_state.rs new file mode 100644 index 00000000..9341045e --- /dev/null +++ b/src/control_plane/consensus/raft/states/transient_state.rs @@ -0,0 +1,186 @@ +use crate::control_plane::NodeId; +use crate::control_plane::consensus::raft::catch_up::CatchUpRepairs; +use crate::control_plane::metadata::MetadataCommand; +use crate::data_plane::SegmentKey; +use std::collections::{BTreeSet, HashMap, HashSet}; +use std::hash::{Hash, Hasher}; + +const ELECTION_JITTER_RANGE: u32 = 20; + +/// Segments whose write leader crashed (sole death). +pub(crate) type LeaderlessSegments = Vec<(SegmentKey, Vec)>; + +pub(crate) struct ElectionJitter { + seed: u64, + counter: u64, +} + +impl ElectionJitter { + fn new(seed: u64) -> Self { + Self { seed, counter: 0 } + } + + pub(crate) fn next(&mut self) -> u32 { + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + self.seed.hash(&mut hasher); + self.counter.hash(&mut hasher); + self.counter += 1; + (hasher.finish() % ELECTION_JITTER_RANGE as u64) as u32 + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Role { + Follower, + Candidate { votes_received: u32 }, + Leader, +} + +// Peer tracking is leader-only. `next_index` is the leader's guess; +// `match_index` is confirmed replication progress. +#[derive(Debug, Clone)] +pub(crate) struct PeerState { + pub(crate) next_index: u64, + pub(crate) match_index: u64, +} + +pub(crate) struct TransientState { + pub(crate) commit_index: u64, + pub(crate) role: Role, + pub(crate) current_leader: Option, + pub(crate) peer_states: HashMap, + pub(crate) learner_states: HashMap, + pub(crate) election_epoch: u64, + pub(crate) election_jitter: ElectionJitter, + + // SegmentKey -> data leader's NodeId: Control plane only actively manages and demands ACKs from the segment's designated Data Leader + pub(crate) confirmed_data_leaders: HashMap, + pub(crate) catch_up: CatchUpRepairs, + pub(crate) pending_proposals: Vec, + pub(crate) leaderless_segments: LeaderlessSegments, + pub(crate) ring_observation_streak: Option<(BTreeSet, u32)>, +} + +impl TransientState { + pub(crate) fn new(election_jitter_seed: u64) -> Self { + Self { + commit_index: 0, + role: Role::Follower, + current_leader: None, + peer_states: HashMap::new(), + learner_states: HashMap::new(), + election_epoch: 0, + election_jitter: ElectionJitter::new(election_jitter_seed), + confirmed_data_leaders: HashMap::new(), + catch_up: CatchUpRepairs::default(), + pending_proposals: Vec::new(), + leaderless_segments: Vec::new(), + ring_observation_streak: None, + } + } + + pub(crate) fn initialize_leader( + &mut self, + node_id: &NodeId, + peers: &HashSet, + next_index: u64, + ) { + self.role = Role::Leader; + self.current_leader = Some(node_id.clone()); + self.peer_states.clear(); + self.learner_states.clear(); + self.ring_observation_streak = None; + self.peer_states.extend(peers.iter().cloned().map(|peer| { + ( + peer, + PeerState { + next_index, + match_index: 0, + }, + ) + })); + } + + pub(crate) fn begin_campaign(&mut self) { + self.role = Role::Candidate { votes_received: 1 }; + } + + pub(crate) fn record_vote(&mut self, quorum: u32) -> bool { + let Role::Candidate { votes_received } = &mut self.role else { + return false; + }; + *votes_received += 1; + *votes_received >= quorum + } + + pub(crate) fn reset_for_follower(&mut self) { + self.role = Role::Follower; + self.current_leader = None; + self.peer_states.clear(); + self.learner_states.clear(); + self.confirmed_data_leaders.clear(); + self.catch_up.clear(); + self.ring_observation_streak = None; + } + + pub(crate) fn advance_election_epoch(&mut self) -> u64 { + self.election_epoch = self.election_epoch.wrapping_add(1); + self.election_epoch + } + + pub(crate) fn record_ring_observation(&mut self, ring: &BTreeSet) -> u32 { + let observations = match self.ring_observation_streak.take() { + Some((previous, count)) if previous == *ring => count.saturating_add(1), + _ => 1, + }; + self.ring_observation_streak = Some((ring.clone(), observations)); + observations + } + + pub(crate) fn is_peer_caught_up(&self, node_id: &NodeId) -> bool { + self.peer_states + .get(node_id) + .is_some_and(|state| state.match_index >= self.commit_index) + } + + pub(crate) fn is_learner_ready_for_promotion(&self, node_id: &NodeId) -> bool { + self.commit_index > 0 + && self + .learner_states + .get(node_id) + .is_some_and(|state| state.match_index >= self.commit_index) + } + + pub(crate) fn replicated_voter_count(&self, index: u64) -> u32 { + self.peer_states + .values() + .filter(|state| state.match_index >= index) + .count() as u32 + + 1 + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn role_transition_reinitializes_leader_only_state() { + let self_id = NodeId::new("node-1"); + let peer = NodeId::new("node-2"); + let mut peers = HashSet::new(); + peers.insert(peer.clone()); + let mut state = TransientState::new(1); + + state.initialize_leader(&self_id, &peers, 7); + assert_eq!(state.role, Role::Leader); + assert_eq!(state.current_leader.as_ref(), Some(&self_id)); + assert_eq!(state.peer_states[&peer].next_index, 7); + + state.reset_for_follower(); + assert_eq!(state.role, Role::Follower); + assert!(state.current_leader.is_none()); + assert!(state.peer_states.is_empty()); + assert!(state.learner_states.is_empty()); + } +} diff --git a/src/control_plane/metadata/topic.rs b/src/control_plane/metadata/topic.rs index a5f66446..e1a594ad 100644 --- a/src/control_plane/metadata/topic.rs +++ b/src/control_plane/metadata/topic.rs @@ -336,7 +336,7 @@ impl TopicMeta { .collect() } - pub(crate) fn find_mergeable_pair(&self, now: u64) -> Option { + pub(crate) fn find_mergeable_range_pair(&self, now: u64) -> Option { if !self.is_merge_eligible() { return None; } From e6e63aefc8e5bf8cd952f4c02027daa947ea27db Mon Sep 17 00:00:00 2001 From: Migorithm Date: Fri, 17 Jul 2026 22:57:00 +0400 Subject: [PATCH 8/9] introduce intermediary layer - consensus --- src/control_plane/consensus/multi_raft.rs | 2 +- src/control_plane/consensus/raft/state.rs | 429 +++++++++--------- .../consensus/raft/states/log_state.rs | 187 -------- .../consensus/raft/states/metadata_state.rs | 7 + .../consensus/raft/states/mod.rs | 23 +- .../consensus/raft/states/transient_state.rs | 186 -------- src/control_plane/consensus/seal_recovery.rs | 2 +- 7 files changed, 222 insertions(+), 614 deletions(-) delete mode 100644 src/control_plane/consensus/raft/states/log_state.rs delete mode 100644 src/control_plane/consensus/raft/states/transient_state.rs diff --git a/src/control_plane/consensus/multi_raft.rs b/src/control_plane/consensus/multi_raft.rs index b837697f..a18d5212 100644 --- a/src/control_plane/consensus/multi_raft.rs +++ b/src/control_plane/consensus/multi_raft.rs @@ -5,7 +5,7 @@ use crate::control_plane::consensus::messages::{ }; use crate::control_plane::consensus::raft::errors::ProposalError; use crate::control_plane::consensus::raft::state::{Raft, TimerSeqs}; -use crate::control_plane::consensus::raft::states::transient_state::LeaderlessSegments; +use crate::control_plane::consensus::raft::states::consensus::LeaderlessSegments; use crate::control_plane::consensus::raft::storage::RaftStorage; use crate::control_plane::consensus::raft::{compute_replacement_replica_set, now_ms}; use crate::control_plane::consensus::seal_recovery::{SealEndRecovery, SealEndStep}; diff --git a/src/control_plane/consensus/raft/state.rs b/src/control_plane/consensus/raft/state.rs index f5bbd14d..54fdef91 100644 --- a/src/control_plane/consensus/raft/state.rs +++ b/src/control_plane/consensus/raft/state.rs @@ -4,11 +4,8 @@ use crate::control_plane::consensus::messages::*; use crate::control_plane::consensus::raft::command::RaftCommand; use crate::control_plane::consensus::raft::errors::{EvictionError, ProposalError}; use crate::control_plane::consensus::raft::log::LogEntry; -use crate::control_plane::consensus::raft::states::log_state::LogState; +use crate::control_plane::consensus::raft::states::consensus::{ConsensusState, PeerState, Role}; use crate::control_plane::consensus::raft::states::metadata_state::MetadataState; -use crate::control_plane::consensus::raft::states::transient_state::{ - PeerState, Role, TransientState, -}; use crate::control_plane::consensus::raft::storage::RaftPersistentState; use crate::control_plane::consensus::raft::{compute_replacement_replica_set, now_ms}; use crate::control_plane::membership::{ShardGroupId, TopologyReader}; @@ -57,9 +54,8 @@ pub struct Raft { pub node_id: NodeId, pub shard_group_id: ShardGroupId, - l_stat: LogState, - t_stat: TransientState, - m_stat: MetadataState, + consensus: ConsensusState, + metadata: MetadataState, peers: HashSet, events: Vec, @@ -86,9 +82,8 @@ impl Raft { node_id, shard_group_id, peers, - l_stat: LogState::from_persistent(persistent), - m_stat: MetadataState::new(shard_group_id), - t_stat: TransientState::new(election_jitter_seed), + consensus: ConsensusState::new(persistent, election_jitter_seed), + metadata: MetadataState::new(shard_group_id), events: Vec::new(), timer_seqs, }; @@ -97,15 +92,15 @@ impl Raft { } pub(crate) fn topic_names(&self) -> Box<[String]> { - self.m_stat.topic_names() + self.metadata.topic_names() } pub(crate) fn topic_stats(&self) -> Box<[TopicStats]> { - self.m_stat.topic_stats() + self.metadata.topic_stats() } pub(crate) fn get_topic_by_name(&self, name: &str) -> Option<&TopicMeta> { - self.m_stat.get_topic_by_name(name) + self.metadata.get_topic_by_name(name) } pub(crate) fn get_consumer_group_assignment( @@ -114,7 +109,7 @@ impl Raft { group_id: &str, member_id: ConsumerMemberId, ) -> Option { - self.m_stat + self.metadata .get_consumer_group_assignment(topic_name, group_id, member_id) } @@ -122,7 +117,7 @@ impl Raft { &self, node_id: &NodeId, ) -> Box<[(SegmentKey, Replicas)]> { - self.m_stat.active_segments_for_node(node_id) + self.metadata.active_segments_for_node(node_id) } /// Full reconciliation against the current topology: assert the ring-assigned @@ -265,9 +260,9 @@ impl Raft { /// is proposed during a check, any `RemovePeer` proposed later in that /// same check is delayed until the add successfully commits. fn has_uncommitted_membership_change(&self) -> bool { - (self.t_stat.commit_index + 1..=self.log_last_index()).any(|i| { + self.consensus.uncommited_log_range().any(|i| { matches!( - self.l_stat.get(i).map(|e| &e.command), + self.consensus.log_entry(i).map(|e| &e.command), Some(RaftCommand::AddPeer(_) | RaftCommand::RemovePeer(_)) ) }) @@ -304,7 +299,7 @@ impl Raft { } let Some(ring_members) = topology_reader.group_ring_members(self.shard_group_id) else { - self.t_stat.ring_observation_streak = None; + self.consensus.clear_ring_observation(); return Err(EvictionError::GroupNotFound); }; let ring: BTreeSet = ring_members.iter().cloned().collect(); @@ -347,7 +342,7 @@ impl Raft { } // It checks if a specific node has a complete, up-to-date copy of all permanently saved data. - let in_sync = self.t_stat.is_peer_caught_up(member); + let in_sync = self.consensus.is_peer_caught_up(member); if !in_sync { return Err(EvictionError::FollowersLagging); } @@ -374,7 +369,7 @@ impl Raft { #[inline(always)] fn record_ring_observation(&mut self, ring: &BTreeSet) -> u32 { - self.t_stat.record_ring_observation(ring) + self.consensus.record_ring_observation(ring) } /// Repair this group's segments whose replica set still names a dead node: @@ -386,7 +381,7 @@ impl Raft { pub(crate) fn reconcile_segments(&mut self, live_set: &HashSet) -> bool { let mut to_roll: Vec<(SegmentKey, Replicas)> = Vec::new(); for (key, rs) in self - .m_stat + .metadata .active_segments_with_dead_members(live_set) .into_vec() { @@ -396,14 +391,14 @@ impl Raft { .filter(|n| live_set.contains(*n)) .cloned() .collect(); - self.t_stat.leaderless_segments.push((key, survivors)); + self.consensus.push_leaderless_segment((key, survivors)); } else { to_roll.push((key, rs)); } } let mut leaderless_segments = vec![]; - for (key, rs) in self.m_stat.boundary_unknown_segments() { + for (key, rs) in self.metadata.boundary_unknown_segments() { let survivors = rs .iter() .filter(|n| live_set.contains(*n)) @@ -411,9 +406,10 @@ impl Raft { .collect(); leaderless_segments.push((key, survivors)); } - self.t_stat.leaderless_segments.extend(leaderless_segments); + self.consensus + .extend_leaderless_segments(leaderless_segments); - let sealed = self.m_stat.sealed_segments_with_dead_members(live_set); + let sealed = self.metadata.sealed_segments_with_dead_members(live_set); let mut changed = false; // Roll the others: seal + reopen with a healthy set. `end_entry_id = None` @@ -480,7 +476,7 @@ impl Raft { /// Re-fill known-end sealed segments left under-replicated by an earlier death fn refill_under_replicated_segments(&mut self, topology: &TopologyReader) -> bool { let targets = self - .m_stat + .metadata .under_replicated_sealed_segments(topology.replication_factor()); let mut changed = false; @@ -520,7 +516,7 @@ impl Raft { /// `DeleteSegments`. Topics with no retention policy contribute nothing. fn reconcile_retention_deletes(&mut self) -> bool { let now = now_ms(); - let targets = self.m_stat.expipred_segments(now); + let targets = self.metadata.expipred_segments(now); let mut changed = false; for (topic_id, range_id, segment_ids) in targets { @@ -587,14 +583,11 @@ impl Raft { } pub(crate) fn has_topic(&self, topic_id: &TopicId) -> bool { - self.m_stat.get_topic(topic_id).is_some() + self.metadata.get_topic(topic_id).is_some() } pub(crate) fn get_replica_set(&self, key: &SegmentKey) -> Option { - let topic = self.m_stat.get_topic(&key.topic_id)?; - let range = topic.ranges.get(&key.range_id)?; - let seg = range.segments.get(&key.segment_id)?; - Some(seg.replica_set.clone()) + Some(self.metadata.get_segment(key)?.replica_set.clone()) } pub(crate) fn take_events(&mut self) -> Vec { @@ -606,25 +599,25 @@ impl Raft { } pub fn take_log_mutations(&mut self) -> Vec { - self.l_stat.take_mutations() + self.consensus.take_log_mutations() } pub(crate) fn take_pending_proposals(&mut self) -> Vec { - std::mem::take(&mut self.t_stat.pending_proposals) + self.consensus.take_pending_proposals() } /// Drain the leaderless segments found by `reconcile_segments` — the actor /// drives seal-end recovery (poll survivors, seal at the recovered end). pub(crate) fn take_leaderless_segments(&mut self) -> Vec<(SegmentKey, Vec)> { - std::mem::take(&mut self.t_stat.leaderless_segments) + self.consensus.take_leaderless_segments() } pub(crate) fn last_applied_index(&self) -> u64 { - self.m_stat.last_applied_index + self.metadata.last_applied_index } pub(crate) fn log_last_index(&self) -> u64 { - self.l_stat.last_index() + self.consensus.last_log_index() } pub fn peers_count(&self) -> usize { @@ -636,11 +629,11 @@ impl Raft { } pub fn current_leader(&self) -> Option<&NodeId> { - self.t_stat.current_leader.as_ref() + self.consensus.current_leader() } pub fn is_leader(&self) -> bool { - self.t_stat.role == Role::Leader + self.consensus.is_leader() } pub fn has_peer(&self, node_id: &NodeId) -> bool { @@ -650,12 +643,12 @@ impl Raft { /// A node the leader is catching up as a non-voting learner (not yet a voter). #[cfg(test)] pub(crate) fn is_learner(&self, node_id: &NodeId) -> bool { - self.t_stat.learner_states.contains_key(node_id) + self.consensus.is_learner(node_id) } #[cfg(test)] pub(crate) fn learner_count(&self) -> usize { - self.t_stat.learner_states.len() + self.consensus.learner_state_count() } /// Minimum number of nodes needed for a majority (strict majority). @@ -666,7 +659,7 @@ impl Raft { } pub(crate) fn stabled_index(&self) -> u64 { - self.l_stat.stabled_index() + self.consensus.stabled_index() } // ------------------------------------------------------------------- @@ -680,7 +673,7 @@ impl Raft { // u64::MAX bypasses the staleness check for direct // invocations in tests; real timers carry the epoch they // were armed with. - if epoch == self.t_stat.election_epoch || epoch == u64::MAX { + if epoch == self.consensus.election_epoch() || epoch == u64::MAX { self.start_election(); return; } @@ -688,7 +681,7 @@ impl Raft { node = %self.node_id, group = self.shard_group_id.0, stale_epoch = epoch, - epoch = self.t_stat.election_epoch, + epoch = self.consensus.election_epoch(), "election: dropped stale election timeout" ); } @@ -724,10 +717,10 @@ impl Raft { fn start_election(&mut self) { // Leaders don't run election timers — they send heartbeats instead. - if self.t_stat.role == Role::Leader { + if self.consensus.is_leader() { return; } - let term = self.l_stat.begin_election(&self.node_id); + let term = self.consensus.begin_campaign(&self.node_id); if self.peers.is_empty() { // Single-node cluster: elect self immediately. @@ -735,7 +728,6 @@ impl Raft { return; } - self.t_stat.begin_campaign(); self.reset_election_timer(); tracing::trace!( node = %self.node_id, @@ -748,7 +740,7 @@ impl Raft { term, candidate_id: self.node_id.clone(), last_log_index: self.log_last_index(), - last_log_term: self.l_stat.last_term(), + last_log_term: self.consensus.last_log_term(), }; for peer_id in self.peers.iter() { self.events.push( @@ -761,11 +753,11 @@ impl Raft { // ! Followers grant or deny - All roles must respond fn handle_request_vote(&mut self, from: NodeId, req: RequestVote) { // If the request term is newer, step down. - if req.term > self.l_stat.current_term() { + if req.term > self.consensus.current_term() { self.step_down(req.term); } - let term_ok = req.term == self.l_stat.current_term(); + let term_ok = req.term == self.consensus.current_term(); let vote_ok = self.vote_available_for(&req.candidate_id); let log_ok = self.log_is_up_to_date(req.last_log_index, req.last_log_term); let vote_granted = term_ok && vote_ok && log_ok; @@ -774,7 +766,7 @@ impl Raft { group = self.shard_group_id.0, from = %req.candidate_id, req_term = req.term, - term = self.l_stat.current_term(), + term = self.consensus.current_term(), granted = vote_granted, term_ok, vote_ok, @@ -783,7 +775,7 @@ impl Raft { ); if vote_granted { - self.l_stat.grant_vote(req.candidate_id); + self.consensus.grant_vote(req.candidate_id); self.reset_election_timer(); } @@ -791,7 +783,7 @@ impl Raft { self.shard_group_id, from, RequestVoteResponse { - term: self.l_stat.current_term(), + term: self.consensus.current_term(), node_id: self.node_id.clone(), vote_granted, }, @@ -804,11 +796,11 @@ impl Raft { group = self.shard_group_id.0, from = %resp.node_id, resp_term = resp.term, - term = self.l_stat.current_term(), + term = self.consensus.current_term(), granted = resp.vote_granted, "election: RequestVoteResponse received" ); - if resp.term > self.l_stat.current_term() { + if resp.term > self.consensus.current_term() { self.step_down(resp.term); return; } @@ -816,22 +808,25 @@ impl Raft { } fn count_vote_if_eligible(&mut self, resp: RequestVoteResponse) { - if resp.term != self.l_stat.current_term() || !resp.vote_granted { + if resp.term != self.consensus.current_term() || !resp.vote_granted { return; } - if self.t_stat.record_vote(self.quorum()) { + if self + .consensus + .record_vote(resp.term, resp.vote_granted, self.quorum()) + { self.become_leader(); } } fn vote_available_for(&self, candidate_id: &NodeId) -> bool { - self.l_stat.vote_available_for(candidate_id) + self.consensus.vote_available_for(candidate_id) } /// §5.4.1: A candidate's log is "at least as up-to-date" if its last /// entry has a higher term, or the same term with a >= index. fn log_is_up_to_date(&self, last_log_index: u64, last_log_term: u64) -> bool { - let my_last_term = self.l_stat.last_term(); + let my_last_term = self.consensus.last_log_term(); let my_last_index = self.log_last_index(); if last_log_term != my_last_term { @@ -845,20 +840,18 @@ impl Raft { // ------------------------------------------------------------------- fn become_leader(&mut self) { - let next = self.log_last_index() + 1; - self.t_stat - .initialize_leader(&self.node_id, &self.peers, next); + self.consensus.initialize_leader(&self.node_id, &self.peers); tracing::debug!( node = %self.node_id, group = self.shard_group_id.0, - term = self.l_stat.current_term(), + term = self.consensus.current_term(), "election: became leader" ); self.raise(LeaderChange { shard_group_id: self.shard_group_id, leader_node_id: self.node_id.clone(), - term: self.l_stat.current_term(), + term: self.consensus.current_term(), }); // Cancel election timer, start heartbeat + merge/ring check timers. @@ -882,10 +875,10 @@ impl Raft { fn step_down(&mut self, new_term: u64) { debug_assert!( - new_term >= self.l_stat.current_term(), + new_term >= self.consensus.current_term(), "step_down must never regress the term" ); - let was_leader = self.t_stat.role == Role::Leader; + let was_leader = self.consensus.is_leader(); tracing::debug!( node = %self.node_id, @@ -901,7 +894,7 @@ impl Raft { // rule `recognize_leader`'s demotion branch follows. All production // callers pass strictly newer terms (guarded `>` at every call // site); tests use equal-term step_down to depose a leader in place. - self.l_stat.advance_term(new_term); + self.consensus.advance_term(new_term); self.cancel_leader_timers(); if was_leader { @@ -913,7 +906,7 @@ impl Raft { // runs no election timer, arms a fresh one. self.reset_election_timer(); } - self.t_stat.reset_for_follower(); + self.consensus.reset_for_follower(); } // ------------------------------------------------------------------- @@ -938,11 +931,9 @@ impl Raft { fn maybe_redrive_segment_assignments(&mut self) { // rederive Metadata <> Datanode segment assignment - let active = self.m_stat.active_segment_assignments(); + let active = self.metadata.active_segment_assignments(); let active_keys: HashSet = active.iter().map(|(k, _, _)| *k).collect(); - self.t_stat - .confirmed_data_leaders - .retain(|k, _| active_keys.contains(k)); + self.consensus.retain_confirmed_data_leaders(&active_keys); let mut redrives = Vec::new(); for (segment_key, replica_set, start_entry_id) in active { @@ -951,7 +942,10 @@ impl Raft { }; // * If data leader acks assignment, it would have been added to confirmed_placement through Raft::handle_segment_placed - if self.t_stat.confirmed_data_leaders.get(&segment_key) == Some(target) { + if self + .consensus + .is_data_leader_confirmed(&segment_key, target) + { continue; } @@ -971,15 +965,14 @@ impl Raft { } pub(crate) fn handle_segment_placed(&mut self, ack: SegmentPlaced) { - self.t_stat - .confirmed_data_leaders - .insert(ack.segment_key, ack.from); + self.consensus + .confirm_data_leader(ack.segment_key, ack.from); } /// Re-drive the catch-up sweep — the sealed-segment analogue of /// `maybe_redrive_segment_assignments`. See `.claude/rules/raft-actor.md` #9. fn maybe_redrive_catch_ups(&mut self) { - let redrives = self.t_stat.catch_up.redrives(self.shard_group_id); + let redrives = self.consensus.catch_up_redrives(self.shard_group_id); if !redrives.is_empty() { self.raise(RaftEvent::RedriveAssignments(redrives)); } @@ -988,7 +981,7 @@ impl Raft { /// A member confirmed it holds a reassigned sealed segment; routed here from /// `MultiRaft`. pub(crate) fn handle_catch_up_ack(&mut self, ack: SegmentCaughtUp) { - self.t_stat.catch_up.confirm(ack.segment_key, ack.from); + self.consensus.confirm_catch_up(ack); } /// Takeover backstop: re-seed catch-up for every known-end sealed segment this @@ -997,11 +990,10 @@ impl Raft { /// heartbeat sweep re-drives the re-seeded set (already-complete members /// full-match-ack cheaply). See `.claude/rules/raft-actor.md` #9. pub(crate) fn reseed_catch_up(&mut self) { - let sealed = self.m_stat.known_end_sealed_segments(); + let sealed = self.metadata.known_end_sealed_segments(); for (segment_key, start, end, replica_set) in sealed { - self.t_stat - .catch_up - .track_sealed(segment_key, start, end, replica_set); + self.consensus + .track_sealed_catch_up(segment_key, start, end, replica_set); } } @@ -1013,19 +1005,14 @@ impl Raft { RaftEvent::ShardLeaderRefresh(LeaderChange { shard_group_id: self.shard_group_id, leader_node_id: self.node_id.clone(), - term: self.l_stat.current_term(), + term: self.consensus.current_term(), }) }) } /// Everyone the leader replicates to: voting peers plus catching-up learners. fn replication_targets(&self) -> Vec { - self.t_stat - .peer_states - .keys() - .chain(self.t_stat.learner_states.keys()) - .cloned() - .collect() + self.consensus.replication_targets() } /// Stage a node as a non-voting learner the leader catches up before promoting it @@ -1037,11 +1024,11 @@ impl Raft { if !self.is_leader() || node == self.node_id || self.peers.contains(&node) - || self.t_stat.learner_states.contains_key(&node) + || self.consensus.is_learner(&node) { return false; } - self.t_stat.learner_states.insert( + self.consensus.stage_learner( node.clone(), PeerState { next_index: self.log_last_index() + 1, @@ -1060,36 +1047,31 @@ impl Raft { if self.has_uncommitted_membership_change() { return; } - if self.t_stat.is_learner_ready_for_promotion(node) { + if self.consensus.is_learner_ready_for_promotion(node) { let _ = self.propose(RaftCommand::AddPeer(node.clone())); } } fn send_append_entries(&mut self, peer_id: NodeId) { - let peer_state = match self - .t_stat - .peer_states - .get(&peer_id) - .or_else(|| self.t_stat.learner_states.get(&peer_id)) - { + let peer_state = match self.consensus.peer_state(&peer_id) { Some(ps) => ps, None => return, }; let prev_log_index = peer_state.next_index.saturating_sub(1); - let prev_log_term = self.l_stat.term_at(prev_log_index); - let entries = self.l_stat.entries_from(peer_state.next_index); + let prev_log_term = self.consensus.log_term_at(prev_log_index); + let entries = self.consensus.log_entries_from(peer_state.next_index); self.raise(OutboundRaftPacket::new( self.shard_group_id, peer_id, AppendEntries { - term: self.l_stat.current_term(), + term: self.consensus.current_term(), leader_id: self.node_id.clone(), prev_log_index, prev_log_term, entries, - leader_commit: self.t_stat.commit_index, + leader_commit: self.consensus.commit_index(), }, )); } @@ -1099,7 +1081,7 @@ impl Raft { // ! - Candidates step down on same term : because receiving entries while being a candidate means another node already won. // ! - Followers process normally. All roles must respond. fn handle_append_entries(&mut self, from: NodeId, req: AppendEntries) { - if req.term < self.l_stat.current_term() { + if req.term < self.consensus.current_term() { self.reject_append_entries(from); return; } @@ -1118,19 +1100,15 @@ impl Raft { } fn recognize_leader(&mut self, req: &AppendEntries) { - if req.term > self.l_stat.current_term() { + if req.term > self.consensus.current_term() { self.step_down(req.term); - } else if self.t_stat.role != Role::Follower { - self.t_stat.role = Role::Follower; - self.t_stat.peer_states.clear(); - self.t_stat.learner_states.clear(); } - self.t_stat.current_leader = Some(req.leader_id.clone()); + self.consensus.recognize_leader(req.leader_id.clone()); self.reset_election_timer(); tracing::debug!( node = %self.node_id, group = self.shard_group_id.0, - term = self.l_stat.current_term(), + term = self.consensus.current_term(), leader = %req.leader_id, "election: leader recognized" ); @@ -1140,25 +1118,25 @@ impl Raft { if prev_log_index == 0 { return true; } - let local_term = self.l_stat.term_at(prev_log_index); + let local_term = self.consensus.log_term_at(prev_log_index); local_term != 0 && local_term == prev_log_term } fn replicate_entries(&mut self, entries: Box<[LogEntry]>) { for entry in entries { - let existing_term = self.l_stat.term_at(entry.index); + let existing_term = self.consensus.log_term_at(entry.index); if existing_term != 0 && existing_term != entry.term { - self.l_stat.truncate_from(entry.index); + self.consensus.truncate_log_from(entry.index); } if entry.index > self.log_last_index() { - self.l_stat.append(entry); + self.consensus.append_log(entry); } } } fn advance_follower_commit(&mut self, leader_commit: u64) { - if leader_commit > self.t_stat.commit_index { - self.t_stat.commit_index = leader_commit.min(self.log_last_index()); + if leader_commit > self.consensus.commit_index() { + self.consensus.advance_follower_commit(leader_commit); self.apply_committed_entries(); } } @@ -1167,7 +1145,7 @@ impl Raft { // ! - term guard // ! - only leaders track peer state fn handle_append_entries_response(&mut self, resp: AppendEntriesResponse) { - if resp.term > self.l_stat.current_term() { + if resp.term > self.consensus.current_term() { self.step_down(resp.term); return; } @@ -1178,12 +1156,8 @@ impl Raft { // The responder may be a voting peer or a catching-up learner. let node_id = resp.node_id.clone(); - let is_voter = self.t_stat.peer_states.contains_key(&node_id); - let peer_state = if is_voter { - self.t_stat.peer_states.get_mut(&node_id) - } else { - self.t_stat.learner_states.get_mut(&node_id) - }; + let is_voter = self.consensus.is_voter(&node_id); + let peer_state = self.consensus.peer_state_mut(&node_id); if let Some(peer_state) = peer_state { if resp.success { peer_state.match_index = resp.last_log_index; @@ -1214,7 +1188,7 @@ impl Raft { self.shard_group_id, target, AppendEntriesResponse { - term: self.l_stat.current_term(), + term: self.consensus.current_term(), node_id: self.node_id.clone(), success, last_log_index: self.log_last_index(), @@ -1239,19 +1213,18 @@ impl Raft { // Instead, it appends a new entry at its own term. Once that entry is committed on a majority, all preceding entries (including index 3) are implicitly committed too. // And that 'implicit commit' does not violate safety because 'new' entry acts as an election shield that physically prevents that overwrite from happening. fn try_advance_commit_index(&mut self) { - let last = self.log_last_index(); let quorum = self.quorum(); // Scan top-down: the highest current-term entry with quorum // implicitly commits everything below it (log matching property). - for n in (self.t_stat.commit_index + 1..=last).rev() { - if self.l_stat.term_at(n) != self.l_stat.current_term() { + for n in (self.consensus.uncommited_log_range()).rev() { + if self.consensus.log_term_at(n) != self.consensus.current_term() { continue; } - let replication_count = self.t_stat.replicated_voter_count(n); + let replication_count = self.consensus.replicated_voter_count(n); if replication_count >= quorum { - self.t_stat.commit_index = n; + self.consensus.set_commit_index(n); self.apply_committed_entries(); return; } @@ -1259,7 +1232,7 @@ impl Raft { } pub(crate) fn advance_stabled_index(&mut self, value: u64) { - self.l_stat.advance_stabled_index(value); + self.consensus.advance_stabled_index(value); self.apply_committed_entries(); #[cfg(any(test, debug_assertions))] @@ -1268,19 +1241,21 @@ impl Raft { #[tracing::instrument(level = "debug", skip_all, fields( group = self.shard_group_id.0, - from = self.m_stat.last_applied_index + 1, - to = self.t_stat.commit_index.min(self.l_stat.stabled_index()), + from = self.metadata.last_applied_index + 1, + to = self.consensus.ready_to_apply_index(), ))] fn apply_committed_entries(&mut self) { - while self.m_stat.last_applied_index - < self.t_stat.commit_index.min(self.l_stat.stabled_index()) - { - self.m_stat.last_applied_index += 1; - let Some(entry) = self.l_stat.get(self.m_stat.last_applied_index).cloned() else { + while self.metadata.last_applied_index < self.consensus.ready_to_apply_index() { + self.metadata.last_applied_index += 1; + let Some(entry) = self + .consensus + .log_entry(self.metadata.last_applied_index) + .cloned() + else { tracing::error!( "[{}] committed entry at index {} missing from log", self.node_id, - self.m_stat.last_applied_index + self.metadata.last_applied_index ); break; }; @@ -1296,7 +1271,7 @@ impl Raft { } fn apply_metadata_entry(&mut self, cmd: MetadataCommand, index: u64) { - match self.m_stat.apply(cmd) { + match self.metadata.apply(cmd) { Ok(result) => { tracing::debug!( "[{}] Applied metadata at index {}: {:?}", @@ -1308,7 +1283,7 @@ impl Raft { if self.is_leader() && let ApplyResult::SegmentReassigned(r) = &result { - self.t_stat.catch_up.track(r); + self.consensus.track_catch_up(r); } self.raise(MetadataCommitted { shard_group_id: self.shard_group_id, @@ -1324,10 +1299,9 @@ impl Raft { e ), } - if self.t_stat.role == Role::Leader { - self.t_stat - .pending_proposals - .extend(self.m_stat.take_pending_proposals()); + if self.consensus.is_leader() { + self.consensus + .extend_pending_proposals(self.metadata.take_pending_proposals()); } } @@ -1340,13 +1314,13 @@ impl Raft { } // Promotion: a learner graduating to a voter carries its catch-up progress, so // the new voter isn't reset to match_index 0 (which would stall commits anew). - let carried = self.t_stat.learner_states.remove(&node_id); - if self.peers.insert(node_id.clone()) && self.t_stat.role == Role::Leader { + let carried = self.consensus.remove_learner(&node_id); + if self.peers.insert(node_id.clone()) && self.consensus.is_leader() { let state = carried.unwrap_or(PeerState { next_index: self.log_last_index() + 1, match_index: 0, }); - self.t_stat.peer_states.insert(node_id, state); + self.consensus.add_voter_state(node_id, state); } } @@ -1354,20 +1328,20 @@ impl Raft { /// `RemovePeer` log entry commits. Never call directly — the peer set is part /// of the replicated state machine and must only mutate through the log. fn apply_remove_peer(&mut self, node_id: NodeId) { - self.t_stat.learner_states.remove(&node_id); + self.consensus.remove_learner(&node_id); if self.peers.remove(&node_id) { - self.t_stat.peer_states.remove(&node_id); + self.consensus.remove_voter_state(&node_id); self.raise(RaftEvent::DisconnectPeer(node_id)); } } fn add_new_entry(&mut self, command: RaftCommand) { let entry = LogEntry { - term: self.l_stat.current_term(), + term: self.consensus.current_term(), index: self.log_last_index() + 1, command, }; - self.l_stat.append(entry); + self.consensus.append_log(entry); } /// Propose a command to the Raft log. Only the leader can accept proposals. @@ -1383,7 +1357,9 @@ impl Raft { #[tracing::instrument(level = "trace", skip_all, fields(group = self.shard_group_id.0, command = ?command))] pub fn propose(&mut self, command: RaftCommand) -> Result { if !self.is_leader() { - return Err(ProposalError::NotLeader(self.t_stat.current_leader.clone())); + return Err(ProposalError::NotLeader( + self.consensus.current_leader().cloned(), + )); } self.add_new_entry(command); @@ -1447,9 +1423,9 @@ impl Raft { } fn reset_election_timer(&mut self) { - let election_epoch = self.t_stat.advance_election_epoch(); + let election_epoch = self.consensus.next_election_epoch(); - let jitter = self.t_stat.election_jitter.next(); + let jitter = self.consensus.next_election_jitter(); tracing::trace!( node = %self.node_id, group = self.shard_group_id.0, @@ -1505,9 +1481,9 @@ impl Raft { return; } - let merge_proposals = self.m_stat.evaluate_merges(now); + let merge_proposals = self.metadata.evaluate_merges(now); for cmd in merge_proposals { - self.t_stat.pending_proposals.push(cmd); + self.consensus.push_pending_proposal(cmd); } self.schedule_merge_check_timer(); @@ -1530,7 +1506,7 @@ impl Raft { } pub(crate) fn cancel_all_timers(&mut self) { - self.t_stat.advance_election_epoch(); + self.consensus.next_election_epoch(); self.raise(RaftEvent::Timer(TimerCommand::CancelSchedule { seq: self.timer_seqs.election, })); @@ -1550,26 +1526,26 @@ impl Raft { impl crate::test_traits::TAssertInvariant for Raft { fn assert_invariants(&self) { assert!( - self.m_stat.last_applied_index <= self.t_stat.commit_index, + self.metadata.last_applied_index <= self.consensus.commit_index(), "last_applied ({}) > commit_index ({})", - self.m_stat.last_applied_index, - self.t_stat.commit_index, + self.metadata.last_applied_index, + self.consensus.commit_index(), ); assert!( - self.m_stat.last_applied_index <= self.l_stat.stabled_index(), + self.metadata.last_applied_index <= self.consensus.stabled_index(), "last_applied ({}) > stabled_index ({}) — applied a non-durable entry", - self.m_stat.last_applied_index, - self.l_stat.stabled_index(), + self.metadata.last_applied_index, + self.consensus.stabled_index(), ); assert!( - self.t_stat.commit_index <= self.log_last_index(), + self.consensus.commit_index() <= self.log_last_index(), "commit_index ({}) > log_last_index ({})", - self.t_stat.commit_index, + self.consensus.commit_index(), self.log_last_index(), ); // Log indices are contiguous and 1-based - for (i, entry) in self.l_stat.entries().iter().enumerate() { + for (i, entry) in self.consensus.log_entries().iter().enumerate() { assert_eq!( entry.index, (i + 1) as u64, @@ -1577,36 +1553,36 @@ impl crate::test_traits::TAssertInvariant for Raft { entry.index, ); assert!( - entry.term <= self.l_stat.current_term(), + entry.term <= self.consensus.current_term(), "log entry at index {} has term {} > current_term {}", entry.index, entry.term, - self.l_stat.current_term(), + self.consensus.current_term(), ); } // Invariant: peer_states exists only on the leader (and matches the peer // set when leader). Followers/candidates carry an empty peer_states. - match self.t_stat.role { + match *self.consensus.role() { Role::Leader => { for peer in &self.peers { assert!( - self.t_stat.peer_states.contains_key(peer), + self.consensus.has_voter_state(peer), "leader missing peer_state for {:?}", peer, ); } assert_eq!( - self.t_stat.peer_states.len(), + self.consensus.voter_state_count(), self.peers.len(), "leader peer_states size ({}) != peers size ({})", - self.t_stat.peer_states.len(), + self.consensus.voter_state_count(), self.peers.len(), ); // Invariant: learners are non-voting and disjoint from voters — a node // is never both — and self is never a learner. Learners are replicated // to but excluded from the commit quorum until promoted via `AddPeer`. - for learner in self.t_stat.learner_states.keys() { + for learner in self.consensus.learner_ids() { assert!( !self.peers.contains(learner), "node {:?} is both a voter and a learner", @@ -1618,23 +1594,23 @@ impl crate::test_traits::TAssertInvariant for Raft { // the local fragment of this: a leader must have voted for itself this // term (and is therefore the only node that could have won this term). assert_eq!( - self.l_stat.voted_for(), + self.consensus.voted_for(), Some(&self.node_id), "leader has voted_for {:?}, expected self ({:?})", - self.l_stat.voted_for(), + self.consensus.voted_for(), self.node_id, ); } Role::Follower | Role::Candidate { .. } => { assert!( - self.t_stat.peer_states.is_empty(), + self.consensus.voter_states_empty(), "non-leader carries peer_states ({} entries)", - self.t_stat.peer_states.len(), + self.consensus.voter_state_count(), ); assert!( - self.t_stat.learner_states.is_empty(), + self.consensus.learner_states_empty(), "non-leader carries learner_states ({} entries)", - self.t_stat.learner_states.len(), + self.consensus.learner_state_count(), ); } } @@ -1667,11 +1643,11 @@ mod tests { } pub(crate) fn current_term(&self) -> u64 { - self.l_stat.current_term() + self.consensus.current_term() } pub(crate) fn voted_for(&self) -> Option { - self.l_stat.voted_for().cloned() + self.consensus.voted_for().cloned() } pub(crate) fn simulate_flush_and_apply(&mut self) { @@ -1680,7 +1656,7 @@ mod tests { } pub(crate) fn state_machine(&self) -> &MetadataState { - &self.m_stat + &self.metadata } } fn node(id: &str) -> NodeId { @@ -1753,17 +1729,17 @@ mod tests { #[test] fn single_node_elects_self_on_timeout() { let mut raft = single_node_raft(); - assert_eq!(raft.t_stat.role, Role::Follower); + assert_eq!(*raft.consensus.role(), Role::Follower); raft.handle_timeout(RaftTimeoutCallback::ElectionTimeout { shard_group_id: TEST_SHARD, epoch: u64::MAX, }); - assert_eq!(raft.t_stat.role, Role::Leader); - assert_eq!(raft.l_stat.current_term(), 1); + assert_eq!(*raft.consensus.role(), Role::Leader); + assert_eq!(raft.consensus.current_term(), 1); assert_eq!( - raft.l_stat.voted_for().cloned(), + raft.consensus.voted_for().cloned(), Some(NodeId::new("node-1")) ); } @@ -1775,7 +1751,7 @@ mod tests { shard_group_id: TEST_SHARD, epoch: u64::MAX, }); - assert_eq!(raft.l_stat.current_term(), 1); + assert_eq!(raft.consensus.current_term(), 1); // Step down and trigger another election raft.step_down(1); @@ -1783,7 +1759,7 @@ mod tests { shard_group_id: TEST_SHARD, epoch: u64::MAX, }); - assert_eq!(raft.l_stat.current_term(), 2); + assert_eq!(raft.consensus.current_term(), 2); } // ------------------------------------------------------------------- @@ -1800,10 +1776,10 @@ mod tests { }); assert!(matches!( - raft.t_stat.role, + *raft.consensus.role(), Role::Candidate { votes_received: 1 } )); - assert_eq!(raft.l_stat.current_term(), 1); + assert_eq!(raft.consensus.current_term(), 1); let out = packets(&mut raft); assert_eq!(out.len(), 2); // one per peer @@ -1834,7 +1810,7 @@ mod tests { _ => panic!("expected RequestVoteResponse"), } assert_eq!( - raft.l_stat.voted_for().cloned(), + raft.consensus.voted_for().cloned(), Some(NodeId::new("node-1")) ); } @@ -1888,7 +1864,7 @@ mod tests { }; raft.handle_rpc(node("node-2"), resp); - assert_eq!(raft.t_stat.role, Role::Leader); + assert_eq!(*raft.consensus.role(), Role::Leader); } #[test] @@ -1907,8 +1883,8 @@ mod tests { }; raft.handle_rpc(node("node-2"), resp); - assert_eq!(raft.t_stat.role, Role::Follower); - assert_eq!(raft.l_stat.current_term(), 5); + assert_eq!(*raft.consensus.role(), Role::Follower); + assert_eq!(raft.consensus.current_term(), 5); } // ------------------------------------------------------------------- @@ -1919,7 +1895,7 @@ mod tests { fn rejects_vote_if_candidate_log_is_stale() { let mut raft = three_node_raft("node-2"); // Give node-2 a log entry at term 2 - raft.l_stat.append(LogEntry { + raft.consensus.append_log(LogEntry { term: 2, index: 1, command: RaftCommand::Noop, @@ -2038,7 +2014,7 @@ mod tests { #[test] fn follower_rejects_append_entries_with_stale_term() { let mut raft = three_node_raft("node-2"); - raft.l_stat.advance_term(5); + raft.consensus.advance_term(5); let ae = AppendEntries { term: 3, @@ -2112,7 +2088,7 @@ mod tests { raft.propose_noop().unwrap(); drain(&mut raft); assert_eq!(raft.log_last_index(), 2); - assert_eq!(raft.t_stat.commit_index, 0); + assert_eq!(raft.consensus.commit_index(), 0); // node-2 acknowledges both entries (noop + proposal) let resp = AppendEntriesResponse { @@ -2124,7 +2100,7 @@ mod tests { raft.handle_rpc(node("node-2"), resp); // Majority achieved (self + node-2 = 2 out of 3) - assert_eq!(raft.t_stat.commit_index, 2); + assert_eq!(raft.consensus.commit_index(), 2); } #[test] @@ -2183,7 +2159,7 @@ mod tests { raft.handle_rpc(node("node-1"), ae); drain(&mut raft); - assert_eq!(raft.t_stat.commit_index, 1); + assert_eq!(raft.consensus.commit_index(), 1); } // ------------------------------------------------------------------- @@ -2252,7 +2228,7 @@ mod tests { }), ); drain(&mut raft); - assert_eq!(raft.t_stat.role, Role::Leader); + assert_eq!(*raft.consensus.role(), Role::Leader); // Receive AppendEntries from a leader with higher term let ae = AppendEntries { @@ -2265,8 +2241,8 @@ mod tests { }; raft.handle_rpc(node("node-3"), ae); - assert_eq!(raft.t_stat.role, Role::Follower); - assert_eq!(raft.l_stat.current_term(), 3); + assert_eq!(*raft.consensus.role(), Role::Follower); + assert_eq!(raft.consensus.current_term(), 3); } // ------------------------------------------------------------------- @@ -2292,8 +2268,8 @@ mod tests { }), ); drain(&mut raft); - assert_eq!(raft.t_stat.role, Role::Leader); - assert_eq!(raft.l_stat.current_term(), 1); + assert_eq!(*raft.consensus.role(), Role::Leader); + assert_eq!(raft.consensus.current_term(), 1); // Stale election timeout arrives — should be ignored raft.handle_timeout(RaftTimeoutCallback::ElectionTimeout { @@ -2302,17 +2278,17 @@ mod tests { }); assert_eq!( - raft.t_stat.role, + *raft.consensus.role(), Role::Leader, "leader must not start a new election" ); - assert_eq!(raft.l_stat.current_term(), 1, "term must not increment"); + assert_eq!(raft.consensus.current_term(), 1, "term must not increment"); } #[test] fn follower_ignores_rpc_timeout() { let mut raft = three_node_raft("node-1"); - assert_eq!(raft.t_stat.role, Role::Follower); + assert_eq!(*raft.consensus.role(), Role::Follower); raft.handle_timeout(RaftTimeoutCallback::RpcTimeout { shard_group_id: TEST_SHARD, @@ -2330,7 +2306,7 @@ mod tests { epoch: u64::MAX, }); drain(&mut raft); - assert!(matches!(raft.t_stat.role, Role::Candidate { .. })); + assert!(matches!(*raft.consensus.role(), Role::Candidate { .. })); raft.handle_timeout(RaftTimeoutCallback::RpcTimeout { shard_group_id: TEST_SHARD, @@ -2369,7 +2345,7 @@ mod tests { ); drain(&mut raft); - assert_eq!(raft.t_stat.role, Role::Leader); + assert_eq!(*raft.consensus.role(), Role::Leader); assert_eq!(raft.current_leader(), Some(&node("node-1"))); } @@ -2659,7 +2635,7 @@ mod tests { let result = raft.propose(cmd.into()); assert!(result.is_ok()); raft.simulate_flush(); - assert!(raft.m_stat.last_applied_index > 0); + assert!(raft.metadata.last_applied_index > 0); } // ------------------------------------------------------------------- @@ -2887,7 +2863,7 @@ mod tests { drain(&mut raft); raft.simulate_flush(); assert_eq!(raft.log_last_index(), 1); - assert_eq!(raft.t_stat.commit_index, 0); + assert_eq!(raft.consensus.commit_index(), 0); // node-1 wins election at term 2 raft.handle_timeout(RaftTimeoutCallback::ElectionTimeout { @@ -2904,12 +2880,12 @@ mod tests { }), ); drain(&mut raft); - assert_eq!(raft.t_stat.role, Role::Leader); - assert_eq!(raft.l_stat.current_term(), 2); + assert_eq!(*raft.consensus.role(), Role::Leader); + assert_eq!(raft.consensus.current_term(), 2); // become_leader appends a Noop at term 2 (index 2) assert_eq!(raft.log_last_index(), 2); - assert_eq!(raft.l_stat.term_at(1), 1); - assert_eq!(raft.l_stat.term_at(2), 2); + assert_eq!(raft.consensus.log_term_at(1), 1); + assert_eq!(raft.consensus.log_term_at(2), 2); // node-3 acks only the old term-1 entry (index 1) but not the term-2 entry raft.handle_rpc( @@ -2923,7 +2899,8 @@ mod tests { ); // commit_index must NOT advance — the replicated entry is term 1, not current term assert_eq!( - raft.t_stat.commit_index, 0, + raft.consensus.commit_index(), + 0, "term-1 entry must not be directly committed even with majority" ); @@ -2939,7 +2916,7 @@ mod tests { ); // Now both entries committed (term-2 entry at index 2 has quorum, // implicitly committing the term-1 entry at index 1) - assert_eq!(raft.t_stat.commit_index, 2); + assert_eq!(raft.consensus.commit_index(), 2); } #[test] @@ -3098,7 +3075,7 @@ mod tests { .next() .expect("three_node_raft must have peers") .clone(); - let term = raft.l_stat.current_term(); + let term = raft.consensus.current_term(); raft.handle_rpc( peer.clone(), RaftRpc::RequestVoteResponse(RequestVoteResponse { @@ -3109,7 +3086,7 @@ mod tests { ); drain(&mut raft); assert_eq!( - raft.t_stat.role, + *raft.consensus.role(), Role::Leader, "must be leader after election" ); @@ -3126,14 +3103,14 @@ mod tests { epoch: u64::MAX, }); drain(&mut raft); - assert_eq!(raft.t_stat.role, Role::Leader); + assert_eq!(*raft.consensus.role(), Role::Leader); raft } /// Proposals in the log after the become_leader noop (index 1). fn proposals_after_become_leader(raft: &Raft) -> Vec { (2..=raft.log_last_index()) - .filter_map(|i| raft.l_stat.get(i).map(|e| e.command.clone())) + .filter_map(|i| raft.consensus.log_entry(i).map(|e| e.command.clone())) .collect() } @@ -3286,7 +3263,7 @@ mod tests { } drain(&mut raft); assert_eq!( - raft.t_stat.role, + *raft.consensus.role(), Role::Leader, "must be leader after election" ); @@ -3689,7 +3666,7 @@ mod tests { epoch: u64::MAX, }); drain(&mut raft); - assert_eq!(raft.t_stat.role, Role::Leader); + assert_eq!(*raft.consensus.role(), Role::Leader); raft } @@ -3904,8 +3881,8 @@ mod tests { fn catch_up_repairs_dropped_on_step_down() { let (mut raft, _) = raft_with_seeded_catch_up(vec![node("node-1"), node("y")]); // A higher term deposes the leader → leader-volatile tracker is cleared. - raft.step_down(raft.l_stat.current_term() + 1); - assert!(raft.t_stat.catch_up.is_empty()); + raft.step_down(raft.consensus.current_term() + 1); + assert!(raft.consensus.catch_up_is_empty()); assert!(drain_catch_up_redrives(&mut raft).is_empty()); } @@ -3916,7 +3893,7 @@ mod tests { // Simulate the takeover gap: the leader-volatile tracker is empty, but the // sealed segment is still under-replicated in the state machine. - raft.t_stat.catch_up.clear(); + raft.consensus.clear_catch_up(); assert!(drain_catch_up_redrives(&mut raft).is_empty()); raft.reseed_catch_up(); diff --git a/src/control_plane/consensus/raft/states/log_state.rs b/src/control_plane/consensus/raft/states/log_state.rs deleted file mode 100644 index c53a522a..00000000 --- a/src/control_plane/consensus/raft/states/log_state.rs +++ /dev/null @@ -1,187 +0,0 @@ -use crate::control_plane::NodeId; -use crate::control_plane::consensus::messages::LogMutation; -use crate::control_plane::consensus::raft::log::LogEntry; -use crate::control_plane::consensus::raft::storage::RaftPersistentState; - -pub(crate) struct LogState { - current_term: u64, - voted_for: Option, - entries: Vec, - stabled_index: u64, - unflushed_mutations: Vec, -} - -impl LogState { - pub(crate) fn from_persistent(persistent: RaftPersistentState) -> Self { - Self { - stabled_index: persistent.stabled_index(), - current_term: persistent.term, - voted_for: persistent.voted_for, - entries: persistent.log, - unflushed_mutations: Vec::new(), - } - } - - pub(crate) fn current_term(&self) -> u64 { - self.current_term - } - - pub(crate) fn voted_for(&self) -> Option<&NodeId> { - self.voted_for.as_ref() - } - - pub(crate) fn vote_available_for(&self, candidate_id: &NodeId) -> bool { - self.voted_for - .as_ref() - .is_none_or(|voted_for| voted_for == candidate_id) - } - - pub(crate) fn begin_election(&mut self, node_id: &NodeId) -> u64 { - self.current_term += 1; - self.voted_for = Some(node_id.clone()); - self.buffer_hard_state(); - self.current_term - } - - pub(crate) fn grant_vote(&mut self, candidate_id: NodeId) { - self.voted_for = Some(candidate_id); - self.buffer_hard_state(); - } - - pub(crate) fn advance_term(&mut self, new_term: u64) -> bool { - if new_term <= self.current_term { - return false; - } - self.current_term = new_term; - self.voted_for = None; - self.buffer_hard_state(); - true - } - - pub(crate) fn stabled_index(&self) -> u64 { - self.stabled_index - } - - pub(crate) fn advance_stabled_index(&mut self, index: u64) { - self.stabled_index = self.stabled_index.max(index); - } - - pub(crate) fn entries(&self) -> &[LogEntry] { - &self.entries - } - - pub(crate) fn last_index(&self) -> u64 { - self.entries.last().map_or(0, |entry| entry.index) - } - - pub(crate) fn last_term(&self) -> u64 { - self.entries.last().map_or(0, |entry| entry.term) - } - - pub(crate) fn term_at(&self, index: u64) -> u64 { - if index == 0 { - return 0; - } - self.get(index).map_or(0, |entry| entry.term) - } - - pub(crate) fn get(&self, index: u64) -> Option<&LogEntry> { - if index == 0 { - return None; - } - self.entries.get((index - 1) as usize) - } - - pub(crate) fn entries_from(&self, start_index: u64) -> Box<[LogEntry]> { - if start_index == 0 || start_index > self.last_index() { - return Box::new([]); - } - self.entries[(start_index - 1) as usize..].into() - } - - pub(crate) fn append(&mut self, entry: LogEntry) { - debug_assert_eq!( - entry.index, - self.last_index() + 1, - "log entry index must be contiguous" - ); - self.unflushed_mutations - .push(LogMutation::Append(entry.clone())); - self.entries.push(entry); - } - - pub(crate) fn truncate_from(&mut self, from_index: u64) { - if from_index == 0 || from_index > self.last_index() + 1 { - return; - } - self.unflushed_mutations - .push(LogMutation::TruncateFrom(from_index)); - self.entries.truncate((from_index - 1) as usize); - } - - fn buffer_hard_state(&mut self) { - self.unflushed_mutations.push(LogMutation::HardState { - term: self.current_term, - voted_for: self.voted_for.clone(), - }); - } - - pub(crate) fn take_mutations(&mut self) -> Vec { - std::mem::take(&mut self.unflushed_mutations) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::control_plane::consensus::raft::command::RaftCommand; - - #[test] - fn append_and_truncate_emit_persistence_mutations() { - let mut state = LogState::from_persistent(RaftPersistentState::default()); - let entry = LogEntry { - term: 1, - index: 1, - command: RaftCommand::Noop, - }; - - state.append(entry.clone()); - state.truncate_from(1); - - assert_eq!(state.last_index(), 0); - let mutations = state.take_mutations(); - assert!(matches!( - &mutations[..], - [ - LogMutation::Append(appended), - LogMutation::TruncateFrom(1) - ] if appended == &entry - )); - } - - #[test] - fn election_transitions_buffer_hard_state() { - let node = NodeId::new("node-1"); - let candidate = NodeId::new("node-2"); - let mut state = LogState::from_persistent(RaftPersistentState::default()); - - assert_eq!(state.begin_election(&node), 1); - assert_eq!(state.voted_for(), Some(&node)); - assert!(!state.vote_available_for(&candidate)); - - assert!(state.advance_term(2)); - assert!(state.voted_for().is_none()); - state.grant_vote(candidate.clone()); - assert_eq!(state.voted_for(), Some(&candidate)); - - let mutations = state.take_mutations(); - assert_eq!(mutations.len(), 3); - assert!(matches!( - mutations.last(), - Some(LogMutation::HardState { - term: 2, - voted_for: Some(voted_node), - }) if voted_node == &candidate - )); - } -} diff --git a/src/control_plane/consensus/raft/states/metadata_state.rs b/src/control_plane/consensus/raft/states/metadata_state.rs index 0b14dc2c..3805f3ed 100644 --- a/src/control_plane/consensus/raft/states/metadata_state.rs +++ b/src/control_plane/consensus/raft/states/metadata_state.rs @@ -1,3 +1,4 @@ +use crate::control_plane::metadata::SegmentMeta; use crate::control_plane::metadata::command::*; use crate::control_plane::metadata::event::*; @@ -43,6 +44,12 @@ impl MetadataState { .and_then(|id| self.topics.get(id)) } + pub(crate) fn get_segment(&self, key: &SegmentKey) -> Option<&SegmentMeta> { + let topic = self.get_topic(&key.topic_id)?; + let range = topic.ranges.get(&key.range_id)?; + range.segments.get(&key.segment_id) + } + pub(crate) fn get_consumer_group_assignment( &self, topic_name: &str, diff --git a/src/control_plane/consensus/raft/states/mod.rs b/src/control_plane/consensus/raft/states/mod.rs index c90a9a9b..a3190469 100644 --- a/src/control_plane/consensus/raft/states/mod.rs +++ b/src/control_plane/consensus/raft/states/mod.rs @@ -1,13 +1,13 @@ /*! Ownership boundaries for one Raft group's state. -[`LogState`] owns the in-memory Raft logs and its persistence bookkeeping: +The log half of [`ConsensusState`] owns the in-memory Raft log and its persistence bookkeeping: [`current_term`], [`voted_for`], the in-memory [`LogState::entries`], and the local durability watermark [`stabled_index`]. Its [`last_index`] is derived from the final in-memory Raft log entry; it is not a data-plane WAL position. -[`TransientState`] owns knowledge that may be reconstructed after restart: +The transient half of [`ConsensusState`] owns knowledge that may be reconstructed after restart: the current role, leader identity, replication progress, and [`commit_index`]. A vote is not transient: forgetting [`voted_for`] after a crash could let one replica vote for two candidates in the same term. @@ -29,18 +29,15 @@ relationship is `applied ⊆ committed ⊆ stable ⊆ in-memory log`. The data-plane WAL is separate: it stores segment records, while these boundaries describe the control-plane Raft metadata log. -[`LogState`]: self::log_state::LogState -[`current_term`]: self::log_state::LogState::current_term -[`voted_for`]: self::log_state::LogState::voted_for -[`LogState::entries`]: self::log_state::LogState::entries -[`LogState::unflushed_mutations`]: self::log_state::LogState::unflushed_mutations -[`stabled_index`]: self::log_state::LogState::stabled_index -[`last_index`]: self::log_state::LogState::last_index -[`TransientState`]: self::transient_state::TransientState -[`commit_index`]: self::transient_state::TransientState::commit_index +[`ConsensusState`]: self::consensus::ConsensusState +[`current_term`]: self::consensus::ConsensusState::current_term +[`voted_for`]: self::consensus::ConsensusState::voted_for +[`LogState::entries`]: self::consensus::ConsensusState::log_entries +[`stabled_index`]: self::consensus::ConsensusState::stabled_index +[`last_index`]: self::consensus::ConsensusState::last_log_index +[`commit_index`]: self::consensus::ConsensusState::commit_index [`MetadataState`]: self::metadata_state::MetadataState [`last_applied_index`]: self::metadata_state::MetadataState::last_applied_index */ -pub(crate) mod log_state; +pub(crate) mod consensus; pub(crate) mod metadata_state; -pub(crate) mod transient_state; diff --git a/src/control_plane/consensus/raft/states/transient_state.rs b/src/control_plane/consensus/raft/states/transient_state.rs deleted file mode 100644 index 9341045e..00000000 --- a/src/control_plane/consensus/raft/states/transient_state.rs +++ /dev/null @@ -1,186 +0,0 @@ -use crate::control_plane::NodeId; -use crate::control_plane::consensus::raft::catch_up::CatchUpRepairs; -use crate::control_plane::metadata::MetadataCommand; -use crate::data_plane::SegmentKey; -use std::collections::{BTreeSet, HashMap, HashSet}; -use std::hash::{Hash, Hasher}; - -const ELECTION_JITTER_RANGE: u32 = 20; - -/// Segments whose write leader crashed (sole death). -pub(crate) type LeaderlessSegments = Vec<(SegmentKey, Vec)>; - -pub(crate) struct ElectionJitter { - seed: u64, - counter: u64, -} - -impl ElectionJitter { - fn new(seed: u64) -> Self { - Self { seed, counter: 0 } - } - - pub(crate) fn next(&mut self) -> u32 { - let mut hasher = std::collections::hash_map::DefaultHasher::new(); - self.seed.hash(&mut hasher); - self.counter.hash(&mut hasher); - self.counter += 1; - (hasher.finish() % ELECTION_JITTER_RANGE as u64) as u32 - } -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum Role { - Follower, - Candidate { votes_received: u32 }, - Leader, -} - -// Peer tracking is leader-only. `next_index` is the leader's guess; -// `match_index` is confirmed replication progress. -#[derive(Debug, Clone)] -pub(crate) struct PeerState { - pub(crate) next_index: u64, - pub(crate) match_index: u64, -} - -pub(crate) struct TransientState { - pub(crate) commit_index: u64, - pub(crate) role: Role, - pub(crate) current_leader: Option, - pub(crate) peer_states: HashMap, - pub(crate) learner_states: HashMap, - pub(crate) election_epoch: u64, - pub(crate) election_jitter: ElectionJitter, - - // SegmentKey -> data leader's NodeId: Control plane only actively manages and demands ACKs from the segment's designated Data Leader - pub(crate) confirmed_data_leaders: HashMap, - pub(crate) catch_up: CatchUpRepairs, - pub(crate) pending_proposals: Vec, - pub(crate) leaderless_segments: LeaderlessSegments, - pub(crate) ring_observation_streak: Option<(BTreeSet, u32)>, -} - -impl TransientState { - pub(crate) fn new(election_jitter_seed: u64) -> Self { - Self { - commit_index: 0, - role: Role::Follower, - current_leader: None, - peer_states: HashMap::new(), - learner_states: HashMap::new(), - election_epoch: 0, - election_jitter: ElectionJitter::new(election_jitter_seed), - confirmed_data_leaders: HashMap::new(), - catch_up: CatchUpRepairs::default(), - pending_proposals: Vec::new(), - leaderless_segments: Vec::new(), - ring_observation_streak: None, - } - } - - pub(crate) fn initialize_leader( - &mut self, - node_id: &NodeId, - peers: &HashSet, - next_index: u64, - ) { - self.role = Role::Leader; - self.current_leader = Some(node_id.clone()); - self.peer_states.clear(); - self.learner_states.clear(); - self.ring_observation_streak = None; - self.peer_states.extend(peers.iter().cloned().map(|peer| { - ( - peer, - PeerState { - next_index, - match_index: 0, - }, - ) - })); - } - - pub(crate) fn begin_campaign(&mut self) { - self.role = Role::Candidate { votes_received: 1 }; - } - - pub(crate) fn record_vote(&mut self, quorum: u32) -> bool { - let Role::Candidate { votes_received } = &mut self.role else { - return false; - }; - *votes_received += 1; - *votes_received >= quorum - } - - pub(crate) fn reset_for_follower(&mut self) { - self.role = Role::Follower; - self.current_leader = None; - self.peer_states.clear(); - self.learner_states.clear(); - self.confirmed_data_leaders.clear(); - self.catch_up.clear(); - self.ring_observation_streak = None; - } - - pub(crate) fn advance_election_epoch(&mut self) -> u64 { - self.election_epoch = self.election_epoch.wrapping_add(1); - self.election_epoch - } - - pub(crate) fn record_ring_observation(&mut self, ring: &BTreeSet) -> u32 { - let observations = match self.ring_observation_streak.take() { - Some((previous, count)) if previous == *ring => count.saturating_add(1), - _ => 1, - }; - self.ring_observation_streak = Some((ring.clone(), observations)); - observations - } - - pub(crate) fn is_peer_caught_up(&self, node_id: &NodeId) -> bool { - self.peer_states - .get(node_id) - .is_some_and(|state| state.match_index >= self.commit_index) - } - - pub(crate) fn is_learner_ready_for_promotion(&self, node_id: &NodeId) -> bool { - self.commit_index > 0 - && self - .learner_states - .get(node_id) - .is_some_and(|state| state.match_index >= self.commit_index) - } - - pub(crate) fn replicated_voter_count(&self, index: u64) -> u32 { - self.peer_states - .values() - .filter(|state| state.match_index >= index) - .count() as u32 - + 1 - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn role_transition_reinitializes_leader_only_state() { - let self_id = NodeId::new("node-1"); - let peer = NodeId::new("node-2"); - let mut peers = HashSet::new(); - peers.insert(peer.clone()); - let mut state = TransientState::new(1); - - state.initialize_leader(&self_id, &peers, 7); - assert_eq!(state.role, Role::Leader); - assert_eq!(state.current_leader.as_ref(), Some(&self_id)); - assert_eq!(state.peer_states[&peer].next_index, 7); - - state.reset_for_follower(); - assert_eq!(state.role, Role::Follower); - assert!(state.current_leader.is_none()); - assert!(state.peer_states.is_empty()); - assert!(state.learner_states.is_empty()); - } -} diff --git a/src/control_plane/consensus/seal_recovery.rs b/src/control_plane/consensus/seal_recovery.rs index 0728629d..f5f8beaf 100644 --- a/src/control_plane/consensus/seal_recovery.rs +++ b/src/control_plane/consensus/seal_recovery.rs @@ -11,7 +11,7 @@ use std::collections::{HashMap, HashSet}; use crate::control_plane::NodeId; -use crate::control_plane::consensus::raft::states::transient_state::LeaderlessSegments; +use crate::control_plane::consensus::raft::states::consensus::LeaderlessSegments; use crate::control_plane::membership::ShardGroupId; use crate::control_plane::metadata::EntryId; use crate::data_plane::SegmentKey; From c32e664e70a59c81ed290fdfc0e0c7ed74f7560b Mon Sep 17 00:00:00 2001 From: Migorithm Date: Fri, 17 Jul 2026 23:00:08 +0400 Subject: [PATCH 9/9] consensus layer --- .../raft/states/consensus/log_state.rs | 187 ++++++++++++ .../consensus/raft/states/consensus/mod.rs | 276 ++++++++++++++++++ .../raft/states/consensus/transient_state.rs | 195 +++++++++++++ 3 files changed, 658 insertions(+) create mode 100644 src/control_plane/consensus/raft/states/consensus/log_state.rs create mode 100644 src/control_plane/consensus/raft/states/consensus/mod.rs create mode 100644 src/control_plane/consensus/raft/states/consensus/transient_state.rs diff --git a/src/control_plane/consensus/raft/states/consensus/log_state.rs b/src/control_plane/consensus/raft/states/consensus/log_state.rs new file mode 100644 index 00000000..c53a522a --- /dev/null +++ b/src/control_plane/consensus/raft/states/consensus/log_state.rs @@ -0,0 +1,187 @@ +use crate::control_plane::NodeId; +use crate::control_plane::consensus::messages::LogMutation; +use crate::control_plane::consensus::raft::log::LogEntry; +use crate::control_plane::consensus::raft::storage::RaftPersistentState; + +pub(crate) struct LogState { + current_term: u64, + voted_for: Option, + entries: Vec, + stabled_index: u64, + unflushed_mutations: Vec, +} + +impl LogState { + pub(crate) fn from_persistent(persistent: RaftPersistentState) -> Self { + Self { + stabled_index: persistent.stabled_index(), + current_term: persistent.term, + voted_for: persistent.voted_for, + entries: persistent.log, + unflushed_mutations: Vec::new(), + } + } + + pub(crate) fn current_term(&self) -> u64 { + self.current_term + } + + pub(crate) fn voted_for(&self) -> Option<&NodeId> { + self.voted_for.as_ref() + } + + pub(crate) fn vote_available_for(&self, candidate_id: &NodeId) -> bool { + self.voted_for + .as_ref() + .is_none_or(|voted_for| voted_for == candidate_id) + } + + pub(crate) fn begin_election(&mut self, node_id: &NodeId) -> u64 { + self.current_term += 1; + self.voted_for = Some(node_id.clone()); + self.buffer_hard_state(); + self.current_term + } + + pub(crate) fn grant_vote(&mut self, candidate_id: NodeId) { + self.voted_for = Some(candidate_id); + self.buffer_hard_state(); + } + + pub(crate) fn advance_term(&mut self, new_term: u64) -> bool { + if new_term <= self.current_term { + return false; + } + self.current_term = new_term; + self.voted_for = None; + self.buffer_hard_state(); + true + } + + pub(crate) fn stabled_index(&self) -> u64 { + self.stabled_index + } + + pub(crate) fn advance_stabled_index(&mut self, index: u64) { + self.stabled_index = self.stabled_index.max(index); + } + + pub(crate) fn entries(&self) -> &[LogEntry] { + &self.entries + } + + pub(crate) fn last_index(&self) -> u64 { + self.entries.last().map_or(0, |entry| entry.index) + } + + pub(crate) fn last_term(&self) -> u64 { + self.entries.last().map_or(0, |entry| entry.term) + } + + pub(crate) fn term_at(&self, index: u64) -> u64 { + if index == 0 { + return 0; + } + self.get(index).map_or(0, |entry| entry.term) + } + + pub(crate) fn get(&self, index: u64) -> Option<&LogEntry> { + if index == 0 { + return None; + } + self.entries.get((index - 1) as usize) + } + + pub(crate) fn entries_from(&self, start_index: u64) -> Box<[LogEntry]> { + if start_index == 0 || start_index > self.last_index() { + return Box::new([]); + } + self.entries[(start_index - 1) as usize..].into() + } + + pub(crate) fn append(&mut self, entry: LogEntry) { + debug_assert_eq!( + entry.index, + self.last_index() + 1, + "log entry index must be contiguous" + ); + self.unflushed_mutations + .push(LogMutation::Append(entry.clone())); + self.entries.push(entry); + } + + pub(crate) fn truncate_from(&mut self, from_index: u64) { + if from_index == 0 || from_index > self.last_index() + 1 { + return; + } + self.unflushed_mutations + .push(LogMutation::TruncateFrom(from_index)); + self.entries.truncate((from_index - 1) as usize); + } + + fn buffer_hard_state(&mut self) { + self.unflushed_mutations.push(LogMutation::HardState { + term: self.current_term, + voted_for: self.voted_for.clone(), + }); + } + + pub(crate) fn take_mutations(&mut self) -> Vec { + std::mem::take(&mut self.unflushed_mutations) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::control_plane::consensus::raft::command::RaftCommand; + + #[test] + fn append_and_truncate_emit_persistence_mutations() { + let mut state = LogState::from_persistent(RaftPersistentState::default()); + let entry = LogEntry { + term: 1, + index: 1, + command: RaftCommand::Noop, + }; + + state.append(entry.clone()); + state.truncate_from(1); + + assert_eq!(state.last_index(), 0); + let mutations = state.take_mutations(); + assert!(matches!( + &mutations[..], + [ + LogMutation::Append(appended), + LogMutation::TruncateFrom(1) + ] if appended == &entry + )); + } + + #[test] + fn election_transitions_buffer_hard_state() { + let node = NodeId::new("node-1"); + let candidate = NodeId::new("node-2"); + let mut state = LogState::from_persistent(RaftPersistentState::default()); + + assert_eq!(state.begin_election(&node), 1); + assert_eq!(state.voted_for(), Some(&node)); + assert!(!state.vote_available_for(&candidate)); + + assert!(state.advance_term(2)); + assert!(state.voted_for().is_none()); + state.grant_vote(candidate.clone()); + assert_eq!(state.voted_for(), Some(&candidate)); + + let mutations = state.take_mutations(); + assert_eq!(mutations.len(), 3); + assert!(matches!( + mutations.last(), + Some(LogMutation::HardState { + term: 2, + voted_for: Some(voted_node), + }) if voted_node == &candidate + )); + } +} diff --git a/src/control_plane/consensus/raft/states/consensus/mod.rs b/src/control_plane/consensus/raft/states/consensus/mod.rs new file mode 100644 index 00000000..524780ec --- /dev/null +++ b/src/control_plane/consensus/raft/states/consensus/mod.rs @@ -0,0 +1,276 @@ +mod log_state; +mod transient_state; + +use crate::control_plane::Replicas; +use crate::control_plane::consensus::messages::LogMutation; +use crate::control_plane::consensus::raft::log::LogEntry; +use crate::control_plane::consensus::raft::storage::RaftPersistentState; +use crate::control_plane::membership::ShardGroupId; +use crate::control_plane::metadata::MetadataCommand; +use crate::control_plane::metadata::event::SegmentReassigned; +use crate::data_plane::SegmentKey; +use crate::data_plane::messages::command::SegmentCaughtUp; +use crate::data_plane::transport::command::DataTransportCommand; +use crate::{client::EntryId, control_plane::NodeId}; +use std::{ + collections::{BTreeSet, HashSet}, + ops::RangeInclusive, +}; + +use log_state::LogState; +use transient_state::TransientState; +pub(crate) use transient_state::{LeaderlessSegments, PeerState, Role}; + +/// Owns Raft consensus state while preserving its durable/volatile boundary. +pub(crate) struct ConsensusState { + log: LogState, + transient: TransientState, +} + +impl ConsensusState { + pub(crate) fn new(persistent: RaftPersistentState, election_jitter_seed: u64) -> Self { + Self { + log: LogState::from_persistent(persistent), + transient: TransientState::new(election_jitter_seed), + } + } + + pub(crate) fn take_log_mutations(&mut self) -> Vec { + self.log.take_mutations() + } + pub(crate) fn last_log_index(&self) -> u64 { + self.log.last_index() + } + pub(crate) fn last_log_term(&self) -> u64 { + self.log.last_term() + } + pub(crate) fn log_term_at(&self, index: u64) -> u64 { + self.log.term_at(index) + } + pub(crate) fn log_entry(&self, index: u64) -> Option<&LogEntry> { + self.log.get(index) + } + pub(crate) fn log_entries_from(&self, index: u64) -> Box<[LogEntry]> { + self.log.entries_from(index) + } + pub(crate) fn append_log(&mut self, entry: LogEntry) { + self.log.append(entry); + } + pub(crate) fn truncate_log_from(&mut self, index: u64) { + self.log.truncate_from(index); + } + pub(crate) fn stabled_index(&self) -> u64 { + self.log.stabled_index() + } + /// Highest committed entry that is also locally durable and therefore safe to apply. + pub(crate) fn ready_to_apply_index(&self) -> u64 { + self.transient.commit_index.min(self.log.stabled_index()) + } + pub(crate) fn advance_stabled_index(&mut self, index: u64) { + self.log.advance_stabled_index(index); + } + pub(crate) fn log_entries(&self) -> &[LogEntry] { + self.log.entries() + } + pub(crate) fn current_term(&self) -> u64 { + self.log.current_term() + } + pub(crate) fn voted_for(&self) -> Option<&NodeId> { + self.log.voted_for() + } + pub(crate) fn vote_available_for(&self, node: &NodeId) -> bool { + self.log.vote_available_for(node) + } + pub(crate) fn grant_vote(&mut self, node: NodeId) { + self.log.grant_vote(node); + } + pub(crate) fn advance_term(&mut self, term: u64) -> bool { + self.log.advance_term(term) + } + + pub(crate) fn begin_campaign(&mut self, node_id: &NodeId) -> u64 { + let term = self.log.begin_election(node_id); + self.transient.begin_campaign(); + term + } + pub(crate) fn record_vote(&mut self, term: u64, granted: bool, quorum: u32) -> bool { + term == self.current_term() && granted && self.transient.record_vote(quorum) + } + pub(crate) fn role(&self) -> &Role { + &self.transient.role + } + pub(crate) fn is_leader(&self) -> bool { + self.transient.role == Role::Leader + } + pub(crate) fn current_leader(&self) -> Option<&NodeId> { + self.transient.current_leader.as_ref() + } + pub(crate) fn initialize_leader(&mut self, node: &NodeId, peers: &HashSet) { + let next_index = self.log.last_index() + 1; + self.transient.initialize_leader(node, peers, next_index); + } + pub(crate) fn reset_for_follower(&mut self) { + self.transient.reset_for_follower(); + } + pub(crate) fn recognize_leader(&mut self, leader: NodeId) { + self.transient.recognize_leader(leader); + } + + pub(crate) fn commit_index(&self) -> u64 { + self.transient.commit_index + } + + pub(crate) fn uncommited_log_range(&self) -> RangeInclusive { + self.commit_index() + 1..=self.last_log_index() + } + pub(crate) fn set_commit_index(&mut self, index: u64) { + self.transient.commit_index = index; + } + pub(crate) fn advance_follower_commit(&mut self, leader_commit: u64) { + if leader_commit > self.commit_index() { + self.transient.commit_index = leader_commit.min(self.last_log_index()); + } + } + pub(crate) fn replicated_voter_count(&self, index: u64) -> u32 { + self.transient.replicated_voter_count(index) + } + pub(crate) fn is_peer_caught_up(&self, node: &NodeId) -> bool { + self.transient.is_peer_caught_up(node) + } + pub(crate) fn is_learner_ready_for_promotion(&self, node: &NodeId) -> bool { + self.transient.is_learner_ready_for_promotion(node) + } + pub(crate) fn peer_state(&self, node: &NodeId) -> Option<&PeerState> { + self.transient + .peer_states + .get(node) + .or_else(|| self.transient.learner_states.get(node)) + } + pub(crate) fn peer_state_mut(&mut self, node: &NodeId) -> Option<&mut PeerState> { + self.transient + .peer_states + .get_mut(node) + .or_else(|| self.transient.learner_states.get_mut(node)) + } + pub(crate) fn is_voter(&self, node: &NodeId) -> bool { + self.transient.peer_states.contains_key(node) + } + pub(crate) fn voter_state_count(&self) -> usize { + self.transient.peer_states.len() + } + pub(crate) fn has_voter_state(&self, node: &NodeId) -> bool { + self.transient.peer_states.contains_key(node) + } + pub(crate) fn voter_states_empty(&self) -> bool { + self.transient.peer_states.is_empty() + } + pub(crate) fn learner_states_empty(&self) -> bool { + self.transient.learner_states.is_empty() + } + pub(crate) fn learner_state_count(&self) -> usize { + self.transient.learner_states.len() + } + pub(crate) fn learner_ids(&self) -> impl Iterator { + self.transient.learner_states.keys() + } + pub(crate) fn replication_targets(&self) -> Vec { + self.transient + .peer_states + .keys() + .chain(self.transient.learner_states.keys()) + .cloned() + .collect() + } + pub(crate) fn is_learner(&self, node: &NodeId) -> bool { + self.transient.learner_states.contains_key(node) + } + pub(crate) fn stage_learner(&mut self, node: NodeId, state: PeerState) { + self.transient.learner_states.insert(node, state); + } + pub(crate) fn remove_learner(&mut self, node: &NodeId) -> Option { + self.transient.learner_states.remove(node) + } + pub(crate) fn add_voter_state(&mut self, node: NodeId, state: PeerState) { + self.transient.peer_states.insert(node, state); + } + pub(crate) fn remove_voter_state(&mut self, node: &NodeId) { + self.transient.peer_states.remove(node); + } + + pub(crate) fn clear_ring_observation(&mut self) { + self.transient.ring_observation_streak = None; + } + pub(crate) fn record_ring_observation(&mut self, ring: &BTreeSet) -> u32 { + self.transient.record_ring_observation(ring) + } + pub(crate) fn next_election_epoch(&mut self) -> u64 { + self.transient.advance_election_epoch() + } + pub(crate) fn election_epoch(&self) -> u64 { + self.transient.election_epoch + } + pub(crate) fn next_election_jitter(&mut self) -> u32 { + self.transient.election_jitter.next() + } + + pub(crate) fn take_pending_proposals(&mut self) -> Vec { + std::mem::take(&mut self.transient.pending_proposals) + } + pub(crate) fn push_pending_proposal(&mut self, command: MetadataCommand) { + self.transient.pending_proposals.push(command); + } + pub(crate) fn extend_pending_proposals( + &mut self, + commands: impl IntoIterator, + ) { + self.transient.pending_proposals.extend(commands); + } + pub(crate) fn take_leaderless_segments(&mut self) -> LeaderlessSegments { + std::mem::take(&mut self.transient.leaderless_segments) + } + pub(crate) fn push_leaderless_segment(&mut self, segment: (SegmentKey, Vec)) { + self.transient.leaderless_segments.push(segment); + } + pub(crate) fn extend_leaderless_segments(&mut self, segments: LeaderlessSegments) { + self.transient.leaderless_segments.extend(segments); + } + pub(crate) fn retain_confirmed_data_leaders(&mut self, active: &HashSet) { + self.transient + .confirmed_data_leaders + .retain(|key, _| active.contains(key)); + } + pub(crate) fn is_data_leader_confirmed(&self, key: &SegmentKey, node: &NodeId) -> bool { + self.transient.confirmed_data_leaders.get(key) == Some(node) + } + pub(crate) fn confirm_data_leader(&mut self, key: SegmentKey, node: NodeId) { + self.transient.confirmed_data_leaders.insert(key, node); + } + pub(crate) fn catch_up_redrives(&self, group_id: ShardGroupId) -> Vec { + self.transient.catch_up.redrives(group_id) + } + pub(crate) fn confirm_catch_up(&mut self, ack: SegmentCaughtUp) { + self.transient.catch_up.confirm(ack.segment_key, ack.from); + } + pub(crate) fn track_catch_up(&mut self, reassigned: &SegmentReassigned) { + self.transient.catch_up.track(reassigned); + } + pub(crate) fn track_sealed_catch_up( + &mut self, + key: SegmentKey, + start: EntryId, + end: EntryId, + replicas: Replicas, + ) { + self.transient + .catch_up + .track_sealed(key, start, end, replicas); + } + #[cfg(test)] + pub(crate) fn clear_catch_up(&mut self) { + self.transient.catch_up.clear(); + } + #[cfg(test)] + pub(crate) fn catch_up_is_empty(&self) -> bool { + self.transient.catch_up.is_empty() + } +} diff --git a/src/control_plane/consensus/raft/states/consensus/transient_state.rs b/src/control_plane/consensus/raft/states/consensus/transient_state.rs new file mode 100644 index 00000000..9b4e99db --- /dev/null +++ b/src/control_plane/consensus/raft/states/consensus/transient_state.rs @@ -0,0 +1,195 @@ +use crate::control_plane::NodeId; +use crate::control_plane::consensus::raft::catch_up::CatchUpRepairs; +use crate::control_plane::metadata::MetadataCommand; +use crate::data_plane::SegmentKey; +use std::collections::{BTreeSet, HashMap, HashSet}; +use std::hash::{Hash, Hasher}; + +const ELECTION_JITTER_RANGE: u32 = 20; + +/// Segments whose write leader crashed (sole death). +pub(crate) type LeaderlessSegments = Vec<(SegmentKey, Vec)>; + +pub(crate) struct ElectionJitter { + seed: u64, + counter: u64, +} + +impl ElectionJitter { + fn new(seed: u64) -> Self { + Self { seed, counter: 0 } + } + + pub(crate) fn next(&mut self) -> u32 { + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + self.seed.hash(&mut hasher); + self.counter.hash(&mut hasher); + self.counter += 1; + (hasher.finish() % ELECTION_JITTER_RANGE as u64) as u32 + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Role { + Follower, + Candidate { votes_received: u32 }, + Leader, +} + +// Peer tracking is leader-only. `next_index` is the leader's guess; +// `match_index` is confirmed replication progress. +#[derive(Debug, Clone)] +pub(crate) struct PeerState { + pub(crate) next_index: u64, + pub(crate) match_index: u64, +} + +pub(crate) struct TransientState { + pub(crate) commit_index: u64, + pub(crate) role: Role, + pub(crate) current_leader: Option, + pub(crate) peer_states: HashMap, + pub(crate) learner_states: HashMap, + pub(crate) election_epoch: u64, + pub(crate) election_jitter: ElectionJitter, + + // SegmentKey -> data leader's NodeId: Control plane only actively manages and demands ACKs from the segment's designated Data Leader + pub(crate) confirmed_data_leaders: HashMap, + pub(crate) catch_up: CatchUpRepairs, + pub(crate) pending_proposals: Vec, + pub(crate) leaderless_segments: LeaderlessSegments, + pub(crate) ring_observation_streak: Option<(BTreeSet, u32)>, +} + +impl TransientState { + pub(crate) fn new(election_jitter_seed: u64) -> Self { + Self { + commit_index: 0, + role: Role::Follower, + current_leader: None, + peer_states: HashMap::new(), + learner_states: HashMap::new(), + election_epoch: 0, + election_jitter: ElectionJitter::new(election_jitter_seed), + confirmed_data_leaders: HashMap::new(), + catch_up: CatchUpRepairs::default(), + pending_proposals: Vec::new(), + leaderless_segments: Vec::new(), + ring_observation_streak: None, + } + } + + pub(crate) fn initialize_leader( + &mut self, + node_id: &NodeId, + peers: &HashSet, + next_index: u64, + ) { + self.role = Role::Leader; + self.current_leader = Some(node_id.clone()); + self.peer_states.clear(); + self.learner_states.clear(); + self.ring_observation_streak = None; + self.peer_states.extend(peers.iter().cloned().map(|peer| { + ( + peer, + PeerState { + next_index, + match_index: 0, + }, + ) + })); + } + + pub(crate) fn begin_campaign(&mut self) { + self.role = Role::Candidate { votes_received: 1 }; + } + + pub(crate) fn record_vote(&mut self, quorum: u32) -> bool { + let Role::Candidate { votes_received } = &mut self.role else { + return false; + }; + *votes_received += 1; + *votes_received >= quorum + } + + pub(crate) fn reset_for_follower(&mut self) { + self.role = Role::Follower; + self.current_leader = None; + self.peer_states.clear(); + self.learner_states.clear(); + self.confirmed_data_leaders.clear(); + self.catch_up.clear(); + self.ring_observation_streak = None; + } + + pub(crate) fn recognize_leader(&mut self, leader: NodeId) { + if self.role != Role::Follower { + self.role = Role::Follower; + self.peer_states.clear(); + self.learner_states.clear(); + } + self.current_leader = Some(leader); + } + + pub(crate) fn advance_election_epoch(&mut self) -> u64 { + self.election_epoch = self.election_epoch.wrapping_add(1); + self.election_epoch + } + + pub(crate) fn record_ring_observation(&mut self, ring: &BTreeSet) -> u32 { + let observations = match self.ring_observation_streak.take() { + Some((previous, count)) if previous == *ring => count.saturating_add(1), + _ => 1, + }; + self.ring_observation_streak = Some((ring.clone(), observations)); + observations + } + + pub(crate) fn is_peer_caught_up(&self, node_id: &NodeId) -> bool { + self.peer_states + .get(node_id) + .is_some_and(|state| state.match_index >= self.commit_index) + } + + pub(crate) fn is_learner_ready_for_promotion(&self, node_id: &NodeId) -> bool { + self.commit_index > 0 + && self + .learner_states + .get(node_id) + .is_some_and(|state| state.match_index >= self.commit_index) + } + + pub(crate) fn replicated_voter_count(&self, index: u64) -> u32 { + self.peer_states + .values() + .filter(|state| state.match_index >= index) + .count() as u32 + + 1 + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn role_transition_reinitializes_leader_only_state() { + let self_id = NodeId::new("node-1"); + let peer = NodeId::new("node-2"); + let mut peers = HashSet::new(); + peers.insert(peer.clone()); + let mut state = TransientState::new(1); + + state.initialize_leader(&self_id, &peers, 7); + assert_eq!(state.role, Role::Leader); + assert_eq!(state.current_leader.as_ref(), Some(&self_id)); + assert_eq!(state.peer_states[&peer].next_index, 7); + + state.reset_for_follower(); + assert_eq!(state.role, Role::Follower); + assert!(state.current_leader.is_none()); + assert!(state.peer_states.is_empty()); + assert!(state.learner_states.is_empty()); + } +}