Skip to main content

nodedb_raft/
state.rs

1/// Persistent state that must survive restarts.
2///
3/// Corresponds to Raft paper Figure 2 "Persistent state on all servers".
4#[derive(
5    Debug,
6    Clone,
7    PartialEq,
8    Eq,
9    serde::Serialize,
10    serde::Deserialize,
11    zerompk::ToMessagePack,
12    zerompk::FromMessagePack,
13)]
14pub struct HardState {
15    /// Latest term this server has seen.
16    pub current_term: u64,
17    /// Candidate that received vote in current term (0 = none).
18    pub voted_for: u64,
19}
20
21impl HardState {
22    pub fn new() -> Self {
23        Self {
24            current_term: 0,
25            voted_for: 0,
26        }
27    }
28}
29
30impl Default for HardState {
31    fn default() -> Self {
32        Self::new()
33    }
34}
35
36/// Role of a Raft node within a group.
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub enum NodeRole {
39    Follower,
40    Candidate,
41    Leader,
42    /// Non-voting member catching up : new node joins as learner first).
43    Learner,
44}
45
46/// Volatile state on all servers.
47#[derive(Debug, Clone)]
48pub struct VolatileState {
49    /// Index of highest log entry known to be committed.
50    pub commit_index: u64,
51    /// Index of highest log entry applied to state machine.
52    pub last_applied: u64,
53}
54
55impl VolatileState {
56    pub fn new() -> Self {
57        Self {
58            commit_index: 0,
59            last_applied: 0,
60        }
61    }
62}
63
64impl Default for VolatileState {
65    fn default() -> Self {
66        Self::new()
67    }
68}
69
70/// Volatile state on leaders (reinitialized after election).
71#[derive(Debug, Clone)]
72pub struct LeaderState {
73    /// For each peer: index of next log entry to send.
74    pub next_index: Vec<(u64, u64)>,
75    /// For each peer: index of highest log entry known to be replicated.
76    pub match_index: Vec<(u64, u64)>,
77}
78
79impl LeaderState {
80    pub fn new(peers: &[u64], last_log_index: u64) -> Self {
81        Self {
82            next_index: peers.iter().map(|&id| (id, last_log_index + 1)).collect(),
83            match_index: peers.iter().map(|&id| (id, 0)).collect(),
84        }
85    }
86
87    pub fn next_index_for(&self, peer: u64) -> u64 {
88        self.next_index
89            .iter()
90            .find(|&&(id, _)| id == peer)
91            .map(|&(_, idx)| idx)
92            .unwrap_or(1)
93    }
94
95    pub fn set_next_index(&mut self, peer: u64, index: u64) {
96        if let Some(entry) = self.next_index.iter_mut().find(|e| e.0 == peer) {
97            entry.1 = index;
98        }
99    }
100
101    pub fn match_index_for(&self, peer: u64) -> u64 {
102        self.match_index
103            .iter()
104            .find(|&&(id, _)| id == peer)
105            .map(|&(_, idx)| idx)
106            .unwrap_or(0)
107    }
108
109    pub fn set_match_index(&mut self, peer: u64, index: u64) {
110        if let Some(entry) = self.match_index.iter_mut().find(|e| e.0 == peer) {
111            entry.1 = index;
112        }
113    }
114
115    /// Add a new peer to leader tracking. Initializes next_index to
116    /// `last_log_index + 1` (the peer needs to catch up from the end of the log).
117    pub fn add_peer(&mut self, peer: u64, last_log_index: u64) {
118        if !self.next_index.iter().any(|&(id, _)| id == peer) {
119            self.next_index.push((peer, last_log_index + 1));
120            self.match_index.push((peer, 0));
121        }
122    }
123
124    /// Remove a peer from leader tracking.
125    pub fn remove_peer(&mut self, peer: u64) {
126        self.next_index.retain(|&(id, _)| id != peer);
127        self.match_index.retain(|&(id, _)| id != peer);
128    }
129
130    /// Current peer list tracked by this leader state.
131    pub fn peers(&self) -> Vec<u64> {
132        self.next_index.iter().map(|&(id, _)| id).collect()
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139
140    #[test]
141    fn hard_state_default() {
142        let hs = HardState::new();
143        assert_eq!(hs.current_term, 0);
144        assert_eq!(hs.voted_for, 0);
145    }
146
147    #[test]
148    fn leader_state_initialization() {
149        let peers = vec![2, 3, 4];
150        let ls = LeaderState::new(&peers, 10);
151        assert_eq!(ls.next_index_for(2), 11);
152        assert_eq!(ls.next_index_for(3), 11);
153        assert_eq!(ls.match_index_for(2), 0);
154    }
155
156    #[test]
157    fn leader_state_update() {
158        let peers = vec![2, 3];
159        let mut ls = LeaderState::new(&peers, 5);
160        ls.set_next_index(2, 8);
161        ls.set_match_index(2, 7);
162        assert_eq!(ls.next_index_for(2), 8);
163        assert_eq!(ls.match_index_for(2), 7);
164        // Peer 3 unchanged.
165        assert_eq!(ls.next_index_for(3), 6);
166    }
167
168    #[test]
169    fn node_role_equality() {
170        assert_eq!(NodeRole::Follower, NodeRole::Follower);
171        assert_ne!(NodeRole::Follower, NodeRole::Leader);
172    }
173}