krabka_raft/kraft/controller/handle.rs
1//! The cloneable [`KraftController`] handle: the accessors that read the
2//! engine's published `watch` channels directly, and the command-sending
3//! methods that round-trip a request through the engine task.
4
5use {
6 super::KraftController,
7 crate::{
8 DelegationTokenMutation, RegistrationAdmission, SubmitChangeResult,
9 error::RaftError,
10 kraft::{
11 event::Event,
12 transport::{Command, Inbound, MetadataFetchSlice, QuorumStateSnapshot},
13 types::{NodeId, ReplicaKey},
14 },
15 reconfig::{ReconfigOutcome, VoterChange},
16 },
17 krabka_metadata::MetadataImage,
18 krabka_units::prelude::ByteSize,
19 std::{sync::Arc, time::SystemTime},
20 tokio::sync::{oneshot, watch},
21};
22
23impl KraftController {
24 /// Send the command `make` builds around a fresh reply sender to the
25 /// engine task and await its reply.
26 ///
27 /// # Errors
28 /// Returns [`RaftError::Shutdown`] if the engine task is gone, before or
29 /// after it took the command.
30 async fn ask<T>(
31 &self,
32 make: impl FnOnce(oneshot::Sender<T>) -> Command,
33 ) -> Result<T, RaftError> {
34 let (reply, rx) = oneshot::channel();
35 self.cmd_tx
36 .send(make(reply))
37 .await
38 .map_err(|_| RaftError::Shutdown)?;
39 rx.await.map_err(|_| RaftError::Shutdown)
40 }
41
42 /// Epoch milliseconds: the clock the server stamps connection expiries
43 /// with, and the engine's clock for delegation-token deadlines and for the
44 /// create-time stamped on every batch the leader appends. `Engine::now`
45 /// is monotonic from this process's own start, and a snapshot header
46 /// timestamp has to mean the same instant on every node that reads it.
47 pub(crate) fn wall_clock_ms() -> i64 { krabka_log::epoch_ms(SystemTime::now()) }
48
49 /// The node id this controller runs as.
50 #[must_use]
51 pub fn node_id(&self) -> NodeId { self.me }
52
53 /// Probe a proposed voter through the controller's configured TLS/SASL
54 /// dialer rather than opening an unauthenticated plaintext connection.
55 pub(crate) async fn probe_kraft_version(
56 &self,
57 address: &str,
58 finalized_version: u16,
59 ) -> Result<bool, RaftError> {
60 self.peers
61 .probe_kraft_version(address, finalized_version)
62 .await
63 }
64
65 /// A snapshot of the latest applied [`MetadataImage`].
66 #[must_use]
67 pub fn current_image(&self) -> Arc<MetadataImage> { self.image_rx.borrow().clone() }
68
69 /// The finalized-features epoch of the published [`MetadataImage`]: the
70 /// offset of the last record of the newest batch or snapshot that carried
71 /// a `FeatureLevelRecord`, or `-1` before the first. Kafka's controller
72 /// listener answers `ApiVersions` with it, from `FeaturesPublisher`.
73 ///
74 /// The engine publishes it after the image, so a caller that reads it
75 /// before [`Self::current_image`] never gets an epoch ahead of the
76 /// features.
77 #[must_use]
78 pub fn finalized_features_epoch(&self) -> i64 { *self.features_epoch_rx.borrow() }
79
80 /// Watch the published [`MetadataImage`].
81 #[must_use]
82 pub fn watch_image(&self) -> watch::Receiver<Arc<MetadataImage>> { self.image_rx.clone() }
83
84 /// Watch the fault that stopped the engine: `None` while it runs, and
85 /// the reason once it stops over a fault Kafka halts the process over.
86 /// That is the I/O error of the first write to the metadata log
87 /// directory that fails (KIP-858), or a controller activation that
88 /// failed.
89 #[must_use]
90 pub fn watch_fault(&self) -> watch::Receiver<Option<String>> { self.fault_rx.clone() }
91
92 /// Watch the current leader id.
93 #[must_use]
94 pub fn watch_leader(&self) -> watch::Receiver<Option<NodeId>> { self.leader_rx.clone() }
95
96 /// A synchronous snapshot of consensus state (the handle's `quorum_state()`
97 /// reads this without an mpsc round-trip — the engine republishes it on
98 /// every event). Cheap `watch` borrow + clone.
99 #[must_use]
100 pub fn quorum_snapshot(&self) -> QuorumStateSnapshot { self.quorum_rx.borrow().clone() }
101
102 /// Submit a metadata change. On the leader, appends the batch at the current
103 /// leader epoch and returns once it is committed (HWM ≥ the appended end
104 /// offset) AND applied, surfacing the first per-record rejection. On a
105 /// follower, returns [`RaftError::NotLeader`] with the leader hint; the
106 /// handle layer forwards via `forward_submit_to`.
107 ///
108 /// # Errors
109 /// - [`RaftError::Metadata`] if a record fails `validate`.
110 /// - [`RaftError::NotLeader`] if this node is not the leader.
111 /// - [`RaftError::Shutdown`] if the engine task is gone.
112 pub async fn submit_change(
113 &self,
114 records: Vec<krabka_metadata::MetadataRecord>,
115 ) -> Result<SubmitChangeResult, RaftError> {
116 let records = self.admit_registrations(records).await?;
117 self.ask(|reply| Command::SubmitChange {
118 records,
119 expected_image: None,
120 reply,
121 })
122 .await?
123 }
124
125 /// Install the hook every batch that registers a broker anew passes through
126 /// on the leader. It can be installed once.
127 ///
128 /// # Errors
129 /// Returns an error if a hook is already installed.
130 pub fn install_registration_admission(
131 &self,
132 admission: Arc<dyn RegistrationAdmission>,
133 ) -> Result<(), &'static str> {
134 self.registration_admission
135 .set(admission)
136 .map_err(|_| "registration admission already installed")
137 }
138
139 /// `records` as the installed [`RegistrationAdmission`] admits them, when
140 /// this node leads and the batch registers a broker anew, and unchanged
141 /// otherwise. A follower forwards the batch as it was submitted, and the
142 /// leader that receives it admits it.
143 async fn admit_registrations(
144 &self,
145 records: Vec<krabka_metadata::MetadataRecord>,
146 ) -> Result<Vec<krabka_metadata::MetadataRecord>, RaftError> {
147 let registers_anew = records.iter().any(|record| {
148 matches!(record, krabka_metadata::MetadataRecord::V1BrokerRegistration(registration)
149 if registration.broker_epoch < 0)
150 });
151 let leads = *self.leader_rx.borrow() == Some(self.me);
152 match self.registration_admission.get() {
153 Some(admission) if registers_anew && leads => admission.admit(records).await,
154 _ => Ok(records),
155 }
156 }
157
158 /// Submit a leader-local mutation only while its admission image is current.
159 /// Pending metadata must commit before a caller can plan a guarded mutation.
160 ///
161 /// # Errors
162 /// Returns a leadership, stale-image, uncommitted-tail, or submission error.
163 pub async fn submit_change_if_image(
164 &self,
165 records: Vec<krabka_metadata::MetadataRecord>,
166 expected_image: Arc<MetadataImage>,
167 ) -> Result<SubmitChangeResult, RaftError> {
168 self.ask(|reply| Command::SubmitChange {
169 records,
170 expected_image: Some(expected_image),
171 reply,
172 })
173 .await?
174 }
175
176 /// Submit generation-bound delegation-token mutations.
177 ///
178 /// # Errors
179 ///
180 /// Returns a leadership, shutdown, or mutation-rejection error.
181 pub async fn submit_delegation_token_mutations(
182 &self,
183 mutations: Vec<DelegationTokenMutation>,
184 ) -> Result<SubmitChangeResult, RaftError> {
185 self.ask(|reply| Command::SubmitDelegationTokenMutations { mutations, reply })
186 .await?
187 }
188
189 /// Submit one KIP-853 voter or kraft-version control operation.
190 pub(crate) async fn reconfigure(
191 &self,
192 change: VoterChange,
193 ) -> Result<ReconfigOutcome, RaftError> {
194 self.ask(|reply| Command::Reconfigure { change, reply })
195 .await?
196 }
197
198 /// Atomically finalize `kraft.version` in the Raft control log.
199 ///
200 /// # Errors
201 /// Returns an unsupported-version, leadership, timeout, or storage error.
202 pub async fn finalize_kraft_version(&self, version: u16) -> Result<ReconfigOutcome, RaftError> {
203 self.reconfigure(VoterChange::FinalizeKraftVersion(version))
204 .await
205 }
206
207 /// Run the leader-side checks of [`Self::finalize_kraft_version`] and
208 /// append nothing.
209 ///
210 /// # Errors
211 /// Returns the error `finalize_kraft_version` would return for the same
212 /// request.
213 pub async fn validate_kraft_version(&self, version: u16) -> Result<ReconfigOutcome, RaftError> {
214 self.reconfigure(VoterChange::ValidateKraftVersion(version))
215 .await
216 }
217
218 /// A structured snapshot of consensus state for the broker's
219 /// `DescribeQuorum` admin view.
220 ///
221 /// # Errors
222 /// Returns [`RaftError::Shutdown`] if the engine task is gone.
223 pub async fn quorum_state(&self) -> Result<QuorumStateSnapshot, RaftError> {
224 self.ask(|reply| Command::QuorumStateSnapshot { reply })
225 .await
226 }
227
228 /// Read a committed `__cluster_metadata` slice for an observer's
229 /// `API_KEY_METADATA_FETCH` (1004). `replica` names the observer that asked,
230 /// when it is a replica, so a leader can track its progress.
231 ///
232 /// # Errors
233 /// Returns [`RaftError::Shutdown`] if the engine task is gone.
234 pub async fn metadata_fetch(
235 &self,
236 fetch_offset: i64,
237 max_size: ByteSize,
238 replica: Option<ReplicaKey>,
239 ) -> Result<MetadataFetchSlice, RaftError> {
240 self.ask(|reply| Command::MetadataFetch {
241 fetch_offset,
242 max_size,
243 replica,
244 reply,
245 })
246 .await
247 }
248
249 /// Serialize the current image to a KIP-630 checkpoint under the data dir.
250 ///
251 /// # Errors
252 /// Returns [`RaftError`] if serialization or the file write fails.
253 pub async fn trigger_snapshot(&self) -> Result<(), RaftError> {
254 self.ask(|reply| Command::TriggerSnapshot { reply }).await?
255 }
256
257 /// Inject a raw core [`Event`] into the loop (test/driver entrypoint and the
258 /// internal feedback path for peer-RPC responses).
259 ///
260 /// # Errors
261 /// Returns [`RaftError::Shutdown`] if the engine task is gone.
262 pub async fn inject_event(&self, event: Event) -> Result<(), RaftError> {
263 self.cmd_tx
264 .send(Command::Event(event))
265 .await
266 .map_err(|_| RaftError::Shutdown)
267 }
268
269 /// Deliver an inbound peer RPC to the engine.
270 ///
271 /// # Errors
272 /// Returns [`RaftError::Shutdown`] if the engine task is gone.
273 pub async fn deliver(&self, inbound: Inbound) -> Result<(), RaftError> {
274 self.cmd_tx
275 .send(Command::Inbound(inbound))
276 .await
277 .map_err(|_| RaftError::Shutdown)
278 }
279
280 /// Stop the engine task gracefully and wait for its log to close.
281 ///
282 /// A leader resigns first and sends `EndQuorumEpoch`, and the engine keeps
283 /// taking part in the quorum until a successor is known or its graceful
284 /// shutdown window runs out, as Kafka's `KafkaRaftClient.shutdown` does. A
285 /// follower or an observer stops at once.
286 pub async fn shutdown(&self) { self.stop(Command::Shutdown).await; }
287
288 /// Stop the engine task at once, handing nothing over, and wait for its
289 /// log to close.
290 ///
291 /// This is the stop of a halted process: a fatal fault, or a test that
292 /// models a crash. The quorum finds out through its fetch watchdog.
293 pub async fn halt(&self) { self.stop(Command::Halt).await; }
294
295 async fn stop(&self, command: Command) {
296 let mut image = self.image_rx.clone();
297 let _ = self.cmd_tx.send(command).await;
298 // The engine owns the image sender. Wait for it to drop before callers
299 // can reopen or remove the log directory.
300 while image.changed().await.is_ok() {}
301 }
302
303 /// Test-only: append `records` as a committed batch and apply them through
304 /// the real pipeline; returns the appended base offset.
305 #[cfg(test)]
306 pub(super) async fn test_append_and_commit(
307 &self,
308 records: Vec<krabka_metadata::MetadataRecord>,
309 ) -> Result<i64, RaftError> {
310 self.ask(|reply| Command::TestAppendAndCommit { records, reply })
311 .await
312 }
313}