krabka_raft/controller/
membership.rs1use {
7 super::ControllerHandle,
8 crate::{
9 error::RaftError,
10 reconfig::{AddVoter, ReconfigOutcome, RemoveVoter, UpdateVoter, VoterChange},
11 types::{Node, NodeId},
12 },
13 std::{
14 collections::BTreeSet,
15 future::{self, Ready},
16 },
17};
18
19impl ControllerHandle {
20 pub async fn change_membership(&self, new_voters: BTreeSet<NodeId>) -> Result<(), RaftError> {
25 let current = self.engine.quorum_snapshot().voters;
26 let current_ids: BTreeSet<NodeId> = current.ids().into_iter().collect();
27 let added: Vec<NodeId> = new_voters.difference(¤t_ids).copied().collect();
28 let removed: Vec<NodeId> = current_ids.difference(&new_voters).copied().collect();
29 if added.len() + removed.len() > 1 {
30 return Err(RaftError::ReconfigRejected(
31 "only one voter change may be submitted at a time".into(),
32 ));
33 }
34 let outcome = if let Some(id) = removed.first() {
35 let voter = current.get(*id).ok_or_else(|| {
36 RaftError::ReconfigRejected(format!(
37 "voter {id} disappeared while preparing removal"
38 ))
39 })?;
40 self.remove_voter(RemoveVoter {
41 id: *id,
42 directory_id: voter.directory_id,
43 })
44 .await?
45 } else if let Some(id) = added.first() {
46 let node = self
47 .staged_learners
48 .lock()
49 .map_err(|_| RaftError::ReconfigRejected("staged learner lock poisoned".into()))?
50 .get(id)
51 .cloned()
52 .ok_or_else(|| {
53 RaftError::ReconfigRejected(format!(
54 "voter {id} must be staged with add_learner first"
55 ))
56 })?;
57 self.add_voter(AddVoter {
58 voter: krabka_metadata::Voter {
59 id: *id,
60 directory_id: node.directory_id,
61 endpoints: node.endpoints,
62 kraft_version: node.kraft_version,
63 },
64 ack_when_committed: true,
65 })
66 .await?
67 } else {
68 return Ok(());
69 };
70 match outcome {
71 ReconfigOutcome::Committed => Ok(()),
72 ReconfigOutcome::NotLeader { leader } => Err(RaftError::NotLeader {
73 current_leader: leader,
74 }),
75 }
76 }
77
78 pub fn add_learner(&self, node_id: NodeId, node: Node) -> Ready<Result<(), RaftError>> {
84 let result = self
85 .staged_learners
86 .lock()
87 .map_err(|_| RaftError::ReconfigRejected("staged learner lock poisoned".into()))
88 .map(|mut learners| {
89 learners.insert(node_id, node);
90 });
91 future::ready(result)
92 }
93
94 pub async fn add_voter(&self, req: AddVoter) -> Result<ReconfigOutcome, RaftError> {
99 self.engine.reconfigure(VoterChange::Add(req)).await
100 }
101
102 pub async fn check_add_voter(&self, req: AddVoter) -> Result<ReconfigOutcome, RaftError> {
110 self.engine.reconfigure(VoterChange::CheckAdd(req)).await
111 }
112
113 pub async fn remove_voter(&self, req: RemoveVoter) -> Result<ReconfigOutcome, RaftError> {
118 self.engine.reconfigure(VoterChange::Remove(req)).await
119 }
120
121 pub async fn update_voter(&self, req: UpdateVoter) -> Result<ReconfigOutcome, RaftError> {
126 self.engine.reconfigure(VoterChange::Update(req)).await
127 }
128
129 pub async fn finalize_kraft_version(&self, version: u16) -> Result<ReconfigOutcome, RaftError> {
134 self.engine.finalize_kraft_version(version).await
135 }
136
137 pub async fn validate_kraft_version(&self, version: u16) -> Result<ReconfigOutcome, RaftError> {
143 self.engine.validate_kraft_version(version).await
144 }
145}
146
147#[cfg(test)]
148mod tests {
149 use {
150 super::*, crate::controller::test_support::bootstrap_controller, std::collections::BTreeSet,
151 };
152
153 #[tokio::test]
154 async fn change_membership_validates_delta_count() {
155 let (_dir, ctrl) = bootstrap_controller("bootstrap").await;
156
157 let res_zero = ctrl.change_membership(BTreeSet::from([NodeId(1)])).await;
159 assert2::assert!(res_zero.is_ok());
160
161 let res_swap = ctrl.change_membership(BTreeSet::from([NodeId(2)])).await;
163 assert2::assert!(matches!(
164 res_swap,
165 Err(RaftError::ReconfigRejected(ref msg)) if msg.contains("only one voter change")
166 ));
167
168 let res_multi_add = ctrl
170 .change_membership(BTreeSet::from([NodeId(1), NodeId(2), NodeId(3)]))
171 .await;
172 assert2::assert!(matches!(
173 res_multi_add,
174 Err(RaftError::ReconfigRejected(ref msg)) if msg.contains("only one voter change")
175 ));
176
177 let res_single_add = ctrl
179 .change_membership(BTreeSet::from([NodeId(1), NodeId(2)]))
180 .await;
181 assert2::assert!(matches!(
182 res_single_add,
183 Err(RaftError::ReconfigRejected(ref msg)) if msg.contains("must be staged with add_learner first")
184 ));
185
186 ctrl.shutdown().await;
187 }
188}