kafrust/
streams.rs

1//! Kafka Streams group membership and heartbeat lifecycle.
2
3use crate::client::Client;
4use crate::config::{ClientConfig, SecurityProtocol};
5use crate::error::{BrokerErrorKind, Error, Result};
6use crate::metrics::ClientMetrics;
7use kafrust_protocol::api::api_versions::ApiVersionsResponseV3;
8use kafrust_protocol::api::find_coordinator::FindCoordinatorResponseV1;
9pub use kafrust_protocol::api::streams_group_heartbeat::{
10    StreamsGroupHeartbeatEndpoint, StreamsGroupHeartbeatEndpointPartitions,
11    StreamsGroupHeartbeatKeyValue, StreamsGroupHeartbeatRequestV0, StreamsGroupHeartbeatResponseV0,
12    StreamsGroupHeartbeatStatus, StreamsGroupHeartbeatSubtopology, StreamsGroupHeartbeatTask,
13    StreamsGroupHeartbeatTaskOffset, StreamsGroupHeartbeatTopic, StreamsGroupHeartbeatTopicConfig,
14    StreamsGroupHeartbeatTopology,
15};
16use std::collections::{BTreeMap, BTreeSet};
17use std::time::Duration;
18use tokio::sync::{mpsc, oneshot, watch};
19use tokio::task::JoinHandle;
20use tokio::time::sleep;
21use tracing::debug;
22
23const STREAMS_GROUP_HEARTBEAT_API_KEY: i16 = 88;
24const DEFAULT_STREAMS_GROUP_MAX_RETRIES: u32 = 5;
25const STREAMS_GROUP_RETRY_BACKOFF: Duration = Duration::from_millis(50);
26const STREAMS_GROUP_MAX_RETRY_BACKOFF: Duration = Duration::from_secs(1);
27const STREAMS_GROUP_COMMAND_CAPACITY: usize = 16;
28
29/// Configuration for a Kafka Streams group membership session.
30///
31/// This API manages the broker-side Streams group protocol. It does not
32/// execute Kafka Streams processors or provide a DSL; applications remain
33/// responsible for processing records and reporting task state.
34pub struct StreamsGroupConfig {
35    client: ClientConfig,
36    group_id: String,
37    topology: StreamsGroupHeartbeatTopology,
38    instance_id: Option<String>,
39    rack_id: Option<String>,
40    rebalance_timeout_ms: i32,
41    process_id: Option<String>,
42    user_endpoint: Option<StreamsGroupHeartbeatEndpoint>,
43    client_tags: Option<Vec<StreamsGroupHeartbeatKeyValue>>,
44    max_retries: u32,
45}
46
47impl StreamsGroupConfig {
48    /// Creates a session configuration with a required initial topology.
49    pub fn new(
50        bootstrap_servers: impl IntoIterator<Item = impl Into<String>>,
51        group_id: impl Into<String>,
52        topology: StreamsGroupHeartbeatTopology,
53    ) -> Self {
54        Self {
55            client: ClientConfig::new(bootstrap_servers),
56            group_id: group_id.into(),
57            topology,
58            instance_id: None,
59            rack_id: None,
60            rebalance_timeout_ms: 30_000,
61            process_id: None,
62            user_endpoint: None,
63            client_tags: None,
64            max_retries: DEFAULT_STREAMS_GROUP_MAX_RETRIES,
65        }
66    }
67
68    /// Replaces the underlying client configuration.
69    ///
70    /// Use this to configure TLS, SASL, bootstrap rotation, decode limits, or
71    /// shared metrics without duplicating those client settings here.
72    pub fn client_config(mut self, client: ClientConfig) -> Self {
73        self.client = client;
74        self
75    }
76
77    /// Replaces the shared client configuration using the common builder
78    /// naming used by the other high-level clients.
79    pub fn with_client_config(self, client: ClientConfig) -> Self {
80        self.client_config(client)
81    }
82
83    /// Sets the Kafka client ID.
84    pub fn client_id(mut self, client_id: impl Into<String>) -> Self {
85        self.client = self.client.client_id(client_id);
86        self
87    }
88
89    /// Sets the static membership instance ID.
90    pub fn instance_id(mut self, instance_id: impl Into<String>) -> Self {
91        self.instance_id = Some(instance_id.into());
92        self
93    }
94
95    /// Sets the rack ID used by the Streams assignment algorithm.
96    pub fn rack_id(mut self, rack_id: impl Into<String>) -> Self {
97        self.rack_id = Some(rack_id.into());
98        self
99    }
100
101    /// Sets the maximum time Kafka may wait for this member during rebalance.
102    pub fn rebalance_timeout_ms(mut self, timeout_ms: i32) -> Self {
103        self.rebalance_timeout_ms = timeout_ms;
104        self
105    }
106
107    /// Sets the Streams process identity used by task assignment.
108    pub fn process_id(mut self, process_id: impl Into<String>) -> Self {
109        self.process_id = Some(process_id.into());
110        self
111    }
112
113    /// Sets the Interactive Queries endpoint advertised by this member.
114    pub fn user_endpoint(mut self, endpoint: StreamsGroupHeartbeatEndpoint) -> Self {
115        self.user_endpoint = Some(endpoint);
116        self
117    }
118
119    /// Sets rack-aware client tags.
120    pub fn client_tags(mut self, tags: Vec<StreamsGroupHeartbeatKeyValue>) -> Self {
121        self.client_tags = Some(tags);
122        self
123    }
124
125    /// Sets the bounded reconnect and rejoin retry count.
126    pub fn max_retries(mut self, max_retries: u32) -> Self {
127        self.max_retries = max_retries;
128        self
129    }
130
131    /// Sets the request timeout.
132    pub fn request_timeout_ms(mut self, timeout_ms: u64) -> Self {
133        self.client = self.client.request_timeout_ms(timeout_ms);
134        self
135    }
136
137    /// Sets the security protocol.
138    pub fn security_protocol(mut self, security_protocol: SecurityProtocol) -> Self {
139        self.client = self.client.security_protocol(security_protocol);
140        self
141    }
142
143    /// Sets the TLS server name used for broker certificate validation.
144    pub fn tls_server_name(mut self, server_name: impl Into<String>) -> Self {
145        self.client = self.client.tls_server_name(server_name);
146        self
147    }
148
149    /// Adds a DER-encoded TLS root certificate for broker validation.
150    pub fn tls_root_certificate_der(mut self, certificate: impl Into<Vec<u8>>) -> Self {
151        self.client = self.client.tls_root_certificate_der(certificate);
152        self
153    }
154
155    /// Adds a DER-encoded client certificate for TLS mutual authentication.
156    pub fn tls_client_certificate_der(mut self, certificate: impl Into<Vec<u8>>) -> Self {
157        self.client = self.client.tls_client_certificate_der(certificate);
158        self
159    }
160
161    /// Sets the DER-encoded private key for TLS mutual authentication.
162    pub fn tls_client_private_key_der(mut self, key: impl Into<Vec<u8>>) -> Self {
163        self.client = self.client.tls_client_private_key_der(key);
164        self
165    }
166
167    /// Sets the shared metrics handle.
168    pub fn metrics(mut self, metrics: ClientMetrics) -> Self {
169        self.client = self.client.metrics(metrics);
170        self
171    }
172
173    /// Returns the configured group ID.
174    pub fn group_id(&self) -> &str {
175        &self.group_id
176    }
177
178    /// Returns the configured topology.
179    pub fn topology(&self) -> &StreamsGroupHeartbeatTopology {
180        &self.topology
181    }
182
183    /// Validates the session configuration without connecting to Kafka.
184    pub fn validate(&self) -> Result<()> {
185        self.client.validate()?;
186        if self.group_id.trim().is_empty() {
187            return Err(Error::InvalidConfiguration {
188                field: "group_id",
189                reason: "must not be empty",
190            });
191        }
192        if self.topology.subtopologies.is_empty() {
193            return Err(Error::InvalidConfiguration {
194                field: "topology.subtopologies",
195                reason: "must contain at least one subtopology",
196            });
197        }
198        if self.rebalance_timeout_ms <= 0 {
199            return Err(Error::InvalidConfiguration {
200                field: "rebalance_timeout_ms",
201                reason: "must be greater than zero",
202            });
203        }
204        Ok(())
205    }
206
207    /// Validates and returns this Streams group configuration without opening
208    /// a broker connection.
209    pub fn build_config(self) -> Result<Self> {
210        self.validate()?;
211        Ok(self)
212    }
213}
214
215/// An active Kafka Streams group membership session.
216///
217/// A session starts by sending API 88 with a client-generated member ID, member
218/// epoch zero, and the configured topology. Subsequent calls send the current
219/// member epoch and only transmit changed membership data according to Kafka's
220/// nullable protocol fields. Call [`Self::close`] to leave gracefully.
221pub struct StreamsGroupSession {
222    config: StreamsGroupConfig,
223    coordinator: Option<Client>,
224    member_id: String,
225    member_epoch: i32,
226    endpoint_information_epoch: i32,
227    heartbeat_interval: Duration,
228    pending_active_tasks: Option<Vec<StreamsGroupHeartbeatTask>>,
229    pending_standby_tasks: Option<Vec<StreamsGroupHeartbeatTask>>,
230    pending_warmup_tasks: Option<Vec<StreamsGroupHeartbeatTask>>,
231    pending_task_offsets: Option<Vec<StreamsGroupHeartbeatTaskOffset>>,
232    pending_task_end_offsets: Option<Vec<StreamsGroupHeartbeatTaskOffset>>,
233    assignment: StreamsGroupSessionAssignment,
234    closed: bool,
235}
236
237enum StreamsGroupCommand {
238    SetTaskState {
239        active_tasks: Vec<StreamsGroupHeartbeatTask>,
240        standby_tasks: Vec<StreamsGroupHeartbeatTask>,
241        warmup_tasks: Vec<StreamsGroupHeartbeatTask>,
242        task_offsets: Option<Vec<StreamsGroupHeartbeatTaskOffset>>,
243        task_end_offsets: Option<Vec<StreamsGroupHeartbeatTaskOffset>>,
244        acknowledged: oneshot::Sender<()>,
245    },
246    Heartbeat {
247        response: oneshot::Sender<Result<StreamsGroupHeartbeatResponseV0>>,
248    },
249    Close,
250}
251
252/// A handle for a Streams session whose heartbeat lifecycle is owned by one
253/// background Tokio task.
254///
255/// The task owns [`StreamsGroupSession`] and is the only code that mutates its
256/// member epoch, coordinator connection, pending task state, and assignment.
257/// Commands use a bounded channel, so task-state updates apply backpressure
258/// instead of growing an unbounded queue. Call [`Self::close`] and await it for
259/// a graceful member-epoch `-1` leave; dropping the handle aborts the task and
260/// cannot provide that broker-side leave guarantee.
261#[must_use = "a StreamsGroupSessionHandle must be closed or kept alive"]
262pub struct StreamsGroupSessionHandle {
263    commands: mpsc::Sender<StreamsGroupCommand>,
264    assignment: watch::Receiver<StreamsGroupSessionAssignment>,
265    task: Option<JoinHandle<Result<()>>>,
266}
267
268struct StreamsHeartbeatPayload {
269    member_epoch: i32,
270    topology: Option<StreamsGroupHeartbeatTopology>,
271    instance_id: Option<String>,
272    rack_id: Option<String>,
273    process_id: Option<String>,
274    user_endpoint: Option<StreamsGroupHeartbeatEndpoint>,
275    client_tags: Option<Vec<StreamsGroupHeartbeatKeyValue>>,
276    active_tasks: Option<Vec<StreamsGroupHeartbeatTask>>,
277    standby_tasks: Option<Vec<StreamsGroupHeartbeatTask>>,
278    warmup_tasks: Option<Vec<StreamsGroupHeartbeatTask>>,
279    task_offsets: Option<Vec<StreamsGroupHeartbeatTaskOffset>>,
280    task_end_offsets: Option<Vec<StreamsGroupHeartbeatTaskOffset>>,
281    shutdown_application: bool,
282}
283
284/// The latest assignment and task-state information returned by Kafka.
285///
286/// The Streams group protocol returns assignment state on every successful
287/// heartbeat. Keeping the last successful snapshot on the session gives an
288/// application a stable reconciliation point after a rebalance or reconnect;
289/// it also preserves nullable broker responses instead of collapsing an
290/// omitted update into an empty assignment.
291#[derive(Debug, Clone, Default, PartialEq, Eq)]
292pub struct StreamsGroupSessionAssignment {
293    /// Broker status entries for the current member, when present.
294    pub status: Option<Vec<StreamsGroupHeartbeatStatus>>,
295    /// Active tasks assigned to this member, if the broker sent the field.
296    pub active_tasks: Option<Vec<StreamsGroupHeartbeatTask>>,
297    /// Standby tasks assigned to this member, if the broker sent the field.
298    pub standby_tasks: Option<Vec<StreamsGroupHeartbeatTask>>,
299    /// Warmup tasks assigned to this member, if the broker sent the field.
300    pub warmup_tasks: Option<Vec<StreamsGroupHeartbeatTask>>,
301    /// Broker-provided recovery lag threshold for the current assignment.
302    pub acceptable_recovery_lag: i32,
303    /// Broker-provided interval for reporting task offsets.
304    pub task_offset_interval_ms: i32,
305    /// Interactive Queries endpoint partition assignments, if present.
306    pub partitions_by_user_endpoint: Option<Vec<StreamsGroupHeartbeatEndpointPartitions>>,
307}
308
309impl StreamsGroupSessionAssignment {
310    fn from_response(response: &StreamsGroupHeartbeatResponseV0) -> Self {
311        Self {
312            status: response.status.clone(),
313            active_tasks: response.active_tasks.clone(),
314            standby_tasks: response.standby_tasks.clone(),
315            warmup_tasks: response.warmup_tasks.clone(),
316            acceptable_recovery_lag: response.acceptable_recovery_lag,
317            task_offset_interval_ms: response.task_offset_interval_ms,
318            partitions_by_user_endpoint: response.partitions_by_user_endpoint.clone(),
319        }
320    }
321}
322
323/// The role a Kafka Streams task has on this member.
324#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
325pub enum StreamsTaskRole {
326    /// The task actively processes its input partitions.
327    Active,
328    /// The task maintains standby state without processing records actively.
329    Standby,
330    /// The task is warming up before it can become active.
331    Warmup,
332}
333
334impl StreamsTaskRole {
335    /// Returns the stable role name used in diagnostics and logs.
336    pub const fn as_str(self) -> &'static str {
337        match self {
338            Self::Active => "active",
339            Self::Standby => "standby",
340            Self::Warmup => "warmup",
341        }
342    }
343}
344
345/// The canonical identity of one Kafka Streams task assignment.
346///
347/// Kafka identifies the task by its subtopology and the input partitions it
348/// processes. Partition order is not semantic, so this type stores partitions
349/// in sorted order and rejects duplicates or negative partition indexes.
350#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
351pub struct StreamsTaskId {
352    subtopology_id: String,
353    partitions: Vec<i32>,
354}
355
356impl StreamsTaskId {
357    /// Creates a canonical task identity from a subtopology and partitions.
358    pub fn new(subtopology_id: impl Into<String>, mut partitions: Vec<i32>) -> Result<Self> {
359        let subtopology_id = subtopology_id.into();
360        if subtopology_id.trim().is_empty() {
361            return Err(Error::StreamsTaskAssignmentInvalid {
362                subtopology_id,
363                reason: "subtopology ID must not be empty",
364            });
365        }
366        if partitions.is_empty() {
367            return Err(Error::StreamsTaskAssignmentInvalid {
368                subtopology_id,
369                reason: "task must contain at least one partition",
370            });
371        }
372        if partitions.iter().any(|partition| *partition < 0) {
373            return Err(Error::StreamsTaskAssignmentInvalid {
374                subtopology_id,
375                reason: "task partitions must not be negative",
376            });
377        }
378        partitions.sort_unstable();
379        if partitions.windows(2).any(|window| window[0] == window[1]) {
380            return Err(Error::StreamsTaskAssignmentInvalid {
381                subtopology_id,
382                reason: "task partitions must not contain duplicates",
383            });
384        }
385        Ok(Self {
386            subtopology_id,
387            partitions,
388        })
389    }
390
391    fn from_heartbeat_task(task: &StreamsGroupHeartbeatTask) -> Result<Self> {
392        Self::new(task.subtopology_id.clone(), task.partitions.clone())
393    }
394
395    /// Returns the subtopology identifier.
396    pub fn subtopology_id(&self) -> &str {
397        &self.subtopology_id
398    }
399
400    /// Returns the canonical, sorted input partitions.
401    pub fn partitions(&self) -> &[i32] {
402        &self.partitions
403    }
404}
405
406/// A task and its currently reconciled role.
407#[derive(Debug, Clone, PartialEq, Eq)]
408pub struct StreamsTaskAssignment {
409    /// Canonical task identity.
410    pub task: StreamsTaskId,
411    /// Role currently held by the task.
412    pub role: StreamsTaskRole,
413}
414
415/// A deterministic lifecycle change produced by [`StreamsTaskRuntime`].
416#[derive(Debug, Clone, PartialEq, Eq)]
417pub enum StreamsTaskTransition {
418    /// A task became assigned to this member.
419    Added {
420        /// Canonical task identity.
421        task: StreamsTaskId,
422        /// New task role.
423        role: StreamsTaskRole,
424    },
425    /// A task was removed from this member.
426    Removed {
427        /// Canonical task identity.
428        task: StreamsTaskId,
429        /// Role held before removal.
430        role: StreamsTaskRole,
431    },
432    /// A task stayed on this member but changed role.
433    RoleChanged {
434        /// Canonical task identity.
435        task: StreamsTaskId,
436        /// Role held before the assignment update.
437        previous_role: StreamsTaskRole,
438        /// Role after the assignment update.
439        role: StreamsTaskRole,
440    },
441}
442
443/// Bounded, deterministic reconciliation state for Streams task assignments.
444///
445/// This type intentionally stops at assignment lifecycle. It does not spawn
446/// processors, own consumer assignments, or manage state stores. Applications
447/// can apply the returned transitions to those components while keeping Kafka's
448/// nullable response semantics and task identity rules in one place.
449#[derive(Debug, Default)]
450pub struct StreamsTaskRuntime {
451    tasks: BTreeMap<StreamsTaskId, StreamsTaskRole>,
452}
453
454impl StreamsTaskRuntime {
455    /// Creates an empty task runtime.
456    pub fn new() -> Self {
457        Self::default()
458    }
459
460    /// Returns the current task assignment in canonical order.
461    pub fn assignment(&self) -> Vec<StreamsTaskAssignment> {
462        self.tasks
463            .iter()
464            .map(|(task, role)| StreamsTaskAssignment {
465                task: task.clone(),
466                role: *role,
467            })
468            .collect()
469    }
470
471    /// Returns the current role for a task, if it is assigned.
472    pub fn role(&self, task: &StreamsTaskId) -> Option<StreamsTaskRole> {
473        self.tasks.get(task).copied()
474    }
475
476    /// Applies the changed task-role fields from a broker assignment.
477    ///
478    /// A `None` role field means that Kafka did not change that role since the
479    /// previous heartbeat and is therefore retained. `Some(Vec::new())` is an
480    /// explicit revocation of that role. The returned transitions are sorted by
481    /// canonical task identity and the runtime is left unchanged if validation
482    /// fails.
483    pub fn reconcile_assignment(
484        &mut self,
485        assignment: &StreamsGroupSessionAssignment,
486    ) -> Result<Vec<StreamsTaskTransition>> {
487        let mut desired = self.tasks.clone();
488        let updates = [
489            (StreamsTaskRole::Active, assignment.active_tasks.as_ref()),
490            (StreamsTaskRole::Standby, assignment.standby_tasks.as_ref()),
491            (StreamsTaskRole::Warmup, assignment.warmup_tasks.as_ref()),
492        ];
493
494        for (role, tasks) in updates {
495            if tasks.is_some() {
496                desired.retain(|_, current_role| *current_role != role);
497            }
498        }
499
500        let mut occupied = BTreeSet::new();
501        for task in desired.keys() {
502            for partition in task.partitions() {
503                if !occupied.insert((task.subtopology_id().to_owned(), *partition)) {
504                    return Err(Error::StreamsTaskAssignmentConflict {
505                        subtopology_id: task.subtopology_id().to_owned(),
506                        partition: *partition,
507                    });
508                }
509            }
510        }
511
512        for (role, tasks) in updates {
513            let Some(tasks) = tasks else {
514                continue;
515            };
516            for task in tasks {
517                let task = StreamsTaskId::from_heartbeat_task(task)?;
518                for partition in task.partitions() {
519                    if !occupied.insert((task.subtopology_id().to_owned(), *partition)) {
520                        return Err(Error::StreamsTaskAssignmentConflict {
521                            subtopology_id: task.subtopology_id().to_owned(),
522                            partition: *partition,
523                        });
524                    }
525                }
526                if desired.insert(task.clone(), role).is_some() {
527                    return Err(Error::StreamsTaskAssignmentConflict {
528                        subtopology_id: task.subtopology_id().to_owned(),
529                        partition: task.partitions()[0],
530                    });
531                }
532            }
533        }
534
535        let mut transitions = Vec::new();
536        for (task, previous_role) in &self.tasks {
537            match desired.get(task) {
538                None => transitions.push(StreamsTaskTransition::Removed {
539                    task: task.clone(),
540                    role: *previous_role,
541                }),
542                Some(role) if role != previous_role => {
543                    transitions.push(StreamsTaskTransition::RoleChanged {
544                        task: task.clone(),
545                        previous_role: *previous_role,
546                        role: *role,
547                    });
548                }
549                Some(_) => {}
550            }
551        }
552        for (task, role) in &desired {
553            if !self.tasks.contains_key(task) {
554                transitions.push(StreamsTaskTransition::Added {
555                    task: task.clone(),
556                    role: *role,
557                });
558            }
559        }
560        self.tasks = desired;
561        Ok(transitions)
562    }
563}
564
565impl StreamsGroupSessionHandle {
566    /// Reconciles the latest broker assignment into an application-owned task runtime.
567    pub fn reconcile_task_runtime(
568        &self,
569        runtime: &mut StreamsTaskRuntime,
570    ) -> Result<Vec<StreamsTaskTransition>> {
571        runtime.reconcile_assignment(&self.assignment())
572    }
573}
574
575impl StreamsGroupSession {
576    /// Joins the configured Kafka Streams group.
577    pub async fn join(config: StreamsGroupConfig) -> Result<Self> {
578        config.validate()?;
579        let mut session = Self {
580            config,
581            coordinator: None,
582            member_id: new_streams_member_id(),
583            member_epoch: 0,
584            endpoint_information_epoch: 0,
585            heartbeat_interval: Duration::from_secs(1),
586            pending_active_tasks: None,
587            pending_standby_tasks: None,
588            pending_warmup_tasks: None,
589            pending_task_offsets: None,
590            pending_task_end_offsets: None,
591            assignment: StreamsGroupSessionAssignment::default(),
592            closed: false,
593        };
594        session.join_with_retry().await?;
595        Ok(session)
596    }
597
598    /// Returns the Kafka Streams group ID.
599    pub fn group_id(&self) -> &str {
600        self.config.group_id()
601    }
602
603    /// Returns the client-generated member ID used for this Streams session.
604    pub fn member_id(&self) -> &str {
605        &self.member_id
606    }
607
608    /// Returns the current member epoch.
609    pub fn member_epoch(&self) -> i32 {
610        self.member_epoch
611    }
612
613    /// Returns the broker-requested heartbeat interval.
614    pub fn heartbeat_interval(&self) -> Duration {
615        self.heartbeat_interval
616    }
617
618    /// Returns the latest successful assignment snapshot.
619    ///
620    /// A `None` collection means that Kafka omitted the nullable field in its
621    /// response; it is distinct from `Some(Vec::new())`, which is an explicit
622    /// empty assignment.
623    pub fn assignment(&self) -> &StreamsGroupSessionAssignment {
624        &self.assignment
625    }
626
627    /// Replaces the task state and changelog offsets reported on the next
628    /// heartbeat.
629    ///
630    /// Empty vectors are meaningful and are sent as empty compact arrays. The
631    /// values remain pending until a heartbeat succeeds, so a reconnect does
632    /// not silently lose a task-state update.
633    pub fn set_task_state(
634        &mut self,
635        active_tasks: Vec<StreamsGroupHeartbeatTask>,
636        standby_tasks: Vec<StreamsGroupHeartbeatTask>,
637        warmup_tasks: Vec<StreamsGroupHeartbeatTask>,
638        task_offsets: Vec<StreamsGroupHeartbeatTaskOffset>,
639        task_end_offsets: Vec<StreamsGroupHeartbeatTaskOffset>,
640    ) {
641        self.set_task_state_with_optional_offsets(
642            active_tasks,
643            standby_tasks,
644            warmup_tasks,
645            Some(task_offsets),
646            Some(task_end_offsets),
647        );
648    }
649
650    /// Replaces task state while optionally omitting changelog offsets.
651    ///
652    /// Kafka 4.3 currently rejects non-null task-offset fields for some Streams
653    /// group configurations. Passing `None` preserves the protocol's nullable
654    /// field semantics and lets the broker request offsets when supported.
655    pub fn set_task_state_with_optional_offsets(
656        &mut self,
657        active_tasks: Vec<StreamsGroupHeartbeatTask>,
658        standby_tasks: Vec<StreamsGroupHeartbeatTask>,
659        warmup_tasks: Vec<StreamsGroupHeartbeatTask>,
660        task_offsets: Option<Vec<StreamsGroupHeartbeatTaskOffset>>,
661        task_end_offsets: Option<Vec<StreamsGroupHeartbeatTaskOffset>>,
662    ) {
663        self.pending_active_tasks = Some(active_tasks);
664        self.pending_standby_tasks = Some(standby_tasks);
665        self.pending_warmup_tasks = Some(warmup_tasks);
666        self.pending_task_offsets = task_offsets;
667        self.pending_task_end_offsets = task_end_offsets;
668    }
669
670    /// Returns whether the session has left the group.
671    pub fn is_closed(&self) -> bool {
672        self.closed
673    }
674
675    /// Sends one Streams heartbeat and returns the broker's assignment state.
676    ///
677    /// Coordinator transport failures and member/coordinator epoch errors are
678    /// recovered within the configured retry budget. A rejoin resends the
679    /// initial topology and updates the session's member identity.
680    pub async fn heartbeat(&mut self) -> Result<StreamsGroupHeartbeatResponseV0> {
681        self.ensure_open()?;
682        let mut retry = 0;
683        loop {
684            match self.send_heartbeat(false).await {
685                Ok(response) => return Ok(response),
686                Err(error)
687                    if retry < self.config.max_retries && is_retryable_streams_error(&error) =>
688                {
689                    retry += 1;
690                    self.config.client.record_retry();
691                    let rejoin = should_rejoin_streams_group(&error);
692                    self.recover_streams_session(rejoin, &mut retry).await?;
693                    sleep(streams_retry_backoff(retry)).await;
694                }
695                Err(error) => return Err(error),
696            }
697        }
698    }
699
700    /// Leaves the Streams group using member epoch `-1`.
701    pub async fn close(&mut self) -> Result<()> {
702        if self.closed {
703            return Ok(());
704        }
705        let mut retry = 0;
706        loop {
707            self.ensure_open()?;
708            match self.send_heartbeat(true).await {
709                Ok(_) => {
710                    self.closed = true;
711                    self.coordinator = None;
712                    return Ok(());
713                }
714                Err(error)
715                    if retry < self.config.max_retries && is_retryable_streams_error(&error) =>
716                {
717                    retry += 1;
718                    self.config.client.record_retry();
719                    self.recover_streams_session(false, &mut retry).await?;
720                    sleep(streams_retry_backoff(retry)).await;
721                }
722                Err(error) => return Err(error),
723            }
724        }
725    }
726
727    /// Moves this session into a background heartbeat task.
728    ///
729    /// The session must already be joined. The task sends heartbeats using the
730    /// interval most recently returned by Kafka, publishes every successful
731    /// assignment through [`StreamsGroupSessionHandle::subscribe_assignment`],
732    /// and preserves the existing bounded retry and rejoin behavior.
733    pub fn spawn_heartbeat_task(self) -> StreamsGroupSessionHandle {
734        let (commands, command_receiver) = mpsc::channel(STREAMS_GROUP_COMMAND_CAPACITY);
735        let (assignment_sender, assignment_receiver) = watch::channel(self.assignment.clone());
736        let task = tokio::spawn(run_streams_heartbeat_task(
737            self,
738            command_receiver,
739            assignment_sender,
740        ));
741        StreamsGroupSessionHandle {
742            commands,
743            assignment: assignment_receiver,
744            task: Some(task),
745        }
746    }
747
748    async fn join_with_retry(&mut self) -> Result<()> {
749        let mut retry = 0;
750        self.recover_streams_session(true, &mut retry).await
751    }
752
753    async fn recover_streams_session(&mut self, rejoin: bool, retry: &mut u32) -> Result<()> {
754        loop {
755            let result = async {
756                self.reconnect().await?;
757                if rejoin {
758                    if self.member_epoch != 0 {
759                        self.member_id = new_streams_member_id();
760                    }
761                    self.member_epoch = 0;
762                    self.send_initial_heartbeat().await?;
763                }
764                Ok::<(), Error>(())
765            }
766            .await;
767
768            match result {
769                Ok(()) => return Ok(()),
770                Err(error)
771                    if *retry < self.config.max_retries && is_retryable_streams_error(&error) =>
772                {
773                    *retry += 1;
774                    self.config.client.record_retry();
775                    sleep(streams_retry_backoff(*retry)).await;
776                }
777                Err(error) => return Err(error),
778            }
779        }
780    }
781
782    async fn reconnect(&mut self) -> Result<()> {
783        self.coordinator = None;
784        let mut bootstrap = self.config.client.clone().connect().await?;
785        let coordinator = bootstrap
786            .find_group_coordinator(self.config.group_id.clone())
787            .await?;
788        if coordinator.error_code != 0 {
789            return Err(self.config.client.broker_error(
790                coordinator.error_code,
791                format!("find Streams group coordinator {}", self.config.group_id),
792            ));
793        }
794        let address = coordinator_addr(&coordinator);
795        let mut client = self.config.client.connect_broker(address).await?;
796        let versions = client
797            .api_versions_v3_cached("kafrust", env!("CARGO_PKG_VERSION"))
798            .await?;
799        ensure_streams_heartbeat_supported(&versions)?;
800        self.coordinator = Some(client);
801        Ok(())
802    }
803
804    async fn send_initial_heartbeat(&mut self) -> Result<StreamsGroupHeartbeatResponseV0> {
805        let response = self
806            .send_heartbeat_request(StreamsHeartbeatPayload {
807                member_epoch: 0,
808                topology: Some(self.config.topology.clone()),
809                instance_id: self.config.instance_id.clone(),
810                rack_id: self.config.rack_id.clone(),
811                process_id: self.config.process_id.clone(),
812                user_endpoint: self.config.user_endpoint.clone(),
813                client_tags: self.config.client_tags.clone(),
814                active_tasks: Some(Vec::new()),
815                standby_tasks: Some(Vec::new()),
816                warmup_tasks: Some(Vec::new()),
817                task_offsets: None,
818                task_end_offsets: None,
819                shutdown_application: false,
820            })
821            .await?;
822        self.apply_response(&response)?;
823        debug!(
824            group_id = self.config.group_id.as_str(),
825            member_id = self.member_id.as_str(),
826            member_epoch = self.member_epoch,
827            "joined Kafka Streams group"
828        );
829        Ok(response)
830    }
831
832    async fn send_heartbeat(&mut self, leave: bool) -> Result<StreamsGroupHeartbeatResponseV0> {
833        let member_epoch = if leave { -1 } else { self.member_epoch };
834        let active_tasks = if leave {
835            None
836        } else {
837            self.pending_active_tasks.clone()
838        };
839        let standby_tasks = if leave {
840            None
841        } else {
842            self.pending_standby_tasks.clone()
843        };
844        let warmup_tasks = if leave {
845            None
846        } else {
847            self.pending_warmup_tasks.clone()
848        };
849        let task_offsets = if leave {
850            None
851        } else {
852            self.pending_task_offsets.clone()
853        };
854        let task_end_offsets = if leave {
855            None
856        } else {
857            self.pending_task_end_offsets.clone()
858        };
859        let response = self
860            .send_heartbeat_request(StreamsHeartbeatPayload {
861                member_epoch,
862                topology: None,
863                instance_id: None,
864                rack_id: None,
865                process_id: None,
866                user_endpoint: None,
867                client_tags: None,
868                active_tasks,
869                standby_tasks,
870                warmup_tasks,
871                task_offsets,
872                task_end_offsets,
873                shutdown_application: leave,
874            })
875            .await?;
876        self.apply_response(&response)?;
877        if !leave {
878            self.pending_active_tasks = None;
879            self.pending_standby_tasks = None;
880            self.pending_warmup_tasks = None;
881            self.pending_task_offsets = None;
882            self.pending_task_end_offsets = None;
883        }
884        Ok(response)
885    }
886
887    async fn send_heartbeat_request(
888        &mut self,
889        payload: StreamsHeartbeatPayload,
890    ) -> Result<StreamsGroupHeartbeatResponseV0> {
891        let coordinator = self.coordinator.as_mut().ok_or(Error::Unsupported(
892            "Streams group coordinator is not connected",
893        ))?;
894        coordinator
895            .streams_group_heartbeat_v0(
896                self.config.group_id.clone(),
897                self.member_id.clone(),
898                payload.member_epoch,
899                self.endpoint_information_epoch,
900                payload.instance_id,
901                payload.rack_id,
902                if payload.member_epoch == -1 {
903                    -1
904                } else {
905                    self.config.rebalance_timeout_ms
906                },
907                payload.topology,
908                payload.active_tasks,
909                payload.standby_tasks,
910                payload.warmup_tasks,
911                payload.process_id,
912                payload.user_endpoint,
913                payload.client_tags,
914                payload.task_offsets,
915                payload.task_end_offsets,
916                payload.shutdown_application,
917            )
918            .await
919    }
920
921    fn apply_response(&mut self, response: &StreamsGroupHeartbeatResponseV0) -> Result<()> {
922        if response.error_code != 0 {
923            return Err(self.config.client.broker_error(
924                response.error_code,
925                format!(
926                    "Streams group heartbeat {}: {}",
927                    self.config.group_id,
928                    response
929                        .error_message
930                        .as_deref()
931                        .unwrap_or("broker returned a Streams group error")
932                ),
933            ));
934        }
935        self.member_id = response.member_id.clone();
936        self.member_epoch = response.member_epoch;
937        self.endpoint_information_epoch = response.endpoint_information_epoch;
938        self.assignment = StreamsGroupSessionAssignment::from_response(response);
939        if response.heartbeat_interval_ms > 0 {
940            self.heartbeat_interval = Duration::from_millis(
941                u64::try_from(response.heartbeat_interval_ms).unwrap_or(u64::MAX),
942            );
943        }
944        Ok(())
945    }
946
947    fn ensure_open(&self) -> Result<()> {
948        if self.closed {
949            return Err(Error::Unsupported("Streams group session is closed"));
950        }
951        Ok(())
952    }
953}
954
955impl StreamsGroupSessionHandle {
956    /// Returns the latest successful broker assignment snapshot.
957    pub fn assignment(&self) -> StreamsGroupSessionAssignment {
958        self.assignment.borrow().clone()
959    }
960
961    /// Subscribes to successful assignment snapshots produced by the heartbeat
962    /// task. The receiver always retains the newest snapshot when the caller is
963    /// temporarily slower than the heartbeat interval.
964    pub fn subscribe_assignment(&self) -> watch::Receiver<StreamsGroupSessionAssignment> {
965        self.assignment.clone()
966    }
967
968    /// Replaces the task state reported by the next background heartbeat.
969    pub async fn set_task_state(
970        &self,
971        active_tasks: Vec<StreamsGroupHeartbeatTask>,
972        standby_tasks: Vec<StreamsGroupHeartbeatTask>,
973        warmup_tasks: Vec<StreamsGroupHeartbeatTask>,
974        task_offsets: Option<Vec<StreamsGroupHeartbeatTaskOffset>>,
975        task_end_offsets: Option<Vec<StreamsGroupHeartbeatTaskOffset>>,
976    ) -> Result<()> {
977        let (acknowledged, completion) = oneshot::channel();
978        self.commands
979            .send(StreamsGroupCommand::SetTaskState {
980                active_tasks,
981                standby_tasks,
982                warmup_tasks,
983                task_offsets,
984                task_end_offsets,
985                acknowledged,
986            })
987            .await
988            .map_err(|_| Error::StreamsGroupBackgroundTaskClosed)?;
989        completion
990            .await
991            .map_err(|_| Error::StreamsGroupBackgroundTaskClosed)
992    }
993
994    /// Sends a heartbeat immediately instead of waiting for the broker's
995    /// advertised interval and returns its response.
996    pub async fn heartbeat_now(&self) -> Result<StreamsGroupHeartbeatResponseV0> {
997        let (response, completion) = oneshot::channel();
998        self.commands
999            .send(StreamsGroupCommand::Heartbeat { response })
1000            .await
1001            .map_err(|_| Error::StreamsGroupBackgroundTaskClosed)?;
1002        completion
1003            .await
1004            .map_err(|_| Error::StreamsGroupBackgroundTaskClosed)?
1005    }
1006
1007    /// Gracefully leaves the Streams group and waits for the background task to
1008    /// finish. This consumes the handle so no later command can race with the
1009    /// leave operation.
1010    pub async fn close(mut self) -> Result<()> {
1011        let Some(task) = self.task.take() else {
1012            return Ok(());
1013        };
1014        if self
1015            .commands
1016            .send(StreamsGroupCommand::Close)
1017            .await
1018            .is_err()
1019        {
1020            return task.await.map_err(Error::from)?;
1021        }
1022        task.await.map_err(Error::from)?
1023    }
1024}
1025
1026impl Drop for StreamsGroupSessionHandle {
1027    fn drop(&mut self) {
1028        if let Some(task) = self.task.as_ref() {
1029            task.abort();
1030        }
1031    }
1032}
1033
1034async fn run_streams_heartbeat_task(
1035    mut session: StreamsGroupSession,
1036    mut commands: mpsc::Receiver<StreamsGroupCommand>,
1037    assignment_sender: watch::Sender<StreamsGroupSessionAssignment>,
1038) -> Result<()> {
1039    loop {
1040        let heartbeat = sleep(session.heartbeat_interval());
1041        tokio::pin!(heartbeat);
1042        tokio::select! {
1043            command = commands.recv() => match command {
1044                Some(StreamsGroupCommand::SetTaskState {
1045                    active_tasks,
1046                    standby_tasks,
1047                    warmup_tasks,
1048                    task_offsets,
1049                    task_end_offsets,
1050                    acknowledged,
1051                }) => {
1052                    session.set_task_state_with_optional_offsets(
1053                        active_tasks,
1054                        standby_tasks,
1055                        warmup_tasks,
1056                        task_offsets,
1057                        task_end_offsets,
1058                    );
1059                    let _ = acknowledged.send(());
1060                }
1061                Some(StreamsGroupCommand::Heartbeat { response }) => {
1062                    let result = session.heartbeat().await;
1063                    if result.is_ok() {
1064                        let _ = assignment_sender.send(session.assignment.clone());
1065                    }
1066                    let _ = response.send(result);
1067                }
1068                Some(StreamsGroupCommand::Close) => return session.close().await,
1069                None => return Ok(()),
1070            },
1071            _ = &mut heartbeat => {
1072                session.heartbeat().await?;
1073                let _ = assignment_sender.send(session.assignment.clone());
1074            }
1075        }
1076    }
1077}
1078
1079fn ensure_streams_heartbeat_supported(versions: &ApiVersionsResponseV3) -> Result<()> {
1080    if versions
1081        .highest_supported_version(STREAMS_GROUP_HEARTBEAT_API_KEY, 0)
1082        .is_none()
1083    {
1084        return Err(Error::Unsupported(
1085            "broker does not advertise StreamsGroupHeartbeat v0",
1086        ));
1087    }
1088    Ok(())
1089}
1090
1091fn coordinator_addr(coordinator: &FindCoordinatorResponseV1) -> String {
1092    format!("{}:{}", coordinator.host, coordinator.port)
1093}
1094
1095fn new_streams_member_id() -> String {
1096    use rand::RngCore;
1097
1098    let mut bytes = [0_u8; 16];
1099    rand::thread_rng().fill_bytes(&mut bytes);
1100    bytes[6] = (bytes[6] & 0x0f) | 0x40;
1101    bytes[8] = (bytes[8] & 0x3f) | 0x80;
1102    let hex = bytes
1103        .iter()
1104        .fold(String::with_capacity(32), |mut hex, byte| {
1105            use std::fmt::Write;
1106
1107            let _ = write!(hex, "{byte:02x}");
1108            hex
1109        });
1110    format!(
1111        "{}-{}-{}-{}-{}",
1112        &hex[..8],
1113        &hex[8..12],
1114        &hex[12..16],
1115        &hex[16..20],
1116        &hex[20..]
1117    )
1118}
1119
1120fn is_retryable_streams_error(error: &Error) -> bool {
1121    if matches!(error, Error::Io(_)) {
1122        return true;
1123    }
1124    matches!(
1125        error.broker_error_kind(),
1126        Some(
1127            BrokerErrorKind::CoordinatorLoadInProgress
1128                | BrokerErrorKind::CoordinatorNotAvailable
1129                | BrokerErrorKind::NotCoordinator
1130                | BrokerErrorKind::RebalanceInProgress
1131                | BrokerErrorKind::UnknownMemberId
1132                | BrokerErrorKind::FencedMemberEpoch
1133                | BrokerErrorKind::StaleMemberEpoch
1134        )
1135    )
1136}
1137
1138fn should_rejoin_streams_group(error: &Error) -> bool {
1139    matches!(
1140        error.broker_error_kind(),
1141        Some(
1142            BrokerErrorKind::UnknownMemberId
1143                | BrokerErrorKind::FencedMemberEpoch
1144                | BrokerErrorKind::StaleMemberEpoch
1145        )
1146    )
1147}
1148
1149fn streams_retry_backoff(attempt: u32) -> Duration {
1150    let factor = u32::pow(2, attempt.saturating_sub(1).min(5));
1151    STREAMS_GROUP_RETRY_BACKOFF
1152        .checked_mul(factor)
1153        .unwrap_or(STREAMS_GROUP_MAX_RETRY_BACKOFF)
1154        .min(STREAMS_GROUP_MAX_RETRY_BACKOFF)
1155}
1156
1157#[cfg(test)]
1158#[allow(clippy::unwrap_used)]
1159mod tests {
1160    use super::{
1161        streams_retry_backoff, StreamsGroupConfig, StreamsGroupHeartbeatEndpoint,
1162        StreamsGroupHeartbeatKeyValue, StreamsGroupHeartbeatTask, StreamsGroupHeartbeatTaskOffset,
1163        StreamsGroupHeartbeatTopology, StreamsGroupSession, StreamsGroupSessionAssignment,
1164        StreamsTaskRole, StreamsTaskRuntime, StreamsTaskTransition,
1165        STREAMS_GROUP_MAX_RETRY_BACKOFF,
1166    };
1167    use crate::client::Client;
1168    use crate::error::Error;
1169    use kafrust_protocol::api::streams_group_heartbeat::StreamsGroupHeartbeatSubtopology;
1170    use kafrust_protocol::codec::{Decoder, Encoder};
1171    use std::time::Duration;
1172    use tokio::io::{AsyncReadExt, AsyncWriteExt};
1173
1174    fn topology() -> StreamsGroupHeartbeatTopology {
1175        StreamsGroupHeartbeatTopology {
1176            epoch: 1,
1177            subtopologies: vec![StreamsGroupHeartbeatSubtopology {
1178                subtopology_id: "subtopology-0".to_owned(),
1179                source_topics: vec!["orders".to_owned()],
1180                source_topic_regex: Vec::new(),
1181                state_changelog_topics: Vec::new(),
1182                repartition_sink_topics: Vec::new(),
1183                repartition_source_topics: Vec::new(),
1184                copartition_groups: Vec::new(),
1185            }],
1186        }
1187    }
1188
1189    trait DecoderLengthExt {
1190        fn read_compact_array_length(&mut self) -> Option<usize>;
1191    }
1192
1193    impl<'a> DecoderLengthExt for Decoder<'a> {
1194        fn read_compact_array_length(&mut self) -> Option<usize> {
1195            match self.read_unsigned_varint().unwrap() {
1196                0 => None,
1197                length => Some(usize::try_from(length - 1).unwrap()),
1198            }
1199        }
1200    }
1201
1202    #[test]
1203    fn validates_required_streams_topology() {
1204        let config = StreamsGroupConfig::new(["localhost:9092"], "orders", topology());
1205        assert_eq!(config.group_id(), "orders");
1206        assert!(config.validate().is_ok());
1207    }
1208
1209    #[test]
1210    fn rejects_empty_streams_topology() {
1211        let config = StreamsGroupConfig::new(
1212            ["localhost:9092"],
1213            "orders",
1214            StreamsGroupHeartbeatTopology {
1215                epoch: 1,
1216                subtopologies: Vec::new(),
1217            },
1218        );
1219        assert!(config.validate().is_err());
1220    }
1221
1222    #[test]
1223    fn caps_streams_retry_backoff() {
1224        assert_eq!(streams_retry_backoff(100), STREAMS_GROUP_MAX_RETRY_BACKOFF);
1225        assert_eq!(streams_retry_backoff(1), Duration::from_millis(50));
1226    }
1227
1228    #[test]
1229    fn generates_uuid_shaped_streams_member_ids() {
1230        let member_id = super::new_streams_member_id();
1231        assert_eq!(member_id.len(), 36);
1232        assert_eq!(&member_id[8..9], "-");
1233        assert_eq!(&member_id[13..14], "-");
1234        assert_eq!(&member_id[18..19], "-");
1235        assert_eq!(&member_id[23..24], "-");
1236        assert_eq!(&member_id[14..15], "4");
1237        assert!(matches!(&member_id[19..20], "8" | "9" | "a" | "b"));
1238    }
1239
1240    #[tokio::test]
1241    async fn preserves_streams_group_lifecycle_on_the_wire() {
1242        let (client_stream, mut broker_stream) = tokio::io::duplex(16 * 1024);
1243        let broker = tokio::spawn(async move {
1244            let request = read_frame(&mut broker_stream).await;
1245            let (correlation_id, mut decoder) = request_header(&request);
1246            assert_eq!(correlation_id, 1);
1247            assert_eq!(decoder.read_compact_string().unwrap(), "orders");
1248            assert_eq!(decoder.read_compact_string().unwrap(), "client-member");
1249            assert_eq!(decoder.read_i32().unwrap(), 0);
1250            assert_eq!(decoder.read_i32().unwrap(), 0);
1251            assert_eq!(
1252                decoder.read_compact_nullable_string().unwrap(),
1253                Some("instance-a".to_owned())
1254            );
1255            assert_eq!(
1256                decoder.read_compact_nullable_string().unwrap(),
1257                Some("rack-a".to_owned())
1258            );
1259            assert_eq!(decoder.read_i32().unwrap(), 30_000);
1260            assert_eq!(decoder.read_i8().unwrap(), 1);
1261            assert_eq!(decoder.read_i32().unwrap(), 1);
1262            assert_eq!(decoder.read_compact_array_length(), Some(1));
1263            assert_eq!(decoder.read_compact_string().unwrap(), "subtopology-0");
1264            assert_eq!(decoder.read_compact_array_length(), Some(1));
1265            assert_eq!(decoder.read_compact_string().unwrap(), "orders");
1266            assert_eq!(decoder.read_compact_array_length(), Some(0));
1267            assert_eq!(decoder.read_compact_array_length(), Some(0));
1268            assert_eq!(decoder.read_compact_array_length(), Some(0));
1269            assert_eq!(decoder.read_compact_array_length(), Some(0));
1270            assert_eq!(decoder.read_compact_array_length(), Some(0));
1271            decoder.read_tagged_fields().unwrap();
1272            decoder.read_tagged_fields().unwrap();
1273            assert_eq!(decoder.read_compact_array_length(), Some(0));
1274            assert_eq!(decoder.read_compact_array_length(), Some(0));
1275            assert_eq!(decoder.read_compact_array_length(), Some(0));
1276            assert_eq!(
1277                decoder.read_compact_nullable_string().unwrap(),
1278                Some("process-a".to_owned())
1279            );
1280            assert_eq!(decoder.read_i8().unwrap(), 1);
1281            assert_eq!(decoder.read_compact_string().unwrap(), "query-host");
1282            assert_eq!(decoder.read_i16().unwrap(), 7_777);
1283            decoder.read_tagged_fields().unwrap();
1284            assert_eq!(decoder.read_compact_array_length(), Some(1));
1285            assert_eq!(decoder.read_compact_string().unwrap(), "zone");
1286            assert_eq!(decoder.read_compact_string().unwrap(), "a");
1287            decoder.read_tagged_fields().unwrap();
1288            assert_eq!(decoder.read_compact_array_length(), None);
1289            assert_eq!(decoder.read_compact_array_length(), None);
1290            assert!(!decoder.read_bool().unwrap());
1291            decoder.read_tagged_fields().unwrap();
1292            write_streams_response(&mut broker_stream, correlation_id, "member-1", 2, 4).await;
1293
1294            let request = read_frame(&mut broker_stream).await;
1295            let (correlation_id, mut decoder) = request_header(&request);
1296            assert_eq!(correlation_id, 2);
1297            assert_eq!(decoder.read_compact_string().unwrap(), "orders");
1298            assert_eq!(decoder.read_compact_string().unwrap(), "member-1");
1299            assert_eq!(decoder.read_i32().unwrap(), 2);
1300            assert_eq!(decoder.read_i32().unwrap(), 4);
1301            assert_eq!(decoder.read_compact_nullable_string().unwrap(), None);
1302            assert_eq!(decoder.read_compact_nullable_string().unwrap(), None);
1303            assert_eq!(decoder.read_i32().unwrap(), 30_000);
1304            assert_eq!(decoder.read_i8().unwrap(), -1);
1305            assert_eq!(decoder.read_compact_array_length(), Some(1));
1306            assert_eq!(decoder.read_compact_string().unwrap(), "subtopology-0");
1307            assert_eq!(decoder.read_compact_array_length(), Some(2));
1308            assert_eq!(decoder.read_i32().unwrap(), 0);
1309            assert_eq!(decoder.read_i32().unwrap(), 1);
1310            decoder.read_tagged_fields().unwrap();
1311            assert_eq!(decoder.read_compact_array_length(), Some(0));
1312            assert_eq!(decoder.read_compact_array_length(), Some(0));
1313            assert_eq!(decoder.read_compact_nullable_string().unwrap(), None);
1314            assert_eq!(decoder.read_i8().unwrap(), -1);
1315            assert_eq!(decoder.read_compact_array_length(), None);
1316            assert_eq!(decoder.read_compact_array_length(), Some(1));
1317            assert_eq!(decoder.read_compact_string().unwrap(), "subtopology-0");
1318            assert_eq!(decoder.read_i32().unwrap(), 0);
1319            assert_eq!(decoder.read_i64().unwrap(), 10);
1320            decoder.read_tagged_fields().unwrap();
1321            assert_eq!(decoder.read_compact_array_length(), Some(1));
1322            assert_eq!(decoder.read_compact_string().unwrap(), "subtopology-0");
1323            assert_eq!(decoder.read_i32().unwrap(), 0);
1324            assert_eq!(decoder.read_i64().unwrap(), 20);
1325            decoder.read_tagged_fields().unwrap();
1326            assert!(!decoder.read_bool().unwrap());
1327            decoder.read_tagged_fields().unwrap();
1328            write_streams_response(&mut broker_stream, correlation_id, "member-1", 3, 5).await;
1329
1330            let request = read_frame(&mut broker_stream).await;
1331            let (correlation_id, mut decoder) = request_header(&request);
1332            assert_eq!(correlation_id, 3);
1333            assert_eq!(decoder.read_compact_string().unwrap(), "orders");
1334            assert_eq!(decoder.read_compact_string().unwrap(), "member-1");
1335            assert_eq!(decoder.read_i32().unwrap(), -1);
1336            assert_eq!(decoder.read_i32().unwrap(), 5);
1337            assert_eq!(decoder.read_compact_nullable_string().unwrap(), None);
1338            assert_eq!(decoder.read_compact_nullable_string().unwrap(), None);
1339            assert_eq!(decoder.read_i32().unwrap(), -1);
1340            assert_eq!(decoder.read_i8().unwrap(), -1);
1341            assert_eq!(decoder.read_compact_array_length(), None);
1342            assert_eq!(decoder.read_compact_array_length(), None);
1343            assert_eq!(decoder.read_compact_array_length(), None);
1344            assert_eq!(decoder.read_compact_nullable_string().unwrap(), None);
1345            assert_eq!(decoder.read_i8().unwrap(), -1);
1346            assert_eq!(decoder.read_compact_array_length(), None);
1347            assert_eq!(decoder.read_compact_array_length(), None);
1348            assert_eq!(decoder.read_compact_array_length(), None);
1349            assert!(decoder.read_bool().unwrap());
1350            decoder.read_tagged_fields().unwrap();
1351            write_streams_response(&mut broker_stream, correlation_id, "member-1", -1, 5).await;
1352        });
1353
1354        let config = StreamsGroupConfig::new(["localhost:9092"], "orders", topology())
1355            .client_id("streams-test")
1356            .instance_id("instance-a")
1357            .rack_id("rack-a")
1358            .process_id("process-a")
1359            .user_endpoint(StreamsGroupHeartbeatEndpoint {
1360                host: "query-host".to_owned(),
1361                port: 7_777,
1362            })
1363            .client_tags(vec![StreamsGroupHeartbeatKeyValue {
1364                key: "zone".to_owned(),
1365                value: "a".to_owned(),
1366            }])
1367            .max_retries(0);
1368        let client = Client::from_stream(
1369            Box::new(client_stream),
1370            Some("streams-test".to_owned()),
1371            Some(Duration::from_secs(1)),
1372        );
1373        let mut session = StreamsGroupSession {
1374            config,
1375            coordinator: Some(client),
1376            member_id: "client-member".to_owned(),
1377            member_epoch: 0,
1378            endpoint_information_epoch: 0,
1379            heartbeat_interval: Duration::from_secs(1),
1380            pending_active_tasks: None,
1381            pending_standby_tasks: None,
1382            pending_warmup_tasks: None,
1383            pending_task_offsets: None,
1384            pending_task_end_offsets: None,
1385            assignment: StreamsGroupSessionAssignment::default(),
1386            closed: false,
1387        };
1388
1389        session.send_initial_heartbeat().await.unwrap();
1390        assert_eq!(session.member_id(), "member-1");
1391        assert_eq!(session.member_epoch(), 2);
1392        assert_eq!(session.assignment().task_offset_interval_ms, 100);
1393        assert_eq!(session.assignment().acceptable_recovery_lag, 0);
1394        assert!(session.assignment().active_tasks.is_none());
1395        session.set_task_state(
1396            vec![StreamsGroupHeartbeatTask {
1397                subtopology_id: "subtopology-0".to_owned(),
1398                partitions: vec![0, 1],
1399            }],
1400            Vec::new(),
1401            Vec::new(),
1402            vec![StreamsGroupHeartbeatTaskOffset {
1403                subtopology_id: "subtopology-0".to_owned(),
1404                partition: 0,
1405                offset: 10,
1406            }],
1407            vec![StreamsGroupHeartbeatTaskOffset {
1408                subtopology_id: "subtopology-0".to_owned(),
1409                partition: 0,
1410                offset: 20,
1411            }],
1412        );
1413        session.heartbeat().await.unwrap();
1414        assert_eq!(session.member_epoch(), 3);
1415        session.close().await.unwrap();
1416        assert!(session.is_closed());
1417        broker.await.unwrap();
1418    }
1419
1420    #[tokio::test]
1421    async fn background_streams_session_owns_heartbeat_and_graceful_close() {
1422        let (client_stream, mut broker_stream) = tokio::io::duplex(16 * 1024);
1423        let broker = tokio::spawn(async move {
1424            let request = read_frame(&mut broker_stream).await;
1425            let (correlation_id, mut decoder) = request_header(&request);
1426            assert_eq!(decoder.read_compact_string().unwrap(), "orders");
1427            assert_eq!(decoder.read_compact_string().unwrap(), "member-1");
1428            assert_eq!(decoder.read_i32().unwrap(), 2);
1429            assert_eq!(decoder.read_i32().unwrap(), 4);
1430            assert_eq!(decoder.read_compact_nullable_string().unwrap(), None);
1431            assert_eq!(decoder.read_compact_nullable_string().unwrap(), None);
1432            assert_eq!(decoder.read_i32().unwrap(), 30_000);
1433            assert_eq!(decoder.read_i8().unwrap(), -1);
1434            assert_eq!(decoder.read_compact_array_length(), Some(0));
1435            assert_eq!(decoder.read_compact_array_length(), Some(0));
1436            assert_eq!(decoder.read_compact_array_length(), Some(0));
1437            assert_eq!(decoder.read_compact_nullable_string().unwrap(), None);
1438            assert_eq!(decoder.read_i8().unwrap(), -1);
1439            assert_eq!(decoder.read_compact_array_length(), None);
1440            assert_eq!(decoder.read_compact_array_length(), Some(0));
1441            assert_eq!(decoder.read_compact_array_length(), Some(0));
1442            assert!(!decoder.read_bool().unwrap());
1443            decoder.read_tagged_fields().unwrap();
1444            write_streams_response(&mut broker_stream, correlation_id, "member-1", 3, 5).await;
1445
1446            let request = read_frame(&mut broker_stream).await;
1447            let (correlation_id, mut decoder) = request_header(&request);
1448            assert_eq!(decoder.read_compact_string().unwrap(), "orders");
1449            assert_eq!(decoder.read_compact_string().unwrap(), "member-1");
1450            assert_eq!(decoder.read_i32().unwrap(), -1);
1451            assert_eq!(decoder.read_i32().unwrap(), 5);
1452            assert_eq!(decoder.read_compact_nullable_string().unwrap(), None);
1453            assert_eq!(decoder.read_compact_nullable_string().unwrap(), None);
1454            assert_eq!(decoder.read_i32().unwrap(), -1);
1455            assert_eq!(decoder.read_i8().unwrap(), -1);
1456            assert_eq!(decoder.read_compact_array_length(), None);
1457            assert_eq!(decoder.read_compact_array_length(), None);
1458            assert_eq!(decoder.read_compact_array_length(), None);
1459            assert_eq!(decoder.read_compact_nullable_string().unwrap(), None);
1460            assert_eq!(decoder.read_i8().unwrap(), -1);
1461            assert_eq!(decoder.read_compact_array_length(), None);
1462            assert_eq!(decoder.read_compact_array_length(), None);
1463            assert_eq!(decoder.read_compact_array_length(), None);
1464            assert!(decoder.read_bool().unwrap());
1465            decoder.read_tagged_fields().unwrap();
1466            write_streams_response(&mut broker_stream, correlation_id, "member-1", -1, 5).await;
1467        });
1468
1469        let config = StreamsGroupConfig::new(["localhost:9092"], "orders", topology())
1470            .client_id("streams-test")
1471            .max_retries(0);
1472        let client = Client::from_stream(
1473            Box::new(client_stream),
1474            Some("streams-test".to_owned()),
1475            Some(Duration::from_secs(1)),
1476        );
1477        let session = StreamsGroupSession {
1478            config,
1479            coordinator: Some(client),
1480            member_id: "member-1".to_owned(),
1481            member_epoch: 2,
1482            endpoint_information_epoch: 4,
1483            heartbeat_interval: Duration::from_millis(100),
1484            pending_active_tasks: None,
1485            pending_standby_tasks: None,
1486            pending_warmup_tasks: None,
1487            pending_task_offsets: None,
1488            pending_task_end_offsets: None,
1489            assignment: StreamsGroupSessionAssignment::default(),
1490            closed: false,
1491        };
1492
1493        let handle = session.spawn_heartbeat_task();
1494        handle
1495            .set_task_state(
1496                Vec::new(),
1497                Vec::new(),
1498                Vec::new(),
1499                Some(Vec::new()),
1500                Some(Vec::new()),
1501            )
1502            .await
1503            .unwrap();
1504        let mut assignments = handle.subscribe_assignment();
1505        tokio::time::timeout(Duration::from_secs(2), assignments.changed())
1506            .await
1507            .unwrap()
1508            .unwrap();
1509        assert_eq!(handle.assignment().task_offset_interval_ms, 100);
1510        handle.close().await.unwrap();
1511        broker.await.unwrap();
1512    }
1513
1514    fn assignment(
1515        active_tasks: Option<Vec<StreamsGroupHeartbeatTask>>,
1516        standby_tasks: Option<Vec<StreamsGroupHeartbeatTask>>,
1517        warmup_tasks: Option<Vec<StreamsGroupHeartbeatTask>>,
1518    ) -> StreamsGroupSessionAssignment {
1519        StreamsGroupSessionAssignment {
1520            active_tasks,
1521            standby_tasks,
1522            warmup_tasks,
1523            ..StreamsGroupSessionAssignment::default()
1524        }
1525    }
1526
1527    fn task(subtopology_id: &str, partitions: &[i32]) -> StreamsGroupHeartbeatTask {
1528        StreamsGroupHeartbeatTask {
1529            subtopology_id: subtopology_id.to_owned(),
1530            partitions: partitions.to_vec(),
1531        }
1532    }
1533
1534    #[test]
1535    fn task_runtime_reconciles_nullable_roles_and_canonical_task_ids() {
1536        let mut runtime = StreamsTaskRuntime::new();
1537        let transitions = runtime
1538            .reconcile_assignment(&assignment(
1539                Some(vec![task("subtopology-0", &[2, 1])]),
1540                Some(vec![task("subtopology-0", &[3])]),
1541                Some(Vec::new()),
1542            ))
1543            .unwrap();
1544        assert_eq!(transitions.len(), 2);
1545        assert!(matches!(
1546            &transitions[0],
1547            StreamsTaskTransition::Added {
1548                role: StreamsTaskRole::Active,
1549                ..
1550            }
1551        ));
1552        assert!(matches!(
1553            &transitions[1],
1554            StreamsTaskTransition::Added {
1555                role: StreamsTaskRole::Standby,
1556                ..
1557            }
1558        ));
1559        assert_eq!(runtime.assignment()[0].task.partitions(), &[1, 2]);
1560
1561        let unchanged = runtime
1562            .reconcile_assignment(&assignment(None, None, None))
1563            .unwrap();
1564        assert!(unchanged.is_empty());
1565
1566        let transitions = runtime
1567            .reconcile_assignment(&assignment(
1568                Some(Vec::new()),
1569                Some(vec![task("subtopology-0", &[2, 1])]),
1570                None,
1571            ))
1572            .unwrap();
1573        assert_eq!(
1574            transitions,
1575            vec![
1576                StreamsTaskTransition::RoleChanged {
1577                    task: super::StreamsTaskId::new("subtopology-0", vec![1, 2]).unwrap(),
1578                    previous_role: StreamsTaskRole::Active,
1579                    role: StreamsTaskRole::Standby,
1580                },
1581                StreamsTaskTransition::Removed {
1582                    task: super::StreamsTaskId::new("subtopology-0", vec![3]).unwrap(),
1583                    role: StreamsTaskRole::Standby,
1584                },
1585            ]
1586        );
1587    }
1588
1589    #[test]
1590    fn task_runtime_rejects_conflicting_assignment_without_mutating_state() {
1591        let mut runtime = StreamsTaskRuntime::new();
1592        runtime
1593            .reconcile_assignment(&assignment(
1594                Some(vec![task("subtopology-0", &[0])]),
1595                Some(Vec::new()),
1596                Some(Vec::new()),
1597            ))
1598            .unwrap();
1599        let before = runtime.assignment();
1600
1601        let error = runtime
1602            .reconcile_assignment(&assignment(
1603                Some(vec![
1604                    task("subtopology-0", &[1]),
1605                    task("subtopology-0", &[1]),
1606                ]),
1607                Some(Vec::new()),
1608                Some(Vec::new()),
1609            ))
1610            .unwrap_err();
1611        assert!(matches!(
1612            error,
1613            Error::StreamsTaskAssignmentConflict {
1614                subtopology_id,
1615                partition: 1,
1616            } if subtopology_id == "subtopology-0"
1617        ));
1618        assert_eq!(runtime.assignment(), before);
1619    }
1620
1621    fn request_header(request: &[u8]) -> (i32, Decoder<'_>) {
1622        let mut decoder = Decoder::new(request);
1623        assert_eq!(decoder.read_i16().unwrap(), 88);
1624        assert_eq!(decoder.read_i16().unwrap(), 0);
1625        let correlation_id = decoder.read_i32().unwrap();
1626        assert_eq!(
1627            decoder.read_nullable_string().unwrap(),
1628            Some("streams-test".to_owned())
1629        );
1630        decoder.read_tagged_fields().unwrap();
1631        (correlation_id, decoder)
1632    }
1633
1634    async fn write_streams_response(
1635        stream: &mut tokio::io::DuplexStream,
1636        correlation_id: i32,
1637        member_id: &str,
1638        member_epoch: i32,
1639        endpoint_information_epoch: i32,
1640    ) {
1641        let mut response = Encoder::new();
1642        response.write_i32(correlation_id);
1643        response.write_empty_tagged_fields();
1644        response.write_i32(0);
1645        response.write_i16(0);
1646        response.write_compact_nullable_string(None).unwrap();
1647        response.write_compact_string(member_id).unwrap();
1648        response.write_i32(member_epoch);
1649        response.write_i32(1_000);
1650        response.write_i32(0);
1651        response.write_i32(100);
1652        response.write_unsigned_varint(0);
1653        response.write_unsigned_varint(0);
1654        response.write_unsigned_varint(0);
1655        response.write_unsigned_varint(0);
1656        response.write_i32(endpoint_information_epoch);
1657        response.write_unsigned_varint(0);
1658        response.write_empty_tagged_fields();
1659        write_frame(stream, &response.into_bytes()).await;
1660    }
1661
1662    async fn read_frame(stream: &mut tokio::io::DuplexStream) -> Vec<u8> {
1663        let mut size = [0; 4];
1664        stream.read_exact(&mut size).await.unwrap();
1665        let size = usize::try_from(i32::from_be_bytes(size)).unwrap();
1666        let mut frame = vec![0; size];
1667        stream.read_exact(&mut frame).await.unwrap();
1668        frame
1669    }
1670
1671    async fn write_frame(stream: &mut tokio::io::DuplexStream, frame: &[u8]) {
1672        stream
1673            .write_all(&(i32::try_from(frame.len()).unwrap()).to_be_bytes())
1674            .await
1675            .unwrap();
1676        stream.write_all(frame).await.unwrap();
1677        stream.flush().await.unwrap();
1678    }
1679}
kafrust/streams.rs

kafrust/
streams.rs
