krabka_client_core/
client.rs

1//! Top-level [`Client`].
2//!
3//! [`Client`] wraps a [`BrokerPool`] and exposes a typed-request `send` API.
4
5use std::sync::{Arc, Mutex, PoisonError};
6
7use krabka_protocol::owned::{metadata_request, metadata_response};
8use krabka_units::{Time, convert::TimeExt as _, minutes};
9use metadata_request::MetadataRequest;
10use refined_type::rule::GreaterEqualI64;
11use tokio::time::Instant;
12use tracing::field;
13
14use crate::{
15    bootstrap,
16    connection::{
17        ClientDnsTimeout, ClientFrameMax, Connection, ConnectionDispatchQueueCapacity,
18        ConnectionOptions,
19    },
20    error::ClientError,
21    fetch,
22    metadata_topics::{MetadataScope, MetadataTopics},
23    offset_for_leader_epoch,
24    pool::{BrokerInfo, BrokerPool},
25    request::ProtocolRequest,
26    security::ClientSecurity,
27    telemetry::{
28        ClientMetrics, ClientTelemetry, ClientTelemetryConfig, NetworkMetrics,
29        TelemetryConnections, TelemetryTiming, client_instance_id,
30    },
31};
32
33/// A Kafka client backed by a [`BrokerPool`].
34///
35/// Construct a `Client` with [`Client::builder`].
36///
37/// A clone of a `Client` is cheap. The clone shares the underlying
38/// [`BrokerPool`] through an `Arc` and copies the connection options by value.
39#[derive(Clone)]
40pub struct Client {
41    bootstrap: String,
42    pool: Arc<BrokerPool>,
43    options: ConnectionOptions,
44    metadata_recovery: Arc<MetadataRecovery>,
45    metadata_topics: Arc<MetadataTopics>,
46    /// The metric registry, for a client that pushes its metrics.
47    metrics: Option<ClientMetrics>,
48    /// The KIP-714 reporter. The last clone that drops it starts its close.
49    /// The reporter itself holds a clone without it.
50    telemetry: Option<Arc<ClientTelemetry>>,
51}
52
53/// Kafka client behavior when its last-known broker metadata can no longer be
54/// refreshed.
55#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
56pub enum MetadataRecoveryStrategy {
57    /// Return the metadata error without contacting the configured bootstrap
58    /// addresses again.
59    None,
60    /// Re-resolve the configured bootstrap addresses and discard stale broker
61    /// connections and addresses.
62    #[default]
63    Rebootstrap,
64}
65
66/// KIP-1102's default interval from the first unsuccessful metadata attempt
67/// until the client repeats bootstrap discovery.
68pub const DEFAULT_METADATA_RECOVERY_REBOOTSTRAP_TRIGGER: Time = minutes(5);
69
70/// Non-negative, whole-millisecond KIP-1102 rebootstrap trigger.
71#[derive(Clone, Copy, Debug, Eq, PartialEq, derive_more::Into)]
72pub struct MetadataRecoveryRebootstrapTrigger(i64);
73
74impl MetadataRecoveryRebootstrapTrigger {
75    /// Validate a rebootstrap trigger.
76    ///
77    /// # Errors
78    ///
79    /// Returns an error when the value is not a non-negative whole number of
80    /// milliseconds.
81    pub fn new(value: Time) -> Result<Self, String> {
82        let milliseconds = GreaterEqualI64::<0>::new(value.millis_i64())
83            .map_err(|error| format!("metadata recovery rebootstrap trigger: {error}"))?
84            .into_value();
85        if !value.secs_f64().is_finite() || Time::from_millis(milliseconds) != value {
86            return Err(
87                "metadata recovery rebootstrap trigger must be a whole number of milliseconds"
88                    .to_owned(),
89            );
90        }
91        Ok(Self(milliseconds))
92    }
93
94    /// Return the validated trigger interval.
95    #[must_use]
96    pub fn time(self) -> Time {
97        Time::from_millis(self.0)
98    }
99}
100
101impl Default for MetadataRecoveryRebootstrapTrigger {
102    fn default() -> Self {
103        Self(DEFAULT_METADATA_RECOVERY_REBOOTSTRAP_TRIGGER.millis_i64())
104    }
105}
106
107#[derive(Debug)]
108struct MetadataRecovery {
109    strategy: MetadataRecoveryStrategy,
110    trigger: MetadataRecoveryRebootstrapTrigger,
111    first_attempt: Mutex<Option<Instant>>,
112}
113
114impl MetadataRecovery {
115    fn begin_attempt(&self) {
116        if self.strategy == MetadataRecoveryStrategy::None {
117            return;
118        }
119        let mut first_attempt = self
120            .first_attempt
121            .lock()
122            .unwrap_or_else(PoisonError::into_inner);
123        first_attempt.get_or_insert_with(Instant::now);
124    }
125
126    fn complete_attempt(&self) {
127        *self
128            .first_attempt
129            .lock()
130            .unwrap_or_else(PoisonError::into_inner) = None;
131    }
132
133    fn timed_out(&self) -> bool {
134        if self.strategy == MetadataRecoveryStrategy::None {
135            return false;
136        }
137        self.first_attempt
138            .lock()
139            .unwrap_or_else(PoisonError::into_inner)
140            .is_some_and(|started| started.elapsed() >= self.trigger.time().to_std())
141    }
142}
143
144const REBOOTSTRAP_REQUIRED: i16 = 129;
145
146#[bon::bon]
147impl Client {
148    /// Build a [`Client`] pointed at the given bootstrap address.
149    #[builder(start_fn = builder, finish_fn = build)]
150    #[tracing::instrument(
151        level = "info",
152        skip_all,
153        fields(bootstrap = %bootstrap, client_id = %client_id),
154        err(level = "info"),
155    )]
156    /// # Errors
157    /// Returns an error when configuration is invalid, protocol encoding fails, the broker rejects the request, or transport I/O fails.
158    pub async fn start(
159        #[builder(into)] bootstrap: String,
160        #[builder(into, default = "krabka".to_owned())] client_id: String,
161        #[builder(default = crate::DEFAULT_CLIENT_DNS_TIMEOUT)] dns_timeout: Time,
162        #[builder(default = crate::DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT)]
163        socket_connection_setup_timeout: Time,
164        #[builder(default = crate::DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT_MAX)]
165        socket_connection_setup_timeout_max: Time,
166        #[builder(default = crate::DEFAULT_RECONNECT_BACKOFF)] reconnect_backoff: Time,
167        #[builder(default = crate::DEFAULT_RECONNECT_BACKOFF_MAX)] reconnect_backoff_max: Time,
168        #[builder(default = crate::DEFAULT_CONNECTIONS_MAX_IDLE)] connections_max_idle: Time,
169        #[builder(default)] client_dns_lookup: crate::ClientDnsLookup,
170        #[builder(required, default = Some(crate::DEFAULT_SEND_BUFFER))] send_buffer: Option<
171            krabka_units::ByteSize,
172        >,
173        #[builder(required, default = Some(crate::DEFAULT_RECEIVE_BUFFER))] receive_buffer: Option<
174            krabka_units::ByteSize,
175        >,
176        #[builder(default = crate::DEFAULT_CLIENT_REQUEST_TIMEOUT)] request_timeout: Time,
177        #[builder(default = crate::DEFAULT_CONNECTION_DISPATCH_QUEUE_CAPACITY)]
178        dispatch_queue_capacity: usize,
179        #[builder(default = crate::DEFAULT_CLIENT_FRAME_MAX)] frame_max: krabka_units::ByteSize,
180        #[builder(default)] metadata_recovery_strategy: MetadataRecoveryStrategy,
181        #[builder(default = crate::DEFAULT_METADATA_RECOVERY_REBOOTSTRAP_TRIGGER)]
182        metadata_recovery_rebootstrap_trigger: Time,
183        security: Option<ClientSecurity>,
184        /// The topics of [`Client::refresh_metadata`]. Kafka's producer and
185        /// consumer name their topics. The default asks for all topics.
186        #[builder(default)]
187        metadata_scope: MetadataScope,
188        /// Push client metrics (KIP-714) as this kind of client: Kafka's
189        /// `enable.metrics.push`. `None`, the default of this transport
190        /// client, sends neither `GetTelemetrySubscriptions` nor
191        /// `PushTelemetry`. The producer, consumer and admin builders set it.
192        telemetry: Option<ClientTelemetryConfig>,
193    ) -> Result<Self, ClientError> {
194        let dns_timeout = ClientDnsTimeout::new(dns_timeout).map_err(ClientError::InvalidConfig)?;
195        let dispatch_queue_capacity = ConnectionDispatchQueueCapacity::new(dispatch_queue_capacity)
196            .map_err(ClientError::InvalidConfig)?;
197        let frame_max = ClientFrameMax::try_from(frame_max).map_err(ClientError::InvalidConfig)?;
198        let metadata_recovery_rebootstrap_trigger =
199            MetadataRecoveryRebootstrapTrigger::new(metadata_recovery_rebootstrap_trigger)
200                .map_err(ClientError::InvalidConfig)?;
201        // Kafka builds the SSL engine factory when it builds the client, so a
202        // store that does not load fails the client before any connection.
203        // The connections reuse this configuration.
204        if let Some(tls) = security.as_ref().and_then(|security| security.tls.as_ref()) {
205            tls.connector()
206                .map_err(|error| ClientError::InvalidConfig(error.to_string()))?;
207        }
208        let options = ConnectionOptions {
209            client_id,
210            dns_timeout,
211            socket_connection_setup_timeout,
212            socket_connection_setup_timeout_max,
213            reconnect_backoff,
214            reconnect_backoff_max,
215            connections_max_idle,
216            client_dns_lookup,
217            send_buffer,
218            receive_buffer,
219            request_timeout,
220            dispatch_queue_capacity,
221            frame_max,
222            security: security.map(Box::new),
223            // `start_with_options` sets it for a client that pushes metrics.
224            network_metrics: None,
225        };
226        options.validate().map_err(ClientError::InvalidConfig)?;
227        Self::start_with_options(
228            bootstrap,
229            options,
230            (
231                metadata_recovery_strategy,
232                metadata_recovery_rebootstrap_trigger,
233            ),
234            metadata_scope,
235            telemetry,
236        )
237        .await
238    }
239}
240
241impl Client {
242    #[tracing::instrument(
243        level = "info",
244        skip_all,
245        fields(bootstrap = %bootstrap, client_id = %options.client_id),
246        err(level = "info"),
247    )]
248    async fn start_with_options(
249        bootstrap: String,
250        mut options: ConnectionOptions,
251        (metadata_recovery_strategy, metadata_recovery_rebootstrap_trigger): (
252            MetadataRecoveryStrategy,
253            MetadataRecoveryRebootstrapTrigger,
254        ),
255        metadata_scope: MetadataScope,
256        telemetry: Option<ClientTelemetryConfig>,
257    ) -> Result<Self, ClientError> {
258        let addrs = bootstrap::resolve_with_server_names(
259            &bootstrap,
260            options.dns_timeout,
261            options.client_dns_lookup,
262        )
263        .await?;
264        let metrics = telemetry.as_ref().map(|_| ClientMetrics::new());
265        options.network_metrics =
266            telemetry
267                .as_ref()
268                .zip(metrics.as_ref())
269                .map(|(config, metrics)| {
270                    NetworkMetrics::register(metrics, config.client_type.metric_group())
271                });
272        let pool = Arc::new(BrokerPool::new_with_server_names(addrs, options.clone()));
273        let client = Self {
274            bootstrap,
275            pool,
276            options,
277            metadata_recovery: Arc::new(MetadataRecovery {
278                strategy: metadata_recovery_strategy,
279                trigger: metadata_recovery_rebootstrap_trigger,
280                first_attempt: Mutex::new(None),
281            }),
282            metadata_topics: Arc::new(MetadataTopics::new(metadata_scope)),
283            metrics,
284            telemetry: None,
285        };
286        let telemetry = telemetry
287            .zip(client.metrics.clone())
288            .map(|(config, metrics)| {
289                let timing = TelemetryTiming {
290                    request_timeout: client.options.request_timeout.to_std(),
291                    reconnect_backoff: client.options.reconnect_backoff.to_std(),
292                };
293                Arc::new(ClientTelemetry::start(
294                    client.clone(),
295                    &config,
296                    metrics,
297                    timing,
298                ))
299            });
300        Ok(Self {
301            telemetry,
302            ..client
303        })
304    }
305
306    /// The options of every connection that this client opens, such as the
307    /// socket buffers (`send.buffer.bytes`, `receive.buffer.bytes`) that
308    /// each broker socket gets before it connects.
309    #[must_use]
310    pub fn connection_options(&self) -> &ConnectionOptions {
311        &self.options
312    }
313
314    /// The metric registry of a client that pushes its metrics (KIP-714).
315    /// A client crate registers its own metrics in it.
316    #[must_use]
317    pub fn metrics(&self) -> Option<&ClientMetrics> {
318        self.metrics.as_ref()
319    }
320
321    /// Send the terminating telemetry push (KIP-714), if the client pushes
322    /// metrics and has a subscription, and wait for it. Kafka's clients
323    /// send it from `close`. Without this call, the push goes out in the
324    /// background once the last clone of the client drops.
325    pub async fn close_telemetry(&self) {
326        if let Some(telemetry) = &self.telemetry {
327            telemetry.close().await;
328        }
329    }
330
331    /// The client instance id that the broker assigned for metrics push
332    /// (KIP-714), waiting up to `timeout` for the first subscription, as
333    /// Kafka's `clientInstanceId` does. `Ok(None)` when no subscription has
334    /// loaded by then, as Kafka returns `null`.
335    ///
336    /// # Errors
337    /// Returns [`ClientError::InvalidArgument`] for a negative `timeout`, and
338    /// [`ClientError::TelemetryDisabled`] for a client that does not push
339    /// metrics.
340    pub async fn client_instance_id(
341        &self,
342        timeout: Time,
343    ) -> Result<Option<uuid::Uuid>, ClientError> {
344        client_instance_id(self.telemetry.as_deref(), timeout).await
345    }
346
347    /// Send a request that has no fixed target broker.
348    ///
349    /// Before the first metadata response, the request goes to a bootstrap
350    /// address. After it, the request goes to the least loaded known broker,
351    /// as Kafka's `NetworkClient.leastLoadedNode` picks the node.
352    #[tracing::instrument(level = "debug", skip_all, fields(api_key = R::API_KEY), err(level = "debug"))]
353    /// # Errors
354    /// Returns an error when configuration is invalid, protocol encoding fails, the broker rejects the request, or transport I/O fails.
355    pub async fn send<R: ProtocolRequest>(&self, req: R) -> Result<R::Response, ClientError> {
356        let conn = self.pool.least_loaded().await?;
357        conn.send(req).await
358    }
359
360    /// Send a request, as [`send`](Self::send) routes it, only when the
361    /// broker and client share a version at or above `min_version`.
362    ///
363    /// # Errors
364    /// Returns [`ClientError::IncompatibleVersion`] before dispatch when the
365    /// shared range is below `min_version`. Other errors match [`send`](Self::send).
366    pub async fn send_at_least<R: ProtocolRequest>(
367        &self,
368        req: R,
369        min_version: i16,
370    ) -> Result<R::Response, ClientError> {
371        let conn = self.pool.least_loaded().await?;
372        let (broker_min, broker_max) = conn.advertised_api_range(R::API_KEY).unwrap_or((0, 0));
373        let client_min = R::MIN_VERSION.max(min_version);
374        let chosen = R::LATEST_STABLE_VERSION.min(broker_max);
375        if chosen < client_min || chosen < broker_min {
376            return Err(ClientError::IncompatibleVersion {
377                api_key: R::API_KEY,
378                broker_min,
379                broker_max,
380                client_min,
381                client_max: R::LATEST_STABLE_VERSION,
382            });
383        }
384        conn.send(req).await
385    }
386
387    /// Send on one fixed pool connection, without selecting another broker.
388    ///
389    /// `Some(id)` selects that broker's connection. `None` selects the distinct
390    /// bootstrap connection even after cluster metadata has been learned.
391    /// Callers that account for per-connection request limits can keep this
392    /// selection fixed until the request completes, then refresh and retry on
393    /// failure themselves.
394    ///
395    /// # Errors
396    ///
397    /// Returns a routing, connection setup, encoding or response error.
398    pub async fn send_to_connection<R: ProtocolRequest>(
399        &self,
400        broker_id: Option<i32>,
401        req: R,
402    ) -> Result<R::Response, ClientError> {
403        let connection = match broker_id {
404            Some(id) => self.pool.get(id).await?,
405            None => self.pool.bootstrap_connection().await?,
406        };
407        connection.send(req).await
408    }
409
410    /// Send without a response on the fixed connection selected by
411    /// [`send_to_connection`](Self::send_to_connection).
412    ///
413    /// # Errors
414    ///
415    /// Returns a routing, connection setup, encoding or writer enqueue error.
416    pub async fn send_to_connection_no_response<R: ProtocolRequest>(
417        &self,
418        broker_id: Option<i32>,
419        req: R,
420    ) -> Result<(), ClientError> {
421        let connection = match broker_id {
422            Some(id) => self.pool.get(id).await?,
423            None => self.pool.bootstrap_connection().await?,
424        };
425        connection.send_no_response(req).await
426    }
427
428    /// Send a request, as [`send`](Self::send) routes it, without registering
429    /// a pending response. Intended for protocol operations such as Produce
430    /// `acks=0`.
431    ///
432    /// # Errors
433    ///
434    /// Returns an error if the connection cannot be opened or cannot accept
435    /// the encoded request.
436    pub async fn send_no_response<R: ProtocolRequest>(&self, req: R) -> Result<(), ClientError> {
437        let conn = self.pool.least_loaded().await?;
438        conn.send_no_response(req).await
439    }
440
441    /// Drop the cached bootstrap connection and refresh the bootstrap address
442    /// list from the original bootstrap string.
443    ///
444    /// Callers that know their bootstrap request is safe to retry can call
445    /// this after a transport error and before they send again.
446    #[tracing::instrument(level = "debug", skip_all, fields(bootstrap = %self.bootstrap))]
447    pub async fn reconnect_bootstrap(&self) {
448        self.pool.evict_bootstrap();
449        // The retry goes to the bootstrap addresses even when the DNS refresh
450        // below fails and the pool keeps the addresses that it has.
451        self.pool.prefer_bootstrap();
452        if let Ok(addrs) = bootstrap::resolve_with_server_names(
453            &self.bootstrap,
454            self.options.dns_timeout,
455            self.options.client_dns_lookup,
456        )
457        .await
458        {
459            self.pool.replace_bootstrap_with_server_names(addrs);
460        }
461    }
462
463    async fn rebootstrap_metadata(&self) -> Result<(), ClientError> {
464        let addrs = bootstrap::resolve_with_server_names(
465            &self.bootstrap,
466            self.options.dns_timeout,
467            self.options.client_dns_lookup,
468        )
469        .await?;
470        self.pool.rebootstrap_with_server_names(addrs);
471        self.metadata_recovery.complete_attempt();
472        Ok(())
473    }
474
475    async fn request_metadata_from_current_cluster(
476        &self,
477        request: &MetadataRequest,
478    ) -> Result<metadata_response::MetadataResponse, ClientError> {
479        let broker_ids = self.pool.broker_ids_by_load();
480        if broker_ids.is_empty() {
481            let connection = self.pool.bootstrap_connection().await?;
482            return connection.send(request.clone()).await;
483        }
484
485        let mut last_error = None;
486        let mut empty_response = None;
487        for broker_id in broker_ids {
488            let connection = match self.pool.get(broker_id).await {
489                Ok(connection) => connection,
490                Err(
491                    error @ (ClientError::Connect { .. }
492                    | ClientError::Tls { .. }
493                    | ClientError::Sasl { .. }
494                    | ClientError::Timeout(_)
495                    | ClientError::Disconnected
496                    | ClientError::Io(_)),
497                ) => {
498                    self.pool.evict(broker_id);
499                    last_error = Some(error);
500                    continue;
501                }
502                Err(error) => return Err(error),
503            };
504            match connection.send(request.clone()).await {
505                Ok(response)
506                    if response.error_code == REBOOTSTRAP_REQUIRED
507                        || !response.brokers.is_empty() =>
508                {
509                    return Ok(response);
510                }
511                Ok(response) => empty_response = Some(response),
512                Err(
513                    error @ (ClientError::Connect { .. }
514                    | ClientError::Timeout(_)
515                    | ClientError::Disconnected
516                    | ClientError::Io(_)),
517                ) => {
518                    self.pool.evict(broker_id);
519                    last_error = Some(error);
520                }
521                Err(error) => return Err(error),
522            }
523        }
524
525        empty_response.map_or_else(|| Err(last_error.unwrap_or(ClientError::Disconnected)), Ok)
526    }
527
528    /// The finalized feature levels (KIP-584) with the highest epoch that a
529    /// connection of this client received in `ApiVersions`. Kafka's producer
530    /// reads `transaction.version` from them. The epoch is
531    /// [`UNKNOWN_FINALIZED_FEATURES_EPOCH`](crate::UNKNOWN_FINALIZED_FEATURES_EPOCH)
532    /// before a connection to a broker that sends features.
533    #[must_use]
534    pub fn finalized_features(&self) -> crate::FinalizedFeatures {
535        self.pool.finalized_features()
536    }
537
538    /// Whether the pool knows a dialable address for `broker_id`.
539    ///
540    /// The pool knows one when
541    /// [`refresh_metadata`](Client::refresh_metadata) learned it and the port
542    /// was not `0`. A caller can then choose between
543    /// [`broker`](Client::broker) routing and the bootstrap
544    /// [`send`](Client::send) without a speculative connect.
545    // cargo-mutants: one-line delegation to BrokerPool::knows_broker
546    #[must_use]
547    #[cfg_attr(test, mutants::skip)]
548    pub fn knows_broker(&self, broker_id: i32) -> bool {
549        self.pool.knows_broker(broker_id)
550    }
551
552    /// Return a [`BrokerHandle`] that routes requests to a specific broker by id.
553    ///
554    /// [`refresh_metadata`] must have registered the broker first.
555    ///
556    /// [`refresh_metadata`]: Client::refresh_metadata
557    #[must_use]
558    pub fn broker(&self, broker_id: i32) -> BrokerHandle<'_> {
559        BrokerHandle {
560            client: self,
561            broker_id,
562        }
563    }
564
565    /// Drop the pooled connection to `broker_id` so the next request to it
566    /// reconnects to its current advertised address.
567    ///
568    /// Call this after a send fails, so a bounced or failed-over broker is not
569    /// retried over a dead socket.
570    // cargo-mutants: one-line delegation to BrokerPool::evict
571    #[cfg_attr(test, mutants::skip)]
572    pub fn evict_broker(&self, broker_id: i32) {
573        self.pool.evict(broker_id);
574    }
575
576    /// The topics that [`refresh_metadata`](Client::refresh_metadata) names.
577    /// The clones of this client share them.
578    #[must_use]
579    pub fn metadata_topics(&self) -> &MetadataTopics {
580        &self.metadata_topics
581    }
582
583    /// Send the `MetadataRequest` of [`metadata_topics`](Client::metadata_topics),
584    /// parse the broker list from the response, refresh the pool's address
585    /// registry, and return the typed response.
586    ///
587    /// With [`MetadataScope::AllTopics`] the request asks for all topics.
588    /// With [`MetadataScope::Topics`] it names the current topics, as Kafka's
589    /// producer and consumer do. Use
590    /// [`refresh_metadata_with`](Client::refresh_metadata_with) for another
591    /// request.
592    ///
593    /// # Errors
594    /// Returns an error when configuration is invalid, protocol encoding fails, the broker rejects the request, or transport I/O fails.
595    pub async fn refresh_metadata(
596        &self,
597    ) -> Result<metadata_response::MetadataResponse, ClientError> {
598        let response = self
599            .refresh_metadata_with(self.metadata_topics.request())
600            .await?;
601        self.metadata_topics.mark_refreshed();
602        Ok(response)
603    }
604
605    /// Send `request`, parse the broker list from the response, refresh the
606    /// pool's address registry, and return the typed response.
607    ///
608    /// Kafka's producer names its topics in the request
609    /// (`ProducerMetadata.newMetadataRequestBuilder`). A request that names a
610    /// topic lets the broker create the topic when
611    /// `allow_auto_topic_creation` is set, and returns a per-topic error such
612    /// as `TOPIC_AUTHORIZATION_FAILED` for it. Every retry and rebootstrap in
613    /// this call sends the same request.
614    ///
615    /// # Errors
616    /// Returns an error when configuration is invalid, protocol encoding fails, the broker rejects the request, or transport I/O fails.
617    // cargo-mutants: live-broker metadata round-trip; not unit-testable
618    #[cfg_attr(test, mutants::skip)]
619    #[tracing::instrument(
620        level = "info",
621        skip_all,
622        fields(brokers = field::Empty),
623        err(level = "info"),
624    )]
625    pub async fn refresh_metadata_with(
626        &self,
627        request: MetadataRequest,
628    ) -> Result<metadata_response::MetadataResponse, ClientError> {
629        self.metadata_recovery.begin_attempt();
630        let first = self.request_metadata_from_current_cluster(&request).await;
631        let resp = match first {
632            Ok(resp)
633                if resp.error_code == REBOOTSTRAP_REQUIRED
634                    && self.metadata_recovery.strategy == MetadataRecoveryStrategy::Rebootstrap =>
635            {
636                self.rebootstrap_metadata().await?;
637                self.metadata_recovery.begin_attempt();
638                self.request_metadata_from_current_cluster(&request).await?
639            }
640            Ok(resp) if resp.error_code == REBOOTSTRAP_REQUIRED => {
641                return Err(ClientError::Server {
642                    error_code: REBOOTSTRAP_REQUIRED,
643                });
644            }
645            Ok(resp) if resp.brokers.is_empty() && self.metadata_recovery.timed_out() => {
646                self.rebootstrap_metadata().await?;
647                self.metadata_recovery.begin_attempt();
648                self.request_metadata_from_current_cluster(&request).await?
649            }
650            Ok(resp) => resp,
651            // `request_metadata_from_current_cluster` has exhausted every
652            // broker in the last-known metadata. KIP-1102 permits immediate
653            // rebootstrap at that boundary, and additionally covers a
654            // reachable cluster that returns no usable metadata until the
655            // configured trigger expires.
656            Err(
657                error @ (ClientError::Connect { .. }
658                | ClientError::Tls { .. }
659                | ClientError::Sasl { .. }
660                | ClientError::Timeout(_)
661                | ClientError::Disconnected
662                | ClientError::Io(_)),
663            ) => {
664                if self.metadata_recovery.strategy == MetadataRecoveryStrategy::None {
665                    return Err(error);
666                }
667                self.rebootstrap_metadata().await?;
668                self.metadata_recovery.begin_attempt();
669                self.request_metadata_from_current_cluster(&request).await?
670            }
671            Err(error) => return Err(error),
672        };
673        if resp.error_code == REBOOTSTRAP_REQUIRED {
674            return Err(ClientError::Server {
675                error_code: REBOOTSTRAP_REQUIRED,
676            });
677        }
678        if !resp.brokers.is_empty() {
679            self.metadata_recovery.complete_attempt();
680        }
681        let brokers: Vec<BrokerInfo> = resp
682            .brokers
683            .iter()
684            .map(|b| BrokerInfo {
685                id: b.node_id,
686                host: b.host.clone(),
687                port: b.port,
688                rack: b.rack.clone(),
689            })
690            .collect();
691        tracing::Span::current().record("brokers", brokers.len());
692        self.pool.refresh_brokers(&brokers).await;
693        Ok(resp)
694    }
695
696    /// Send a single-partition `OffsetForLeaderEpoch` (`api_key=23`) over the
697    /// bootstrap connection.
698    ///
699    /// This is a thin wrapper over the free
700    /// [`offset_for_leader_epoch`](crate::offset_for_leader_epoch) helper that
701    /// the consumer's KIP-320 position-validation pass uses. `Client` does not
702    /// otherwise expose its connection, so this method borrows the same
703    /// bootstrap connection that `send` uses.
704    ///
705    /// # Errors
706    /// Returns an error on transport or version-negotiation failure, or when
707    /// the response does not contain the partition.
708    #[tracing::instrument(
709        level = "debug",
710        skip_all,
711        fields(topic = %topic, partition, current_leader_epoch, leader_epoch),
712        err(level = "debug"),
713    )]
714    pub async fn offset_for_leader_epoch(
715        &self,
716        topic: &str,
717        partition: i32,
718        current_leader_epoch: i32,
719        leader_epoch: i32,
720    ) -> Result<offset_for_leader_epoch::EpochEndOffset, ClientError> {
721        let conn = match self.pool.bootstrap_connection().await {
722            Ok(conn) => conn,
723            Err(
724                ClientError::Connect { .. }
725                | ClientError::Tls { .. }
726                | ClientError::Sasl { .. }
727                | ClientError::Timeout(_)
728                | ClientError::Disconnected
729                | ClientError::Io(_),
730            ) => {
731                self.reconnect_bootstrap().await;
732                self.pool.bootstrap_connection().await?
733            }
734            Err(e) => return Err(e),
735        };
736        offset_for_leader_epoch::offset_for_leader_epoch(
737            &conn,
738            topic,
739            partition,
740            current_leader_epoch,
741            leader_epoch,
742        )
743        .await
744    }
745
746    /// Send a single-partition `OffsetForLeaderEpoch` (`api_key=23`) to a
747    /// *specific* broker by id, through [`BrokerPool::get`].
748    ///
749    /// This method mirrors
750    /// [`offset_for_leader_epoch`](Client::offset_for_leader_epoch) but
751    /// targets the partition leader instead of the bootstrap connection.
752    /// KIP-320 requires the validation RPC to reach the partition leader,
753    /// which is the only replica with the authoritative epoch→end-offset
754    /// history.
755    ///
756    /// The broker must already be in the pool's registry, which
757    /// [`refresh_metadata`](Client::refresh_metadata) populates.
758    ///
759    /// # Errors
760    /// Returns `Disconnected` if `broker_id` is not in the registry. Returns
761    /// an error on transport or version-negotiation failure. Returns an error
762    /// when the response does not contain the partition.
763    #[tracing::instrument(
764        level = "debug",
765        skip_all,
766        fields(broker_id, topic = %topic, partition, current_leader_epoch, leader_epoch),
767        err(level = "debug"),
768    )]
769    pub async fn offset_for_leader_epoch_on(
770        &self,
771        broker_id: i32,
772        topic: &str,
773        partition: i32,
774        current_leader_epoch: i32,
775        leader_epoch: i32,
776    ) -> Result<offset_for_leader_epoch::EpochEndOffset, ClientError> {
777        let conn = match self.pool.get(broker_id).await {
778            Ok(conn) => conn,
779            Err(
780                ClientError::Connect { .. }
781                | ClientError::Tls { .. }
782                | ClientError::Sasl { .. }
783                | ClientError::Timeout(_)
784                | ClientError::Io(_),
785            ) if self.pool.knows_broker(broker_id) => {
786                self.pool.evict(broker_id);
787                self.refresh_metadata().await?;
788                self.pool.get(broker_id).await?
789            }
790            Err(e) => return Err(e),
791        };
792        offset_for_leader_epoch::offset_for_leader_epoch(
793            &conn,
794            topic,
795            partition,
796            current_leader_epoch,
797            leader_epoch,
798        )
799        .await
800    }
801
802    /// Fetch one partition from a specific broker with the requested isolation
803    /// level. The broker address must have been learned through
804    /// [`refresh_metadata`](Client::refresh_metadata).
805    ///
806    /// # Errors
807    ///
808    /// Returns a connection, protocol, or broker error. A failed connection to
809    /// a known broker is evicted, metadata is refreshed, and the request is
810    /// retried once against the broker's current advertised address.
811    #[tracing::instrument(
812        level = "debug",
813        skip_all,
814        fields(
815            broker_id,
816            topic = %fetch.topic,
817            partition = fetch.partition,
818            fetch_offset = fetch.fetch_offset,
819            isolation_level = fetch.isolation_level,
820        ),
821        err(level = "debug"),
822    )]
823    pub async fn fetch_partition_with_isolation_on(
824        &self,
825        broker_id: i32,
826        fetch: fetch::IsolatedFetch<'_>,
827    ) -> Result<Vec<fetch::FetchedRecord>, ClientError> {
828        let conn = match self.pool.get(broker_id).await {
829            Ok(conn) => conn,
830            Err(ClientError::Disconnected) if !self.pool.knows_broker(broker_id) => {
831                self.pool.bootstrap_connection().await?
832            }
833            Err(
834                ClientError::Connect { .. }
835                | ClientError::Tls { .. }
836                | ClientError::Sasl { .. }
837                | ClientError::Timeout(_)
838                | ClientError::Io(_),
839            ) if self.pool.knows_broker(broker_id) => {
840                self.pool.evict(broker_id);
841                self.refresh_metadata().await?;
842                match self.pool.get(broker_id).await {
843                    Ok(conn) => conn,
844                    Err(ClientError::Disconnected) if !self.pool.knows_broker(broker_id) => {
845                        self.pool.bootstrap_connection().await?
846                    }
847                    Err(error) => return Err(error),
848                }
849            }
850            Err(error) => return Err(error),
851        };
852        fetch::fetch_partition_with_isolation(&conn, fetch).await
853    }
854
855    /// Close the client and all pooled connections.
856    // cargo-mutants: teardown; delegates to BrokerPool::close_all
857    #[cfg_attr(test, mutants::skip)]
858    pub fn close(self) {
859        if let Some(pool) = Arc::into_inner(self.pool) {
860            pool.close_all();
861        }
862    }
863}
864
865impl TelemetryConnections for Client {
866    fn open_connection(&self, preferred: Option<i32>) -> Option<(i32, Arc<Connection>)> {
867        self.pool.open_connection(preferred)
868    }
869}
870
871/// A handle to a specific broker within a [`Client`]'s pool.
872///
873/// [`Client::broker`] returns this handle.
874pub struct BrokerHandle<'a> {
875    client: &'a Client,
876    broker_id: i32,
877}
878
879impl BrokerHandle<'_> {
880    /// Send a request to this specific broker.
881    ///
882    /// When the pool has no dialable address for `broker_id`, this method
883    /// falls back to the bootstrap connection. That happens when the broker
884    /// advertises port `0`, which occurs on a single-broker cluster whose
885    /// OS-assigned port never got rewritten in metadata, so
886    /// [`BrokerPool::refresh_brokers`] deliberately skipped it. On such a
887    /// cluster the bootstrap broker *is* this broker, for example the group
888    /// coordinator a consumer routes to. The request therefore still reaches
889    /// its intended target instead of failing with `Disconnected`.
890    ///
891    /// This fallback does not mask a *known* broker whose connect fails. It
892    /// only triggers when the id was never in the registry.
893    // cargo-mutants: live-broker send path; not unit-testable
894    #[cfg_attr(test, mutants::skip)]
895    #[tracing::instrument(
896        level = "debug",
897        skip_all,
898        fields(broker_id = self.broker_id, api_key = R::API_KEY),
899        err(level = "debug"),
900    )]
901    /// # Errors
902    /// Returns an error when configuration is invalid, protocol encoding fails, the broker rejects the request, or transport I/O fails.
903    pub async fn send<R: ProtocolRequest>(&self, req: R) -> Result<R::Response, ClientError> {
904        let conn = match self.client.pool.get(self.broker_id).await {
905            Ok(conn) => conn,
906            Err(ClientError::Disconnected) if !self.client.pool.knows_broker(self.broker_id) => {
907                self.client.pool.bootstrap_connection().await?
908            }
909            Err(
910                ClientError::Connect { .. }
911                | ClientError::Tls { .. }
912                | ClientError::Sasl { .. }
913                | ClientError::Timeout(_)
914                | ClientError::Io(_),
915            ) if self.client.pool.knows_broker(self.broker_id) => {
916                self.client.pool.evict(self.broker_id);
917                self.client.refresh_metadata().await?;
918                match self.client.pool.get(self.broker_id).await {
919                    Ok(conn) => conn,
920                    Err(ClientError::Disconnected)
921                        if !self.client.pool.knows_broker(self.broker_id) =>
922                    {
923                        self.client.pool.bootstrap_connection().await?
924                    }
925                    Err(e) => return Err(e),
926                }
927            }
928            Err(e) => return Err(e),
929        };
930        conn.send(req).await
931    }
932
933    /// Send a request to this specific broker without waiting for a response.
934    ///
935    /// # Errors
936    ///
937    /// Returns an error if routing, connection setup, encoding, or writer
938    /// enqueue fails.
939    pub async fn send_no_response<R: ProtocolRequest>(&self, req: R) -> Result<(), ClientError> {
940        let conn = match self.client.pool.get(self.broker_id).await {
941            Ok(conn) => conn,
942            Err(ClientError::Disconnected) if !self.client.pool.knows_broker(self.broker_id) => {
943                self.client.pool.bootstrap_connection().await?
944            }
945            Err(
946                ClientError::Connect { .. }
947                | ClientError::Tls { .. }
948                | ClientError::Sasl { .. }
949                | ClientError::Timeout(_)
950                | ClientError::Io(_),
951            ) if self.client.pool.knows_broker(self.broker_id) => {
952                self.client.pool.evict(self.broker_id);
953                self.client.refresh_metadata().await?;
954                match self.client.pool.get(self.broker_id).await {
955                    Ok(conn) => conn,
956                    Err(ClientError::Disconnected)
957                        if !self.client.pool.knows_broker(self.broker_id) =>
958                    {
959                        self.client.pool.bootstrap_connection().await?
960                    }
961                    Err(error) => return Err(error),
962                }
963            }
964            Err(error) => return Err(error),
965        };
966        conn.send_no_response(req).await
967    }
968}
969
970#[cfg(test)]
971mod bootstrap_failover_tests {
972    use std::{
973        net::SocketAddr,
974        sync::{
975            Arc,
976            atomic::{AtomicU16, AtomicUsize, Ordering},
977        },
978        time::Duration,
979    };
980
981    use bytes::BytesMut;
982    use krabka_protocol::{
983        Decode, Encode,
984        owned::{
985            api_versions_request,
986            api_versions_response::{ApiVersion, ApiVersionsResponse},
987            metadata_request::{self, MetadataRequestTopic},
988            metadata_response::{
989                FLEXIBLE_MIN as META_FLEXIBLE_MIN, MetadataResponse, MetadataResponseBroker,
990            },
991        },
992    };
993    use krabka_units::{bytes, mebibytes, millis};
994    use tokio::{
995        net::TcpStream,
996        task::yield_now,
997        time::{sleep, timeout},
998    };
999
1000    use super::*;
1001    use crate::mock::MockBroker;
1002
1003    #[tokio::test]
1004    async fn zero_dns_timeout_is_rejected_before_resolution() {
1005        let result = Client::builder()
1006            .bootstrap("unused.invalid:9092")
1007            .dns_timeout(millis(0))
1008            .build()
1009            .await;
1010        assert2::assert!(matches!(result, Err(ClientError::InvalidConfig(_))));
1011    }
1012
1013    #[tokio::test]
1014    async fn invalid_connection_resource_policy_fails_before_resolution() {
1015        let queue_result = Client::builder()
1016            .bootstrap("unused.invalid:9092")
1017            .dispatch_queue_capacity(0)
1018            .build()
1019            .await;
1020        let Err(queue_error) = queue_result else {
1021            panic!("zero queue capacity must fail");
1022        };
1023        assert2::assert!(queue_error.to_string().contains("dispatch queue capacity"));
1024
1025        let frame_result = Client::builder()
1026            .bootstrap("unused.invalid:9092")
1027            .frame_max(mebibytes(100) + bytes(1))
1028            .build()
1029            .await;
1030        let Err(frame_error) = frame_result else {
1031            panic!("frame limit above fixed ceiling must fail");
1032        };
1033        assert2::assert!(frame_error.to_string().contains("client frame max"));
1034
1035        assert2::assert!(
1036            MetadataRecoveryRebootstrapTrigger::new(Time::ZERO)
1037                .expect("zero is an immediate KIP-1102 trigger")
1038                .time()
1039                == Time::ZERO
1040        );
1041        assert2::assert!(
1042            MetadataRecoveryRebootstrapTrigger::default()
1043                == MetadataRecoveryRebootstrapTrigger::new(
1044                    DEFAULT_METADATA_RECOVERY_REBOOTSTRAP_TRIGGER
1045                )
1046                .unwrap()
1047        );
1048        let recovery_result = Client::builder()
1049            .bootstrap("unused.invalid:9092")
1050            .metadata_recovery_rebootstrap_trigger(Time::from_millis(-1))
1051            .build()
1052            .await;
1053        let Err(recovery_error) = recovery_result else {
1054            panic!("negative metadata recovery trigger must fail");
1055        };
1056        assert2::assert!(recovery_error.to_string().contains("metadata recovery"));
1057    }
1058
1059    fn api_versions_v0() -> Vec<u8> {
1060        let resp = ApiVersionsResponse {
1061            error_code: 0,
1062            api_keys: vec![
1063                ApiVersion {
1064                    api_key: api_versions_request::API_KEY,
1065                    min_version: 0,
1066                    max_version: 3,
1067                    ..Default::default()
1068                },
1069                ApiVersion {
1070                    api_key: metadata_request::API_KEY,
1071                    min_version: 0,
1072                    max_version: 13,
1073                    ..Default::default()
1074                },
1075            ],
1076            ..Default::default()
1077        };
1078        let mut buf = BytesMut::new();
1079        resp.encode(&mut buf, 0).unwrap();
1080        buf.to_vec()
1081    }
1082
1083    fn metadata_v(version: i16, node_id: i32) -> Vec<u8> {
1084        metadata_v_with_port(version, node_id, 9092)
1085    }
1086
1087    fn metadata_v_with_port(version: i16, node_id: i32, port: u16) -> Vec<u8> {
1088        let resp = MetadataResponse {
1089            brokers: vec![MetadataResponseBroker {
1090                node_id,
1091                host: "127.0.0.1".into(),
1092                port: i32::from(port),
1093                ..Default::default()
1094            }],
1095            ..Default::default()
1096        };
1097        let mut buf = BytesMut::new();
1098        if version >= META_FLEXIBLE_MIN {
1099            buf.extend_from_slice(&[0x00u8]); // empty tagged fields
1100        }
1101        resp.encode(&mut buf, version).unwrap();
1102        buf.to_vec()
1103    }
1104
1105    fn empty_metadata_v(version: i16, error_code: i16) -> Vec<u8> {
1106        let resp = MetadataResponse {
1107            error_code,
1108            ..Default::default()
1109        };
1110        let mut buf = BytesMut::new();
1111        if version >= META_FLEXIBLE_MIN {
1112            buf.extend_from_slice(&[0x00u8]);
1113        }
1114        resp.encode(&mut buf, version).unwrap();
1115        buf.to_vec()
1116    }
1117
1118    fn handler(
1119        node_id: i32,
1120    ) -> impl FnMut(i16, i16, i32, &[u8]) -> Option<Vec<u8>> + Send + 'static {
1121        move |api_key, version, _corr, _body| {
1122            if api_key == api_versions_request::API_KEY {
1123                return Some(api_versions_v0());
1124            }
1125            if api_key == metadata_request::API_KEY {
1126                return Some(metadata_v(version, node_id));
1127            }
1128            None
1129        }
1130    }
1131
1132    fn dynamic_port_handler(
1133        node_id: i32,
1134        port: Arc<AtomicU16>,
1135    ) -> impl FnMut(i16, i16, i32, &[u8]) -> Option<Vec<u8>> + Send + 'static {
1136        move |api_key, version, _corr, _body| {
1137            if api_key == api_versions_request::API_KEY {
1138                return Some(api_versions_v0());
1139            }
1140            if api_key == metadata_request::API_KEY {
1141                return Some(metadata_v_with_port(
1142                    version,
1143                    node_id,
1144                    port.load(Ordering::SeqCst),
1145                ));
1146            }
1147            None
1148        }
1149    }
1150
1151    /// Bounded poll: wait until `addr` refuses connections.
1152    ///
1153    /// At that point a stopped `MockBroker`'s listener has torn down, together
1154    /// with its per-connection handlers, which share the same cancelled token.
1155    /// This poll replaces a fixed settle after `stop()`. That sleep only
1156    /// waited for the socket teardown, so this helper polls the teardown
1157    /// directly and then runs the unchanged failover assertion. The timeout is
1158    /// a hang-guard.
1159    async fn wait_for_listener_closed(addr: SocketAddr) {
1160        timeout(Duration::from_secs(10), async {
1161            while TcpStream::connect(addr).await.is_ok() {
1162                yield_now().await;
1163            }
1164        })
1165        .await
1166        .expect("broker listener must close after stop()");
1167    }
1168
1169    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1170    async fn refresh_metadata_fails_over_when_bootstrap_broker_dies() {
1171        // Two bootstrap brokers. The client pins its bootstrap connection to the
1172        // first; when that broker is killed, `refresh_metadata` must evict the
1173        // dead bootstrap connection and reconnect to the second live broker
1174        // instead of failing forever on the dead socket — the client-side cause
1175        // of the on-cluster producer stall where a failover whose killed leader
1176        // was also the bootstrap-pinned broker left the producer stranded.
1177        let a = MockBroker::start(handler(0)).await.unwrap();
1178        let b = MockBroker::start(handler(1)).await.unwrap();
1179        let bootstrap = format!("{},{}", a.addr, b.addr);
1180        let client = Client::builder()
1181            .bootstrap(bootstrap)
1182            .request_timeout(millis(500))
1183            .build()
1184            .await
1185            .expect("client connects to bootstrap A");
1186
1187        // First refresh pins the bootstrap connection to broker A.
1188        client
1189            .refresh_metadata()
1190            .await
1191            .expect("first refresh succeeds via A");
1192
1193        // Broker A is killed. Poll until its listener is actually gone so the
1194        // pinned bootstrap connection is torn down before the failover refresh.
1195        let a_addr = a.addr;
1196        a.stop();
1197        wait_for_listener_closed(a_addr).await;
1198
1199        // The next refresh must transparently fail over to the live broker B,
1200        // not hang/error on the dead A socket.
1201        let md = client
1202            .refresh_metadata()
1203            .await
1204            .expect("refresh must fail over to live bootstrap B after A dies");
1205        assert2::assert!(!md.brokers.is_empty());
1206    }
1207
1208    /// `refresh_metadata_with` sends the given request, and the pool learns
1209    /// the brokers of the response as it does for `refresh_metadata`.
1210    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1211    async fn refresh_metadata_with_sends_the_request_and_learns_the_brokers() {
1212        let orders = MetadataRequest {
1213            topics: Some(vec![MetadataRequestTopic {
1214                name: Some("orders".into()),
1215                ..Default::default()
1216            }]),
1217            allow_auto_topic_creation: true,
1218            ..Default::default()
1219        };
1220        for (name, request) in [
1221            ("all topics", crate::all_topics_request()),
1222            ("the topic orders", orders),
1223        ] {
1224            let seen = Arc::new(Mutex::new(Vec::new()));
1225            let handler_seen = Arc::clone(&seen);
1226            let broker = MockBroker::start(move |api_key, version, _corr, body| {
1227                if api_key == api_versions_request::API_KEY {
1228                    return Some(api_versions_v0());
1229                }
1230                if api_key == metadata_request::API_KEY {
1231                    let header_len =
1232                        2 + "krabka".len() + usize::from(version >= metadata_request::FLEXIBLE_MIN);
1233                    let mut request_body = &body[header_len..];
1234                    let decoded = MetadataRequest::decode(&mut request_body, version)
1235                        .expect("decode Metadata");
1236                    handler_seen.lock().expect("seen").push(decoded);
1237                    return Some(metadata_v(version, 7));
1238                }
1239                None
1240            })
1241            .await
1242            .unwrap();
1243            let client = Client::builder()
1244                .bootstrap(broker.addr.to_string())
1245                .request_timeout(millis(500))
1246                .build()
1247                .await
1248                .expect("client connects");
1249
1250            let refreshed = client.refresh_metadata_with(request.clone()).await.is_ok();
1251            let observed = (
1252                refreshed,
1253                seen.lock().expect("seen").clone(),
1254                client.knows_broker(7),
1255            );
1256            broker.stop();
1257            assert2::assert!(observed == (true, vec![request], true), "{name}");
1258        }
1259    }
1260
1261    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1262    async fn refresh_metadata_uses_a_learned_broker_after_the_only_seed_retires() {
1263        let healthy_port = Arc::new(AtomicU16::new(0));
1264        let healthy_calls = Arc::new(AtomicUsize::new(0));
1265        let h_port = Arc::clone(&healthy_port);
1266        let h_calls = Arc::clone(&healthy_calls);
1267        let healthy = MockBroker::start(move |api_key, version, _corr, _body| {
1268            if api_key == api_versions_request::API_KEY {
1269                return Some(api_versions_v0());
1270            }
1271            if api_key == metadata_request::API_KEY {
1272                h_calls.fetch_add(1, Ordering::SeqCst);
1273                return Some(metadata_v_with_port(
1274                    version,
1275                    7,
1276                    h_port.load(Ordering::SeqCst),
1277                ));
1278            }
1279            None
1280        })
1281        .await
1282        .unwrap();
1283        healthy_port.store(healthy.addr.port(), Ordering::SeqCst);
1284
1285        let seed_port = Arc::clone(&healthy_port);
1286        let seed = MockBroker::start(dynamic_port_handler(7, seed_port))
1287            .await
1288            .unwrap();
1289        let client = Client::builder()
1290            .bootstrap(seed.addr.to_string())
1291            .request_timeout(millis(500))
1292            .build()
1293            .await
1294            .expect("client resolves its only seed");
1295        client
1296            .refresh_metadata()
1297            .await
1298            .expect("seed supplies the current broker metadata");
1299
1300        let seed_addr = seed.addr;
1301        seed.stop();
1302        wait_for_listener_closed(seed_addr).await;
1303
1304        let metadata = client
1305            .refresh_metadata()
1306            .await
1307            .expect("learned broker remains usable after seed retirement");
1308        assert2::assert!(metadata.brokers[0].node_id == 7);
1309        assert2::assert!(healthy_calls.load(Ordering::SeqCst) == 1);
1310    }
1311
1312    /// Kafka's `NetworkClient.leastLoadedNode`: after the first metadata
1313    /// response, a request with no fixed target goes to a known broker, not
1314    /// to the bootstrap address.
1315    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1316    async fn send_goes_to_a_known_broker_after_the_first_metadata() {
1317        let counting_handler = |calls: Arc<AtomicUsize>, port: Arc<AtomicU16>| {
1318            move |api_key, version, _corr, _body: &[u8]| {
1319                if api_key == api_versions_request::API_KEY {
1320                    return Some(api_versions_v0());
1321                }
1322                if api_key == metadata_request::API_KEY {
1323                    calls.fetch_add(1, Ordering::SeqCst);
1324                    return Some(metadata_v_with_port(
1325                        version,
1326                        7,
1327                        port.load(Ordering::SeqCst),
1328                    ));
1329                }
1330                None
1331            }
1332        };
1333        let port = Arc::new(AtomicU16::new(0));
1334        let seed_calls = Arc::new(AtomicUsize::new(0));
1335        let broker_calls = Arc::new(AtomicUsize::new(0));
1336        let broker = MockBroker::start(counting_handler(
1337            Arc::clone(&broker_calls),
1338            Arc::clone(&port),
1339        ))
1340        .await
1341        .unwrap();
1342        port.store(broker.addr.port(), Ordering::SeqCst);
1343        let seed = MockBroker::start(counting_handler(Arc::clone(&seed_calls), Arc::clone(&port)))
1344            .await
1345            .unwrap();
1346        let client = Client::builder()
1347            .bootstrap(seed.addr.to_string())
1348            .request_timeout(millis(500))
1349            .build()
1350            .await
1351            .expect("client builds");
1352
1353        client.send(MetadataRequest::default()).await.unwrap();
1354        client.refresh_metadata().await.unwrap();
1355        let before = (
1356            seed_calls.load(Ordering::SeqCst),
1357            broker_calls.load(Ordering::SeqCst),
1358        );
1359        for _ in 0..3 {
1360            client.send(MetadataRequest::default()).await.unwrap();
1361        }
1362        let after = (
1363            seed_calls.load(Ordering::SeqCst),
1364            broker_calls.load(Ordering::SeqCst),
1365        );
1366        seed.stop();
1367        broker.stop();
1368        assert2::assert!(before == (2, 0));
1369        assert2::assert!(after == (2, 3));
1370    }
1371
1372    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1373    async fn reconnect_bootstrap_forces_next_send_to_redial() {
1374        let a = MockBroker::start(handler(0)).await.unwrap();
1375        let b = MockBroker::start(handler(1)).await.unwrap();
1376        let bootstrap = format!("{},{}", a.addr, b.addr);
1377        let client = Client::builder()
1378            .bootstrap(bootstrap)
1379            .request_timeout(millis(500))
1380            .build()
1381            .await
1382            .expect("client builds");
1383
1384        let _ = client
1385            .send(metadata_request::MetadataRequest::default())
1386            .await
1387            .expect("first send succeeds via A");
1388
1389        let a_addr = a.addr;
1390        a.stop();
1391        wait_for_listener_closed(a_addr).await;
1392        client.reconnect_bootstrap().await;
1393
1394        let md = client
1395            .send(metadata_request::MetadataRequest::default())
1396            .await
1397            .expect("send after reconnect_bootstrap reaches B");
1398        assert2::assert!(!md.brokers.is_empty());
1399    }
1400
1401    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1402    async fn broker_send_refreshes_metadata_when_learned_address_is_stale() {
1403        let a_port = Arc::new(AtomicU16::new(0));
1404        let b_port = Arc::new(AtomicU16::new(0));
1405        let a = MockBroker::start(dynamic_port_handler(1, Arc::clone(&a_port)))
1406            .await
1407            .unwrap();
1408        a_port.store(a.addr.port(), Ordering::SeqCst);
1409        let b = MockBroker::start(dynamic_port_handler(1, Arc::clone(&b_port)))
1410            .await
1411            .unwrap();
1412        b_port.store(b.addr.port(), Ordering::SeqCst);
1413
1414        let bootstrap = format!("{},{}", a.addr, b.addr);
1415        let client = Client::builder()
1416            .bootstrap(bootstrap)
1417            .socket_connection_setup_timeout(millis(500))
1418            .request_timeout(millis(500))
1419            .build()
1420            .await
1421            .expect("client builds");
1422
1423        client
1424            .refresh_metadata()
1425            .await
1426            .expect("first metadata refresh learns broker 1 at A");
1427        let a_addr = a.addr;
1428        a.stop();
1429        wait_for_listener_closed(a_addr).await;
1430
1431        let md = client
1432            .broker(1)
1433            .send(metadata_request::MetadataRequest::default())
1434            .await
1435            .expect("broker send should refresh metadata and redial broker 1 at B");
1436        assert2::assert!(!md.brokers.is_empty());
1437    }
1438
1439    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1440    async fn metadata_error_129_rebootstraps_immediately() {
1441        let handshakes = Arc::new(AtomicUsize::new(0));
1442        let metadata_calls = Arc::new(AtomicUsize::new(0));
1443        let broker_port = Arc::new(AtomicU16::new(0));
1444        let h_handshakes = Arc::clone(&handshakes);
1445        let h_metadata_calls = Arc::clone(&metadata_calls);
1446        let h_broker_port = Arc::clone(&broker_port);
1447        let broker = MockBroker::start(move |api_key, version, _corr, _body| {
1448            if api_key == api_versions_request::API_KEY {
1449                h_handshakes.fetch_add(1, Ordering::SeqCst);
1450                return Some(api_versions_v0());
1451            }
1452            if api_key == metadata_request::API_KEY {
1453                let call = h_metadata_calls.fetch_add(1, Ordering::SeqCst);
1454                return Some(if call == 0 {
1455                    empty_metadata_v(version, REBOOTSTRAP_REQUIRED)
1456                } else {
1457                    metadata_v_with_port(version, 7, h_broker_port.load(Ordering::SeqCst))
1458                });
1459            }
1460            None
1461        })
1462        .await
1463        .unwrap();
1464        broker_port.store(broker.addr.port(), Ordering::SeqCst);
1465        let client = Client::builder()
1466            .bootstrap(broker.addr.to_string())
1467            .build()
1468            .await
1469            .expect("client builds");
1470
1471        let metadata = client
1472            .refresh_metadata()
1473            .await
1474            .expect("error 129 triggers one rebootstrap and retry");
1475
1476        assert2::assert!(metadata.brokers[0].node_id == 7);
1477        assert2::assert!(metadata_calls.load(Ordering::SeqCst) == 2);
1478        assert2::assert!(handshakes.load(Ordering::SeqCst) == 2);
1479    }
1480
1481    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1482    async fn stale_metadata_timeout_rebootstraps_after_first_bad_attempt() {
1483        let handshakes = Arc::new(AtomicUsize::new(0));
1484        let metadata_calls = Arc::new(AtomicUsize::new(0));
1485        let broker_port = Arc::new(AtomicU16::new(0));
1486        let h_handshakes = Arc::clone(&handshakes);
1487        let h_metadata_calls = Arc::clone(&metadata_calls);
1488        let h_broker_port = Arc::clone(&broker_port);
1489        let broker = MockBroker::start(move |api_key, version, _corr, _body| {
1490            if api_key == api_versions_request::API_KEY {
1491                h_handshakes.fetch_add(1, Ordering::SeqCst);
1492                return Some(api_versions_v0());
1493            }
1494            if api_key == metadata_request::API_KEY {
1495                let call = h_metadata_calls.fetch_add(1, Ordering::SeqCst);
1496                return Some(if call < 2 {
1497                    empty_metadata_v(version, 0)
1498                } else {
1499                    metadata_v_with_port(version, 8, h_broker_port.load(Ordering::SeqCst))
1500                });
1501            }
1502            None
1503        })
1504        .await
1505        .unwrap();
1506        broker_port.store(broker.addr.port(), Ordering::SeqCst);
1507        let client = Client::builder()
1508            .bootstrap(broker.addr.to_string())
1509            .metadata_recovery_rebootstrap_trigger(millis(20))
1510            .build()
1511            .await
1512            .expect("client builds");
1513
1514        let first = client
1515            .refresh_metadata()
1516            .await
1517            .expect("first empty metadata response remains observable");
1518        assert2::assert!(first.brokers.is_empty());
1519        sleep(Duration::from_millis(25)).await;
1520        let recovered = client
1521            .refresh_metadata()
1522            .await
1523            .expect("expired metadata attempt triggers rebootstrap");
1524
1525        assert2::assert!(recovered.brokers[0].node_id == 8);
1526        assert2::assert!(metadata_calls.load(Ordering::SeqCst) == 3);
1527        assert2::assert!(handshakes.load(Ordering::SeqCst) == 2);
1528    }
1529
1530    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1531    async fn metadata_recovery_none_surfaces_error_129_without_retry() {
1532        let metadata_calls = Arc::new(AtomicUsize::new(0));
1533        let h_metadata_calls = Arc::clone(&metadata_calls);
1534        let broker = MockBroker::start(move |api_key, version, _corr, _body| {
1535            if api_key == api_versions_request::API_KEY {
1536                return Some(api_versions_v0());
1537            }
1538            if api_key == metadata_request::API_KEY {
1539                h_metadata_calls.fetch_add(1, Ordering::SeqCst);
1540                return Some(empty_metadata_v(version, REBOOTSTRAP_REQUIRED));
1541            }
1542            None
1543        })
1544        .await
1545        .unwrap();
1546        let client = Client::builder()
1547            .bootstrap(broker.addr.to_string())
1548            .metadata_recovery_strategy(MetadataRecoveryStrategy::None)
1549            .build()
1550            .await
1551            .expect("client builds");
1552
1553        let error = client
1554            .refresh_metadata()
1555            .await
1556            .expect_err("disabled recovery surfaces the proxy request");
1557
1558        assert2::assert!(matches!(
1559            error,
1560            ClientError::Server {
1561                error_code: REBOOTSTRAP_REQUIRED
1562            }
1563        ));
1564        assert2::assert!(metadata_calls.load(Ordering::SeqCst) == 1);
1565    }
1566}
1567
1568#[cfg(test)]
1569mod authentication_failure_tests {
1570    use std::sync::{
1571        Arc,
1572        atomic::{AtomicU16, AtomicUsize, Ordering},
1573    };
1574
1575    use bytes::BytesMut;
1576    use krabka_protocol::{
1577        Encode,
1578        owned::{
1579            api_versions_request,
1580            api_versions_response::{ApiVersion, ApiVersionsResponse},
1581            metadata_request::{self, MetadataRequest},
1582            metadata_response::{MetadataResponse, MetadataResponseBroker},
1583            sasl_handshake_request,
1584        },
1585    };
1586    use krabka_security::ListenerProtocol;
1587
1588    use super::*;
1589    use crate::{
1590        OutboundSaslError, SaslAuthenticationError,
1591        error::AuthenticationError,
1592        mock::{MockBroker, MockReply, MockSaslAnswer},
1593        security::{ClientSecurity, SaslCredentials},
1594    };
1595
1596    const UNSUPPORTED_SASL_MECHANISM: i16 = 33;
1597    const ILLEGAL_SASL_STATE: i16 = 34;
1598    const UNSUPPORTED_VERSION: i16 = 35;
1599    const SASL_AUTHENTICATION_FAILED: i16 = 58;
1600    const BROKER_ID: i32 = 1;
1601
1602    /// The request that the test sends after the SASL listener starts.
1603    #[derive(Clone, Copy, Debug)]
1604    enum Route {
1605        /// `Client::refresh_metadata` over the bootstrap connection.
1606        Bootstrap,
1607        /// `BrokerHandle::send` to a broker that a first metadata refresh
1608        /// registered.
1609        Broker,
1610    }
1611
1612    /// How a call ended, as a caller classifies it.
1613    #[derive(Debug, PartialEq, Eq)]
1614    enum Outcome {
1615        Ok,
1616        Rejected(SaslAuthenticationError),
1617        SaslServerError(i16),
1618        Other(String),
1619    }
1620
1621    #[derive(Debug, PartialEq, Eq)]
1622    struct Observed {
1623        outcome: Outcome,
1624        handshakes: usize,
1625    }
1626
1627    fn classify<T>(result: &Result<T, ClientError>) -> Outcome {
1628        match result {
1629            Ok(_) => Outcome::Ok,
1630            Err(ClientError::Authentication {
1631                source: AuthenticationError::Sasl(error),
1632                ..
1633            }) => Outcome::Rejected(error.clone()),
1634            Err(ClientError::Sasl {
1635                source: OutboundSaslError::Server { error_code, .. },
1636                ..
1637            }) => Outcome::SaslServerError(*error_code),
1638            Err(error) => Outcome::Other(error.to_string()),
1639        }
1640    }
1641
1642    fn encode(response: &impl Encode, version: i16) -> Vec<u8> {
1643        let mut body = BytesMut::new();
1644        response.encode(&mut body, version).unwrap();
1645        body.to_vec()
1646    }
1647
1648    /// A SASL listener that answers the exchange of connection `n` with
1649    /// `script[n]`, or with the last entry once the script runs out. It also
1650    /// answers `ApiVersions` and a `Metadata` that names itself as broker 1.
1651    async fn sasl_broker(script: Vec<MockSaslAnswer>, handshakes: Arc<AtomicUsize>) -> MockBroker {
1652        let port = Arc::new(AtomicU16::new(0));
1653        let handler_port = Arc::clone(&port);
1654        let mut answer = *script.last().unwrap();
1655        let broker = MockBroker::start_with_replies(move |api_key, version, _corr, _body| {
1656            if api_key == sasl_handshake_request::API_KEY {
1657                let connection = handshakes.fetch_add(1, Ordering::SeqCst);
1658                answer = script
1659                    .get(connection)
1660                    .copied()
1661                    .unwrap_or_else(|| *script.last().unwrap());
1662            }
1663            if let Some(reply) = answer.reply(api_key, version).unwrap() {
1664                return reply;
1665            }
1666            match api_key {
1667                api_versions_request::API_KEY => MockReply::Respond(encode(
1668                    &ApiVersionsResponse {
1669                        api_keys: vec![
1670                            ApiVersion {
1671                                api_key: api_versions_request::API_KEY,
1672                                min_version: 0,
1673                                max_version: 3,
1674                                ..Default::default()
1675                            },
1676                            ApiVersion {
1677                                api_key: metadata_request::API_KEY,
1678                                min_version: 0,
1679                                max_version: 8,
1680                                ..Default::default()
1681                            },
1682                        ],
1683                        ..Default::default()
1684                    },
1685                    0,
1686                )),
1687                metadata_request::API_KEY => MockReply::Respond(encode(
1688                    &MetadataResponse {
1689                        brokers: vec![MetadataResponseBroker {
1690                            node_id: BROKER_ID,
1691                            host: "127.0.0.1".into(),
1692                            port: i32::from(handler_port.load(Ordering::SeqCst)),
1693                            ..Default::default()
1694                        }],
1695                        ..Default::default()
1696                    },
1697                    version,
1698                )),
1699                _ => MockReply::Silent,
1700            }
1701        })
1702        .await
1703        .unwrap();
1704        port.store(broker.addr.port(), Ordering::SeqCst);
1705        broker
1706    }
1707
1708    fn plain_security() -> ClientSecurity {
1709        ClientSecurity {
1710            protocol: ListenerProtocol::SaslPlaintext,
1711            tls: None,
1712            sasl: Some(SaslCredentials::Plain {
1713                username: "alice".into(),
1714                password: "secret".into(),
1715            }),
1716            sasl_host: None,
1717        }
1718    }
1719
1720    fn authenticate_failed() -> SaslAuthenticationError {
1721        SaslAuthenticationError::Failed(
1722            "SaslAuthenticate(PLAIN) error_code=58 error_message=Some(\"rejected by mock broker\")"
1723                .into(),
1724        )
1725    }
1726
1727    fn unsupported_mechanism() -> SaslAuthenticationError {
1728        SaslAuthenticationError::UnsupportedMechanism(
1729            "client SASL mechanism 'PLAIN' not enabled in the server, enabled mechanisms are \
1730             [\"PLAIN\"]"
1731                .into(),
1732        )
1733    }
1734
1735    fn illegal_state() -> SaslAuthenticationError {
1736        SaslAuthenticationError::IllegalState(
1737            "unexpected handshake request with client mechanism PLAIN, enabled mechanisms are \
1738             [\"PLAIN\"]"
1739                .into(),
1740        )
1741    }
1742
1743    /// Kafka's `NetworkClient.processDisconnection` stores an
1744    /// `AuthenticationException` for the node, and the next call raises it
1745    /// with no retry. A disconnect during authentication with no broker error
1746    /// stays retriable. The client must not evict and redial on a rejection.
1747    #[tokio::test]
1748    async fn sasl_rejection_is_raised_without_a_retry() {
1749        use MockSaslAnswer::{Accept, AuthenticateError, CloseAfterHandshake, HandshakeError};
1750
1751        let cases = [
1752            (
1753                "bootstrap: SaslAuthenticate 58",
1754                Route::Bootstrap,
1755                vec![AuthenticateError(SASL_AUTHENTICATION_FAILED)],
1756                Observed {
1757                    outcome: Outcome::Rejected(authenticate_failed()),
1758                    handshakes: 1,
1759                },
1760            ),
1761            (
1762                "bootstrap: SaslHandshake 33",
1763                Route::Bootstrap,
1764                vec![HandshakeError(UNSUPPORTED_SASL_MECHANISM)],
1765                Observed {
1766                    outcome: Outcome::Rejected(unsupported_mechanism()),
1767                    handshakes: 1,
1768                },
1769            ),
1770            (
1771                "bootstrap: SaslHandshake 34",
1772                Route::Bootstrap,
1773                vec![HandshakeError(ILLEGAL_SASL_STATE)],
1774                Observed {
1775                    outcome: Outcome::Rejected(illegal_state()),
1776                    handshakes: 1,
1777                },
1778            ),
1779            (
1780                "bootstrap: SaslAuthenticate code that is not an authentication error",
1781                Route::Bootstrap,
1782                vec![AuthenticateError(UNSUPPORTED_VERSION)],
1783                Observed {
1784                    outcome: Outcome::SaslServerError(UNSUPPORTED_VERSION),
1785                    handshakes: 2,
1786                },
1787            ),
1788            (
1789                "bootstrap: close after SaslHandshake, then accept",
1790                Route::Bootstrap,
1791                vec![CloseAfterHandshake, Accept],
1792                Observed {
1793                    outcome: Outcome::Ok,
1794                    handshakes: 2,
1795                },
1796            ),
1797            (
1798                "broker: SaslAuthenticate 58",
1799                Route::Broker,
1800                vec![Accept, AuthenticateError(SASL_AUTHENTICATION_FAILED)],
1801                Observed {
1802                    outcome: Outcome::Rejected(authenticate_failed()),
1803                    handshakes: 2,
1804                },
1805            ),
1806            (
1807                "broker: SaslHandshake 33",
1808                Route::Broker,
1809                vec![Accept, HandshakeError(UNSUPPORTED_SASL_MECHANISM)],
1810                Observed {
1811                    outcome: Outcome::Rejected(unsupported_mechanism()),
1812                    handshakes: 2,
1813                },
1814            ),
1815            (
1816                "broker: close after SaslHandshake, then accept",
1817                Route::Broker,
1818                vec![Accept, CloseAfterHandshake, Accept],
1819                Observed {
1820                    outcome: Outcome::Ok,
1821                    handshakes: 3,
1822                },
1823            ),
1824        ];
1825
1826        for (name, route, script, expected) in cases {
1827            let handshakes = Arc::new(AtomicUsize::new(0));
1828            let broker = sasl_broker(script, Arc::clone(&handshakes)).await;
1829            let client = Client::builder()
1830                .bootstrap(broker.addr.to_string())
1831                .security(plain_security())
1832                .build()
1833                .await
1834                .unwrap();
1835
1836            let outcome = match route {
1837                Route::Bootstrap => classify(&client.refresh_metadata().await),
1838                Route::Broker => {
1839                    client.refresh_metadata().await.unwrap();
1840                    classify(
1841                        &client
1842                            .broker(BROKER_ID)
1843                            .send(MetadataRequest::default())
1844                            .await,
1845                    )
1846                }
1847            };
1848            let observed = Observed {
1849                outcome,
1850                handshakes: handshakes.load(Ordering::SeqCst),
1851            };
1852
1853            broker.stop();
1854            assert2::check!(observed == expected, "case {name}");
1855        }
1856    }
1857}
krabka_client_core/client.rs

krabka_client_core/
client.rs
