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