krabka_client_admin/
brokers.rs

1//! Broker and controller registration administration.
2
3use krabka_protocol::owned::{
4    unregister_broker_request::UnregisterBrokerRequest,
5    unregister_controller_request::UnregisterControllerRequest,
6};
7
8use crate::{
9    AdminClient, AdminError, kafka_error_name,
10    retry::{ControllerRetry, REQUEST_TIMED_OUT, RetryPolicy},
11};
12
13impl AdminClient {
14    /// The call behind [`AdminClientLike::unregister_broker`](crate::AdminClientLike::unregister_broker).
15    pub(crate) async fn unregister_broker_impl(&self, broker_id: i32) -> Result<(), AdminError> {
16        self.unregister_broker_with_retry(broker_id, self.retry)
17            .await
18    }
19
20    /// Unregisters one controller from the active `KRaft` controller, with
21    /// an `UnregisterController` request (API key 94).
22    ///
23    /// This is Kafka's `Admin.unregisterController`. Like
24    /// `KafkaAdminClient.unregisterBroker`, it goes to a broker of a broker
25    /// connection or to the active controller of a controller connection. It
26    /// retries `REQUEST_TIMED_OUT` (7) with the retry backoff until the call
27    /// deadline, and fails at once on every other code: `NOT_CONTROLLER`
28    /// (41), `INVALID_REQUEST` (42) for a controller in the voter set, and
29    /// `CONTROLLER_ID_NOT_REGISTERED` (136).
30    ///
31    /// # Errors
32    /// Returns a transport, protocol, or broker error. At the deadline it
33    /// gives `REQUEST_TIMED_OUT` (7).
34    pub async fn unregister_controller(&mut self, controller_id: i32) -> Result<(), AdminError> {
35        self.unregister_controller_with_retry(controller_id, self.retry)
36            .await
37    }
38
39    async fn unregister_broker_with_retry(
40        &self,
41        broker_id: i32,
42        policy: RetryPolicy,
43    ) -> Result<(), AdminError> {
44        let request = UnregisterBrokerRequest {
45            broker_id,
46            ..Default::default()
47        };
48        let mut retry = ControllerRetry::new("UnregisterBroker", policy);
49        loop {
50            let response = retry.bounded(self.conn.send(request.clone())).await?;
51            if response.error_code != REQUEST_TIMED_OUT {
52                return unregister_error(
53                    "UnregisterBroker",
54                    response.error_code,
55                    response.error_message,
56                );
57            }
58            retry.after_retriable("REQUEST_TIMED_OUT").await?;
59        }
60    }
61
62    async fn unregister_controller_with_retry(
63        &self,
64        controller_id: i32,
65        policy: RetryPolicy,
66    ) -> Result<(), AdminError> {
67        let request = UnregisterControllerRequest {
68            controller_id,
69            ..Default::default()
70        };
71        let mut retry = ControllerRetry::new("UnregisterController", policy);
72        loop {
73            let response = retry.bounded(self.conn.send(request.clone())).await?;
74            if response.error_code != REQUEST_TIMED_OUT {
75                return unregister_error(
76                    "UnregisterController",
77                    response.error_code,
78                    response.error_message,
79                );
80            }
81            retry.after_retriable("REQUEST_TIMED_OUT").await?;
82        }
83    }
84}
85
86fn unregister_error(
87    api: &'static str,
88    code: i16,
89    message: Option<String>,
90) -> Result<(), AdminError> {
91    if code == 0 {
92        Ok(())
93    } else {
94        Err(AdminError::Broker {
95            api,
96            code,
97            name: kafka_error_name(code),
98            message,
99        })
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use std::{
106        sync::{
107            Arc, Mutex,
108            atomic::{AtomicUsize, Ordering},
109        },
110        time::Duration,
111    };
112
113    use assert2::assert;
114    use bytes::BytesMut;
115    use krabka_client_core::MockBroker;
116    use krabka_protocol::{
117        Encode,
118        owned::{
119            api_versions_request,
120            api_versions_response::{ApiVersion, ApiVersionsResponse},
121            unregister_broker_request,
122            unregister_broker_response::UnregisterBrokerResponse,
123            unregister_controller_request,
124            unregister_controller_response::UnregisterControllerResponse,
125        },
126    };
127
128    use super::*;
129    use crate::partition_leaders::test_support::decode_request;
130
131    #[test]
132    fn unregister_error_names_the_api_and_the_code() {
133        for (api, code, message, expected) in [
134            ("UnregisterBroker", 0, None, None),
135            ("UnregisterController", 0, None, None),
136            (
137                "UnregisterBroker",
138                42,
139                Some("stale"),
140                Some(("UnregisterBroker", 42, "INVALID_REQUEST", Some("stale"))),
141            ),
142            (
143                "UnregisterController",
144                136,
145                None,
146                Some((
147                    "UnregisterController",
148                    136,
149                    "CONTROLLER_ID_NOT_REGISTERED",
150                    None,
151                )),
152            ),
153        ] {
154            let actual = match unregister_error(api, code, message.map(str::to_owned)) {
155                Ok(()) => None,
156                Err(AdminError::Broker {
157                    api,
158                    code,
159                    name,
160                    message,
161                }) => Some((api, code, name, message)),
162                Err(other) => panic!("unexpected error {other:?}"),
163            };
164            let actual = actual
165                .as_ref()
166                .map(|(api, code, name, message)| (*api, *code, *name, message.as_deref()));
167            assert!(actual == expected, "{api} {code}");
168        }
169    }
170
171    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
172    enum Call {
173        Broker,
174        Controller,
175    }
176
177    /// One unregister request that the mock broker decoded.
178    #[derive(Debug, Clone, PartialEq, Eq)]
179    enum Sent {
180        Broker(UnregisterBrokerRequest),
181        Controller(UnregisterControllerRequest),
182    }
183
184    /// A mock broker that advertises only `ApiVersions` and the API of
185    /// `call`, answers each request of that API with the next code of
186    /// `codes`, repeating the last, and records each request in `sent`.
187    async fn unregister_mock(
188        call: Call,
189        codes: Vec<i16>,
190        sent: Arc<Mutex<Vec<Sent>>>,
191    ) -> MockBroker {
192        let requests = AtomicUsize::new(0);
193        let api_key = match call {
194            Call::Broker => unregister_broker_request::API_KEY,
195            Call::Controller => unregister_controller_request::API_KEY,
196        };
197        MockBroker::start(move |key, version, _, request_body| {
198            let mut body = BytesMut::new();
199            if key == api_versions_request::API_KEY {
200                ApiVersionsResponse {
201                    api_keys: vec![
202                        ApiVersion {
203                            api_key: api_versions_request::API_KEY,
204                            ..Default::default()
205                        },
206                        ApiVersion {
207                            api_key,
208                            ..Default::default()
209                        },
210                    ],
211                    ..Default::default()
212                }
213                .encode(&mut body, 0)
214                .unwrap();
215                return Some(body.to_vec());
216            }
217            if key != api_key {
218                return None;
219            }
220            let error_code = codes[requests.fetch_add(1, Ordering::SeqCst).min(codes.len() - 1)];
221            // Both APIs are flexible from v0: a tagged-field byte follows the
222            // correlation id of the response header.
223            body.extend_from_slice(&[0]);
224            match call {
225                Call::Broker => {
226                    sent.lock().unwrap().push(Sent::Broker(decode_request(
227                        request_body,
228                        version,
229                        true,
230                    )));
231                    UnregisterBrokerResponse {
232                        error_code,
233                        ..Default::default()
234                    }
235                    .encode(&mut body, version)
236                    .unwrap();
237                }
238                Call::Controller => {
239                    sent.lock().unwrap().push(Sent::Controller(decode_request(
240                        request_body,
241                        version,
242                        true,
243                    )));
244                    UnregisterControllerResponse {
245                        error_code,
246                        ..Default::default()
247                    }
248                    .encode(&mut body, version)
249                    .unwrap();
250                }
251            }
252            Some(body.to_vec())
253        })
254        .await
255        .unwrap()
256    }
257
258    /// Kafka's `unregisterBroker` and `unregisterController` retry
259    /// `REQUEST_TIMED_OUT` (7) until the deadline and fail at once on
260    /// `NOT_CONTROLLER` (41) and every other code.
261    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
262    async fn unregister_retries_request_timed_out_as_kafka_does() {
263        const LONG: Duration = Duration::from_secs(5);
264        const SHORT: Duration = Duration::from_millis(300);
265        for call in [Call::Broker, Call::Controller] {
266            for (name, codes, timeout, expected_result, expected_requests) in [
267                ("success", vec![0], LONG, Ok(()), 1),
268                (
269                    "request timed out, then success",
270                    vec![7, 7, 0],
271                    LONG,
272                    Ok(()),
273                    3,
274                ),
275                ("not controller is final", vec![41], LONG, Err(41), 1),
276                ("invalid request is final", vec![42], LONG, Err(42), 1),
277                (
278                    "controller id not registered is final",
279                    vec![136],
280                    LONG,
281                    Err(136),
282                    1,
283                ),
284                (
285                    "request timed out past the deadline",
286                    vec![7],
287                    SHORT,
288                    Err(7),
289                    1,
290                ),
291            ] {
292                let sent = Arc::new(Mutex::new(Vec::new()));
293                let broker = unregister_mock(call, codes, Arc::clone(&sent)).await;
294                let admin = AdminClient::connect(&[broker.addr.to_string()])
295                    .await
296                    .expect("admin connects");
297                let backoff = if timeout == LONG {
298                    Duration::from_millis(1)
299                } else {
300                    timeout
301                };
302                let policy = RetryPolicy {
303                    timeout,
304                    initial_backoff: backoff,
305                    max_backoff: backoff,
306                    jitter: 0.0,
307                    max_retries: u32::MAX,
308                };
309
310                let (result, expected_request) = match call {
311                    Call::Broker => (
312                        admin.unregister_broker_with_retry(9, policy).await,
313                        Sent::Broker(UnregisterBrokerRequest {
314                            broker_id: 9,
315                            ..Default::default()
316                        }),
317                    ),
318                    Call::Controller => (
319                        admin.unregister_controller_with_retry(9, policy).await,
320                        Sent::Controller(UnregisterControllerRequest {
321                            controller_id: 9,
322                            ..Default::default()
323                        }),
324                    ),
325                };
326                let result = result.map_err(|error| match error {
327                    AdminError::Broker { code, .. } => code,
328                    other => panic!("case {call:?} {name}: unexpected error {other:?}"),
329                });
330
331                broker.stop();
332                let sent = sent.lock().unwrap().clone();
333                assert!(
334                    (result, sent) == (expected_result, vec![expected_request; expected_requests]),
335                    "case {call:?} {name}"
336                );
337            }
338        }
339    }
340}
krabka_client_admin/brokers.rs

krabka_client_admin/
brokers.rs
