Skip to main content

krabka_client_consumer/share/
consumer.rs

1//! `ShareConsumer`, the public lifecycle handle for a KIP-932 share group.
2//!
3//! Build it with [`ShareConsumer::builder`]. The constructor joins the share
4//! group with one `ShareGroupHeartbeat` that carries a client-generated member
5//! UUID, epoch 0, and the subscription. It captures the returned member id, epoch,
6//! heartbeat interval, and assignment. It resolves the assignment's topic ids
7//! to names with Metadata, then spawns the background heartbeat loop.
8
9use std::{collections::HashMap, sync::Arc};
10
11use krabka_client_core::{Client, security, telemetry};
12use krabka_protocol::{
13    owned::share_group_heartbeat_request::ShareGroupHeartbeatRequest,
14    primitives::uuid::Uuid as WireUuid,
15};
16use krabka_units::{
17    ByteSize, Time, bytes,
18    convert::{ByteSizeExt as _, TimeExt as _},
19    mebibytes, secs,
20};
21use refined_type::rule::GreaterI64;
22use tokio::{sync::Mutex, task::JoinHandle};
23use tokio_util::sync::CancellationToken;
24
25use super::{
26    coordinator::ShareCoordinatorState,
27    types::{ShareAckMode, ShareAcquireMode},
28};
29use crate::{
30    consumer::{client_instance_id, protocol_bytes, validated_positive_i32},
31    coordinator,
32    error::ConsumerError,
33};
34
35/// Default deadline for the final best-effort `ShareGroup` leave heartbeat.
36pub const DEFAULT_SHARE_CONSUMER_LEAVE_HEARTBEAT_TIMEOUT: Time = secs(5);
37
38/// Default minimum response size for a `ShareFetch`.
39pub const DEFAULT_SHARE_CONSUMER_FETCH_MIN: ByteSize = bytes(1);
40/// Default maximum response size for a `ShareFetch`.
41pub const DEFAULT_SHARE_CONSUMER_FETCH_MAX: ByteSize = mebibytes(50);
42/// Default maximum records acquired by a `ShareFetch`.
43pub const DEFAULT_SHARE_CONSUMER_FETCH_MAX_RECORDS: i32 = 500;
44
45/// Validated minimum response bytes for a `ShareFetch`.
46#[derive(Clone, Copy, Debug, Eq, PartialEq, derive_more::Into)]
47pub struct ShareConsumerFetchMinBytes(i32);
48
49impl ShareConsumerFetchMinBytes {
50    /// Validate a positive minimum byte count.
51    ///
52    /// # Errors
53    ///
54    /// Returns an error when `value` is zero or negative.
55    pub fn new(value: i32) -> Result<Self, String> {
56        validated_positive_i32(value, "share consumer fetch min bytes").map(Self)
57    }
58
59    /// Return the validated byte count.
60    #[must_use]
61    pub const fn bytes(self) -> i32 {
62        self.0
63    }
64}
65
66impl TryFrom<ByteSize> for ShareConsumerFetchMinBytes {
67    type Error = String;
68
69    fn try_from(value: ByteSize) -> Result<Self, Self::Error> {
70        Self::new(protocol_bytes("share consumer fetch min", value)?)
71    }
72}
73
74impl Default for ShareConsumerFetchMinBytes {
75    fn default() -> Self {
76        Self::new(DEFAULT_SHARE_CONSUMER_FETCH_MIN.bytes_i32())
77            .expect("default share consumer fetch min bytes is valid")
78    }
79}
80
81/// Validated maximum response bytes for a `ShareFetch`.
82#[derive(Clone, Copy, Debug, Eq, PartialEq, derive_more::Into)]
83pub struct ShareConsumerFetchMaxBytes(i32);
84
85impl ShareConsumerFetchMaxBytes {
86    /// Validate a positive maximum byte count.
87    ///
88    /// # Errors
89    ///
90    /// Returns an error when `value` is zero or negative.
91    pub fn new(value: i32) -> Result<Self, String> {
92        validated_positive_i32(value, "share consumer fetch max bytes").map(Self)
93    }
94
95    /// Return the validated byte count.
96    #[must_use]
97    pub const fn bytes(self) -> i32 {
98        self.0
99    }
100}
101
102impl TryFrom<ByteSize> for ShareConsumerFetchMaxBytes {
103    type Error = String;
104
105    fn try_from(value: ByteSize) -> Result<Self, Self::Error> {
106        Self::new(protocol_bytes("share consumer fetch max", value)?)
107    }
108}
109
110impl Default for ShareConsumerFetchMaxBytes {
111    fn default() -> Self {
112        Self::new(DEFAULT_SHARE_CONSUMER_FETCH_MAX.bytes_i32())
113            .expect("default share consumer fetch max bytes is valid")
114    }
115}
116
117/// Validated maximum records acquired by a `ShareFetch`.
118#[derive(Clone, Copy, Debug, Eq, PartialEq, derive_more::Into)]
119pub struct ShareConsumerFetchMaxRecords(i32);
120
121impl ShareConsumerFetchMaxRecords {
122    /// Validate a positive maximum record count.
123    ///
124    /// # Errors
125    ///
126    /// Returns an error when `value` is zero or negative.
127    pub fn new(value: i32) -> Result<Self, String> {
128        validated_positive_i32(value, "share consumer fetch max records").map(Self)
129    }
130
131    /// Return the validated record count.
132    #[must_use]
133    pub const fn records(self) -> i32 {
134        self.0
135    }
136}
137
138impl Default for ShareConsumerFetchMaxRecords {
139    fn default() -> Self {
140        Self::new(DEFAULT_SHARE_CONSUMER_FETCH_MAX_RECORDS)
141            .expect("default share consumer fetch max records is valid")
142    }
143}
144
145/// Validated deadline for the final best-effort `ShareGroup` leave heartbeat.
146#[derive(Clone, Copy, Debug, PartialEq, derive_more::Into)]
147pub struct ShareConsumerLeaveHeartbeatTimeout(Time);
148
149impl ShareConsumerLeaveHeartbeatTimeout {
150    /// Validate a positive, whole-millisecond timeout.
151    ///
152    /// # Errors
153    ///
154    /// Returns an error for nonfinite, zero, negative, or fractional values.
155    pub fn new(value: Time) -> Result<Self, String> {
156        validated_time("share consumer leave-heartbeat timeout", value).map(Self)
157    }
158
159    /// Return the validated timeout.
160    #[must_use]
161    pub const fn time(self) -> Time {
162        self.0
163    }
164}
165
166impl Default for ShareConsumerLeaveHeartbeatTimeout {
167    fn default() -> Self {
168        Self::new(DEFAULT_SHARE_CONSUMER_LEAVE_HEARTBEAT_TIMEOUT)
169            .expect("default share consumer leave-heartbeat timeout is valid")
170    }
171}
172
173#[derive(bon::Builder)]
174#[builder(start_fn = builder, finish_fn = build)]
175struct ShareClientStart {
176    #[builder(into)]
177    bootstrap: String,
178    #[builder(into)]
179    client_id: String,
180    dispatch_queue_capacity: usize,
181    frame_max: ByteSize,
182    metadata_recovery_strategy: krabka_client_core::MetadataRecoveryStrategy,
183    metadata_recovery_rebootstrap_trigger: Time,
184    security: Option<security::ClientSecurity>,
185    allow_auto_create_topics: bool,
186}
187
188#[derive(bon::Builder, Default)]
189struct ShareClientOptions {
190    telemetry: Option<telemetry::ClientTelemetryConfig>,
191}
192
193async fn build_share_client(
194    start: &ShareClientStart,
195    options: ShareClientOptions,
196) -> Result<Client, ConsumerError> {
197    Ok(Client::builder()
198        .bootstrap(&start.bootstrap)
199        .client_id(start.client_id.clone())
200        .dispatch_queue_capacity(start.dispatch_queue_capacity)
201        .frame_max(start.frame_max)
202        .metadata_recovery_strategy(start.metadata_recovery_strategy)
203        .metadata_recovery_rebootstrap_trigger(start.metadata_recovery_rebootstrap_trigger)
204        .maybe_security(start.security.clone())
205        .metadata_scope(krabka_client_core::MetadataScope::Topics {
206            allow_auto_topic_creation: start.allow_auto_create_topics,
207        })
208        .maybe_telemetry(options.telemetry)
209        .build()
210        .await?)
211}
212
213fn build_join_heartbeat_request(
214    group_id: String,
215    subscribe: Vec<String>,
216) -> ShareGroupHeartbeatRequest {
217    ShareGroupHeartbeatRequest {
218        group_id,
219        member_id: uuid::Uuid::new_v4().to_string(),
220        subscribed_topic_names: Some(subscribe),
221        ..Default::default()
222    }
223}
224
225fn response_has_error(error_code: i16) -> bool {
226    error_code != 0
227}
228
229fn validated_time(name: &str, value: Time) -> Result<Time, String> {
230    let milliseconds = GreaterI64::<0>::new(value.millis_i64())
231        .map_err(|error| format!("{name}: {error}"))?
232        .into_value();
233    if value.secs_f64().is_finite() && Time::from_millis(milliseconds) == value {
234        Ok(value)
235    } else {
236        Err(format!("{name} must be a whole number of milliseconds"))
237    }
238}
239
240fn heartbeat_interval_from_response(heartbeat_interval_ms: i32, configured: Time) -> Time {
241    if heartbeat_interval_ms > 0 {
242        Time::from_millis(i64::from(heartbeat_interval_ms))
243    } else {
244        configured
245    }
246}
247
248fn has_assignment_partitions(partitions_len: usize) -> bool {
249    partitions_len > 0
250}
251
252fn should_stage_implicit_accepts(ack_mode: ShareAckMode) -> bool {
253    ack_mode == ShareAckMode::Implicit
254}
255
256fn stage_implicit_accepts(
257    prev_delivered: &mut Vec<(WireUuid, i32, i64, i64)>,
258    pending_acks: &mut Vec<(WireUuid, i32, i64, i64, i8)>,
259) {
260    for (tid, partition, first, last) in std::mem::take(prev_delivered) {
261        pending_acks.push((
262            tid,
263            partition,
264            first,
265            last,
266            super::types::ShareAckType::Accept.wire(),
267        ));
268    }
269}
270
271/// A share-group consumer. Construct via [`ShareConsumer::builder`].
272///
273/// It joins the group and keeps the membership alive with a background
274/// heartbeat. [`poll`](ShareConsumer::poll) issues `ShareFetch` over the live
275/// assignment and returns the acquired records. Acknowledgement, implicit or
276/// explicit per [`ShareAckMode`], travels back to the broker on the next
277/// `ShareFetch` as a piggyback, or in a standalone `ShareAcknowledge` from
278/// [`commit`](ShareConsumer::commit).
279pub struct ShareConsumer {
280    pub(crate) client: Client,
281    pub(crate) group_id: String,
282    pub(crate) member_id: String,
283    /// The live member epoch. The background heartbeat loop owns and advances
284    /// it, and holds the other `Arc`. The consumer keeps this clone so the
285    /// shared cell outlives the heartbeat task. `poll()` does not read it,
286    /// because the data path keys off the share-session epoch and not the
287    /// member epoch.
288    #[allow(dead_code)]
289    pub(crate) member_epoch: Arc<Mutex<i32>>,
290    /// Live assignment as `(topic_id, topic_name, partition)`. The heartbeat
291    /// loop updates it.
292    pub(crate) assignment: Arc<Mutex<Vec<(WireUuid, String, i32)>>>,
293    pub(crate) topic_names: Arc<Mutex<HashMap<WireUuid, String>>>,
294    /// `ShareFetch` session epoch: 0 opens the session, then 1, 2, … per
295    /// successful fetch. `poll()` owns it.
296    pub(crate) share_session_epoch: i32,
297    pub(crate) fetch_min: ByteSize,
298    pub(crate) fetch_max: ByteSize,
299    pub(crate) fetch_max_records: i32,
300    pub(crate) acquire_mode: ShareAcquireMode,
301    pub(crate) ack_mode: ShareAckMode,
302    /// Acks staged for the next `ShareFetch` / `ShareAcknowledge` as
303    /// `(topic_id, partition, first_offset, last_offset, ack_type_wire)`.
304    pub(crate) pending_acks: Vec<(WireUuid, i32, i64, i64, i8)>,
305    /// Ranges delivered by the previous `poll()` as
306    /// `(topic_id, partition, first_offset, last_offset)`, for implicit-accept.
307    pub(crate) prev_delivered: Vec<(WireUuid, i32, i64, i64)>,
308    pub(crate) shutdown: CancellationToken,
309    pub(crate) hb_handle: Option<JoinHandle<()>>,
310    /// A fatal heartbeat error that the next `poll()` returns. The heartbeat
311    /// loop holds the other `Arc`.
312    pub(crate) poll_error: coordinator::PollErrorSlot,
313}
314
315#[bon::bon]
316impl ShareConsumer {
317    /// Join a share group and start heartbeating.
318    ///
319    /// This method sends one `ShareGroupHeartbeat` that carries a client-generated
320    /// member UUID, epoch 0, and `subscribe`. It captures the returned member id, epoch,
321    /// heartbeat interval, and assignment. It resolves the assignment topic ids
322    /// → names with Metadata, then spawns the heartbeat loop.
323    #[builder(start_fn = builder, finish_fn = build)]
324    #[tracing::instrument(
325        name = "share_consumer.start",
326        level = "info",
327        skip_all,
328        fields(
329            group_id = %group_id,
330            member_id = tracing::field::Empty,
331            member_epoch = tracing::field::Empty,
332            assigned_partitions = tracing::field::Empty,
333        ),
334        err(level = "info")
335    )]
336    /// # Errors
337    /// Returns an error when configuration is invalid, protocol encoding fails, the broker rejects the request, or transport I/O fails.
338    pub async fn start(
339        #[builder(into)] bootstrap: String,
340        #[builder(into, default = "krabka-share-consumer".to_string())] client_id: String,
341        #[builder(into)] group_id: String,
342        #[builder(into)] subscribe: Vec<String>,
343        #[builder(default = ShareAckMode::Implicit)] ack_mode: ShareAckMode,
344        #[builder(default = ShareAcquireMode::BatchOptimized)] acquire_mode: ShareAcquireMode,
345        /// The minimum `ShareFetch` response bytes, as a raw protocol count.
346        #[builder(default = DEFAULT_SHARE_CONSUMER_FETCH_MIN.bytes_i32())]
347        fetch_min: i32,
348        /// The maximum `ShareFetch` response bytes, as a raw protocol count.
349        #[builder(default = DEFAULT_SHARE_CONSUMER_FETCH_MAX.bytes_i32())]
350        fetch_max: i32,
351        /// The maximum records acquired by one `ShareFetch`.
352        #[builder(default = DEFAULT_SHARE_CONSUMER_FETCH_MAX_RECORDS)]
353        fetch_max_records: i32,
354        #[builder(default = secs(3))] heartbeat_interval: Time,
355        #[builder(default = DEFAULT_SHARE_CONSUMER_LEAVE_HEARTBEAT_TIMEOUT)]
356        leave_heartbeat_timeout: Time,
357        #[builder(default = krabka_client_core::DEFAULT_CONNECTION_DISPATCH_QUEUE_CAPACITY)]
358        dispatch_queue_capacity: usize,
359        #[builder(default = krabka_client_core::DEFAULT_CLIENT_FRAME_MAX)] frame_max: ByteSize,
360        #[builder(default)]
361        metadata_recovery_strategy: krabka_client_core::MetadataRecoveryStrategy,
362        #[builder(default = krabka_client_core::DEFAULT_METADATA_RECOVERY_REBOOTSTRAP_TRIGGER)]
363        metadata_recovery_rebootstrap_trigger: Time,
364        security: Option<security::ClientSecurity>,
365        /// Kafka's `allow.auto.create.topics`, which
366        /// `ShareConsumerMetadata` reads.
367        #[builder(default = true)]
368        allow_auto_create_topics: bool,
369        /// Kafka's `enable.metrics.push` (default `true`): push the client
370        /// metrics that a broker's client metrics subscription names
371        /// (KIP-714), with a terminating push on close. When it is `false`,
372        /// the consumer sends neither `GetTelemetrySubscriptions` nor
373        /// `PushTelemetry`.
374        #[builder(default = true)]
375        enable_metrics_push: bool,
376    ) -> Result<Self, ConsumerError> {
377        if subscribe.is_empty() {
378            return Err(ConsumerError::NotSubscribed);
379        }
380        if group_id.is_empty() {
381            return Err(ConsumerError::RebalanceFailed("group_id required".into()));
382        }
383        // The newtypes validate the raw protocol values before turning them
384        // into quantities for the fetch implementation.
385        let fetch_min = ByteSize::from_bytes_i64(i64::from(
386            ShareConsumerFetchMinBytes::new(fetch_min)
387                .map_err(ConsumerError::RebalanceFailed)?
388                .bytes(),
389        ));
390        let fetch_max = ByteSize::from_bytes_i64(i64::from(
391            ShareConsumerFetchMaxBytes::new(fetch_max)
392                .map_err(ConsumerError::RebalanceFailed)?
393                .bytes(),
394        ));
395        let fetch_max_records = ShareConsumerFetchMaxRecords::new(fetch_max_records)
396            .map_err(ConsumerError::RebalanceFailed)?
397            .records();
398        if fetch_min > fetch_max {
399            return Err(ConsumerError::RebalanceFailed(
400                "share consumer fetch min bytes must not exceed fetch max bytes".to_owned(),
401            ));
402        }
403        let leave_heartbeat_timeout =
404            ShareConsumerLeaveHeartbeatTimeout::new(leave_heartbeat_timeout)
405                .map_err(ConsumerError::RebalanceFailed)?
406                .time();
407        let heartbeat_interval =
408            validated_time("share consumer heartbeat interval", heartbeat_interval)
409                .map_err(ConsumerError::RebalanceFailed)?;
410        let dispatch_queue_capacity =
411            krabka_client_core::ConnectionDispatchQueueCapacity::new(dispatch_queue_capacity)
412                .map_err(ConsumerError::RebalanceFailed)?;
413        let frame_max = krabka_client_core::ClientFrameMax::try_from(frame_max)
414            .map_err(ConsumerError::RebalanceFailed)?;
415        let metadata_recovery_rebootstrap_trigger =
416            krabka_client_core::MetadataRecoveryRebootstrapTrigger::new(
417                metadata_recovery_rebootstrap_trigger,
418            )
419            .map_err(ConsumerError::RebalanceFailed)?;
420
421        let client_start = ShareClientStart::builder()
422            .bootstrap(bootstrap.clone())
423            .client_id(client_id.clone())
424            .dispatch_queue_capacity(dispatch_queue_capacity.get())
425            .frame_max(frame_max.size())
426            .metadata_recovery_strategy(metadata_recovery_strategy)
427            .metadata_recovery_rebootstrap_trigger(metadata_recovery_rebootstrap_trigger.time())
428            .maybe_security(security.clone())
429            .allow_auto_create_topics(allow_auto_create_topics)
430            .build();
431        let telemetry = enable_metrics_push.then(|| {
432            // `ClientTelemetryProvider` labels a consumer's metrics with
433            // its `group.id`. A share consumer has no
434            // `group.instance.id` and no `client.rack`.
435            telemetry::ClientTelemetryConfig::consumer(Some(&group_id), None, None)
436        });
437        let client = build_share_client(
438            &client_start,
439            ShareClientOptions::builder()
440                .maybe_telemetry(telemetry)
441                .build(),
442        )
443        .await?;
444        client.metadata_topics().set(subscribe.iter().cloned());
445
446        // 1. Join: client-generated member UUID + epoch 0 + the subscription.
447        //    The broker returns the member id and bumps us to a live epoch.
448        let join = client
449            .send(build_join_heartbeat_request(
450                group_id.clone(),
451                subscribe.clone(),
452            ))
453            .await?;
454        if response_has_error(join.error_code) {
455            return Err(ConsumerError::Server(join.error_code));
456        }
457        let member_id = join.member_id.clone().unwrap_or_default();
458        if member_id.is_empty() {
459            return Err(ConsumerError::RebalanceFailed(
460                "broker did not return a member_id".into(),
461            ));
462        }
463        let member_epoch_val = join.member_epoch;
464        {
465            let span = tracing::Span::current();
466            span.record("member_id", member_id.as_str());
467            span.record("member_epoch", member_epoch_val);
468        }
469        // Honor the broker's heartbeat interval when it supplies one; else keep
470        // the configured default.
471        let hb_interval =
472            heartbeat_interval_from_response(join.heartbeat_interval_ms, heartbeat_interval);
473
474        // 2. Resolve assignment topic ids → names via Metadata.
475        let md = client.refresh_metadata().await?;
476        let mut topic_names: HashMap<WireUuid, String> = HashMap::new();
477        for t in &md.topics {
478            if let Some(name) = &t.name {
479                topic_names.insert(t.topic_id, name.clone());
480            }
481        }
482
483        // 3. Decode the initial assignment (if the broker placed us already).
484        let mut assignment_vec: Vec<(WireUuid, String, i32)> = Vec::new();
485        if let Some(assignment) = join.assignment {
486            for tp in &assignment.topic_partitions {
487                let name = topic_names.get(&tp.topic_id).cloned().unwrap_or_default();
488                if has_assignment_partitions(tp.partitions.len()) {
489                    for &partition in &tp.partitions {
490                        assignment_vec.push((tp.topic_id, name.clone(), partition));
491                    }
492                }
493            }
494        }
495
496        tracing::Span::current().record("assigned_partitions", assignment_vec.len());
497        let member_epoch = Arc::new(Mutex::new(member_epoch_val));
498        let assignment = Arc::new(Mutex::new(assignment_vec));
499        let topic_names = Arc::new(Mutex::new(topic_names));
500        let shutdown = CancellationToken::new();
501        let poll_error = coordinator::PollErrorSlot::default();
502
503        // 4. Spawn the heartbeat loop on its own connection so a parked
504        //    request on the data path can't head-of-line-block heartbeats
505        //    (the broker serves a connection's requests serially).
506        let coordinator_client =
507            build_share_client(&client_start, ShareClientOptions::default()).await?;
508        coordinator_client
509            .metadata_topics()
510            .set(subscribe.iter().cloned());
511        let state = ShareCoordinatorState {
512            client: coordinator_client,
513            group_id: group_id.clone(),
514            member_id: member_id.clone(),
515            member_epoch: Arc::clone(&member_epoch),
516            assignment: Arc::clone(&assignment),
517            topic_names: Arc::clone(&topic_names),
518            subscribe,
519            heartbeat_interval: hb_interval,
520            leave_heartbeat_timeout,
521            poll_error: Arc::clone(&poll_error),
522        };
523        let hb_handle = tokio::spawn(super::coordinator::run(state, shutdown.clone()));
524
525        Ok(ShareConsumer {
526            client,
527            group_id,
528            member_id,
529            member_epoch,
530            assignment,
531            topic_names,
532            share_session_epoch: 0,
533            fetch_min,
534            fetch_max,
535            fetch_max_records,
536            acquire_mode,
537            ack_mode,
538            pending_acks: Vec::new(),
539            prev_delivered: Vec::new(),
540            shutdown,
541            hb_handle: Some(hb_handle),
542            poll_error,
543        })
544    }
545}
546
547impl ShareConsumer {
548    /// The share group id.
549    #[must_use]
550    pub fn group_id(&self) -> &str {
551        &self.group_id
552    }
553
554    /// The client-generated member id returned by the broker at join time.
555    #[must_use]
556    pub fn member_id(&self) -> &str {
557        &self.member_id
558    }
559
560    /// Snapshot of the currently assigned `(topic, partition)` pairs.
561    pub async fn assignment(&self) -> Vec<(String, i32)> {
562        self.assignment
563            .lock()
564            .await
565            .iter()
566            .map(|(_, name, p)| (name.clone(), *p))
567            .collect()
568    }
569
570    /// Stop heartbeating, acknowledge outstanding records, and leave the group.
571    ///
572    /// This method first flushes any outstanding acknowledgements in a
573    /// standalone `ShareAcknowledge`. In Implicit mode it auto-`Accept`s the
574    /// previous poll's delivered ranges. In Explicit mode it sends any staged
575    /// `acknowledge()` calls. It then cancels the heartbeat task and awaits it.
576    /// The task sends a best-effort leave heartbeat with `member_epoch = -1` on
577    /// its way out, so the broker evicts this member promptly and does not wait
578    /// out the session timeout. A flush failure is best-effort and logged, so
579    /// close still leaves the group. Last, it sends the terminating
580    /// telemetry push (KIP-714) when the consumer pushes metrics and has a
581    /// subscription.
582    #[tracing::instrument(
583        name = "share_consumer.close",
584        level = "info",
585        skip_all,
586        fields(group_id = %self.group_id, member_id = %self.member_id),
587        err(level = "info")
588    )]
589    /// # Errors
590    /// Returns an error when configuration is invalid, protocol encoding fails, the broker rejects the request, or transport I/O fails.
591    pub async fn close(&mut self) -> Result<(), ConsumerError> {
592        // Roll the previous poll's implicit Accepts into the explicit ack queue
593        // so the final flush below covers both modes in one ShareAcknowledge.
594        if should_stage_implicit_accepts(self.ack_mode) {
595            stage_implicit_accepts(&mut self.prev_delivered, &mut self.pending_acks);
596        }
597        if let Err(e) = self.flush_pending_acks().await {
598            tracing::warn!(error = %e, "share consumer close: final acknowledge failed");
599        }
600
601        self.shutdown.cancel();
602        if let Some(h) = self.hb_handle.take() {
603            let _ = h.await;
604        }
605        // The terminating telemetry push (KIP-714).
606        self.client.close_telemetry().await;
607        Ok(())
608    }
609
610    /// The client instance id that the broker assigned for metrics push
611    /// (KIP-714), as Kafka's `KafkaShareConsumer.clientInstanceId` returns
612    /// it.
613    ///
614    /// The call waits up to `timeout` for the first
615    /// `GetTelemetrySubscriptions` response. It returns `Ok(None)` when none
616    /// has come by then, as Kafka returns `null`. A zero `timeout` does not
617    /// wait.
618    ///
619    /// # Errors
620    /// Returns [`ClientError::TelemetryDisabled`] when `enable_metrics_push`
621    /// is `false`, and [`ClientError::InvalidArgument`] for a negative
622    /// `timeout`, as Kafka throws an `IllegalStateException` and an
623    /// `IllegalArgumentException`.
624    ///
625    /// [`ClientError::TelemetryDisabled`]: ClientError::TelemetryDisabled
626    /// [`ClientError::InvalidArgument`]: ClientError::InvalidArgument
627    pub async fn client_instance_id(
628        &self,
629        timeout: Time,
630    ) -> Result<Option<uuid::Uuid>, ConsumerError> {
631        client_instance_id(&self.client, timeout).await
632    }
633}
634
635#[cfg(test)]
636mod tests {
637    use assert2::check;
638    use krabka_protocol::tagged_fields::UnknownTaggedFields;
639    use krabka_units::millis;
640
641    use super::*;
642    use crate::test_support::wire_uuid;
643
644    #[test]
645    fn leave_heartbeat_timeout_uses_default_and_valid_override() {
646        let default = ShareConsumerLeaveHeartbeatTimeout::default();
647        check!(default.time() == DEFAULT_SHARE_CONSUMER_LEAVE_HEARTBEAT_TIMEOUT);
648
649        let configured = ShareConsumerLeaveHeartbeatTimeout::new(millis(37)).unwrap();
650        check!(configured.time() == millis(37));
651    }
652
653    #[test]
654    fn leave_heartbeat_timeout_validates_millisecond_boundaries() {
655        assert2::assert!(
656            ShareConsumerLeaveHeartbeatTimeout::new(secs(0))
657                .unwrap_err()
658                .contains("share consumer leave-heartbeat timeout")
659        );
660        assert2::assert!(
661            ShareConsumerLeaveHeartbeatTimeout::new(Time::from_micros(1_001))
662                .unwrap_err()
663                .contains("whole number of milliseconds")
664        );
665        assert2::assert!(
666            ShareConsumerLeaveHeartbeatTimeout::new(Time::from_secs_f64(f64::INFINITY))
667                .unwrap_err()
668                .contains("share consumer leave-heartbeat timeout")
669        );
670    }
671
672    /// A refused share-consumer start returns its error to the caller, so it
673    /// is recorded at the span's own level, info, not at ERROR
674    /// (krabka-io/krabka-broker#1515).
675    #[tokio::test]
676    async fn a_refused_start_is_recorded_at_its_span_level() {
677        let (result, levels) = krabka_client_core::mock_test_support::recorded_levels(Box::pin(
678            ShareConsumer::builder()
679                .bootstrap("invalid.invalid:9092")
680                .group_id("log-level")
681                .subscribe(["topic".to_owned()])
682                .leave_heartbeat_timeout(krabka_units::secs(0))
683                .build(),
684        ))
685        .await;
686        assert2::assert!(result.is_err());
687        assert2::check!(levels == vec![tracing::Level::INFO]);
688    }
689
690    #[tokio::test]
691    async fn invalid_leave_heartbeat_timeout_fails_before_broker_lookup() {
692        let error = ShareConsumer::builder()
693            .bootstrap("invalid.invalid:9092")
694            .group_id("leave-validation")
695            .subscribe(["topic".to_owned()])
696            .leave_heartbeat_timeout(krabka_units::secs(0))
697            .build()
698            .await
699            .err()
700            .expect("zero leave-heartbeat timeout must fail");
701
702        assert2::assert!(
703            error
704                .to_string()
705                .contains("share consumer leave-heartbeat timeout")
706        );
707    }
708
709    #[test]
710    fn share_fetch_limits_use_defaults_and_valid_overrides() {
711        check!(ShareConsumerFetchMinBytes::default().bytes() == 1);
712        check!(ShareConsumerFetchMaxBytes::default().bytes() == 52_428_800);
713        check!(DEFAULT_SHARE_CONSUMER_FETCH_MIN.bytes_i32() == 1);
714        check!(DEFAULT_SHARE_CONSUMER_FETCH_MAX.bytes_i32() == 52_428_800);
715        check!(
716            ShareConsumerFetchMaxRecords::default().records()
717                == DEFAULT_SHARE_CONSUMER_FETCH_MAX_RECORDS
718        );
719
720        check!(ShareConsumerFetchMinBytes::new(7).unwrap().bytes() == 7);
721        check!(ShareConsumerFetchMaxBytes::new(65_536).unwrap().bytes() == 65_536);
722        check!(ShareConsumerFetchMaxRecords::new(37).unwrap().records() == 37);
723    }
724
725    #[test]
726    fn share_fetch_limits_validate_boundaries() {
727        for invalid in [-1, 0] {
728            check!(
729                ShareConsumerFetchMinBytes::new(invalid)
730                    .unwrap_err()
731                    .contains("share consumer fetch min bytes")
732            );
733            check!(
734                ShareConsumerFetchMaxBytes::new(invalid)
735                    .unwrap_err()
736                    .contains("share consumer fetch max bytes")
737            );
738            check!(
739                ShareConsumerFetchMaxRecords::new(invalid)
740                    .unwrap_err()
741                    .contains("share consumer fetch max records")
742            );
743        }
744
745        check!(ShareConsumerFetchMinBytes::new(i32::MAX).unwrap().bytes() == i32::MAX);
746        check!(ShareConsumerFetchMaxBytes::new(i32::MAX).unwrap().bytes() == i32::MAX);
747        check!(
748            ShareConsumerFetchMaxRecords::new(i32::MAX)
749                .unwrap()
750                .records()
751                == i32::MAX
752        );
753    }
754
755    #[tokio::test]
756    async fn invalid_share_fetch_limits_fail_before_broker_lookup() {
757        let error = ShareConsumer::builder()
758            .bootstrap("invalid.invalid:9092")
759            .group_id("fetch-limit-validation")
760            .subscribe(["topic".to_owned()])
761            .fetch_min(2)
762            .fetch_max(1)
763            .build()
764            .await
765            .err()
766            .expect("minimum above maximum must fail");
767
768        check!(
769            error
770                .to_string()
771                .contains("share consumer fetch min bytes must not exceed fetch max bytes")
772        );
773    }
774
775    #[tokio::test]
776    async fn non_positive_share_fetch_limit_fails_before_broker_lookup() {
777        let error = ShareConsumer::builder()
778            .bootstrap("invalid.invalid:9092")
779            .group_id("fetch-limit-validation")
780            .subscribe(["topic".to_owned()])
781            .fetch_min(0)
782            .build()
783            .await
784            .err()
785            .expect("non-positive fetch limit must fail");
786
787        check!(error.to_string().contains("must be greater than 0"));
788    }
789
790    #[tokio::test]
791    async fn invalid_client_resource_policy_fails_before_broker_lookup() {
792        let error = ShareConsumer::builder()
793            .bootstrap("invalid.invalid:9092")
794            .group_id("client-policy-validation")
795            .subscribe(["topic".to_owned()])
796            .dispatch_queue_capacity(0)
797            .build()
798            .await
799            .err()
800            .expect("invalid client policy");
801
802        check!(error.to_string().contains("client dispatch queue capacity"));
803    }
804
805    async fn test_consumer() -> ShareConsumer {
806        ShareConsumer {
807            client: Client::builder()
808                .bootstrap("127.0.0.1:1")
809                .client_id("share-test")
810                .build()
811                .await
812                .unwrap(),
813            group_id: "group-a".into(),
814            member_id: "member-a".into(),
815            member_epoch: Arc::new(Mutex::new(3)),
816            assignment: Arc::new(Mutex::new(vec![(wire_uuid(7), "topic-a".into(), 2)])),
817            topic_names: Arc::new(Mutex::new(HashMap::new())),
818            share_session_epoch: 0,
819            fetch_min: DEFAULT_SHARE_CONSUMER_FETCH_MIN,
820            fetch_max: DEFAULT_SHARE_CONSUMER_FETCH_MAX,
821            fetch_max_records: DEFAULT_SHARE_CONSUMER_FETCH_MAX_RECORDS,
822            acquire_mode: ShareAcquireMode::BatchOptimized,
823            ack_mode: ShareAckMode::Explicit,
824            pending_acks: Vec::new(),
825            prev_delivered: Vec::new(),
826            shutdown: CancellationToken::new(),
827            hb_handle: None,
828            poll_error: coordinator::PollErrorSlot::default(),
829        }
830    }
831
832    #[test]
833    fn join_heartbeat_request_mints_uuid_and_preserves_group_epoch_and_subscription() {
834        let req = build_join_heartbeat_request("group-a".into(), vec!["topic-a".into()]);
835        let second = build_join_heartbeat_request("group-a".into(), vec!["topic-a".into()]);
836
837        for member_id in [&req.member_id, &second.member_id] {
838            assert2::assert!(uuid::Uuid::parse_str(member_id).is_ok());
839            assert2::assert!(member_id.len() == 36);
840        }
841        assert2::assert!(req.member_id != second.member_id);
842
843        assert2::assert!(
844            req == ShareGroupHeartbeatRequest {
845                group_id: "group-a".into(),
846                member_id: req.member_id.clone(),
847                member_epoch: 0,
848                rack_id: None,
849                subscribed_topic_names: Some(vec!["topic-a".into()]),
850                unknown_tagged_fields: UnknownTaggedFields(Vec::new()),
851            }
852        );
853    }
854
855    #[test]
856    fn join_response_helpers_preserve_error_interval_and_assignment_boundaries() {
857        for (name, code, expected) in [("success", 0, false), ("error", 17, true)] {
858            check!(response_has_error(code) == expected, "case {name}");
859        }
860        for (name, broker_ms, expected) in [
861            ("broker interval", 2500, krabka_units::millis(2500)),
862            ("fallback interval", 0, krabka_units::secs(3)),
863        ] {
864            check!(
865                heartbeat_interval_from_response(broker_ms, krabka_units::secs(3)) == expected,
866                "case {name}"
867            );
868        }
869        for (name, count, expected) in [("empty", 0, false), ("present", 1, true)] {
870            check!(has_assignment_partitions(count) == expected, "case {name}");
871        }
872        for (name, mode, expected) in [
873            ("implicit", ShareAckMode::Implicit, true),
874            ("explicit", ShareAckMode::Explicit, false),
875        ] {
876            check!(
877                should_stage_implicit_accepts(mode) == expected,
878                "case {name}"
879            );
880        }
881    }
882
883    #[test]
884    fn stage_implicit_accepts_moves_delivered_ranges_to_pending_acks() {
885        let mut prev = vec![(wire_uuid(7), 2, 10, 12)];
886        let mut pending = Vec::new();
887
888        stage_implicit_accepts(&mut prev, &mut pending);
889
890        assert2::assert!(prev.is_empty());
891        assert2::assert!(
892            pending
893                == vec![(
894                    wire_uuid(7),
895                    2,
896                    10,
897                    12,
898                    crate::share::types::ShareAckType::Accept.wire()
899                )]
900        );
901    }
902
903    #[tokio::test]
904    async fn accessors_return_share_identity_and_assignment() {
905        let consumer = test_consumer().await;
906
907        check!(
908            (
909                consumer.group_id(),
910                consumer.member_id(),
911                consumer.assignment().await
912            ) == ("group-a", "member-a", vec![("topic-a".into(), 2)])
913        );
914    }
915
916    #[tokio::test]
917    async fn close_cancels_shutdown_token_without_spawned_handle() {
918        let mut consumer = test_consumer().await;
919
920        assert2::assert!(!consumer.shutdown.is_cancelled());
921        consumer.close().await.unwrap();
922        assert2::assert!(consumer.shutdown.is_cancelled());
923    }
924}