Skip to main content

krabka_raft/controller/
membership.rs

1//! The voter-administration surface of [`ControllerHandle`]: the KIP-853 add,
2//! remove and update operations, the `kraft.version` finalization, the staging
3//! area for an observer awaiting promotion, and the whole-set membership delta
4//! that is reconciled into a single one of those engine commands.
5
6use {
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    /// Reconcile a single-node voter-set delta through the KIP-853 engine.
21    ///
22    /// # Errors
23    /// Rejects multi-node batch changes; KIP-642 is a separate operation.
24    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(&current_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    /// Stage a caught-up observer identity for later voter promotion.
79    ///
80    /// # Errors
81    /// The observer catches up by fetching from the leader; this call does not
82    /// alter quorum membership.
83    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    /// Add a caught-up controller voter through the Raft control log.
95    ///
96    /// # Errors
97    /// Returns a validation, leadership, timeout, or storage error from Raft.
98    pub async fn add_voter(&self, req: AddVoter) -> Result<ReconfigOutcome, RaftError> {
99        self.engine.reconfigure(VoterChange::Add(req)).await
100    }
101
102    /// Run the leader-local admission checks of `AddVoterHandler` for `req`
103    /// without contacting the candidate or appending anything.
104    ///
105    /// # Errors
106    /// Returns the refusal the same request would get from [`Self::add_voter`]
107    /// for a pending change, an unset high watermark, `kraft.version` below 1,
108    /// an uncommitted voters record or a duplicate voter id.
109    pub async fn check_add_voter(&self, req: AddVoter) -> Result<ReconfigOutcome, RaftError> {
110        self.engine.reconfigure(VoterChange::CheckAdd(req)).await
111    }
112
113    /// Remove the exact node/directory pair through the Raft control log.
114    ///
115    /// # Errors
116    /// Returns a validation, leadership, timeout, or storage error from Raft.
117    pub async fn remove_voter(&self, req: RemoveVoter) -> Result<ReconfigOutcome, RaftError> {
118        self.engine.reconfigure(VoterChange::Remove(req)).await
119    }
120
121    /// Update the exact voter's endpoint and supported feature range.
122    ///
123    /// # Errors
124    /// Returns a validation, leadership, timeout, or storage error from Raft.
125    pub async fn update_voter(&self, req: UpdateVoter) -> Result<ReconfigOutcome, RaftError> {
126        self.engine.reconfigure(VoterChange::Update(req)).await
127    }
128
129    /// Atomically append `KRaftVersionRecord` and the initial `VotersRecord`.
130    ///
131    /// # Errors
132    /// Returns an unsupported-version, leadership, timeout, or storage error.
133    pub async fn finalize_kraft_version(&self, version: u16) -> Result<ReconfigOutcome, RaftError> {
134        self.engine.finalize_kraft_version(version).await
135    }
136
137    /// Run the checks of [`Self::finalize_kraft_version`] and write nothing.
138    ///
139    /// # Errors
140    /// Returns the error `finalize_kraft_version` would return for the same
141    /// request.
142    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        // 0 changes: current is {1}, target is {1} -> Ok(())
158        let res_zero = ctrl.change_membership(BTreeSet::from([NodeId(1)])).await;
159        assert2::assert!(res_zero.is_ok());
160
161        // 1 added + 1 removed = 2 changes: {1} -> {2} -> ReconfigRejected
162        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        // 2 added = 2 changes: {1} -> {1, 2, 3} -> ReconfigRejected
169        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        // 1 change: {1} -> {1, 2} -> fails because NodeId(2) is not staged as learner
178        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}