1#[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 pub current_term: u64,
17 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub enum NodeRole {
39 Follower,
40 Candidate,
41 Leader,
42 Learner,
44}
45
46#[derive(Debug, Clone)]
48pub struct VolatileState {
49 pub commit_index: u64,
51 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#[derive(Debug, Clone)]
72pub struct LeaderState {
73 pub next_index: Vec<(u64, u64)>,
75 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 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 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 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 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}