barnabas_client/cluster.rs
1//! Connections to brokers, and the routing that decides which one to use.
2//!
3//! A [`Cluster`] owns one connection per broker plus the metadata that says
4//! which broker leads which partition. It is the piece that turns "send this
5//! fetch" into "send it to node 3, and if node 3 says it is no longer the
6//! leader, find out who is and try again".
7//!
8//! # Locality is the binding's choice
9//!
10//! The connection map is a plain `HashMap` behind `&mut self`, with no locking
11//! anywhere. Under a per-core binding that is simply correct — the sockets
12//! belong to the core that opened them and there is no other thread to contend
13//! with. Under a work-stealing binding the handle is what carries the
14//! synchronisation, not this map.
15//!
16//! One consequence is worth stating, because it is the interesting cost of the
17//! design: **each core keeps its own connections**, so a node with C cores and
18//! B brokers holds C×B connections rather than B. In exchange every partition a
19//! core owns shares one connection to each broker, and their fetches batch into
20//! one request — which is strictly better than a client per *partition*, the
21//! shape a wrapper around a threaded C client forces.
22
23use std::collections::{BTreeMap, HashMap};
24use std::time::{Duration, Instant};
25
26use barnabas_core::{BrokerAddr, Connection, Metadata};
27use bytes::Bytes;
28use kafka_protocol::messages::{
29 metadata_request::MetadataRequestTopic, ApiKey, ApiVersionsRequest, ApiVersionsResponse,
30 MetadataRequest, MetadataResponse, SaslAuthenticateRequest, SaslAuthenticateResponse,
31 SaslHandshakeRequest, SaslHandshakeResponse, TopicName,
32};
33use kafka_protocol::protocol::{Decodable, Encodable, StrBytes};
34
35use crate::sasl::{plain_message, Credentials, SaslMechanism, ScramExchange};
36use crate::timeout::with_timeout;
37use crate::{check, Error, Result, Transport};
38
39/// How long a single request may take before its connection is considered
40/// broken.
41///
42/// A broker that accepts a connection and then never answers is
43/// indistinguishable from a slow one, so a deadline is the only way out. It is
44/// deliberately generous: this is a liveness backstop, not a latency budget.
45pub const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
46
47/// Largest single read. Bounds the per-connection buffer while still being big
48/// enough that a multi-megabyte fetch response arrives in a few reads.
49const MAX_READ: usize = 1024 * 1024;
50
51/// One broker connection.
52pub(crate) struct Broker<T: Transport> {
53 stream: T::Stream,
54 conn: Connection,
55 /// Reused across reads, sized to the frame being read rather than to a
56 /// fixed chunk. See [`Self::recv`].
57 read_buf: Vec<u8>,
58 /// What this broker said it speaks, per API key: `(min, max)`.
59 ///
60 /// **The point of the `ApiVersions` handshake, which this client used to
61 /// perform and then ignore.** Every request went out at a hardcoded
62 /// version, which works against Apache Kafka because those versions happen
63 /// to be the ones it supports — and fails immediately against anything
64 /// else. Redpanda, for one, answers a version it does not know by *closing
65 /// the connection*, so the symptom is an unexplained EOF at connect rather
66 /// than an error code.
67 versions: BTreeMap<i16, (i16, i16)>,
68}
69
70impl<T: Transport> Broker<T> {
71 pub(crate) async fn connect(
72 transport: &T,
73 addr: &str,
74 client_id: &str,
75 credentials: Option<&Credentials>,
76 ) -> Result<Self> {
77 let stream = transport
78 .connect(addr)
79 .await
80 .map_err(|source| Error::Connect {
81 addr: addr.to_owned(),
82 source,
83 })?;
84 let mut me = Self {
85 stream,
86 conn: Connection::new(StrBytes::from_string(client_id.to_owned())),
87 read_buf: Vec::new(),
88 versions: BTreeMap::new(),
89 };
90
91 // `ApiVersions` first, as every client does: an unsupported version
92 // then fails at connect rather than in the middle of a fetch loop.
93 let mut req = ApiVersionsRequest::default();
94 req.client_software_name = StrBytes::from_string(client_id.to_owned());
95 req.client_software_version = StrBytes::from_static_str("0.0.1");
96 let resp: ApiVersionsResponse = me.call(ApiKey::ApiVersions, 3, &req).await?;
97 check("ApiVersions", resp.error_code)?;
98 for api in &resp.api_keys {
99 me.versions
100 .insert(api.api_key, (api.min_version, api.max_version));
101 }
102
103 // Authentication comes after `ApiVersions` and before anything else: a
104 // broker on a SASL listener rejects every other request until it has
105 // happened.
106 if let Some(credentials) = credentials {
107 me.authenticate(credentials).await?;
108 }
109
110 Ok(me)
111 }
112
113 /// `SaslHandshake` then one or more `SaslAuthenticate` round trips.
114 ///
115 /// Version 1 of `SaslAuthenticate`, which wraps the SASL bytes in a Kafka
116 /// request. The older form wrote raw SASL frames onto the socket ahead of
117 /// the protocol, which is unframed, undiagnosable, and not supported here.
118 async fn authenticate(&mut self, credentials: &Credentials) -> Result<()> {
119 let mut handshake = SaslHandshakeRequest::default();
120 handshake.mechanism = StrBytes::from_string(credentials.mechanism.as_str().to_owned());
121
122 let resp: SaslHandshakeResponse = self.call(ApiKey::SaslHandshake, 1, &handshake).await?;
123 if resp.error_code != 0 {
124 // The broker lists what it *would* accept, which is the single most
125 // useful thing to say when authentication fails at this stage.
126 let offered: Vec<String> = resp.mechanisms.iter().map(|m| m.to_string()).collect();
127 return Err(Error::Sasl(format!(
128 "broker rejected mechanism {}; it offers {}",
129 credentials.mechanism.as_str(),
130 if offered.is_empty() {
131 "nothing".to_owned()
132 } else {
133 offered.join(", ")
134 }
135 )));
136 }
137
138 match credentials.mechanism {
139 SaslMechanism::Plain => {
140 self.sasl_exchange(plain_message(credentials)).await?;
141 }
142 SaslMechanism::ScramSha256 | SaslMechanism::ScramSha512 => {
143 let nonce = self.scram_nonce();
144 let (mut exchange, client_first) = ScramExchange::start(credentials, &nonce);
145 let server_first = self.sasl_exchange(client_first).await?;
146 let client_final = exchange
147 .client_final(&server_first)
148 .map_err(|e| Error::Sasl(e.to_string()))?;
149 let server_final = self.sasl_exchange(client_final).await?;
150 // Verifying is what proves the peer knows the password. A
151 // client that skips it authenticates to anyone.
152 exchange
153 .verify(&server_final)
154 .map_err(|e| Error::Sasl(e.to_string()))?;
155 }
156 }
157 Ok(())
158 }
159
160 async fn sasl_exchange(&mut self, bytes: Vec<u8>) -> Result<Bytes> {
161 let mut req = SaslAuthenticateRequest::default();
162 req.auth_bytes = Bytes::from(bytes);
163 let resp: SaslAuthenticateResponse = self.call(ApiKey::SaslAuthenticate, 1, &req).await?;
164 if resp.error_code != 0 {
165 return Err(Error::Sasl(resp.error_message.map_or_else(
166 || "authentication failed".to_owned(),
167 |m| m.to_string(),
168 )));
169 }
170 Ok(resp.auth_bytes)
171 }
172
173 /// A nonce for SCRAM: unique per exchange, and unpredictable enough that a
174 /// replayed server-first cannot be matched to a future exchange.
175 ///
176 /// Built from the address of a stack local and a counter rather than a
177 /// `rand` dependency — SCRAM needs the client nonce to be *fresh*, not
178 /// cryptographically random, since the security rests on the shared secret.
179 fn scram_nonce(&self) -> String {
180 use std::sync::atomic::{AtomicU64, Ordering};
181 static COUNTER: AtomicU64 = AtomicU64::new(0);
182 let counter = COUNTER.fetch_add(1, Ordering::Relaxed);
183 let local = 0u8;
184 let entropy = std::ptr::addr_of!(local) as usize;
185 let nanos = std::time::SystemTime::now()
186 .duration_since(std::time::UNIX_EPOCH)
187 .map_or(0, |d| d.subsec_nanos());
188 format!("{counter:x}{entropy:x}{nanos:x}")
189 }
190
191 /// Send one request and read its response.
192 ///
193 /// **Requires an idle connection, and says so rather than guessing.** Kafka
194 /// answers in order, so calling this while something else is in flight
195 /// reads that something else's response. Three separate bugs in this client
196 /// were exactly that — a prefetched `Fetch` decoded as a `Metadata`, a
197 /// `ListOffsets`, and finally an `OffsetCommit`, the last of which only
198 /// surfaced because it failed to parse. A response that *did* parse would
199 /// have been silent.
200 ///
201 /// Callers that pipeline on purpose use [`Self::send`] and [`Self::recv`]
202 /// and own the ordering themselves.
203 pub(crate) async fn call<Req, Resp>(
204 &mut self,
205 api_key: ApiKey,
206 version: i16,
207 req: &Req,
208 ) -> Result<Resp>
209 where
210 Req: Encodable,
211 Resp: Decodable,
212 {
213 if self.conn.in_flight() != 0 {
214 return Err(Error::ConnectionBusy {
215 op: api_key,
216 addr: String::new(),
217 in_flight: self.conn.in_flight(),
218 });
219 }
220 self.send(api_key, version, req).await?;
221 self.recv().await
222 }
223
224 /// The version to actually send, given what this broker supports.
225 ///
226 /// Callers name the version they *prefer* — the newest this client knows
227 /// how to read. If the broker caps lower, the request goes out at its
228 /// maximum; if the broker requires newer, at its minimum. An API the broker
229 /// did not mention is sent as asked, which keeps behaviour unchanged
230 /// against a broker whose `ApiVersions` says nothing useful.
231 fn negotiated(&self, api_key: ApiKey, preferred: i16) -> i16 {
232 match self.versions.get(&(api_key as i16)) {
233 Some((min, max)) => preferred.clamp(*min, *max),
234 None => preferred,
235 }
236 }
237
238 /// Write a request and **do not wait for it**.
239 ///
240 /// The half of `call` that makes a request outstanding. Kafka answers a
241 /// connection's requests in the order they arrived and the core matches
242 /// them by position ([`Connection::in_flight`]), so several may be in
243 /// flight at once — which is what lets the consumer keep a fetch permanently
244 /// outstanding instead of issuing one only when its caller asks.
245 pub(crate) async fn send<Req: Encodable>(
246 &mut self,
247 api_key: ApiKey,
248 version: i16,
249 req: &Req,
250 ) -> Result<()> {
251 let version = self.negotiated(api_key, version);
252 let wire = self.conn.request(api_key, version, req)?;
253 T::write_all(&mut self.stream, &wire).await?;
254 Ok(())
255 }
256
257 /// Read the answer to the **oldest** outstanding request.
258 ///
259 /// **The read is sized to the frame, not to a fixed chunk.** This used to
260 /// read into a 16 KiB stack buffer, which for a 10 MiB fetch response meant
261 /// more than six hundred reads and six hundred copies into the decoder's
262 /// buffer. Produce responses are a few hundred bytes and never noticed;
263 /// fetch responses are megabytes, which is why the cost showed up on the
264 /// consume side of the benchmark and nowhere else.
265 ///
266 /// The length prefix says how much is coming, so after the first small read
267 /// the rest arrives in a handful of large ones.
268 pub(crate) async fn recv<Resp: Decodable>(&mut self) -> Result<Resp> {
269 loop {
270 if let Some(resp) = self.conn.next_response()? {
271 return Ok(Connection::decode(&resp)?);
272 }
273
274 // Capped so a large `fetch.max.bytes` cannot turn into one
275 // enormous buffer, and floored so the prefix read is not a
276 // four-byte syscall.
277 let want = self.conn.needed().clamp(16 * 1024, MAX_READ);
278 if self.read_buf.len() < want {
279 self.read_buf.resize(want, 0);
280 }
281
282 let n = T::read(&mut self.stream, &mut self.read_buf[..want]).await?;
283 if n == 0 {
284 return Err(Error::Io(std::io::Error::new(
285 std::io::ErrorKind::UnexpectedEof,
286 "broker closed the connection",
287 )));
288 }
289 self.conn.push_bytes(&self.read_buf[..n]);
290 }
291 }
292
293 /// How many requests are awaiting an answer.
294 pub(crate) fn in_flight(&self) -> usize {
295 self.conn.in_flight()
296 }
297}
298
299/// Connections plus the cluster map.
300pub struct Cluster<T: Transport> {
301 transport: T,
302 client_id: String,
303 bootstrap: Vec<String>,
304 /// Whether a metadata request may create the topic.
305 ///
306 /// **True by default, which is what librdkafka and the Java client do.** A
307 /// producer's first write to a new topic otherwise fails with error 3
308 /// rather than creating it — and a caller that expected the usual
309 /// behaviour has no way to tell that from a genuinely missing topic.
310 allow_auto_topic_creation: bool,
311 /// Keyed by `host:port` rather than node id: a broker that restarts with a
312 /// new id at the same address should reuse the socket, and a bootstrap
313 /// address has no node id until the first `Metadata` reply.
314 conns: HashMap<String, Broker<T>>,
315 /// Connections used **only** for group-coordinator traffic.
316 ///
317 /// A consumer keeps a fetch permanently in flight, and its group
318 /// coordinator is usually one of the brokers it fetches from. Sharing the
319 /// connection would put a `Heartbeat` behind a `Fetch` — which Kafka
320 /// answers in order, so the heartbeat reads the fetch's response. Draining
321 /// the fetch first would work and would cost the prefetch on every poll;
322 /// a second connection costs one socket per coordinator and nothing else.
323 coordinator_conns: HashMap<String, Broker<T>>,
324 metadata: Metadata,
325 /// When each topic's metadata was last read from a broker.
326 ///
327 /// Metadata is otherwise only refreshed when something is *missing* — an
328 /// unknown leader, an unknown count. That never notices metadata that is
329 /// merely **stale**, which is exactly what a topic that grew partitions
330 /// looks like: every leader still known, every answer still wrong.
331 refreshed_at: HashMap<String, Instant>,
332 /// How long a topic's metadata may go unrefreshed.
333 ///
334 /// Five minutes, which is `metadata.max.age.ms`'s default in the Java
335 /// client and librdkafka's `topic.metadata.refresh.interval.ms` in spirit.
336 /// It bounds how long a partition expansion goes unnoticed.
337 metadata_max_age: Duration,
338 request_timeout: Duration,
339 credentials: Option<Credentials>,
340}
341
342impl<T: Transport> Cluster<T> {
343 /// Connect to the first reachable bootstrap address and load metadata.
344 ///
345 /// # Errors
346 /// If no bootstrap address answers.
347 pub async fn connect(transport: T, bootstrap: &[String], client_id: &str) -> Result<Self> {
348 let mut me = Self {
349 transport,
350 client_id: client_id.to_owned(),
351 bootstrap: bootstrap.to_vec(),
352 conns: HashMap::new(),
353 coordinator_conns: HashMap::new(),
354 metadata: Metadata::new(),
355 request_timeout: DEFAULT_REQUEST_TIMEOUT,
356 credentials: None,
357 allow_auto_topic_creation: true,
358 refreshed_at: HashMap::new(),
359 metadata_max_age: Duration::from_secs(300),
360 };
361 // Touch one bootstrap address so a bad configuration fails here rather
362 // than at the first fetch.
363 me.any_broker().await?;
364 Ok(me)
365 }
366
367 /// Authenticate every connection with `credentials`.
368 ///
369 /// Must be set before the first request; connections already open are not
370 /// re-authenticated, because Kafka has no way to do so.
371 pub fn set_credentials(&mut self, credentials: Credentials) {
372 self.credentials = Some(credentials);
373 }
374
375 /// How long a request may take before its connection is dropped as broken.
376 pub fn set_request_timeout(&mut self, timeout: Duration) {
377 self.request_timeout = timeout;
378 }
379
380 /// Send a request to a specific broker, reconnecting once if the
381 /// connection turns out to be dead.
382 ///
383 /// **This is what makes a broker restart survivable.** A pooled connection
384 /// can be closed at any time — a rolling upgrade, an idle reaper, a network
385 /// blip — and the client only finds out when it writes to it. Without this,
386 /// the dead socket stays in the pool and every subsequent request fails
387 /// forever, which is a client that dies the first time its cluster is
388 /// maintained.
389 ///
390 /// **A timed-out request also drops the connection**, and that is not
391 /// caution: the response may still arrive later, and reading it as the
392 /// answer to the *next* request would desynchronise the stream. Kafka
393 /// answers in order, so a request abandoned mid-flight poisons everything
394 /// behind it.
395 ///
396 /// Retried exactly once: a second failure is the peer, not the socket.
397 pub(crate) async fn call_at<Req, Resp>(
398 &mut self,
399 addr: &str,
400 api_key: ApiKey,
401 version: i16,
402 req: &Req,
403 ) -> Result<Resp>
404 where
405 Req: Encodable,
406 Resp: Decodable,
407 {
408 let timeout = self.request_timeout;
409 self.call_at_with_timeout(addr, api_key, version, req, timeout)
410 .await
411 }
412
413 /// As [`Self::call_at`], with a deadline of its own.
414 ///
415 /// **Some requests are long-polls and the general deadline is wrong for
416 /// them.** `JoinGroup` is held by the coordinator until every member of the
417 /// group has joined — up to `rebalance_timeout_ms`, which is minutes — so a
418 /// 30-second request timeout cancels a request that is behaving correctly,
419 /// and takes the connection with it.
420 pub(crate) async fn call_at_with_timeout<Req, Resp>(
421 &mut self,
422 addr: &str,
423 api_key: ApiKey,
424 version: i16,
425 req: &Req,
426 timeout: Duration,
427 ) -> Result<Resp>
428 where
429 Req: Encodable,
430 Resp: Decodable,
431 {
432 for attempt in 0..2 {
433 let broker = self.broker_at(addr).await?;
434 if broker.in_flight() != 0 {
435 return Err(Error::ConnectionBusy {
436 op: api_key,
437 addr: addr.to_owned(),
438 in_flight: broker.in_flight(),
439 });
440 }
441 let outcome = with_timeout::<T, _>(timeout, broker.call(api_key, version, req)).await;
442
443 match outcome {
444 Some(Ok(resp)) => return Ok(resp),
445 Some(Err(Error::Io(e))) => {
446 self.conns.remove(addr);
447 if attempt == 1 {
448 return Err(Error::Io(e));
449 }
450 }
451 Some(Err(e)) => return Err(e),
452 None => {
453 self.conns.remove(addr);
454 return Err(Error::Timeout {
455 op: api_key,
456 addr: addr.to_owned(),
457 });
458 }
459 }
460 }
461 unreachable!("the loop returns on its last attempt")
462 }
463
464 /// Write a request to `addr` and leave it outstanding.
465 ///
466 /// Reconnects once if the pooled connection turns out to be dead, like
467 /// [`Self::call_at`]. Pair with [`Self::recv_many`].
468 pub(crate) async fn send_at<Req: Encodable>(
469 &mut self,
470 api_key: ApiKey,
471 version: i16,
472 addr: &str,
473 req: &Req,
474 ) -> Result<()> {
475 for attempt in 0..2 {
476 let broker = self.broker_at(addr).await?;
477 match broker.send(api_key, version, req).await {
478 Ok(()) => return Ok(()),
479 Err(Error::Io(e)) => {
480 self.conns.remove(addr);
481 if attempt == 1 {
482 return Err(Error::Io(e));
483 }
484 }
485 Err(e) => return Err(e),
486 }
487 }
488 unreachable!("the loop returns on its last attempt")
489 }
490
491 /// Collect one outstanding answer from each of `addrs`, **all at once**.
492 ///
493 /// Connections are taken out of the pool for the duration for the same
494 /// reason [`Self::call_many`] does it. A broker whose read failed or timed
495 /// out is dropped rather than returned: its stream position is no longer
496 /// known, and reading a late response as the answer to the next request
497 /// would desynchronise it.
498 ///
499 /// No reconnect-once here, deliberately — the request this would retry was
500 /// already sent and its answer lost, so re-sending is the caller's decision.
501 /// Both callers re-issue on the next round, which for a fetch is free.
502 pub(crate) async fn recv_many<Resp: Decodable>(
503 &mut self,
504 op: ApiKey,
505 addrs: &[String],
506 ) -> Vec<Result<Resp>> {
507 let timeout = self.request_timeout;
508 let mut taken = Vec::with_capacity(addrs.len());
509 let mut outcomes: Vec<Option<Result<Resp>>> = (0..addrs.len()).map(|_| None).collect();
510
511 for (index, addr) in addrs.iter().enumerate() {
512 match self.conns.remove(addr) {
513 Some(broker) => taken.push((index, addr.clone(), broker)),
514 // Dropped between send and receive — nothing outstanding to
515 // read, so say so rather than block forever.
516 None => {
517 outcomes[index] = Some(Err(Error::Io(std::io::Error::new(
518 std::io::ErrorKind::NotConnected,
519 "connection dropped before its response was read",
520 ))))
521 }
522 }
523 }
524
525 let reads: Vec<_> = taken
526 .into_iter()
527 .map(|(index, addr, mut broker)| async move {
528 let outcome = with_timeout::<T, _>(timeout, broker.recv()).await;
529 (index, addr, broker, outcome)
530 })
531 .collect();
532
533 for (index, addr, broker, outcome) in crate::join::join_all(reads).await {
534 outcomes[index] = Some(match outcome {
535 Some(Ok(resp)) => {
536 self.conns.insert(addr, broker);
537 Ok(resp)
538 }
539 Some(Err(e @ Error::Io(_))) => Err(e),
540 Some(Err(e)) => {
541 self.conns.insert(addr, broker);
542 Err(e)
543 }
544 None => Err(Error::Timeout { op, addr }),
545 });
546 }
547
548 outcomes
549 .into_iter()
550 .map(|o| o.expect("every address produced an outcome"))
551 .collect()
552 }
553
554 /// Read and throw away one outstanding answer per address.
555 ///
556 /// Used when a request in flight has been made irrelevant — the consumer's
557 /// assignment changed under it. The bytes must still be consumed or the
558 /// connection is left pointing at a response nobody expects; a connection
559 /// that cannot be drained is dropped instead.
560 pub(crate) async fn discard_many<Resp: Decodable>(&mut self, op: ApiKey, addrs: &[String]) {
561 let _ = self.recv_many::<Resp>(op, addrs).await;
562 }
563
564 /// Test-only doors onto the pipelining primitives, so the invariant above
565 /// can be checked without a broker.
566 #[doc(hidden)]
567 pub async fn send_at_for_test<Req: Encodable>(
568 &mut self,
569 api_key: ApiKey,
570 version: i16,
571 addr: &str,
572 req: &Req,
573 ) -> Result<()> {
574 self.send_at(api_key, version, addr, req).await
575 }
576
577 #[doc(hidden)]
578 pub async fn call_at_for_test<Req: Encodable, Resp: Decodable>(
579 &mut self,
580 addr: &str,
581 api_key: ApiKey,
582 version: i16,
583 req: &Req,
584 ) -> Result<Resp> {
585 self.call_at(addr, api_key, version, req).await
586 }
587
588 /// As [`Self::call_at`], on the connection reserved for coordinator
589 /// traffic. See [`Self::coordinator_conns`].
590 pub(crate) async fn call_coordinator<Req, Resp>(
591 &mut self,
592 addr: &str,
593 api_key: ApiKey,
594 version: i16,
595 req: &Req,
596 timeout: Duration,
597 ) -> Result<Resp>
598 where
599 Req: Encodable,
600 Resp: Decodable,
601 {
602 for attempt in 0..2 {
603 if !self.coordinator_conns.contains_key(addr) {
604 let broker = Broker::<T>::connect(
605 &self.transport,
606 addr,
607 &self.client_id,
608 self.credentials.as_ref(),
609 )
610 .await?;
611 self.coordinator_conns.insert(addr.to_owned(), broker);
612 }
613 let broker = self.coordinator_conns.get_mut(addr).expect("just inserted");
614
615 match with_timeout::<T, _>(timeout, broker.call(api_key, version, req)).await {
616 Some(Ok(resp)) => return Ok(resp),
617 Some(Err(Error::Io(e))) => {
618 self.coordinator_conns.remove(addr);
619 if attempt == 1 {
620 return Err(Error::Io(e));
621 }
622 }
623 Some(Err(e)) => return Err(e),
624 None => {
625 self.coordinator_conns.remove(addr);
626 return Err(Error::Timeout {
627 op: api_key,
628 addr: addr.to_owned(),
629 });
630 }
631 }
632 }
633 unreachable!("the loop returns on its last attempt")
634 }
635
636 /// As [`Self::call_at`], for requests any broker can serve.
637 pub(crate) async fn call_any<Req, Resp>(
638 &mut self,
639 api_key: ApiKey,
640 version: i16,
641 req: &Req,
642 ) -> Result<Resp>
643 where
644 Req: Encodable,
645 Resp: Decodable,
646 {
647 let addr = {
648 let broker_addr = self.any_broker_addr().await?;
649 broker_addr
650 };
651 self.call_at(&addr, api_key, version, req).await
652 }
653
654 /// The address of some live-looking broker, connecting if the pool is
655 /// empty.
656 async fn any_broker_addr(&mut self) -> Result<String> {
657 if let Some(addr) = self.conns.keys().next().cloned() {
658 return Ok(addr);
659 }
660 let mut last_err = None;
661 for addr in self.bootstrap.clone() {
662 match Broker::<T>::connect(
663 &self.transport,
664 &addr,
665 &self.client_id,
666 self.credentials.as_ref(),
667 )
668 .await
669 {
670 Ok(broker) => {
671 self.conns.insert(addr.clone(), broker);
672 return Ok(addr);
673 }
674 Err(e) => last_err = Some(e),
675 }
676 }
677 Err(last_err.unwrap_or(Error::Missing("bootstrap address")))
678 }
679
680 /// Whether metadata requests may create missing topics. See the field.
681 pub fn set_allow_auto_topic_creation(&mut self, allow: bool) {
682 self.allow_auto_topic_creation = allow;
683 }
684
685 /// The cluster map, for callers that want to inspect leadership.
686 #[must_use]
687 pub fn metadata(&self) -> &Metadata {
688 &self.metadata
689 }
690
691 /// A connection to *some* broker: an existing one if there is one, else the
692 /// first bootstrap address that answers.
693 ///
694 /// Used for requests that any broker can serve — `Metadata` above all,
695 /// which is what makes bootstrapping work at all.
696 pub(crate) async fn any_broker(&mut self) -> Result<&mut Broker<T>> {
697 if let Some(addr) = self.conns.keys().next().cloned() {
698 return Ok(self.conns.get_mut(&addr).expect("just found"));
699 }
700
701 let mut last_err = None;
702 for addr in self.bootstrap.clone() {
703 match Broker::<T>::connect(
704 &self.transport,
705 &addr,
706 &self.client_id,
707 self.credentials.as_ref(),
708 )
709 .await
710 {
711 Ok(broker) => {
712 self.conns.insert(addr.clone(), broker);
713 return Ok(self.conns.get_mut(&addr).expect("just inserted"));
714 }
715 Err(e) => last_err = Some(e),
716 }
717 }
718 Err(last_err.unwrap_or(Error::Missing("bootstrap address")))
719 }
720
721 /// Connect to `addr` if not already connected, and return it.
722 pub(crate) async fn broker_at(&mut self, addr: &str) -> Result<&mut Broker<T>> {
723 if !self.conns.contains_key(addr) {
724 let broker = Broker::<T>::connect(
725 &self.transport,
726 addr,
727 &self.client_id,
728 self.credentials.as_ref(),
729 )
730 .await?;
731 self.conns.insert(addr.to_owned(), broker);
732 }
733 Ok(self.conns.get_mut(addr).expect("just inserted"))
734 }
735
736 /// Ask any broker for `topic`'s metadata and merge it in.
737 ///
738 /// # Errors
739 /// If no broker answers, or the topic carries an error code.
740 pub async fn refresh_metadata(&mut self, topic: &str) -> Result<()> {
741 let mut req_topic = MetadataRequestTopic::default();
742 req_topic.name = Some(TopicName(StrBytes::from_string(topic.to_owned())));
743 let mut req = MetadataRequest::default();
744 req.topics = Some(vec![req_topic]);
745 req.allow_auto_topic_creation = self.allow_auto_topic_creation;
746
747 let resp: MetadataResponse = self.call_any(ApiKey::Metadata, 12, &req).await?;
748 for t in &resp.topics {
749 let code = barnabas_core::ErrorCode(t.error_code);
750 // An invalid-metadata code is "not yet", not "no": a topic being
751 // auto-created reports 3 until it exists. The response is still
752 // merged — its brokers are real — and the caller retries.
753 if !code.is_ok() && code.disposition() != barnabas_core::Disposition::RefreshMetadata {
754 check("Metadata", t.error_code)?;
755 }
756 }
757 self.metadata.update(&resp);
758 self.refreshed_at.insert(topic.to_owned(), Instant::now());
759 Ok(())
760 }
761
762 /// How long a topic's metadata may go unrefreshed before a lookup re-reads
763 /// it. See the `refreshed_at` field.
764 pub fn set_metadata_max_age(&mut self, age: Duration) {
765 self.metadata_max_age = age;
766 }
767
768 /// Whether this topic's metadata is older than the maximum age. A topic
769 /// never read is stale.
770 #[must_use]
771 pub fn is_metadata_stale(&self, topic: &str) -> bool {
772 self.refreshed_at
773 .get(topic)
774 .is_none_or(|at| at.elapsed() >= self.metadata_max_age)
775 }
776
777 /// Re-read this topic if it is stale, and report `(before, after)` if its
778 /// partition count **grew**.
779 ///
780 /// Only growth: Kafka has no operation that removes partitions from a live
781 /// topic, so a smaller number is a transient answer — a broker that has not
782 /// caught up, a topic mid-creation — and acting on it would move every key
783 /// twice. A topic seen for the first time reports no growth either; there
784 /// is nothing to have grown from.
785 ///
786 /// # Errors
787 /// If the refresh fails.
788 pub async fn refresh_if_stale(&mut self, topic: &str) -> Result<Option<(i32, i32)>> {
789 if !self.is_metadata_stale(topic) {
790 return Ok(None);
791 }
792 let before = self.metadata.partition_count(topic);
793 self.refresh_metadata(topic).await?;
794 let after = self.metadata.partition_count(topic);
795 Ok((before > 0 && after > before).then_some((before, after)))
796 }
797
798 /// Forget a topic that has been deleted — its count and every leader.
799 pub fn forget_topic(&mut self, topic: &str) {
800 self.metadata.forget_topic(topic);
801 self.refreshed_at.remove(topic);
802 }
803
804 /// Refresh brokers and the controller **without** naming a topic.
805 ///
806 /// An empty topic list is not the same as no topic list: `None` asks for
807 /// every topic in the cluster, which on a large cluster is a large answer
808 /// for information this does not want.
809 ///
810 /// # Errors
811 /// If no broker answers.
812 pub async fn refresh_cluster(&mut self) -> Result<()> {
813 let mut req = MetadataRequest::default();
814 req.topics = Some(Vec::new());
815 req.allow_auto_topic_creation = false;
816 let resp: MetadataResponse = self.call_any(ApiKey::Metadata, 12, &req).await?;
817 self.metadata.update(&resp);
818 Ok(())
819 }
820
821 /// The controller's address, refreshing if it is unknown.
822 ///
823 /// # Errors
824 /// If metadata cannot be refreshed, or no controller is elected — which
825 /// happens briefly during a controller election and is a wait, not a
826 /// failure, so callers retry it.
827 pub async fn controller_addr(&mut self) -> Result<String> {
828 if let Some(broker) = self.metadata.controller() {
829 return Ok(broker.addr());
830 }
831 self.refresh_cluster().await?;
832 self.metadata
833 .controller()
834 .map(barnabas_core::metadata::BrokerAddr::addr)
835 .ok_or(Error::Missing("a controller"))
836 }
837
838 /// Forget which broker is the controller, after it said it is not.
839 pub fn invalidate_controller(&mut self) {
840 self.metadata.invalidate_controller();
841 }
842
843 /// The address of `topic`/`partition`'s leader, refreshing if unknown.
844 ///
845 /// # Errors
846 /// If metadata cannot be refreshed, or the partition still has no leader
847 /// afterwards — which is what a partition mid-election looks like, and is
848 /// [`Error::NoLeader`] so a caller can back off and try again rather than
849 /// treat it as fatal.
850 pub async fn leader_addr(&mut self, topic: &str, partition: i32) -> Result<String> {
851 if let Some(addr) = self.metadata.leader_for(topic, partition) {
852 return Ok(addr.addr());
853 }
854 self.refresh_metadata(topic).await?;
855 self.metadata
856 .leader_for(topic, partition)
857 .map(BrokerAddr::addr)
858 .ok_or_else(|| Error::NoLeader {
859 topic: topic.to_owned(),
860 partition,
861 })
862 }
863
864 /// How many partitions `topic` has, refreshing metadata if the client has
865 /// not seen it yet.
866 ///
867 /// # Errors
868 /// If metadata cannot be refreshed, or the topic has no partitions after it.
869 pub async fn partition_count(&mut self, topic: &str) -> Result<i32> {
870 // **Stale counts as unknown.** Everything that places a key by hash
871 // reaches this — the producer's partitioner, the group leader's
872 // assignment — so a count that has quietly gone out of date here is a
873 // producer writing to the wrong partitions and a group ignoring the new
874 // ones. It is the one lookup where age matters more than a round trip.
875 if self.metadata.partition_count(topic) == 0 || self.is_metadata_stale(topic) {
876 self.refresh_metadata(topic).await?;
877 }
878 let count = self.metadata.partition_count(topic);
879 if count == 0 {
880 return Err(Error::NoLeader {
881 topic: topic.to_owned(),
882 partition: -1,
883 });
884 }
885 Ok(count)
886 }
887
888 /// Forget one partition's leader after the broker said it is not the
889 /// leader. Per partition on purpose — see [`barnabas_core::metadata`].
890 pub fn invalidate(&mut self, topic: &str, partition: i32) {
891 self.metadata.invalidate_partition(topic, partition);
892 }
893
894 /// Number of open connections. Exposed because connection count is a real
895 /// cost of the per-core design, and something a caller may want to watch.
896 #[must_use]
897 pub fn connection_count(&self) -> usize {
898 self.conns.len()
899 }
900}