krabka_client_streams/membership/
client.rs

1//! `StreamsMembership`: the public handle for a KIP-1071 streams group.
2//!
3//! `start` generates a member id and sends the join heartbeat, which is epoch 0
4//! with the topology. It captures the broker-assigned epoch, the heartbeat
5//! interval, and the initial assignment. It then spawns the background heartbeat
6//! loop on its own connection, because the broker serves a connection serially.
7//!
8//! `next_event` drains coordinator events. `close` leaves the group.
9
10use std::{sync::Arc, time::Duration};
11
12use krabka_client_core::{
13    Client, ClientDnsTimeout, ClientError, ClientFrameMax, ConnectionDispatchQueueCapacity,
14    CoordinatorKeyType, build_find_coordinator, coordinator_endpoint,
15};
16use krabka_protocol::owned::streams_group_heartbeat_request::StreamsGroupHeartbeatRequest;
17use krabka_units::prelude::*;
18use refined_type::rule::MinMaxU128;
19use tokio::{
20    sync::{Mutex, mpsc},
21    task::JoinHandle,
22};
23use tokio_util::sync::CancellationToken;
24
25use super::{
26    coordinator::{self, CoordinatorState},
27    status::map_status,
28    types::{StreamsAssignment, StreamsEvent, TaskOffsetTracker},
29};
30use crate::{error::StreamsClientError, membership::assignment::resolve};
31
32const COORDINATOR_LOAD_IN_PROGRESS: i16 = 14;
33const COORDINATOR_NOT_AVAILABLE: i16 = 15;
34const NOT_COORDINATOR: i16 = 16;
35
36struct RoutingHeartbeatTransport {
37    bootstrap: Client,
38    coordinator: Mutex<Client>,
39    group_id: String,
40    client_id: String,
41    broker_dns_timeout: ClientDnsTimeout,
42    dispatch_queue_capacity: ConnectionDispatchQueueCapacity,
43    frame_max: ClientFrameMax,
44    security: Option<krabka_client_core::security::ClientSecurity>,
45}
46
47impl RoutingHeartbeatTransport {
48    async fn new(
49        bootstrap: &str,
50        group_id: &str,
51        client_id: &str,
52        broker_dns_timeout: ClientDnsTimeout,
53        dispatch_queue_capacity: ConnectionDispatchQueueCapacity,
54        frame_max: ClientFrameMax,
55        security: Option<krabka_client_core::security::ClientSecurity>,
56    ) -> Result<Self, ClientError> {
57        let bootstrap = Client::builder()
58            .bootstrap(bootstrap)
59            .client_id(client_id)
60            .dns_timeout(broker_dns_timeout.time())
61            .dispatch_queue_capacity(dispatch_queue_capacity.get())
62            .frame_max(frame_max.size())
63            .maybe_security(security.clone())
64            .build()
65            .await?;
66        let _ = bootstrap.refresh_metadata().await;
67        let coordinator = Self::discover(
68            &bootstrap,
69            group_id,
70            client_id,
71            broker_dns_timeout,
72            dispatch_queue_capacity,
73            frame_max,
74            security.clone(),
75        )
76        .await?;
77        Ok(Self {
78            bootstrap,
79            coordinator: Mutex::new(coordinator),
80            group_id: group_id.to_owned(),
81            client_id: client_id.to_owned(),
82            broker_dns_timeout,
83            dispatch_queue_capacity,
84            frame_max,
85            security,
86        })
87    }
88
89    async fn discover(
90        bootstrap: &Client,
91        group_id: &str,
92        client_id: &str,
93        broker_dns_timeout: ClientDnsTimeout,
94        dispatch_queue_capacity: ConnectionDispatchQueueCapacity,
95        frame_max: ClientFrameMax,
96        security: Option<krabka_client_core::security::ClientSecurity>,
97    ) -> Result<Client, ClientError> {
98        let response = bootstrap
99            .send(build_find_coordinator(group_id, CoordinatorKeyType::Group))
100            .await?;
101        let endpoint = coordinator_endpoint(group_id, response)?;
102        Client::builder()
103            .bootstrap(endpoint.address())
104            .client_id(client_id)
105            .dns_timeout(broker_dns_timeout.time())
106            .dispatch_queue_capacity(dispatch_queue_capacity.get())
107            .frame_max(frame_max.size())
108            .maybe_security(security)
109            .build()
110            .await
111    }
112
113    async fn rediscover(&self) -> Result<Client, ClientError> {
114        Self::discover(
115            &self.bootstrap,
116            &self.group_id,
117            &self.client_id,
118            self.broker_dns_timeout,
119            self.dispatch_queue_capacity,
120            self.frame_max,
121            self.security.clone(),
122        )
123        .await
124    }
125}
126
127#[async_trait::async_trait]
128impl coordinator::HeartbeatTransport for RoutingHeartbeatTransport {
129    async fn send_heartbeat(
130        &self,
131        request: StreamsGroupHeartbeatRequest,
132    ) -> Result<
133        krabka_protocol::owned::streams_group_heartbeat_response::StreamsGroupHeartbeatResponse,
134        ClientError,
135    > {
136        let first = self.coordinator.lock().await.send(request.clone()).await;
137        if !should_rediscover(&first) {
138            return first;
139        }
140        let coordinator = self.rediscover().await?;
141        let response = coordinator.send(request).await;
142        *self.coordinator.lock().await = coordinator;
143        response
144    }
145}
146
147fn should_rediscover(
148    result: &Result<
149        krabka_protocol::owned::streams_group_heartbeat_response::StreamsGroupHeartbeatResponse,
150        ClientError,
151    >,
152) -> bool {
153    matches!(
154        result,
155        Ok(response)
156            if response.error_code == COORDINATOR_NOT_AVAILABLE
157                || response.error_code == NOT_COORDINATOR
158    ) || matches!(
159        result,
160        Err(ClientError::Connect { .. }
161            | ClientError::Disconnected
162            | ClientError::Timeout(_)
163            | ClientError::Io(_))
164    )
165}
166
167fn should_retry_coordinator_discovery(error: &ClientError) -> bool {
168    matches!(
169        error,
170        ClientError::Server {
171            error_code: COORDINATOR_LOAD_IN_PROGRESS | COORDINATOR_NOT_AVAILABLE | NOT_COORDINATOR
172        } | ClientError::NoCoordinator { .. }
173            | ClientError::Connect { .. }
174            | ClientError::Disconnected
175            | ClientError::Timeout(_)
176            | ClientError::Io(_)
177    )
178}
179
180/// Default Client Streams rebalance timeout.
181pub const DEFAULT_STREAMS_REBALANCE_TIMEOUT: Duration = Duration::from_secs(30);
182
183/// Default delay between Client Streams initial join retries.
184pub const DEFAULT_STREAMS_JOIN_RETRY_BACKOFF: Duration = Duration::from_millis(200);
185
186/// Default deadline for the final Client Streams leave heartbeat.
187pub const DEFAULT_STREAMS_LEAVE_HEARTBEAT_TIMEOUT: Duration = Duration::from_secs(5);
188
189fn validate_positive_whole_milliseconds(field: &str, value: Duration) -> Result<u64, String> {
190    let milliseconds = MinMaxU128::<1, { u64::MAX as u128 }>::new(value.as_millis())
191        .map_err(|error| format!("{field}: {error}"))?
192        .into_value();
193    let milliseconds = u64::try_from(milliseconds).map_err(|error| format!("{field}: {error}"))?;
194    if Duration::from_millis(milliseconds) != value {
195        return Err(format!("{field} must be a whole number of milliseconds"));
196    }
197    Ok(milliseconds)
198}
199
200/// Positive, whole-millisecond delay between Client Streams initial join retries.
201#[derive(Clone, Copy, Debug, Eq, PartialEq)]
202pub struct StreamsJoinRetryBackoff(Duration);
203
204impl StreamsJoinRetryBackoff {
205    /// Validate a Client Streams initial join retry backoff.
206    ///
207    /// # Errors
208    ///
209    /// Returns an error for zero, for fractional milliseconds, or for a value
210    /// whose milliseconds do not fit in `u64`.
211    pub fn new(value: Duration) -> Result<Self, String> {
212        validate_positive_whole_milliseconds("streams join retry backoff", value)?;
213        Ok(Self(value))
214    }
215
216    /// Return the validated duration.
217    #[must_use]
218    pub const fn duration(self) -> Duration {
219        self.0
220    }
221
222    /// Return the validated whole milliseconds.
223    ///
224    /// # Panics
225    ///
226    /// Panics if the validated duration no longer fits in `u64` milliseconds.
227    #[must_use]
228    pub fn milliseconds(self) -> u64 {
229        u64::try_from(self.0.as_millis()).expect("validated streams join retry backoff fits u64")
230    }
231}
232
233impl Default for StreamsJoinRetryBackoff {
234    fn default() -> Self {
235        Self::new(DEFAULT_STREAMS_JOIN_RETRY_BACKOFF)
236            .expect("default streams join retry backoff is valid")
237    }
238}
239
240/// Positive, whole-millisecond deadline for the final leave heartbeat.
241#[derive(Clone, Copy, Debug, Eq, PartialEq)]
242pub struct StreamsLeaveHeartbeatTimeout(Duration);
243
244impl StreamsLeaveHeartbeatTimeout {
245    /// Validate a final leave-heartbeat timeout.
246    ///
247    /// # Errors
248    ///
249    /// Returns an error for zero, for fractional milliseconds, or for a value
250    /// whose milliseconds do not fit in `u64`.
251    pub fn new(value: Duration) -> Result<Self, String> {
252        validate_positive_whole_milliseconds("streams leave heartbeat timeout", value)?;
253        Ok(Self(value))
254    }
255
256    /// Return the validated duration.
257    #[must_use]
258    pub const fn duration(self) -> Duration {
259        self.0
260    }
261
262    /// Return the validated whole milliseconds.
263    ///
264    /// # Panics
265    ///
266    /// Panics if the validated timeout no longer fits in `u64` milliseconds.
267    #[must_use]
268    pub fn milliseconds(self) -> u64 {
269        u64::try_from(self.0.as_millis())
270            .expect("validated streams leave heartbeat timeout fits u64")
271    }
272}
273
274impl Default for StreamsLeaveHeartbeatTimeout {
275    fn default() -> Self {
276        Self::new(DEFAULT_STREAMS_LEAVE_HEARTBEAT_TIMEOUT)
277            .expect("default streams leave heartbeat timeout is valid")
278    }
279}
280
281fn join_retry_delay(error_code: i16, backoff: StreamsJoinRetryBackoff) -> Option<Duration> {
282    matches!(
283        error_code,
284        COORDINATOR_LOAD_IN_PROGRESS | COORDINATOR_NOT_AVAILABLE | NOT_COORDINATOR
285    )
286    .then(|| backoff.duration())
287}
288
289/// Positive, whole-millisecond rebalance timeout representable on the Kafka wire.
290#[derive(Clone, Copy, Debug, Eq, PartialEq)]
291pub struct StreamsRebalanceTimeout(Duration);
292
293impl StreamsRebalanceTimeout {
294    /// Validate a Client Streams rebalance timeout.
295    ///
296    /// # Errors
297    ///
298    /// Returns an error for zero, fractional milliseconds, or a value greater
299    /// than `i32::MAX` milliseconds.
300    ///
301    /// # Panics
302    ///
303    /// Panics if a positive `i32` cannot be represented as `u64`.
304    pub fn new(value: Duration) -> Result<Self, String> {
305        let milliseconds = MinMaxU128::<1, { i32::MAX as u128 }>::new(value.as_millis())
306            .map_err(|error| format!("streams rebalance timeout: {error}"))?
307            .into_value();
308        let milliseconds = i32::try_from(milliseconds)
309            .map_err(|error| format!("streams rebalance timeout: {error}"))?;
310        let whole = Duration::from_millis(
311            u64::try_from(milliseconds).expect("positive i32 milliseconds fit u64"),
312        );
313        if whole != value {
314            return Err(
315                "streams rebalance timeout must be a whole number of milliseconds".to_owned(),
316            );
317        }
318        Ok(Self(value))
319    }
320
321    #[must_use]
322    /// Return the validated duration.
323    pub const fn duration(self) -> Duration {
324        self.0
325    }
326
327    #[must_use]
328    /// Return the validated signed wire milliseconds.
329    ///
330    /// # Panics
331    ///
332    /// Panics if the validated timeout no longer fits in `i32`.
333    pub fn milliseconds(self) -> i32 {
334        i32::try_from(self.0.as_millis()).expect("validated streams rebalance timeout fits i32")
335    }
336}
337
338impl Default for StreamsRebalanceTimeout {
339    fn default() -> Self {
340        Self::new(DEFAULT_STREAMS_REBALANCE_TIMEOUT)
341            .expect("default streams rebalance timeout is valid")
342    }
343}
344
345/// Hook that runs once at membership start to resolve schema ids before
346/// processing. `SchemaCache` implements it under the `schema-serde` feature.
347#[async_trait::async_trait]
348pub trait SchemaPrewarm: Send + Sync {
349    async fn prewarm(&self) -> Result<(), StreamsClientError>;
350}
351
352/// A live streams-group membership. Construct via [`StreamsMembership::builder`].
353pub struct StreamsMembership {
354    member_id: String,
355    group_id: String,
356    /// Shared with the coordinator loop. It reads the live member epoch for
357    /// [`group_metadata`](Self::group_metadata), which EOS uses in
358    /// `send_offsets_to_transaction`.
359    member_epoch: Arc<Mutex<i32>>,
360    events: mpsc::UnboundedReceiver<StreamsEvent>,
361    shutdown: CancellationToken,
362    hb_handle: Option<JoinHandle<()>>,
363    tracker: Arc<Mutex<TaskOffsetTracker>>,
364}
365
366#[bon::bon]
367impl StreamsMembership {
368    /// Join a streams group and start the heartbeat loop.
369    #[builder(start_fn = builder, finish_fn = build)]
370    #[tracing::instrument(
371        name = "streams.membership.start",
372        level = "info",
373        skip_all,
374        fields(group_id = %group_id, member_id = tracing::field::Empty),
375        err,
376    )]
377    /// # Errors
378    /// Returns an error when configuration is invalid, protocol encoding fails, the broker rejects the request, or transport I/O fails.
379    /// # Panics
380    /// Panics if synchronized client state is poisoned or a response violates an invariant established by protocol validation.
381    pub async fn start(
382        #[builder(into)] bootstrap: String,
383        #[builder(into, default = "krabka-streams".to_string())] client_id: String,
384        #[builder(into)] group_id: String,
385        topology: std::sync::Arc<crate::topology::BuiltTopology>,
386        #[builder(into)] process_id: Option<String>,
387        #[builder(into)] instance_id: Option<String>,
388        #[builder(default = DEFAULT_STREAMS_REBALANCE_TIMEOUT)] rebalance_timeout: Duration,
389        #[builder(default = DEFAULT_STREAMS_JOIN_RETRY_BACKOFF)] join_retry_backoff: Duration,
390        #[builder(default = DEFAULT_STREAMS_LEAVE_HEARTBEAT_TIMEOUT)]
391        leave_heartbeat_timeout: Duration,
392        #[builder(default)] broker_dns_timeout: ClientDnsTimeout,
393        #[builder(default)] dispatch_queue_capacity: ConnectionDispatchQueueCapacity,
394        #[builder(default)] frame_max: ClientFrameMax,
395        security: Option<krabka_client_core::security::ClientSecurity>,
396        schema_prewarm: Option<std::sync::Arc<dyn SchemaPrewarm>>,
397    ) -> Result<Self, StreamsClientError> {
398        if group_id.is_empty() {
399            return Err(StreamsClientError::Server(0));
400        }
401        let rebalance_timeout =
402            StreamsRebalanceTimeout::new(rebalance_timeout).map_err(StreamsClientError::Runtime)?;
403        let join_retry_backoff = StreamsJoinRetryBackoff::new(join_retry_backoff)
404            .map_err(StreamsClientError::Runtime)?;
405        let leave_heartbeat_timeout = StreamsLeaveHeartbeatTimeout::new(leave_heartbeat_timeout)
406            .map_err(StreamsClientError::Runtime)?;
407        if let Some(prewarm) = &schema_prewarm {
408            prewarm.prewarm().await?;
409        }
410        let process_id = process_id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
411        let member_id = uuid::Uuid::new_v4().to_string();
412        tracing::Span::current().record("member_id", tracing::field::display(&member_id));
413        let rebalance_timeout = Time::from_millis(i64::from(rebalance_timeout.milliseconds()));
414
415        let client = loop {
416            match RoutingHeartbeatTransport::new(
417                &bootstrap,
418                &group_id,
419                &client_id,
420                broker_dns_timeout,
421                dispatch_queue_capacity,
422                frame_max,
423                security.clone(),
424            )
425            .await
426            {
427                Ok(client) => break client,
428                Err(error) if should_retry_coordinator_discovery(&error) => {
429                    tokio::time::sleep(join_retry_backoff.duration()).await;
430                }
431                Err(error) => return Err(error.into()),
432            }
433        };
434
435        let (events_tx, events_rx) = mpsc::unbounded_channel();
436        let join = loop {
437            let response = coordinator::HeartbeatTransport::send_heartbeat(
438                &client,
439                coordinator::join_heartbeat(
440                    &group_id,
441                    &member_id,
442                    &process_id,
443                    instance_id.clone(),
444                    rebalance_timeout,
445                    &topology,
446                ),
447            )
448            .await;
449            let resp = match response {
450                Ok(response) => response,
451                Err(error) if should_retry_coordinator_discovery(&error) => {
452                    tokio::time::sleep(join_retry_backoff.duration()).await;
453                    continue;
454                }
455                Err(error) => return Err(error.into()),
456            };
457            if let Some(delay) = join_retry_delay(resp.error_code, join_retry_backoff) {
458                tokio::time::sleep(delay).await;
459                continue;
460            }
461            break map_error(resp)?;
462        };
463
464        let member_epoch_val = join.member_epoch;
465        let hb_interval = heartbeat_interval(join.heartbeat_interval_ms);
466
467        if should_emit_statuses(join.status.as_ref()) {
468            let statuses = join.status.as_ref().expect("checked above");
469            let _ = events_tx.send(StreamsEvent::NotReady(
470                statuses.iter().map(map_status).collect(),
471            ));
472        }
473        let owned_active = Arc::new(Mutex::new(join.active_tasks.clone().unwrap_or_default()));
474        let owned_standby = Arc::new(Mutex::new(join.standby_tasks.clone().unwrap_or_default()));
475        let owned_warmup = Arc::new(Mutex::new(join.warmup_tasks.clone().unwrap_or_default()));
476        let tracker = Arc::new(Mutex::new(TaskOffsetTracker::default()));
477        let initial = StreamsAssignment {
478            active: resolve(join.active_tasks.as_ref(), &topology),
479            standby: resolve(join.standby_tasks.as_ref(), &topology),
480            warmup: resolve(join.warmup_tasks.as_ref(), &topology),
481        };
482        if initial != StreamsAssignment::default() {
483            let _ = events_tx.send(StreamsEvent::Assigned(initial.clone()));
484        }
485
486        let shutdown = CancellationToken::new();
487        // Shared epoch handle: the coordinator advances it each heartbeat; the
488        // membership reads it for EOS `group_metadata()`.
489        let member_epoch = Arc::new(Mutex::new(member_epoch_val));
490        let state = CoordinatorState {
491            client,
492            group_id: group_id.clone(),
493            member_id: member_id.clone(),
494            process_id,
495            instance_id,
496            rebalance_timeout,
497            topology: Arc::clone(&topology),
498            member_epoch: Arc::clone(&member_epoch),
499            owned_active,
500            owned_standby,
501            owned_warmup,
502            last_sent_tasks: Mutex::new(None),
503            tracker: tracker.clone(),
504            heartbeat_interval: hb_interval,
505            leave_heartbeat_timeout: leave_heartbeat_timeout.duration().as_time(),
506            events: events_tx,
507            last_assignment: tokio::sync::Mutex::new(initial),
508        };
509        let hb_handle = tokio::spawn(coordinator::run(state, shutdown.clone()));
510
511        Ok(Self {
512            member_id,
513            group_id,
514            member_epoch,
515            events: events_rx,
516            shutdown,
517            hb_handle: Some(hb_handle),
518            tracker,
519        })
520    }
521}
522
523impl Drop for StreamsMembership {
524    fn drop(&mut self) {
525        if let Some(handle) = self.hb_handle.take() {
526            handle.abort();
527        }
528        self.shutdown.cancel();
529    }
530}
531
532impl StreamsMembership {
533    /// The client-generated member id.
534    #[must_use]
535    pub fn member_id(&self) -> &str {
536        &self.member_id
537    }
538
539    /// Get the shared task offset tracker.
540    #[must_use]
541    pub fn tracker(&self) -> Arc<Mutex<TaskOffsetTracker>> {
542        self.tracker.clone()
543    }
544
545    /// The streams group id.
546    #[must_use]
547    pub fn group_id(&self) -> &str {
548        &self.group_id
549    }
550
551    /// Streams group metadata for the EOS `send_offsets_to_transaction` call.
552    ///
553    /// The `generation_id` maps to the live member epoch, the next-generation
554    /// "generation". The epoch lives behind the coordinator's async `Mutex`, so
555    /// this reader is `async`. A sync accessor would have to `blocking_lock`,
556    /// which panics inside the runtime's async supervisor.
557    #[tracing::instrument(
558        name = "streams.membership.group_metadata",
559        level = "debug",
560        skip_all,
561        fields(group_id = %self.group_id, member_id = %self.member_id),
562    )]
563    pub async fn group_metadata(&self) -> crate::runtime::eos::StreamsGroupMeta {
564        let epoch = *self.member_epoch.lock().await;
565        crate::runtime::eos::StreamsGroupMeta {
566            group: self.group_id.clone(),
567            generation: epoch,
568            member: self.member_id.clone(),
569            group_instance: None,
570        }
571    }
572
573    /// Await the next membership event: assignment, not-ready, or fenced.
574    ///
575    /// This method returns [`StreamsClientError::Closed`] once the heartbeat loop
576    /// has ended.
577    ///
578    /// # Errors
579    /// Returns an error when configuration is invalid, protocol encoding fails, the broker rejects the request, or transport I/O fails.
580    pub async fn next_event(&mut self) -> Result<StreamsEvent, StreamsClientError> {
581        self.events.recv().await.ok_or(StreamsClientError::Closed)
582    }
583
584    /// Leave the group and stop the heartbeat loop.
585    #[tracing::instrument(
586        name = "streams.membership.close",
587        level = "info",
588        skip_all,
589        fields(group_id = %self.group_id, member_id = %self.member_id),
590        err,
591    )]
592    /// # Errors
593    /// Returns an error when configuration is invalid, protocol encoding fails, the broker rejects the request, or transport I/O fails.
594    pub async fn close(&mut self) -> Result<(), StreamsClientError> {
595        self.shutdown.cancel();
596        if let Some(h) = self.hb_handle.take() {
597            let _ = h.await;
598        }
599        Ok(())
600    }
601}
602
603/// The coordinator's advertised heartbeat cadence, or the JVM default when the
604/// broker sends a non-positive value. `raw` is the response's raw `int32`
605/// milliseconds.
606fn heartbeat_interval(raw: i32) -> Time {
607    if raw > 0 {
608        Time::from_millis(i64::from(raw))
609    } else {
610        secs(3)
611    }
612}
613
614fn should_emit_statuses<T>(statuses: Option<&Vec<T>>) -> bool {
615    statuses.is_some_and(|statuses| !statuses.is_empty())
616}
617
618/// Map a join-response error code to a typed error (0 = ok).
619fn map_error(
620    resp: krabka_protocol::owned::streams_group_heartbeat_response::StreamsGroupHeartbeatResponse,
621) -> Result<
622    krabka_protocol::owned::streams_group_heartbeat_response::StreamsGroupHeartbeatResponse,
623    StreamsClientError,
624> {
625    // Kafka error codes for the STREAMS_INVALID_TOPOLOGY family (KIP-1071).
626    // The broker coordinator surfaces topology problems via the response Status
627    // list, but these are still valid top-level response codes per the
628    // StreamsGroupHeartbeatResponse schema.
629    const STREAMS_INVALID_TOPOLOGY: i16 = 130;
630    const STREAMS_INVALID_TOPOLOGY_EPOCH: i16 = 131;
631    const STREAMS_TOPOLOGY_FENCED: i16 = 132;
632    // Verified against crates/broker/src/codes.rs:
633    const GROUP_AUTHORIZATION_FAILED: i16 = 30; // codes::GROUP_AUTHORIZATION_FAILED
634    const TOPIC_AUTHORIZATION_FAILED: i16 = 29; // codes::TOPIC_AUTHORIZATION_FAILED
635    const GROUP_ID_NOT_FOUND: i16 = 69; // codes::GROUP_ID_NOT_FOUND
636    match resp.error_code {
637        0 => Ok(resp),
638        c @ (STREAMS_INVALID_TOPOLOGY
639        | STREAMS_INVALID_TOPOLOGY_EPOCH
640        | STREAMS_TOPOLOGY_FENCED) => Err(StreamsClientError::InvalidTopology {
641            code: c,
642            message: resp.error_message.unwrap_or_default(),
643        }),
644        c @ (GROUP_AUTHORIZATION_FAILED | TOPIC_AUTHORIZATION_FAILED) => {
645            Err(StreamsClientError::Authorization(c))
646        }
647        GROUP_ID_NOT_FOUND => Err(StreamsClientError::GroupIdNotFound),
648        other => Err(StreamsClientError::Server(other)),
649    }
650}
651
652#[cfg(test)]
653mod tests {
654    use std::{sync::Arc, time::Duration};
655
656    use assert2::check;
657    use krabka_client_core::ClientError;
658    use krabka_protocol::owned::streams_group_heartbeat_response::StreamsGroupHeartbeatResponse;
659    use krabka_units::prelude::*;
660    use tokio::sync::{Mutex, mpsc};
661    use tokio_util::sync::CancellationToken;
662
663    use super::{
664        COORDINATOR_LOAD_IN_PROGRESS, COORDINATOR_NOT_AVAILABLE, NOT_COORDINATOR,
665        StreamsJoinRetryBackoff, StreamsLeaveHeartbeatTimeout, StreamsRebalanceTimeout,
666        heartbeat_interval, join_retry_delay, map_error, should_emit_statuses, should_rediscover,
667        should_retry_coordinator_discovery,
668    };
669    use crate::{
670        error::StreamsClientError, membership::types::TaskOffsetTracker, topology::Topology,
671    };
672
673    fn resp(code: i16) -> StreamsGroupHeartbeatResponse {
674        StreamsGroupHeartbeatResponse {
675            error_code: code,
676            error_message: Some("detail".into()),
677            ..Default::default()
678        }
679    }
680
681    #[test]
682    fn ok_code_passes_through() {
683        check!(map_error(resp(0)).is_ok());
684    }
685
686    #[test]
687    #[allow(clippy::duration_suboptimal_units)]
688    fn rebalance_timeout_uses_default_and_valid_override() {
689        let default = StreamsRebalanceTimeout::default();
690        check!(default.duration() == Duration::from_secs(30));
691        check!(default.milliseconds() == 30_000);
692
693        let timeout = StreamsRebalanceTimeout::new(Duration::from_millis(45_000))
694            .expect("valid rebalance timeout");
695        check!(timeout.duration() == Duration::from_secs(45));
696        check!(timeout.milliseconds() == 45_000);
697    }
698
699    #[test]
700    fn rebalance_timeout_rejects_invalid_wire_values() {
701        check!(StreamsRebalanceTimeout::new(Duration::ZERO).is_err());
702        check!(
703            StreamsRebalanceTimeout::new(Duration::from_millis(1) + Duration::from_nanos(1))
704                .is_err()
705        );
706        check!(
707            StreamsRebalanceTimeout::new(Duration::from_millis(
708                u64::try_from(i32::MAX).expect("i32 max fits u64") + 1,
709            ))
710            .is_err()
711        );
712    }
713
714    #[test]
715    fn join_retry_backoff_uses_default_and_valid_override() {
716        let default = StreamsJoinRetryBackoff::default();
717        check!(default.duration() == Duration::from_millis(200));
718        check!(default.milliseconds() == 200);
719
720        let backoff = StreamsJoinRetryBackoff::new(Duration::from_millis(37))
721            .expect("positive whole milliseconds");
722        check!(backoff.duration() == Duration::from_millis(37));
723        check!(backoff.milliseconds() == 37);
724    }
725
726    #[test]
727    fn join_retry_backoff_validates_millisecond_boundaries() {
728        check!(StreamsJoinRetryBackoff::new(Duration::ZERO).is_err());
729        check!(
730            StreamsJoinRetryBackoff::new(Duration::from_millis(1) + Duration::from_nanos(1))
731                .is_err()
732        );
733        check!(StreamsJoinRetryBackoff::new(Duration::from_millis(u64::MAX)).is_ok());
734        check!(
735            StreamsJoinRetryBackoff::new(
736                Duration::from_millis(u64::MAX) + Duration::from_millis(1)
737            )
738            .is_err()
739        );
740    }
741
742    #[test]
743    fn leave_heartbeat_timeout_uses_default_and_valid_override() {
744        let default = StreamsLeaveHeartbeatTimeout::default();
745        check!(default.duration() == Duration::from_secs(5));
746        check!(default.milliseconds() == 5_000);
747
748        let timeout = StreamsLeaveHeartbeatTimeout::new(Duration::from_millis(37))
749            .expect("positive whole milliseconds");
750        check!(timeout.duration() == Duration::from_millis(37));
751        check!(timeout.milliseconds() == 37);
752    }
753
754    #[test]
755    fn leave_heartbeat_timeout_validates_millisecond_boundaries() {
756        check!(StreamsLeaveHeartbeatTimeout::new(Duration::ZERO).is_err());
757        check!(
758            StreamsLeaveHeartbeatTimeout::new(Duration::from_millis(1) + Duration::from_nanos(1))
759                .is_err()
760        );
761        check!(StreamsLeaveHeartbeatTimeout::new(Duration::from_millis(u64::MAX)).is_ok());
762        check!(StreamsLeaveHeartbeatTimeout::new(Duration::from_secs(u64::MAX)).is_err());
763    }
764
765    struct CountingPrewarm(Arc<std::sync::atomic::AtomicUsize>);
766
767    #[async_trait::async_trait]
768    impl super::SchemaPrewarm for CountingPrewarm {
769        async fn prewarm(&self) -> Result<(), StreamsClientError> {
770            self.0.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
771            Ok(())
772        }
773    }
774
775    #[tokio::test]
776    async fn invalid_leave_heartbeat_timeout_fails_before_prewarm_or_broker_lookup() {
777        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
778        let mut topology = Topology::new();
779        let source = topology.add_source::<String, String>("source", ["input"]);
780        topology.add_sink("sink", "output", [&source]);
781        let topology = Arc::new(topology.build("leave-validation").expect("topology"));
782
783        let error = super::StreamsMembership::builder()
784            .bootstrap("invalid.invalid:9092")
785            .group_id("leave-validation")
786            .topology(topology)
787            .leave_heartbeat_timeout(Duration::ZERO)
788            .schema_prewarm(Arc::new(CountingPrewarm(Arc::clone(&calls))))
789            .build()
790            .await
791            .err()
792            .expect("invalid configuration");
793
794        check!(
795            error
796                .to_string()
797                .contains("streams leave heartbeat timeout")
798        );
799        check!(calls.load(std::sync::atomic::Ordering::Relaxed) == 0);
800    }
801
802    #[test]
803    fn join_retry_path_uses_backoff_while_coordinator_loads_or_moves() {
804        let backoff = StreamsJoinRetryBackoff::new(Duration::from_millis(37))
805            .expect("positive whole milliseconds");
806        for code in [
807            COORDINATOR_LOAD_IN_PROGRESS,
808            COORDINATOR_NOT_AVAILABLE,
809            NOT_COORDINATOR,
810        ] {
811            check!(join_retry_delay(code, backoff) == Some(Duration::from_millis(37)));
812        }
813        check!(join_retry_delay(0, backoff).is_none());
814        check!(join_retry_delay(17, backoff).is_none());
815    }
816
817    #[test]
818    fn coordinator_movement_and_transport_errors_trigger_rediscovery() {
819        for code in [COORDINATOR_NOT_AVAILABLE, NOT_COORDINATOR] {
820            check!(should_rediscover(&Ok(resp(code))));
821        }
822        check!(should_rediscover(&Err(ClientError::Disconnected)));
823        check!(!should_rediscover(&Ok(resp(0))));
824        check!(!should_rediscover(&Ok(resp(30))));
825        check!(!should_rediscover(&Err(ClientError::Server {
826            error_code: 30,
827        })));
828    }
829
830    #[test]
831    fn unavailable_coordinator_discovery_retries() {
832        for error in [
833            ClientError::Server {
834                error_code: COORDINATOR_LOAD_IN_PROGRESS,
835            },
836            ClientError::Server {
837                error_code: COORDINATOR_NOT_AVAILABLE,
838            },
839            ClientError::Server {
840                error_code: NOT_COORDINATOR,
841            },
842            ClientError::NoCoordinator {
843                key: "streams-group".into(),
844            },
845        ] {
846            check!(should_retry_coordinator_discovery(&error));
847        }
848        check!(!should_retry_coordinator_discovery(&ClientError::Server {
849            error_code: 30,
850        }));
851        check!(should_retry_coordinator_discovery(
852            &ClientError::Disconnected
853        ));
854    }
855
856    #[test]
857    fn heartbeat_interval_uses_positive_broker_value_or_default() {
858        check!(heartbeat_interval(1) == millis(1));
859        check!(heartbeat_interval(3_000) == secs(3));
860        check!(heartbeat_interval(0) == secs(3));
861        check!(heartbeat_interval(-1) == secs(3));
862    }
863
864    #[test]
865    fn should_emit_statuses_only_for_non_empty_status_list() {
866        check!(!should_emit_statuses::<i32>(None));
867        check!(!should_emit_statuses(Some(&Vec::<i32>::new())));
868        check!(should_emit_statuses(Some(&vec![1])));
869    }
870
871    #[tokio::test]
872    async fn accessors_return_membership_identity_and_tracker() {
873        let (_tx, rx) = mpsc::unbounded_channel();
874        let member_epoch = Arc::new(Mutex::new(42));
875        let tracker = Arc::new(Mutex::new(TaskOffsetTracker::default()));
876        let membership = super::StreamsMembership {
877            member_id: "member-1".into(),
878            group_id: "group-1".into(),
879            member_epoch,
880            events: rx,
881            shutdown: CancellationToken::new(),
882            hb_handle: None,
883            tracker: tracker.clone(),
884        };
885
886        check!(membership.member_id() == "member-1");
887        check!(membership.group_id() == "group-1");
888        check!(Arc::ptr_eq(&membership.tracker(), &tracker));
889
890        let meta = membership.group_metadata().await;
891        check!(meta.member == "member-1");
892        check!(meta.group == "group-1");
893        check!(meta.generation == 42);
894        check!(meta.group_instance.is_none());
895    }
896
897    #[test]
898    fn invalid_topology_family_maps() {
899        for code in [130i16, 131, 132] {
900            check!(matches!(
901                map_error(resp(code)),
902                Err(StreamsClientError::InvalidTopology { code: c, .. }) if c == code
903            ));
904        }
905    }
906
907    #[test]
908    fn auth_not_found_and_unknown_codes_map() {
909        check!(matches!(
910            map_error(resp(30)),
911            Err(StreamsClientError::Authorization(30))
912        ));
913        check!(matches!(
914            map_error(resp(29)),
915            Err(StreamsClientError::Authorization(29))
916        ));
917        check!(matches!(
918            map_error(resp(69)),
919            Err(StreamsClientError::GroupIdNotFound)
920        ));
921        check!(matches!(
922            map_error(resp(99)),
923            Err(StreamsClientError::Server(99))
924        ));
925    }
926}
krabka_client_streams/membership/client.rs

krabka_client_streams/membership/
client.rs
