krabka_client_consumer/
control.rs

1//! Consumer control calls: Kafka's `KafkaConsumer.wakeup`, `enforceRebalance`
2//! and `close(CloseOptions)`.
3
4use std::sync::{
5    Arc,
6    atomic::{AtomicBool, Ordering},
7};
8
9use krabka_units::{Time, convert::TimeExt, secs};
10use tokio::sync::Notify;
11
12use crate::{
13    consumer::{Consumer, GroupMembershipOperation},
14    error::ConsumerError,
15};
16
17/// Kafka's `ConsumerUtils.DEFAULT_CLOSE_TIMEOUT_MS`.
18pub const DEFAULT_CONSUMER_CLOSE_TIMEOUT: Time = secs(30);
19
20/// The `JoinGroup` reason of [`Consumer::enforce_rebalance`] without a reason.
21/// Kafka's `ClassicKafkaConsumer.DEFAULT_REASON`.
22pub const ENFORCED_REBALANCE_REASON: &str = "rebalance enforced by user";
23
24/// How [`Consumer::close_with`](Consumer::close_with) closes the
25/// consumer. Kafka's `CloseOptions` (KIP-1092).
26#[derive(Clone, Copy, Debug, Default, PartialEq, bon::Builder)]
27pub struct CloseOptions {
28    /// What the consumer does with its group membership.
29    #[builder(default)]
30    pub group_membership_operation: GroupMembershipOperation,
31    /// The time that the close can take. `None` is
32    /// [`DEFAULT_CONSUMER_CLOSE_TIMEOUT`].
33    #[builder(with = |millis: i64| Time::from_millis(millis))]
34    pub timeout: Option<Time>,
35}
36
37impl CloseOptions {
38    /// The options with `timeout`. Kafka's `CloseOptions.timeout`.
39    #[must_use]
40    pub fn timeout(timeout: Time) -> Self {
41        Self {
42            timeout: Some(timeout),
43            ..Self::default()
44        }
45    }
46
47    /// The options with `operation`. Kafka's
48    /// `CloseOptions.groupMembershipOperation`.
49    #[must_use]
50    pub fn group_membership_operation(operation: GroupMembershipOperation) -> Self {
51        Self::builder()
52            .group_membership_operation(operation)
53            .build()
54    }
55}
56
57/// What the coordinator task needs when it stops.
58#[derive(Clone, Copy, Debug, Default)]
59pub struct CloseRequest {
60    pub operation: GroupMembershipOperation,
61}
62
63/// The close timeout of `options`, or an error for a negative timeout.
64pub fn close_timeout(options: CloseOptions) -> Result<Time, ConsumerError> {
65    let timeout = options.timeout.unwrap_or(DEFAULT_CONSUMER_CLOSE_TIMEOUT);
66    if timeout.secs_f64() < 0.0 || !timeout.secs_f64().is_finite() {
67        return Err(ConsumerError::InvalidArgument(
68            "The timeout cannot be negative.".to_owned(),
69        ));
70    }
71    Ok(timeout)
72}
73
74/// A handle that wakes a [`Consumer::poll`](Consumer::poll) from
75/// another task. Kafka's `KafkaConsumer.wakeup`.
76///
77/// Get it with [`Consumer::wakeup_handle`](Consumer::wakeup_handle).
78/// The poll that runs, or the next poll, returns [`ConsumerError::Wakeup`]
79/// once.
80#[derive(Clone, Debug, Default)]
81pub struct WakeupHandle {
82    inner: Arc<WakeupState>,
83}
84
85#[derive(Debug, Default)]
86struct WakeupState {
87    pending: AtomicBool,
88    notify: Notify,
89}
90
91impl WakeupHandle {
92    /// Wake the poll that runs, or the next poll.
93    pub fn wakeup(&self) {
94        self.inner.pending.store(true, Ordering::SeqCst);
95        self.inner.notify.notify_waiters();
96    }
97
98    /// Take a pending wakeup. Kafka's `ConsumerNetworkClient.maybeTriggerWakeup`.
99    pub(crate) fn take(&self) -> bool {
100        self.inner.pending.swap(false, Ordering::SeqCst)
101    }
102
103    /// Wait until a wakeup comes, and take it.
104    pub(crate) async fn woken(&self) {
105        loop {
106            let notified = self.inner.notify.notified();
107            tokio::pin!(notified);
108            notified.as_mut().enable();
109            if self.take() {
110                return;
111            }
112            notified.await;
113        }
114    }
115}
116
117impl Consumer {
118    /// A handle that wakes the `poll` of this consumer from another task.
119    /// Kafka's `KafkaConsumer.wakeup` is thread safe; `poll` takes `&mut self`
120    /// here, so another task needs its own handle.
121    #[must_use]
122    pub fn wakeup_handle(&self) -> WakeupHandle {
123        self.wakeup.clone()
124    }
125
126    /// Make the next `poll` return [`ConsumerError::Wakeup`]. Kafka's
127    /// `KafkaConsumer.wakeup`.
128    pub fn wakeup(&self) {
129        self.wakeup.wakeup();
130    }
131
132    /// Ask the group to rebalance. The next `poll` sends `JoinGroup` with
133    /// `reason`, or with Kafka's default reason `rebalance enforced by user`
134    /// when `reason` is `None` or empty. Kafka's
135    /// `KafkaConsumer.enforceRebalance`.
136    ///
137    /// # Errors
138    ///
139    /// Returns [`ConsumerError::InvalidGroupId`] without a group id, and
140    /// [`ConsumerError::IllegalState`] for a consumer without group
141    /// membership (Kafka: `Tried to force a rebalance but consumer does not
142    /// have a group.`). A member of the consumer group protocol logs a
143    /// warning and does nothing, as Kafka's `AsyncKafkaConsumer` does.
144    pub fn enforce_rebalance(&self, reason: Option<&str>) -> Result<(), ConsumerError> {
145        self.require_group_membership()?;
146        if self.group_protocol == crate::GroupProtocol::Consumer {
147            // Kafka's `AsyncKafkaConsumer.enforceRebalance` logs
148            // `Operation not supported in new consumer group protocol` and
149            // does nothing: the group coordinator owns the assignment.
150            tracing::warn!("enforce_rebalance is not supported in the consumer group protocol");
151            return Ok(());
152        }
153        if self.subscription.borrow().manual_assignment {
154            return Err(ConsumerError::IllegalState(
155                "Tried to force a rebalance but the consumer has a manual assignment.".to_owned(),
156            ));
157        }
158        let reason = reason
159            .filter(|reason| !reason.is_empty())
160            .unwrap_or(ENFORCED_REBALANCE_REASON);
161        // A closed channel means that the coordinator task stopped, and no
162        // rebalance can come.
163        // The `poll` count now: the rebalance starts with the first `poll`
164        // after this call, also when the task sees the request after that
165        // `poll`.
166        let polls = *self.poll_signal.borrow();
167        let _ = self.enforced_rebalances.send((reason.to_owned(), polls));
168        Ok(())
169    }
170}
171
172#[cfg(test)]
173mod tests {
174    use std::time::Duration;
175
176    use krabka_units::millis;
177    use tokio::time::timeout;
178
179    use super::*;
180
181    #[test]
182    fn close_timeout_follows_kafkas_close() {
183        let cases = [
184            (
185                "default timeout",
186                CloseOptions::builder().build(),
187                Ok(DEFAULT_CONSUMER_CLOSE_TIMEOUT),
188            ),
189            (
190                "finite timeout",
191                CloseOptions::timeout(millis(250)),
192                Ok(millis(250)),
193            ),
194            (
195                "negative timeout",
196                CloseOptions::timeout(Time::from_secs_f64(-1.0)),
197                Err("invalid argument: The timeout cannot be negative.".to_owned()),
198            ),
199            (
200                "NaN timeout",
201                CloseOptions::timeout(Time::from_secs_f64(f64::NAN)),
202                Err("invalid argument: The timeout cannot be negative.".to_owned()),
203            ),
204            (
205                "infinite timeout",
206                CloseOptions::timeout(Time::from_secs_f64(f64::INFINITY)),
207                Err("invalid argument: The timeout cannot be negative.".to_owned()),
208            ),
209        ];
210
211        for (name, options, expected) in cases {
212            let actual = close_timeout(options).map_err(|error| error.to_string());
213            assert2::assert!(actual == expected, "case: {name}");
214        }
215    }
216
217    #[tokio::test]
218    async fn wakeup_is_taken_once() {
219        let handle = WakeupHandle::default();
220        handle.wakeup();
221        let first = timeout(Duration::from_millis(50), handle.woken()).await;
222        let second = timeout(Duration::from_millis(50), handle.woken()).await;
223        assert2::assert!((first.is_ok(), second.is_ok()) == (true, false));
224    }
225}
krabka_client_consumer/control.rs

krabka_client_consumer/
control.rs
