krabka_client_consumer/
partition_state.rs

1//! Per-partition state that the application reads and changes: pause and
2//! resume, the fetch position and the committed offset.
3//!
4//! These are Kafka's `KafkaConsumer.pause`, `resume`, `paused`, `position` and
5//! `committed`.
6
7use std::{
8    collections::{HashMap, HashSet},
9    sync::{Mutex, PoisonError},
10};
11
12use krabka_units::convert::TimeExt as _;
13use tokio::time::{Instant, sleep, timeout_at};
14
15use crate::{
16    commit::OffsetAndMetadata,
17    consumer::Consumer,
18    coordinator::{CoordinatorRetryPolicy, send_offset_fetch},
19    error::ConsumerError,
20    offset_wire,
21    poll::is_reset_sentinel,
22    position::with_offsets_and_positions,
23};
24
25/// Paused partitions, with the ownership of each partition at the pause.
26///
27/// Kafka keeps the pause in the `TopicPartitionState` of the partition. A
28/// rebalance that keeps the partition keeps its state, and an assignment that
29/// adds the partition again creates a new state that is not paused
30/// (`SubscriptionState.assignFromSubscribed`). The ownership id changes in the
31/// same cases, so a pause is in effect only while the ownership id is the one
32/// of the pause.
33pub type PausedPartitions = Mutex<HashMap<(String, i32), u64>>;
34
35impl Consumer {
36    /// Stop fetching from `partitions`. Kafka's `KafkaConsumer.pause`.
37    ///
38    /// `poll` does not return records of a paused partition and does not fetch
39    /// it. The consumer stays in the group. Records that `poll` fetched before
40    /// the pause come after [`resume`](Self::resume). A rebalance that keeps
41    /// the partition keeps the pause.
42    ///
43    /// # Errors
44    ///
45    /// Returns [`ConsumerError::NoCurrentAssignment`] if the consumer does not
46    /// own a partition. The call then pauses no partition.
47    pub async fn pause(&self, partitions: &[(String, i32)]) -> Result<(), ConsumerError> {
48        let ownership_ids = self.commit_identity.lock().await.ownership_ids.clone();
49        let mut paused = Vec::with_capacity(partitions.len());
50        for partition in partitions {
51            let Some(ownership_id) = ownership_ids.get(partition) else {
52                return Err(no_current_assignment(partition));
53            };
54            paused.push((partition.clone(), *ownership_id));
55        }
56        self.paused
57            .lock()
58            .unwrap_or_else(PoisonError::into_inner)
59            .extend(paused);
60        Ok(())
61    }
62
63    /// Fetch from `partitions` again. Kafka's `KafkaConsumer.resume`.
64    ///
65    /// # Errors
66    ///
67    /// Returns [`ConsumerError::NoCurrentAssignment`] if the consumer does not
68    /// own a partition. The call then resumes no partition.
69    pub async fn resume(&self, partitions: &[(String, i32)]) -> Result<(), ConsumerError> {
70        let ownership_ids = self.commit_identity.lock().await.ownership_ids.clone();
71        if let Some(partition) = partitions
72            .iter()
73            .find(|partition| !ownership_ids.contains_key(*partition))
74        {
75            return Err(no_current_assignment(partition));
76        }
77        self.paused
78            .lock()
79            .unwrap_or_else(PoisonError::into_inner)
80            .retain(|partition, _| !partitions.contains(partition));
81        Ok(())
82    }
83
84    /// The paused partitions that the consumer owns. Kafka's
85    /// `KafkaConsumer.paused`.
86    pub async fn paused(&self) -> HashSet<(String, i32)> {
87        let ownership_ids = self.commit_identity.lock().await.ownership_ids.clone();
88        self.paused_of(&ownership_ids)
89    }
90
91    /// The partitions whose pause is in effect for `ownership_ids`.
92    pub(crate) fn paused_of(
93        &self,
94        ownership_ids: &HashMap<(String, i32), u64>,
95    ) -> HashSet<(String, i32)> {
96        let mut paused = self.paused.lock().unwrap_or_else(PoisonError::into_inner);
97        paused.retain(|partition, ownership_id| ownership_ids.get(partition) == Some(ownership_id));
98        paused.keys().cloned().collect()
99    }
100
101    /// The offset of the next record that `poll` fetches for `(topic,
102    /// partition)`. Kafka's `KafkaConsumer.position`.
103    ///
104    /// When the partition waits for an offset reset or for a leader epoch
105    /// validation, the call sends `ListOffsets` or `OffsetForLeaderEpoch` and
106    /// waits for the position, for at most `default_api_timeout`.
107    ///
108    /// # Errors
109    ///
110    /// Returns [`ConsumerError::NoCurrentAssignment`] if the consumer does not
111    /// own the partition, [`ConsumerError::Timeout`] if the position is not
112    /// known before `default_api_timeout`, or the error of a request.
113    pub async fn position(
114        &self,
115        topic: impl Into<String>,
116        partition: i32,
117    ) -> Result<i64, ConsumerError> {
118        let key = (topic.into(), partition);
119        let deadline = Instant::now() + self.default_api_timeout.to_std();
120        let wait = async {
121            loop {
122                if !self.assigned.lock().await.contains(&key) {
123                    return Err(no_current_assignment(&key));
124                }
125                if let Some(offset) = self.valid_position(&key).await {
126                    return Ok(offset);
127                }
128                self.update_fetch_positions(None, Some(deadline)).await?;
129                if self.valid_position(&key).await.is_none() {
130                    sleep(self.retry_policy.initial_backoff).await;
131                }
132            }
133        };
134        // The timeout bounds the requests too, as Kafka's `position` passes
135        // its timer to `updateFetchPositions` and `client.poll`.
136        timeout_at(deadline, wait).await.unwrap_or_else(|_| {
137            Err(ConsumerError::timeout(format!(
138                "the position for partition {}-{} could not be determined",
139                key.0, key.1
140            )))
141        })
142    }
143
144    /// The position of `key` when it is valid: Kafka's
145    /// `SubscriptionState.validPosition`. A position that waits for a reset or
146    /// for validation is not valid.
147    async fn valid_position(&self, key: &(String, i32)) -> Option<i64> {
148        with_offsets_and_positions(&self.next_offsets, &self.positions, |offsets, positions| {
149            let offset = *offsets.get(key)?;
150            let awaiting_validation = positions
151                .get(key)
152                .is_some_and(|position| position.awaiting_validation);
153            (!is_reset_sentinel(offset) && !awaiting_validation).then_some(offset)
154        })
155        .await
156    }
157
158    /// The committed offsets of `partitions` in the group. Kafka's
159    /// `KafkaConsumer.committed`.
160    ///
161    /// A partition without a committed offset maps to `None`. The partitions
162    /// do not need to be assigned to this consumer.
163    ///
164    /// # Errors
165    ///
166    /// Returns the error of the `OffsetFetch` request after the retries that
167    /// Kafka's consumer does, for at most `default_api_timeout`.
168    pub async fn committed(
169        &self,
170        partitions: &[(String, i32)],
171    ) -> Result<HashMap<(String, i32), Option<OffsetAndMetadata>>, ConsumerError> {
172        let deadline = Instant::now() + self.default_api_timeout.to_std();
173        self.committed_until(partitions, deadline).await
174    }
175
176    /// The committed offsets of `partitions`, within `deadline`.
177    pub(crate) async fn committed_until(
178        &self,
179        partitions: &[(String, i32)],
180        deadline: Instant,
181    ) -> Result<HashMap<(String, i32), Option<OffsetAndMetadata>>, ConsumerError> {
182        self.require_group_id()?;
183        if partitions.is_empty() {
184            return Ok(HashMap::new());
185        }
186        let mut by_topic: HashMap<String, Vec<i32>> = HashMap::new();
187        for (topic, partition) in partitions {
188            by_topic.entry(topic.clone()).or_default().push(*partition);
189        }
190        let request = offset_wire::build_offset_fetch(&self.group_id, &by_topic);
191        let fetch = send_offset_fetch(
192            &self.client,
193            &self.group_id,
194            &self.coordinator_id,
195            &request,
196            CoordinatorRetryPolicy {
197                timeout: deadline.saturating_duration_since(Instant::now()),
198                ..self.retry_policy
199            },
200        );
201        // The deadline bounds each request too. Kafka's
202        // `fetchCommittedOffsets` gives up when its timer expires.
203        let response = timeout_at(deadline, fetch).await.map_err(|_| {
204            ConsumerError::timeout("the last committed offsets could not be determined".to_owned())
205        })??;
206        let mut committed = offset_wire::parse_committed_offsets(&response);
207        committed.retain(|partition, _| partitions.contains(partition));
208        for partition in partitions {
209            committed.entry(partition.clone()).or_insert(None);
210        }
211        Ok(committed)
212    }
213}
214
215fn no_current_assignment((topic, partition): &(String, i32)) -> ConsumerError {
216    ConsumerError::NoCurrentAssignment {
217        topic: topic.clone(),
218        partition: *partition,
219    }
220}
221
222#[cfg(test)]
223mod tests {
224    use std::{collections::HashSet, time::Duration};
225
226    use krabka_client_core::{Client, MockBroker};
227    use krabka_protocol::{
228        Encode,
229        owned::{
230            api_versions_request,
231            api_versions_response::{ApiVersion, ApiVersionsResponse},
232            list_offsets_request, offset_fetch_request,
233            offset_fetch_response::{
234                OffsetFetchResponse, OffsetFetchResponsePartition, OffsetFetchResponseTopic,
235            },
236        },
237    };
238    use krabka_units::secs;
239    use tokio::time::Instant;
240
241    use crate::{
242        commit::OffsetAndMetadata,
243        poll::{BEGINNING_SENTINEL, partition_error_tests::consumer_with_client},
244        subscription::shared,
245    };
246
247    fn encode(response: &impl Encode, version: i16) -> Vec<u8> {
248        let mut buf = bytes::BytesMut::new();
249        response.encode(&mut buf, version).expect("encode");
250        buf.to_vec()
251    }
252
253    /// A coordinator that answers `OffsetFetch` v5 with `rows`:
254    /// `(partition, committed offset, leader epoch, metadata, error code)` of
255    /// `orders`.
256    async fn offset_fetch_coordinator(rows: Vec<(i32, i64, i32, &'static str, i16)>) -> MockBroker {
257        MockBroker::start(move |api_key, version, _corr_id, _body| {
258            if api_key == api_versions_request::API_KEY {
259                let versions = ApiVersionsResponse {
260                    api_keys: [
261                        (api_versions_request::API_KEY, 0, 3),
262                        (offset_fetch_request::API_KEY, 5, 5),
263                    ]
264                    .into_iter()
265                    .map(|(api_key, min_version, max_version)| ApiVersion {
266                        api_key,
267                        min_version,
268                        max_version,
269                        ..Default::default()
270                    })
271                    .collect(),
272                    ..Default::default()
273                };
274                return Some(encode(&versions, 0));
275            }
276            if api_key != offset_fetch_request::API_KEY {
277                return None;
278            }
279            let response = OffsetFetchResponse {
280                topics: vec![OffsetFetchResponseTopic {
281                    name: "orders".into(),
282                    partitions: rows
283                        .iter()
284                        .map(
285                            |(
286                                partition_index,
287                                committed_offset,
288                                committed_leader_epoch,
289                                metadata,
290                                error_code,
291                            )| {
292                                OffsetFetchResponsePartition {
293                                    partition_index: *partition_index,
294                                    committed_offset: *committed_offset,
295                                    committed_leader_epoch: *committed_leader_epoch,
296                                    metadata: Some((*metadata).into()),
297                                    error_code: *error_code,
298                                    ..Default::default()
299                                }
300                            },
301                        )
302                        .collect(),
303                    ..Default::default()
304                }],
305                ..Default::default()
306            };
307            Some(encode(&response, version))
308        })
309        .await
310        .unwrap()
311    }
312
313    /// Kafka's `committed` maps a partition without a committed offset to
314    /// `null`, and keeps the leader epoch and metadata of a committed one
315    /// (`ConsumerCoordinator.OffsetFetchResponseHandler`).
316    #[tokio::test]
317    async fn committed_returns_kafkas_offsets_and_metadata() {
318        let mock = offset_fetch_coordinator(vec![
319            (0, 10, 3, "note", 0),
320            (1, -1, -1, "", 0),
321            (3, 7, -1, "", 0),
322        ])
323        .await;
324        let client = Client::builder()
325            .bootstrap(mock.addr.to_string())
326            .request_timeout(secs(30))
327            .build()
328            .await
329            .expect("client");
330        let consumer = consumer_with_client(client);
331        let partitions = [
332            ("orders".to_owned(), 0),
333            ("orders".to_owned(), 1),
334            ("orders".to_owned(), 2),
335        ];
336        let committed = consumer
337            .committed(&partitions)
338            .await
339            .map_err(|error| error.to_string());
340        mock.stop();
341        assert2::assert!(
342            committed
343                == Ok(
344                    maplit::hashmap! {("orders".to_owned(), 0) => Some(OffsetAndMetadata {
345                        offset: 10,
346                        leader_epoch: Some(3),
347                        metadata: "note".into(),
348                    }), ("orders".to_owned(), 1) => None, ("orders".to_owned(), 2) => None}
349                )
350        );
351    }
352
353    /// Kafka's `pause`, `resume` and `paused` act on assigned partitions only,
354    /// and a partition that the consumer owns again after a rebalance is not
355    /// paused (`SubscriptionState.assignFromSubscribed` gives it a new state).
356    #[tokio::test]
357    async fn pause_follows_the_ownership_of_the_partition() {
358        let key = ("orders".to_owned(), 0);
359        let mut actual = Vec::new();
360        let mut wanted = Vec::new();
361        for (name, pause, resume, new_ownership, expected_error, expected_paused) in [
362            (
363                "pause",
364                vec![key.clone()],
365                vec![],
366                false,
367                None,
368                vec![key.clone()],
369            ),
370            (
371                "pause and resume",
372                vec![key.clone()],
373                vec![key.clone()],
374                false,
375                None,
376                vec![],
377            ),
378            ("owned again", vec![key.clone()], vec![], true, None, vec![]),
379            (
380                "unowned partition",
381                vec![("orders".to_owned(), 9)],
382                vec![],
383                false,
384                Some("no current assignment for partition orders-9"),
385                vec![],
386            ),
387        ] {
388            let mock = offset_fetch_coordinator(vec![]).await;
389            let client = Client::builder()
390                .bootstrap(mock.addr.to_string())
391                .build()
392                .await
393                .expect("client");
394            let consumer = consumer_with_client(client);
395            let mut error = consumer
396                .pause(&pause)
397                .await
398                .err()
399                .map(|error| error.to_string());
400            if let Err(resume_error) = consumer.resume(&resume).await {
401                error.get_or_insert_with(|| resume_error.to_string());
402            }
403            if new_ownership {
404                consumer
405                    .commit_identity
406                    .lock()
407                    .await
408                    .ownership_ids
409                    .insert(key.clone(), 2);
410            }
411            let paused = consumer.paused().await;
412            mock.stop();
413            actual.push((name, error, paused));
414            wanted.push((
415                name,
416                expected_error.map(str::to_owned),
417                expected_paused.into_iter().collect::<HashSet<_>>(),
418            ));
419        }
420        assert2::assert!(actual == wanted);
421    }
422
423    /// Kafka's `committed(partitions, timeout)` gives up when the timeout
424    /// expires, also while a request waits for its response.
425    #[tokio::test]
426    async fn committed_returns_by_the_api_timeout_when_the_coordinator_is_silent() {
427        let mock = MockBroker::start(|api_key, _version, _corr_id, _body| {
428            (api_key == api_versions_request::API_KEY).then(|| {
429                encode(
430                    &ApiVersionsResponse {
431                        api_keys: [
432                            (api_versions_request::API_KEY, 0, 3),
433                            (offset_fetch_request::API_KEY, 5, 5),
434                        ]
435                        .into_iter()
436                        .map(|(api_key, min_version, max_version)| ApiVersion {
437                            api_key,
438                            min_version,
439                            max_version,
440                            ..Default::default()
441                        })
442                        .collect(),
443                        ..Default::default()
444                    },
445                    0,
446                )
447            })
448        })
449        .await
450        .unwrap();
451        let client = Client::builder()
452            .bootstrap(mock.addr.to_string())
453            .request_timeout(secs(30))
454            .build()
455            .await
456            .expect("client");
457        let mut consumer = consumer_with_client(client);
458        consumer.default_api_timeout = krabka_units::millis(200);
459        let started = Instant::now();
460        let committed = consumer
461            .committed(&[("orders".to_owned(), 0)])
462            .await
463            .map_err(|error| error.to_string());
464        let in_time = started.elapsed() < Duration::from_secs(5);
465        mock.stop();
466        assert2::assert!(
467            (committed, in_time)
468                == (
469                    Err("timeout: the last committed offsets could not be determined".to_owned()),
470                    true
471                )
472        );
473    }
474
475    /// Kafka's `position` gives up at `default.api.timeout.ms` when the reset
476    /// of the partition gets no `ListOffsets` answer.
477    #[tokio::test]
478    async fn position_returns_by_the_api_timeout_when_the_reset_gets_no_answer() {
479        let mock = MockBroker::start(|api_key, _version, _corr_id, _body| {
480            (api_key == api_versions_request::API_KEY).then(|| {
481                encode(
482                    &ApiVersionsResponse {
483                        api_keys: [
484                            (api_versions_request::API_KEY, 0, 3),
485                            (list_offsets_request::API_KEY, 1, 7),
486                        ]
487                        .into_iter()
488                        .map(|(api_key, min_version, max_version)| ApiVersion {
489                            api_key,
490                            min_version,
491                            max_version,
492                            ..Default::default()
493                        })
494                        .collect(),
495                        ..Default::default()
496                    },
497                    0,
498                )
499            })
500        })
501        .await
502        .unwrap();
503        let client = Client::builder()
504            .bootstrap(mock.addr.to_string())
505            .request_timeout(secs(30))
506            .build()
507            .await
508            .expect("client");
509        let mut consumer = consumer_with_client(client);
510        consumer.default_api_timeout = krabka_units::millis(200);
511        consumer
512            .next_offsets
513            .lock()
514            .await
515            .insert(("orders".to_owned(), 0), BEGINNING_SENTINEL);
516        let started = Instant::now();
517
518        let position = consumer
519            .position("orders", 0)
520            .await
521            .map_err(|error| error.to_string());
522
523        let in_time = started.elapsed() < Duration::from_secs(5);
524        mock.stop();
525        assert2::assert!(
526            (position, in_time)
527                == (
528                    Err(
529                        "timeout: the position for partition orders-0 could not be determined"
530                            .to_owned()
531                    ),
532                    true
533                )
534        );
535    }
536
537    /// Kafka's `seekToEnd` requests a reset and keeps the high watermark of
538    /// the last fetch. A partition that waits for the reset is not at the end.
539    #[tokio::test]
540    async fn seek_to_end_keeps_the_end_offset_and_leaves_the_log_end() {
541        let mock = offset_fetch_coordinator(vec![]).await;
542        let client = Client::builder()
543            .bootstrap(mock.addr.to_string())
544            .build()
545            .await
546            .expect("client");
547        let consumer = consumer_with_client(client);
548        let key = ("orders".to_owned(), 0);
549        consumer.next_offsets.lock().await.insert(key.clone(), 12);
550        consumer.end_offsets.lock().await.insert(key.clone(), 12);
551        let before = consumer.at_log_end().await;
552        consumer.seek_to_end(&[]).await.expect("seek to end");
553        let after = consumer.at_log_end().await;
554        let end = consumer.end_offsets.lock().await.get(&key).copied();
555        mock.stop();
556        assert2::assert!((before, after, end) == (true, false, Some(12)));
557    }
558
559    /// Kafka's `assign` keeps the position of a partition that stays
560    /// assigned, and a new partition starts at the committed offset of the
561    /// group, or by `auto.offset.reset` without one
562    /// (`refreshCommittedOffsetsIfNeeded`, `resetInitializingPositions`).
563    #[tokio::test]
564    async fn assign_starts_new_partitions_at_the_committed_offset_or_the_reset() {
565        let mock = offset_fetch_coordinator(vec![(1, 10, 2, "", 0), (2, -1, -1, "", 0)]).await;
566        let client = Client::builder()
567            .bootstrap(mock.addr.to_string())
568            .request_timeout(secs(30))
569            .build()
570            .await
571            .expect("client");
572        let mut consumer = consumer_with_client(client);
573        consumer.subscription = shared(Vec::new(), None, true);
574        consumer.auto_offset_reset = crate::AutoOffsetReset::Earliest;
575        let key = |partition: i32| ("orders".to_owned(), partition);
576        consumer
577            .assign(&[key(0), key(1), key(2)])
578            .await
579            .expect("assign");
580        // The metadata of the poll gave the partition the current leader
581        // epoch before the committed offset came.
582        consumer
583            .positions
584            .lock()
585            .await
586            .entry(key(1))
587            .or_default()
588            .leader_epoch = krabka_ids::LeaderEpoch(5);
589        consumer
590            .resolve_committed_sentinels(None)
591            .await
592            .expect("resolve");
593        let mut offsets: Vec<_> = consumer
594            .next_offsets
595            .lock()
596            .await
597            .clone()
598            .into_iter()
599            .collect();
600        offsets.sort();
601        let identity = consumer.commit_identity.lock().await.clone();
602        let awaiting_validation = consumer
603            .positions
604            .lock()
605            .await
606            .get(&key(1))
607            .is_some_and(|position| position.awaiting_validation);
608        let rejected = consumer
609            .subscribe(["orders"])
610            .await
611            .map_err(|error| error.to_string());
612        mock.stop();
613        assert2::assert!(
614            (
615                offsets,
616                awaiting_validation,
617                identity.generation,
618                identity.member_id,
619                consumer.assignment().await,
620                rejected
621            ) == (
622                vec![(key(0), 5), (key(1), 10), (key(2), BEGINNING_SENTINEL)],
623                true,
624                -1,
625                String::new(),
626                vec![key(0), key(1), key(2)],
627                Err("illegal state: the consumer was built without a subscription, so it has no group membership; build it with subscribe or subscribe_pattern to join a group".to_owned())
628            )
629        );
630    }
631}
krabka_client_consumer/partition_state.rs

krabka_client_consumer/
partition_state.rs
