krabka_raft/
controller.rs

1//! `Controller` is the public entry point. It owns the async
2//! [`KraftController`] consensus engine, the controller TCP listener, and the
3//! `submit_change` leader-aware forwarding logic, behind a stable
4//! [`ControllerHandle`] API the broker depends on.
5//!
6//! Cluster formation is driven by `BootstrapMode`: a fresh `Bootstrap`/`Join`
7//! node seeds its quorum state from configured static voters or a dynamic
8//! bootstrap snapshot; a restarted `Rejoin` node recovers from its on-disk
9//! metadata log, checkpoint, and quorum-state file (handled inside
10//! [`KraftController::open`]).
11//!
12//! KIP-853 voter changes are serialized by the same single-owner engine that
13//! appends, commits, truncates, and snapshots the metadata log.
14
15use {
16    krabka_metadata::MetadataImage,
17    std::{
18        collections::BTreeMap,
19        net::SocketAddr,
20        path::PathBuf,
21        sync::{self, Arc},
22    },
23    tokio::{
24        sync::{Mutex, watch},
25        task::JoinHandle,
26    },
27    tokio_util::sync::CancellationToken,
28    uuid::Uuid,
29};
30
31mod checkpoint;
32mod feature_check;
33mod membership;
34mod metadata_fetch;
35mod startup;
36mod submit;
37#[cfg(test)]
38mod test_support;
39
40use crate::{
41    RegistrationAdmission,
42    kraft::KraftController,
43    network::OutboundDialer,
44    types::{Node, NodeId},
45};
46pub use {
47    self::{
48        checkpoint::{SnapshotRange, SnapshotSlice},
49        startup::{Controller, metadata_log_nonempty},
50    },
51    crate::kraft::transport::QuorumStateSnapshot,
52};
53
54fn connection_options(
55    client_id: &str,
56    dispatch_queue_capacity: krabka_client_core::ConnectionDispatchQueueCapacity,
57    frame_max: krabka_client_core::ClientFrameMax,
58) -> krabka_client_core::ConnectionOptions {
59    krabka_client_core::ConnectionOptions {
60        client_id: client_id.to_owned(),
61        dispatch_queue_capacity,
62        frame_max,
63        ..krabka_client_core::ConnectionOptions::default()
64    }
65}
66
67/// Krabka-native view of the controller's current quorum state.
68///
69/// Surfaced by [`ControllerHandle::quorum_state`] for the broker's
70/// `DescribeQuorum` admin handler so callers don't depend on engine internals
71/// directly.
72#[derive(Debug, Clone, Default)]
73pub struct QuorumState {
74    /// `KRaft` leader epoch (on the wire as `leader_epoch`).
75    pub current_term: u64,
76    /// High watermark — the last committed/applied offset on this node. `0`
77    /// until the first commit.
78    pub last_applied_index: u64,
79    /// Current cluster leader. `None` mid-election.
80    pub current_leader: Option<NodeId>,
81    /// Voter ids in the current (static) membership.
82    pub voters: Vec<NodeId>,
83    /// Full voter node identities (directory id + endpoints + kraft.version)
84    /// keyed by node id. Mirrors `voters`; carries the KIP-853 voter metadata
85    /// the `DescribeQuorum` path needs.
86    pub voter_nodes: BTreeMap<NodeId, Node>,
87    /// Per-voter fetch offset (matched index). Empty on a follower; callers use
88    /// Kafka's unknown sentinel.
89    pub per_voter_matched_index: BTreeMap<NodeId, u64>,
90    /// Per-replica last fetch timestamp in wall-clock milliseconds.
91    pub per_replica_last_fetch_ms: BTreeMap<NodeId, i64>,
92    /// Per-replica last caught-up timestamp in wall-clock milliseconds.
93    pub per_replica_last_caught_up_ms: BTreeMap<NodeId, i64>,
94    /// Discovered directory identities for observer replicas.
95    pub observer_directory_ids: BTreeMap<NodeId, uuid::Uuid>,
96    /// Whether this node currently leads the metadata quorum.
97    pub is_leader: bool,
98}
99
100/// Handle returned by [`Controller::start`].
101///
102/// Owns the live [`KraftController`] engine and the listener task. Drop is NOT
103/// a clean shutdown — call [`Self::shutdown`] (or [`Self::cancel`]) to drain
104/// the listener + stop the engine before the runtime is torn down.
105pub struct ControllerHandle {
106    engine: KraftController,
107    leader: watch::Receiver<Option<NodeId>>,
108    shutdown: CancellationToken,
109    /// The reason this controller stopped itself over a fatal fault, if it did.
110    fatal: watch::Receiver<Option<String>>,
111    listener_task: Mutex<Option<JoinHandle<()>>>,
112    /// Directory holding the metadata log + KIP-630 `.checkpoint` artifacts.
113    data_dir: PathBuf,
114    client_id: String,
115    client_dispatch_queue_capacity: krabka_client_core::ConnectionDispatchQueueCapacity,
116    client_frame_max: krabka_client_core::ClientFrameMax,
117    /// This node's own id, used for leader and membership checks.
118    self_node_id: NodeId,
119    /// Configured bootstrap voter set. Dynamic membership comes from the
120    /// engine snapshot; this remains only as an address fallback during
121    /// initial discovery at kraft.version 0.
122    voters: krabka_metadata::VoterSet,
123    /// Compatibility staging area for callers that still separate observer
124    /// registration from promotion. Membership itself is changed only by an
125    /// engine-owned KIP-853 command.
126    staged_learners: sync::Mutex<BTreeMap<NodeId, Node>>,
127    /// Outbound dialer; `forward_submit_to`/`fetch_metadata_from` reach a peer's
128    /// controller listener with the same TLS/SASL handshake the engine's RPCs
129    /// ride on.
130    dialer: Arc<dyn OutboundDialer>,
131    /// The address the controller listener actually bound to (resolved port when
132    /// `controller_listen_addr` requested port 0).
133    controller_bound_addr: SocketAddr,
134}
135
136impl ControllerHandle {
137    /// Current metadata snapshot (cheap; `Arc` clone).
138    #[must_use]
139    pub fn current_image(&self) -> Arc<MetadataImage> { self.engine.current_image() }
140
141    /// Install the hook every batch that registers a broker anew passes
142    /// through while this node leads. See [`RegistrationAdmission`].
143    ///
144    /// # Errors
145    /// Returns an error if a hook is already installed.
146    pub fn install_registration_admission(
147        &self,
148        admission: Arc<dyn RegistrationAdmission>,
149    ) -> Result<(), &'static str> {
150        self.engine.install_registration_admission(admission)
151    }
152
153    /// The address the controller listener actually bound to.
154    #[must_use]
155    pub fn controller_bound_addr(&self) -> SocketAddr { self.controller_bound_addr }
156
157    /// Subscribe to leader-id changes.
158    #[must_use]
159    pub fn watch_leader(&self) -> watch::Receiver<Option<NodeId>> { self.leader.clone() }
160
161    /// Subscribe to the controller's fatal fault: the message of the fault
162    /// that made it stop itself, and `None` until then.
163    ///
164    /// Kafka handles such a fault (a `FeatureControlManager` replay of a
165    /// feature level that the controller does not support) with a
166    /// `ProcessTerminatingFaultHandler`, which halts the process. A library
167    /// cannot halt its host, so the controller stops itself, publishes the
168    /// message here before it does, and leaves the process to its embedder:
169    /// the broker turns it into its own shutdown. The channel closes without a
170    /// value when the controller is shut down for any other reason.
171    #[must_use]
172    pub fn watch_fatal(&self) -> watch::Receiver<Option<String>> { self.fatal.clone() }
173
174    /// Subscribe to metadata-image changes.
175    #[must_use]
176    pub fn watch_image(&self) -> watch::Receiver<Arc<MetadataImage>> { self.engine.watch_image() }
177
178    /// Snapshot the controller's current quorum state. Used by the broker's
179    /// `DescribeQuorum` (`api_key=55`, KIP-595) handler. Cheap — a `watch`
180    /// borrow of the engine's published snapshot.
181    #[must_use]
182    pub fn quorum_state(&self) -> QuorumState {
183        let snap = self.engine.quorum_snapshot();
184        let voter_nodes: BTreeMap<NodeId, Node> = snap
185            .voters
186            .iter()
187            .map(|v| {
188                (v.id, Node {
189                    directory_id: v.directory_id,
190                    endpoints: v.endpoints.clone(),
191                    kraft_version: v.kraft_version,
192                })
193            })
194            .collect();
195        let per_voter_matched_index: BTreeMap<NodeId, u64> = snap
196            .per_replica_fetch_offset
197            .iter()
198            .map(|(id, off)| (*id, u64::try_from((*off).max(0)).unwrap_or(0)))
199            .collect();
200        QuorumState {
201            current_term: u64::from(snap.leader_epoch),
202            last_applied_index: u64::try_from(snap.high_watermark.max(0)).unwrap_or(0),
203            current_leader: snap.leader_id,
204            voters: snap.voters.ids().into_iter().collect(),
205            voter_nodes,
206            per_voter_matched_index,
207            per_replica_last_fetch_ms: snap.per_replica_last_fetch_ms,
208            per_replica_last_caught_up_ms: snap.per_replica_last_caught_up_ms,
209            observer_directory_ids: snap
210                .observers
211                .iter()
212                .map(|observer| (observer.id, observer.directory_id))
213                .collect(),
214            is_leader: snap.is_leader,
215        }
216    }
217
218    /// Snapshot the full consensus state.
219    #[must_use]
220    pub fn quorum_snapshot(&self) -> QuorumStateSnapshot { self.engine.quorum_snapshot() }
221
222    /// Highest `__cluster_metadata` offset the quorum has committed, as this
223    /// node last heard it, or `-1` before the first committed record.
224    ///
225    /// This is the offset a readiness probe compares against. It differs from
226    /// [`QuorumState::last_applied_index`] on a follower that is still
227    /// replaying: `last_applied_index` is clamped to what this node has
228    /// applied, while this is what the leader said it had committed.
229    #[must_use]
230    pub fn quorum_committed_offset(&self) -> i64 {
231        self.engine
232            .quorum_snapshot()
233            .quorum_high_watermark
234            .saturating_sub(1)
235    }
236
237    /// Directory identity voted for in the current leader epoch, if any.
238    #[must_use]
239    pub fn voted_directory_id(&self) -> Option<Uuid> {
240        self.engine.quorum_snapshot().voted_directory_id
241    }
242
243    /// Drain the listener and stop the engine. Idempotent in practice.
244    pub async fn shutdown(self) { self.cancel().await; }
245
246    /// Stop the engine and cancel the controller listener without consuming
247    /// `self`. Used by `BrokerHandle::shutdown` where the controller is behind
248    /// an `Arc`. Idempotent. Awaits the listener task so the OS port is released
249    /// before returning.
250    ///
251    /// The engine stops gracefully, before the listener does: a leader resigns,
252    /// and its successor's `BeginQuorumEpoch` still reaches the listener, which
253    /// is how the engine learns that the handover is complete.
254    pub async fn cancel(&self) {
255        self.engine.shutdown().await;
256        self.stop_listener().await;
257    }
258
259    /// Stop the engine at once, handing nothing over, and cancel the listener,
260    /// as a halted process stops. The quorum finds out through its fetch
261    /// watchdog. Idempotent.
262    pub async fn halt(&self) {
263        self.shutdown.cancel();
264        self.engine.halt().await;
265        self.stop_listener().await;
266    }
267
268    async fn stop_listener(&self) {
269        self.shutdown.cancel();
270        // The lock stays held until the listener has stopped, so a second
271        // caller waits for it too.
272        let mut task = self.listener_task.lock().await;
273        if let Some(handle) = task.as_mut() {
274            let _ = handle.await;
275            *task = None;
276        }
277    }
278}
279
280#[cfg(test)]
281mod tests {
282    use {
283        super::*,
284        crate::{
285            config::ControllerConfig,
286            controller::test_support::{
287                FAST_ELECTION_TIMEOUT, TEST_OP_TIMEOUT, bind_eventually, bootstrap_controller,
288                committable_topic_record, joining_controller, wait_for_leader,
289            },
290        },
291        assert2::check,
292        krabka_units::prelude::TimeExt as _,
293        tempfile::TempDir,
294        tokio::time,
295    };
296
297    #[tokio::test]
298    async fn quorum_view_reflects_live_single_voter_state_and_submitted_records() {
299        let dir = TempDir::new().unwrap();
300        let mut cfg = ControllerConfig::for_tests(NodeId(1), dir.path().to_path_buf());
301        cfg.election_timeout = FAST_ELECTION_TIMEOUT;
302        let ctrl = Controller::start(cfg).await.expect("bootstrap");
303        wait_for_leader(&ctrl).await;
304
305        let quorum = ctrl.quorum_state();
306        check!(
307            (
308                quorum.voter_nodes.contains_key(&NodeId(1)),
309                quorum.current_leader,
310                quorum.current_leader == Some(NodeId(1)),
311            ) == (true, Some(NodeId(1)), true)
312        );
313
314        time::timeout(
315            TEST_OP_TIMEOUT.to_std(),
316            ctrl.submit_change(vec![committable_topic_record("ops-a")]),
317        )
318        .await
319        .expect("submit ops-a timed out")
320        .expect("submit ops-a");
321        time::timeout(
322            TEST_OP_TIMEOUT.to_std(),
323            ctrl.submit_change(vec![committable_topic_record("ops-b")]),
324        )
325        .await
326        .expect("submit ops-b timed out")
327        .expect("submit ops-b");
328
329        assert2::assert!(
330            (
331                ctrl.current_image().topic("ops-a").is_some(),
332                ctrl.current_image().topic("ops-b").is_some(),
333            ) == (true, true)
334        );
335        let quorum = ctrl.quorum_state();
336        let leader_last = quorum.last_applied_index;
337        assert2::assert!(leader_last >= 2);
338        assert2::assert!(quorum.per_voter_matched_index.get(&NodeId(1)) == Some(&leader_last));
339        assert2::assert!(ctrl.quorum_committed_offset() == 2);
340        assert2::assert!(ctrl.voted_directory_id() == Some(Uuid::nil()));
341        ctrl.shutdown().await;
342    }
343
344    #[tokio::test]
345    async fn quorum_view_reports_join_node_is_not_leader() {
346        let (_dir, ctrl) = joining_controller("join start").await;
347
348        assert2::assert!(ctrl.quorum_state().current_leader.is_none());
349        assert2::assert!(ctrl.voted_directory_id().is_none());
350        ctrl.shutdown().await;
351    }
352
353    #[tokio::test]
354    async fn shutdown_releases_bound_listener_addr() {
355        let (_dir, ctrl) = bootstrap_controller("bootstrap").await;
356        let addr = ctrl.controller_bound_addr();
357
358        ctrl.shutdown().await;
359
360        let rebound = bind_eventually(addr).await;
361        drop(rebound);
362    }
363
364    #[tokio::test]
365    async fn cancel_releases_bound_listener_addr_without_consuming_handle() {
366        let (_dir, ctrl) = bootstrap_controller("bootstrap").await;
367        let addr = ctrl.controller_bound_addr();
368
369        ctrl.cancel().await;
370
371        let rebound = bind_eventually(addr).await;
372        drop(rebound);
373        ctrl.cancel().await;
374    }
375}
krabka_raft/controller.rs

krabka_raft/
controller.rs
