Skip to main content

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}