krabka_client_core/
pool.rs

1//! `BrokerPool`: one connection per broker id, opened on first use.
2//!
3//! The pool follows Kafka's `ClusterConnectionStates`: a reconnect backoff
4//! after each failure or disconnect, a connection setup timeout that grows
5//! per failed attempt, and a turn through every resolved address of a broker.
6
7use std::{
8    net::SocketAddr,
9    sync::{
10        Arc, Mutex, MutexGuard, PoisonError, RwLock,
11        atomic::{AtomicBool, Ordering},
12    },
13};
14
15use dashmap::DashMap;
16use krabka_units::{Time, convert::TimeExt as _};
17use tokio::{
18    sync,
19    time::{Instant, sleep_until},
20};
21
22use crate::{
23    backoff::{ExponentialBackoff, random_below},
24    bootstrap::{bounded_lookup, filter_preferred_addresses},
25    connection::{Connection, ConnectionOptions},
26    error::ClientError,
27    transport::resolve_host,
28    version::FinalizedFeatures,
29};
30
31/// Information about a single Kafka broker, as reported by a `MetadataResponse`.
32#[derive(Debug, Clone)]
33pub struct BrokerInfo {
34    pub id: i32,
35    pub host: String,
36    pub port: i32,
37    pub rack: Option<String>,
38}
39
40/// The live-IO dependency of [`BrokerPool`]: resolve a host and dial an
41/// address.
42///
43/// A trait hides this dependency, so the pool's caching, backoff, address
44/// rotation and eviction logic is testable without a socket. The test
45/// connector hands back a stand-in connection type and records each dial.
46#[async_trait::async_trait]
47pub trait BrokerConnector: Send + Sync {
48    /// Connection handle this connector produces. It is `Connection` in
49    /// production and a cheap stand-in in tests.
50    type Conn: Send + Sync;
51
52    /// Dial `addr` with its pre-DNS `server_name`. `setup_timeout` is the TCP
53    /// connect deadline of this attempt.
54    async fn dial(
55        &self,
56        addr: SocketAddr,
57        server_name: &str,
58        setup_timeout: Time,
59    ) -> Result<Self::Conn, ClientError>;
60
61    /// Resolve `host` to the addresses to try, in order.
62    async fn resolve(&self, host: &str, port: u16) -> Result<Vec<SocketAddr>, ClientError>;
63
64    /// Whether `connection` can still carry requests.
65    fn is_open(connection: &Self::Conn) -> bool;
66
67    /// The number of requests that wait for a response on `connection`.
68    fn in_flight(connection: &Self::Conn) -> usize;
69
70    /// The finalized feature levels that the broker of `connection` sent.
71    fn finalized_features(connection: &Self::Conn) -> FinalizedFeatures;
72}
73
74/// Production [`BrokerConnector`]: opens a real [`Connection`] that honours
75/// the pool's TLS/SASL policy.
76#[derive(Debug)]
77pub struct TcpConnector {
78    options: ConnectionOptions,
79}
80
81#[async_trait::async_trait]
82impl BrokerConnector for TcpConnector {
83    type Conn = Connection;
84
85    #[tracing::instrument(level = "debug", skip_all, fields(addr = %addr), err(level = "debug"))]
86    async fn dial(
87        &self,
88        addr: SocketAddr,
89        server_name: &str,
90        setup_timeout: Time,
91    ) -> Result<Connection, ClientError> {
92        let mut options = self.options.clone();
93        options.socket_connection_setup_timeout = setup_timeout;
94        if let Some(security) = options.security.as_mut() {
95            **security = security.for_target_host(server_name);
96        }
97        Connection::connect_with_options(addr, options).await
98    }
99
100    async fn resolve(&self, host: &str, port: u16) -> Result<Vec<SocketAddr>, ClientError> {
101        let timeout = self.options.dns_timeout;
102        let addresses = bounded_lookup(timeout, resolve_host(host, port))
103            .await
104            .map_err(|_| ClientError::Timeout(timeout.time()))??;
105        Ok(filter_preferred_addresses(addresses))
106    }
107
108    fn is_open(connection: &Connection) -> bool {
109        !connection.is_closed()
110    }
111
112    fn in_flight(connection: &Connection) -> usize {
113        connection.in_flight()
114    }
115
116    fn finalized_features(connection: &Connection) -> FinalizedFeatures {
117        connection.versions().finalized_features().clone()
118    }
119}
120
121/// Synthetic broker id under which the shared bootstrap connection is cached.
122/// Never a real Kafka node id (those are `>= 0`).
123const BOOTSTRAP_ID: i32 = -1;
124
125/// Kafka's reconnect backoff and connection setup timeout for one pool.
126#[derive(Clone, Copy, Debug)]
127struct ConnectPolicy {
128    reconnect: ExponentialBackoff,
129    setup_timeout: ExponentialBackoff,
130}
131
132impl ConnectPolicy {
133    fn new(options: &ConnectionOptions) -> Self {
134        Self {
135            reconnect: ExponentialBackoff::kafka(
136                options.reconnect_backoff,
137                options.reconnect_backoff_max,
138            ),
139            setup_timeout: ExponentialBackoff::kafka(
140                options.socket_connection_setup_timeout,
141                options.socket_connection_setup_timeout_max,
142            ),
143        }
144    }
145}
146
147/// The connection state of one broker id, as Kafka's `NodeConnectionState`
148/// holds it.
149struct NodeState<C> {
150    connection: Option<Arc<C>>,
151    /// The resolved addresses with their TLS server names. The pool resolves
152    /// them again when the list is empty.
153    addresses: Vec<(SocketAddr, String)>,
154    address_index: usize,
155    last_attempted: Option<SocketAddr>,
156    /// When the last connection attempt started.
157    last_attempt_at: Option<Instant>,
158    /// Failures in a row, the exponent of the reconnect backoff.
159    failed_attempts: u32,
160    /// Failed connection attempts in a row, the exponent of the setup timeout.
161    failed_connect_attempts: u32,
162    /// The pool does not dial before this instant.
163    retry_at: Option<Instant>,
164}
165
166impl<C> Default for NodeState<C> {
167    fn default() -> Self {
168        Self {
169            connection: None,
170            addresses: Vec::new(),
171            address_index: 0,
172            last_attempted: None,
173            last_attempt_at: None,
174            failed_attempts: 0,
175            failed_connect_attempts: 0,
176            retry_at: None,
177        }
178    }
179}
180
181impl<C> NodeState<C> {
182    /// A connection became ready (`ClusterConnectionStates.ready`).
183    fn ready(&mut self, connection: Arc<C>) {
184        self.connection = Some(connection);
185        self.failed_attempts = 0;
186        self.failed_connect_attempts = 0;
187        self.retry_at = None;
188    }
189
190    /// An open connection closed (`ClusterConnectionStates.disconnected` for
191    /// a connected node). The next connection resolves the host again.
192    fn disconnected(&mut self, policy: ConnectPolicy) {
193        self.connection = None;
194        self.addresses.clear();
195        self.back_off(policy);
196    }
197
198    /// A connection attempt failed (`ClusterConnectionStates.disconnected`
199    /// for a connecting node). The next attempt uses the next address, and
200    /// resolves the host again after the last one.
201    fn connect_failed(&mut self, policy: ConnectPolicy) {
202        self.failed_connect_attempts = self.failed_connect_attempts.saturating_add(1);
203        self.address_index += 1;
204        if self.address_index >= self.addresses.len() {
205            self.addresses.clear();
206        }
207        self.back_off(policy);
208    }
209
210    fn back_off(&mut self, policy: ConnectPolicy) {
211        let backoff = policy.reconnect.backoff(self.failed_attempts);
212        self.failed_attempts = self.failed_attempts.saturating_add(1);
213        self.retry_at = Some(Instant::now() + backoff);
214    }
215
216    /// Use `addresses` as the new address list. As Kafka's
217    /// `NodeConnectionState.resolveAddresses` does, skip the first address
218    /// when it is the one that the last attempt used.
219    fn set_addresses(&mut self, addresses: Vec<(SocketAddr, String)>) {
220        self.address_index = usize::from(
221            addresses.len() > 1 && addresses.first().map(|(addr, _)| *addr) == self.last_attempted,
222        );
223        self.addresses = addresses;
224    }
225
226    /// The cached connection when it is still open. A closed one counts as a
227    /// disconnect.
228    fn open_connection<K: BrokerConnector<Conn = C>>(
229        &mut self,
230        policy: ConnectPolicy,
231    ) -> Option<Arc<C>> {
232        match &self.connection {
233            Some(connection) if K::is_open(connection) => Some(Arc::clone(connection)),
234            Some(_) => {
235                self.disconnected(policy);
236                None
237            }
238            None => None,
239        }
240    }
241}
242
243/// One broker id: its state, and a gate that lets one task at a time make a
244/// connection attempt. The state lock is never held across an `.await`.
245struct Node<C> {
246    state: Mutex<NodeState<C>>,
247    connect: sync::Mutex<()>,
248}
249
250impl<C> Default for Node<C> {
251    fn default() -> Self {
252        Self {
253            state: Mutex::new(NodeState::default()),
254            connect: sync::Mutex::new(()),
255        }
256    }
257}
258
259impl<C> Node<C> {
260    fn state(&self) -> MutexGuard<'_, NodeState<C>> {
261        self.state.lock().unwrap_or_else(PoisonError::into_inner)
262    }
263}
264
265/// Pool of `Arc<Connection>` keyed by broker id.
266///
267/// The pool opens a connection lazily on first use and caches it afterwards.
268/// A closed connection is replaced on the next use. After a failure or a
269/// disconnect, the pool waits Kafka's reconnect backoff before it dials that
270/// broker again. Concurrent users of one broker share one connection attempt.
271pub struct BrokerPool<C: BrokerConnector = TcpConnector> {
272    nodes: DashMap<i32, Arc<Node<C::Conn>>>,
273    by_endpoint: DashMap<i32, (String, u16)>,
274    bootstrap: RwLock<Vec<(SocketAddr, String)>>,
275    policy: ConnectPolicy,
276    connector: C,
277    /// Set by a bootstrap reconnect: untargeted requests use the bootstrap
278    /// connection until the pool learns brokers again.
279    prefer_bootstrap: AtomicBool,
280    /// The finalized features with the highest epoch that a connection saw,
281    /// as Kafka's `ApiVersions.update` keeps them.
282    finalized_features: Mutex<FinalizedFeatures>,
283}
284
285impl BrokerPool<TcpConnector> {
286    /// Create a new pool with the given bootstrap addresses and connection options.
287    #[must_use]
288    pub fn new(bootstrap: Vec<SocketAddr>, options: ConnectionOptions) -> Self {
289        let bootstrap = bootstrap
290            .into_iter()
291            .map(|address| (address, address.ip().to_string()))
292            .collect();
293        Self::new_with_server_names(bootstrap, options)
294    }
295
296    /// Create a pool whose bootstrap TLS names retain their pre-DNS hostnames.
297    #[must_use]
298    pub fn new_with_server_names(
299        bootstrap: Vec<(SocketAddr, String)>,
300        options: ConnectionOptions,
301    ) -> Self {
302        let policy = ConnectPolicy::new(&options);
303        Self::with_connector_and_names(bootstrap, TcpConnector { options }, policy)
304    }
305}
306
307/// A broker to dial for [`BrokerPool::least_loaded`]: (in backoff, backoff
308/// end, last attempt, id). The smallest of the first three wins.
309type Candidate = (bool, Option<Instant>, Option<Instant>, i32);
310
311/// What a connection attempt needs, read from the node state.
312struct Attempt {
313    addr: SocketAddr,
314    server_name: String,
315    setup_timeout: Time,
316}
317
318impl<C: BrokerConnector> BrokerPool<C> {
319    fn with_connector_and_names(
320        bootstrap: Vec<(SocketAddr, String)>,
321        connector: C,
322        policy: ConnectPolicy,
323    ) -> Self {
324        Self {
325            nodes: DashMap::new(),
326            by_endpoint: DashMap::new(),
327            bootstrap: RwLock::new(bootstrap),
328            policy,
329            connector,
330            prefer_bootstrap: AtomicBool::new(false),
331            finalized_features: Mutex::default(),
332        }
333    }
334
335    fn node(&self, broker_id: i32) -> Arc<Node<C::Conn>> {
336        Arc::clone(self.nodes.entry(broker_id).or_default().value())
337    }
338
339    /// The cached open connection of `node`.
340    fn cached(&self, node: &Node<C::Conn>) -> Option<Arc<C::Conn>> {
341        node.state().open_connection::<C>(self.policy)
342    }
343
344    /// Wait for the reconnect backoff of `node` to end.
345    async fn wait_for_backoff(node: &Node<C::Conn>) {
346        let retry_at = node.state().retry_at;
347        if let Some(retry_at) = retry_at {
348            sleep_until(retry_at).await;
349        }
350    }
351
352    /// The address and setup timeout of the next attempt, or `None` when the
353    /// node has no address.
354    fn next_attempt(&self, node: &Node<C::Conn>) -> Option<Attempt> {
355        let mut state = node.state();
356        let (addr, server_name) = state.addresses.get(state.address_index).cloned()?;
357        state.last_attempted = Some(addr);
358        state.last_attempt_at = Some(Instant::now());
359        Some(Attempt {
360            addr,
361            server_name,
362            setup_timeout: Time::from_std(
363                self.policy
364                    .setup_timeout
365                    .backoff(state.failed_connect_attempts),
366            ),
367        })
368    }
369
370    /// Dial one attempt and store a new connection.
371    async fn dial(
372        &self,
373        node: &Node<C::Conn>,
374        attempt: Attempt,
375    ) -> Result<Arc<C::Conn>, ClientError> {
376        let connection = self
377            .connector
378            .dial(attempt.addr, &attempt.server_name, attempt.setup_timeout)
379            .await?;
380        let connection = Arc::new(connection);
381        node.state().ready(Arc::clone(&connection));
382        let features = C::finalized_features(&connection);
383        let mut known = self
384            .finalized_features
385            .lock()
386            .unwrap_or_else(PoisonError::into_inner);
387        if features.epoch > known.epoch {
388            *known = features;
389        }
390        drop(known);
391        Ok(connection)
392    }
393
394    /// Get-or-connect to a specific broker id. The pool must have already
395    /// learned the (id, address) mapping with [`refresh_brokers`].
396    ///
397    /// When the broker is in its reconnect backoff, this method waits for the
398    /// backoff to end before it dials. One call makes at most one connection
399    /// attempt, to the next address of the broker.
400    ///
401    /// [`refresh_brokers`]: BrokerPool::refresh_brokers
402    ///
403    /// # Errors
404    /// Returns [`ClientError::Disconnected`] for a broker id that the pool
405    /// does not know, the resolve error, or the dial error.
406    #[tracing::instrument(level = "debug", skip_all, fields(broker_id), err(level = "debug"))]
407    pub async fn get(&self, broker_id: i32) -> Result<Arc<C::Conn>, ClientError> {
408        let (host, port) = self
409            .by_endpoint
410            .get(&broker_id)
411            .map(|entry| entry.value().clone())
412            .ok_or(ClientError::Disconnected)?;
413        let node = self.node(broker_id);
414        if let Some(connection) = self.cached(&node) {
415            return Ok(connection);
416        }
417        let _gate = node.connect.lock().await;
418        if let Some(connection) = self.cached(&node) {
419            return Ok(connection);
420        }
421        Self::wait_for_backoff(&node).await;
422        if node.state().addresses.is_empty() {
423            match self.connector.resolve(&host, port).await {
424                Ok(addresses) => node.state().set_addresses(
425                    addresses
426                        .into_iter()
427                        .map(|addr| (addr, host.clone()))
428                        .collect(),
429                ),
430                Err(error) => {
431                    node.state().connect_failed(self.policy);
432                    return Err(error);
433                }
434            }
435        }
436        let Some(attempt) = self.next_attempt(&node) else {
437            node.state().connect_failed(self.policy);
438            return Err(ClientError::Disconnected);
439        };
440        let result = self.dial(&node, attempt).await;
441        if result.is_err() {
442            node.state().connect_failed(self.policy);
443        }
444        result
445    }
446
447    /// Drop the cached connection to `broker_id`, if there is one, so the next
448    /// [`get`](BrokerPool::get) reconnects.
449    ///
450    /// Call this after a send fails. A bounced or failed-over broker must not
451    /// be retried over its dead, cached socket. As Kafka does after a
452    /// disconnect, the next connection waits the reconnect backoff and
453    /// resolves the host again.
454    pub fn evict(&self, broker_id: i32) {
455        if let Some(node) = self.nodes.get(&broker_id) {
456            let mut state = node.state();
457            if state.connection.is_some() {
458                state.disconnected(self.policy);
459            }
460        }
461    }
462
463    /// Drop the cached bootstrap connection so the next
464    /// [`bootstrap_connection`](BrokerPool::bootstrap_connection) re-iterates the
465    /// bootstrap addresses and reconnects to a live broker.
466    ///
467    /// The pool keys the bootstrap connection by the synthetic id `-1`, which
468    /// no real broker id matches, so [`evict`](BrokerPool::evict) can never
469    /// reach it.
470    pub fn evict_bootstrap(&self) {
471        self.evict(BOOTSTRAP_ID);
472    }
473
474    /// Replace the bootstrap address list and drop the cached bootstrap
475    /// connection so the next bootstrap send dials the fresh addresses.
476    pub fn replace_bootstrap(&self, bootstrap: Vec<SocketAddr>) {
477        let bootstrap = bootstrap
478            .into_iter()
479            .map(|address| (address, address.ip().to_string()))
480            .collect();
481        self.replace_bootstrap_with_server_names(bootstrap);
482    }
483
484    /// Replace bootstrap addresses while retaining their TLS server names.
485    ///
486    /// The bootstrap node keeps its reconnect backoff and failure counts, so
487    /// a retry loop that calls this after each failure still backs off. Until
488    /// the next [`refresh_brokers`](Self::refresh_brokers) with brokers,
489    /// [`least_loaded`](Self::least_loaded) uses the bootstrap connection.
490    pub fn replace_bootstrap_with_server_names(&self, bootstrap: Vec<(SocketAddr, String)>) {
491        self.prefer_bootstrap();
492        match self.bootstrap.write() {
493            Ok(mut guard) => *guard = bootstrap,
494            Err(poisoned) => *poisoned.into_inner() = bootstrap,
495        }
496        if let Some(node) = self.nodes.get(&BOOTSTRAP_ID) {
497            let mut state = node.state();
498            state.connection = None;
499            state.addresses.clear();
500            state.address_index = 0;
501        }
502    }
503
504    /// Make [`least_loaded`](Self::least_loaded) use the bootstrap connection
505    /// until the next [`refresh_brokers`](Self::refresh_brokers) with brokers.
506    pub fn prefer_bootstrap(&self) {
507        self.prefer_bootstrap.store(true, Ordering::Relaxed);
508    }
509
510    /// Replace bootstrap addresses and discard every connection and advertised
511    /// broker address learned from stale metadata.
512    pub fn rebootstrap(&self, bootstrap: Vec<SocketAddr>) {
513        let bootstrap = bootstrap
514            .into_iter()
515            .map(|address| (address, address.ip().to_string()))
516            .collect();
517        self.rebootstrap_with_server_names(bootstrap);
518    }
519
520    /// Replace all broker state while retaining bootstrap TLS server names.
521    pub fn rebootstrap_with_server_names(&self, bootstrap: Vec<(SocketAddr, String)>) {
522        self.nodes.clear();
523        self.by_endpoint.clear();
524        match self.bootstrap.write() {
525            Ok(mut guard) => *guard = bootstrap,
526            Err(poisoned) => *poisoned.into_inner() = bootstrap,
527        }
528    }
529
530    /// Get-or-connect to the first reachable bootstrap address. The bootstrap
531    /// connection is cached under the synthetic broker id `-1`.
532    ///
533    /// One call dials each bootstrap address at most once, starting after the
534    /// address that failed last. When every address fails, the next call
535    /// waits the reconnect backoff first.
536    ///
537    /// # Errors
538    /// Returns the error of the last address, an authentication failure at
539    /// once, or [`ClientError::Disconnected`] for an empty address list.
540    #[tracing::instrument(level = "debug", skip_all, err(level = "debug"))]
541    pub async fn bootstrap_connection(&self) -> Result<Arc<C::Conn>, ClientError> {
542        let node = self.node(BOOTSTRAP_ID);
543        if let Some(connection) = self.cached(&node) {
544            return Ok(connection);
545        }
546        let _gate = node.connect.lock().await;
547        if let Some(connection) = self.cached(&node) {
548            return Ok(connection);
549        }
550        Self::wait_for_backoff(&node).await;
551        let bootstrap = match self.bootstrap.read() {
552            Ok(guard) => guard.clone(),
553            Err(poisoned) => poisoned.into_inner().clone(),
554        };
555        let rounds = {
556            let mut state = node.state();
557            state.address_index %= bootstrap.len().max(1);
558            state.addresses = bootstrap;
559            state.addresses.len()
560        };
561        let mut last_error = ClientError::Disconnected;
562        for _ in 0..rounds {
563            let Some(attempt) = self.next_attempt(&node) else {
564                break;
565            };
566            match self.dial(&node, attempt).await {
567                Ok(connection) => return Ok(connection),
568                // Kafka raises an authentication failure from the next call
569                // and does not try another node (`Metadata.fatalError`).
570                Err(error) if error.is_authentication_failure() => {
571                    node.state().connect_failed(self.policy);
572                    return Err(error);
573                }
574                Err(error) => last_error = error,
575            }
576            let mut state = node.state();
577            state.address_index = (state.address_index + 1) % rounds;
578        }
579        {
580            let mut state = node.state();
581            state.failed_connect_attempts = state.failed_connect_attempts.saturating_add(1);
582            state.back_off(self.policy);
583        }
584        Err(last_error)
585    }
586
587    /// Get-or-connect to the least loaded known broker, as Kafka's
588    /// `NetworkClient.leastLoadedNode` picks the node for a request with no
589    /// fixed target. Before the pool knows a broker, this is the bootstrap
590    /// connection.
591    ///
592    /// The order of preference:
593    ///
594    /// 1. An open connection with no request in flight.
595    /// 2. The open connection with the fewest requests in flight.
596    /// 3. A broker out of its reconnect backoff, the one with the oldest
597    ///    connection attempt first, and a broker never dialed before all.
598    /// 4. The broker whose reconnect backoff ends first.
599    ///
600    /// # Errors
601    /// Returns the error of [`get`](Self::get) or
602    /// [`bootstrap_connection`](Self::bootstrap_connection).
603    pub async fn least_loaded(&self) -> Result<Arc<C::Conn>, ClientError> {
604        let ids = self.broker_ids();
605        if ids.is_empty() || self.prefer_bootstrap.load(Ordering::Relaxed) {
606            return self.bootstrap_connection().await;
607        }
608        let offset = random_below(ids.len());
609        let now = Instant::now();
610        let mut fewest: Option<(usize, Arc<C::Conn>)> = None;
611        let mut candidate: Option<Candidate> = None;
612        for &id in ids.iter().cycle().skip(offset).take(ids.len()) {
613            let node = self.node(id);
614            let state = node.state();
615            if let Some(connection) = state.connection.as_ref().filter(|c| C::is_open(c)) {
616                let in_flight = C::in_flight(connection);
617                if in_flight == 0 {
618                    return Ok(Arc::clone(connection));
619                }
620                if fewest.as_ref().is_none_or(|(least, _)| in_flight < *least) {
621                    fewest = Some((in_flight, Arc::clone(connection)));
622                }
623                continue;
624            }
625            let backoff_end = state.retry_at.filter(|retry_at| *retry_at > now);
626            let key = (
627                backoff_end.is_some(),
628                backoff_end,
629                state.last_attempt_at,
630                id,
631            );
632            drop(state);
633            if candidate
634                .as_ref()
635                .is_none_or(|best| (key.0, key.1, key.2) < (best.0, best.1, best.2))
636            {
637                candidate = Some(key);
638            }
639        }
640        if let Some((_, connection)) = fewest {
641            return Ok(connection);
642        }
643        match candidate {
644            // Every known broker is in its reconnect backoff. Kafka's
645            // `leastLoadedNode` then has no node, and the metadata recovery
646            // goes back to the bootstrap servers.
647            Some((true, ..)) | None => self.bootstrap_connection().await,
648            Some((false, _, _, id)) => self.get(id).await,
649        }
650    }
651
652    /// An open connection with its id, without a new connection attempt: the
653    /// one of `preferred` while it is open, or else the one with the fewest
654    /// requests in flight. The bootstrap connection has id `-1`.
655    pub fn open_connection(&self, preferred: Option<i32>) -> Option<(i32, Arc<C::Conn>)> {
656        let open = |id: i32, node: &Node<C::Conn>| {
657            node.state()
658                .connection
659                .as_ref()
660                .filter(|connection| C::is_open(connection))
661                .map(|connection| (id, Arc::clone(connection)))
662        };
663        if let Some((id, node)) = preferred.and_then(|id| Some((id, self.nodes.get(&id)?)))
664            && let Some(found) = open(id, &node)
665        {
666            return Some(found);
667        }
668        let nodes = self
669            .nodes
670            .iter()
671            .map(|entry| (*entry.key(), Arc::clone(entry.value())))
672            .collect::<Vec<_>>();
673        nodes
674            .into_iter()
675            .filter_map(|(id, node)| open(id, &node))
676            .min_by_key(|(id, connection)| (C::in_flight(connection), *id))
677    }
678
679    /// Update the (id, host, port) registry from a list of brokers, which
680    /// usually comes from a `MetadataResponse`. This method opens no new
681    /// connections, and it does not resolve the hosts: the pool resolves a
682    /// host each time it needs new addresses, as Kafka does.
683    ///
684    /// This method skips brokers that advertise port `0`, because that is not
685    /// a dialable address. It shows up for in-process test brokers whose
686    /// advertised port never got rewritten to the real bound port. When the
687    /// registry leaves such an entry out, [`get`](BrokerPool::get) reports
688    /// `Disconnected` for that id. The caller can then fall back to the
689    /// bootstrap connection instead of trying a doomed `host:0` connect.
690    ///
691    /// A broker whose host or port changed loses its cached connection and
692    /// addresses.
693    #[tracing::instrument(level = "debug", skip_all, fields(brokers = brokers.len()))]
694    pub async fn refresh_brokers(&self, brokers: &[BrokerInfo]) {
695        if !brokers.is_empty() {
696            self.prefer_bootstrap.store(false, Ordering::Relaxed);
697            // Kafka's `Metadata.update` replaces the node list, so a broker
698            // that the new metadata does not name leaves the registry and its
699            // connection closes once no request uses it.
700            let named = brokers.iter().map(|broker| broker.id).collect::<Vec<_>>();
701            self.by_endpoint.retain(|id, _| named.contains(id));
702            self.nodes
703                .retain(|id, _| *id == BOOTSTRAP_ID || named.contains(id));
704        }
705        for b in brokers {
706            let Ok(port) = u16::try_from(b.port) else {
707                continue;
708            };
709            if port == 0 {
710                continue;
711            }
712            let endpoint = (b.host.clone(), port);
713            let previous = self.by_endpoint.insert(b.id, endpoint.clone());
714            if previous.is_some_and(|previous| previous != endpoint) {
715                tracing::info!(broker_id = b.id, host = %b.host, port, "broker address changed");
716                self.nodes.remove(&b.id);
717            }
718        }
719    }
720
721    /// Whether the registry knows a dialable address for this broker id. It
722    /// knows one when [`refresh_brokers`](BrokerPool::refresh_brokers)
723    /// learned it and the port was not `0`. A caller can use this to choose
724    /// between a route to a specific broker and a fallback to the bootstrap
725    /// connection, without a speculative connect.
726    #[must_use]
727    pub fn knows_broker(&self, broker_id: i32) -> bool {
728        self.by_endpoint.contains_key(&broker_id)
729    }
730
731    /// Return the broker ids from the most recent usable metadata.
732    ///
733    /// The deterministic order keeps metadata failover predictable. Callers
734    /// use this list to exhaust the last-known cluster before deciding that a
735    /// fresh bootstrap is required.
736    #[must_use]
737    pub fn broker_ids(&self) -> Vec<i32> {
738        let mut ids = self
739            .by_endpoint
740            .iter()
741            .map(|entry| *entry.key())
742            .collect::<Vec<_>>();
743        ids.sort_unstable();
744        ids
745    }
746
747    /// [`broker_ids`](Self::broker_ids), with the brokers that have an open
748    /// connection first, fewest requests in flight first, then the others in
749    /// id order.
750    ///
751    /// Kafka sends a metadata request to `NetworkClient.leastLoadedNode`. A
752    /// broker answers the requests of one connection in order, so a request
753    /// queued behind a parked long-poll `Fetch` waits for up to the fetch's
754    /// `max_wait_ms`. A caller that walks every broker until one answers, as
755    /// the metadata refresh does, walks them in this order.
756    #[must_use]
757    pub fn broker_ids_by_load(&self) -> Vec<i32> {
758        let mut ids = self.broker_ids();
759        ids.sort_by_key(|id| {
760            self.node(*id)
761                .state()
762                .connection
763                .as_ref()
764                .filter(|connection| C::is_open(connection))
765                .map_or((1, 0), |connection| (0, C::in_flight(connection)))
766        });
767        ids
768    }
769
770    /// The finalized feature levels with the highest epoch that any
771    /// connection of the pool received, as Kafka's `ApiVersions.getMaxFinalizedFeaturesInfo`
772    /// gives them.
773    #[must_use]
774    pub fn finalized_features(&self) -> FinalizedFeatures {
775        self.finalized_features
776            .lock()
777            .unwrap_or_else(PoisonError::into_inner)
778            .clone()
779    }
780
781    /// Close every open connection in the pool. Consumes the pool.
782    // cargo-mutants: teardown; no observable return to assert against
783    #[cfg_attr(test, mutants::skip)]
784    pub fn close_all(self) {
785        // Dropping the pool drops each node state and its `Arc`. When the
786        // last reference goes away the background tasks shut down through
787        // the `CancellationToken` in `ConnectionInner`.
788        drop(self);
789    }
790}
791
792#[cfg(test)]
793mod tests;
krabka_client_core/pool.rs

krabka_client_core/
pool.rs
