krabka_client_consumer/seek.rs
1//! `Consumer::seek` sets the fetch position of an assigned partition.
2//!
3//! This is Kafka's `KafkaConsumer.seek` and `SubscriptionState.seekUnvalidated`.
4//! The call changes the position at once. The next `poll` fetches from it.
5//! A seek on a partition that the consumer does not own fails, as Kafka's
6//! `SubscriptionState.assignedState` does.
7//!
8//! The coordinator task sets the position of an added partition before it
9//! publishes the partition in `assigned`. A seek that finds its partition in
10//! `assigned` therefore comes after that position, and the position does not
11//! overwrite the seek. The seek holds the `assigned` lock while it writes the
12//! position, so a rebalance cannot remove the partition during the write.
13//!
14//! `seek_to_beginning` and `seek_to_end` plant a reset sentinel in
15//! `next_offsets`. The next position update resolves it with `ListOffsets`.
16//!
17//! Records that `poll` fetched from the earlier position stay in the fetch
18//! buffer, but `poll` drops them because their position is not the current
19//! one. Kafka's `FetchCollector` does the same check.
20
21use krabka_ids::LeaderEpoch;
22
23use crate::{consumer::Consumer, error::ConsumerError, poll, position};
24
25impl Consumer {
26 /// Set the next offset that `poll` fetches for `(topic, partition)`.
27 ///
28 /// `offset` is the next offset to read, that is the last consumed offset
29 /// plus 1, the same value as a committed group offset. Pass `0` to read a
30 /// partition from the beginning. The consumer does not validate the
31 /// position with `OffsetForLeaderEpoch`, as Kafka's `seek(TopicPartition,
32 /// long)` does not.
33 ///
34 /// # Errors
35 ///
36 /// Returns [`ConsumerError::InvalidOffset`] if `offset` is negative, and
37 /// [`ConsumerError::NoCurrentAssignment`] if the consumer does not own the
38 /// partition.
39 #[tracing::instrument(
40 name = "consumer.seek",
41 level = "debug",
42 skip_all,
43 fields(group_id = %self.group_id, topic = tracing::field::Empty, partition, offset),
44 err(level = "debug")
45 )]
46 pub async fn seek(
47 &self,
48 topic: impl Into<String>,
49 partition: i32,
50 offset: i64,
51 ) -> Result<(), ConsumerError> {
52 self.seek_to_position(topic.into(), partition, offset, None)
53 .await
54 }
55
56 /// Set the next offset that `poll` fetches for `(topic, partition)`, with
57 /// the leader epoch of the record before `offset`.
58 ///
59 /// This is Kafka's `seek(TopicPartition, OffsetAndMetadata)` with a leader
60 /// epoch. When the consumer knows the current leader epoch of the
61 /// partition, it validates the position with `OffsetForLeaderEpoch` before
62 /// the next fetch, as Kafka's `SubscriptionState.seekUnvalidated` does. A
63 /// truncated log then resets the position by `auto.offset.reset`.
64 ///
65 /// # Errors
66 ///
67 /// Returns [`ConsumerError::InvalidOffset`] if `offset` is negative, and
68 /// [`ConsumerError::NoCurrentAssignment`] if the consumer does not own the
69 /// partition.
70 #[tracing::instrument(
71 name = "consumer.seek_with_leader_epoch",
72 level = "debug",
73 skip_all,
74 fields(group_id = %self.group_id, topic = tracing::field::Empty, partition, offset, leader_epoch),
75 err(level = "debug")
76 )]
77 pub async fn seek_with_leader_epoch(
78 &self,
79 topic: impl Into<String>,
80 partition: i32,
81 offset: i64,
82 leader_epoch: i32,
83 ) -> Result<(), ConsumerError> {
84 self.seek_to_position(topic.into(), partition, offset, Some(leader_epoch))
85 .await
86 }
87
88 /// Reset the position of `partitions` to the log start. Kafka's
89 /// `KafkaConsumer.seekToBeginning`.
90 ///
91 /// An empty slice resets every assigned partition. The next `poll` or
92 /// [`position`](Self::position) sends `ListOffsets` with timestamp `-2`
93 /// and fetches from the offset that it returns.
94 ///
95 /// # Errors
96 ///
97 /// Returns [`ConsumerError::NoCurrentAssignment`] if the consumer does not
98 /// own a partition. The call then resets no partition.
99 pub async fn seek_to_beginning(
100 &self,
101 partitions: &[(String, i32)],
102 ) -> Result<(), ConsumerError> {
103 self.request_offset_reset(partitions, poll::BEGINNING_SENTINEL)
104 .await
105 }
106
107 /// Reset the position of `partitions` to the log end. Kafka's
108 /// `KafkaConsumer.seekToEnd`.
109 ///
110 /// An empty slice resets every assigned partition. The next `poll` or
111 /// [`position`](Self::position) sends `ListOffsets` with timestamp `-1`.
112 /// With `read_committed`, that is the last stable offset.
113 ///
114 /// # Errors
115 ///
116 /// Returns [`ConsumerError::NoCurrentAssignment`] if the consumer does not
117 /// own a partition. The call then resets no partition.
118 pub async fn seek_to_end(&self, partitions: &[(String, i32)]) -> Result<(), ConsumerError> {
119 self.request_offset_reset(partitions, poll::END_SENTINEL)
120 .await
121 }
122
123 /// Kafka's `SubscriptionState.requestOffsetReset`: plant `sentinel` for each
124 /// partition, so that the next position update resolves it.
125 async fn request_offset_reset(
126 &self,
127 partitions: &[(String, i32)],
128 sentinel: i64,
129 ) -> Result<(), ConsumerError> {
130 let assigned = self.assigned.lock().await;
131 let targets = if partitions.is_empty() {
132 assigned.clone()
133 } else {
134 if let Some((topic, partition)) = partitions
135 .iter()
136 .find(|partition| !assigned.contains(*partition))
137 {
138 return Err(ConsumerError::NoCurrentAssignment {
139 topic: topic.clone(),
140 partition: *partition,
141 });
142 }
143 partitions.to_vec()
144 };
145 let mut offsets = self.next_offsets.lock().await;
146 let mut positions = self.positions.lock().await;
147 for key in &targets {
148 offsets.insert(key.clone(), sentinel);
149 seek_position(positions.entry(key.clone()).or_default(), None);
150 }
151 drop(positions);
152 drop(offsets);
153 drop(assigned);
154 Ok(())
155 }
156
157 async fn seek_to_position(
158 &self,
159 topic: String,
160 partition: i32,
161 offset: i64,
162 leader_epoch: Option<i32>,
163 ) -> Result<(), ConsumerError> {
164 if offset < 0 {
165 return Err(ConsumerError::InvalidOffset(offset));
166 }
167 tracing::Span::current().record("topic", tracing::field::display(&topic));
168 let key = (topic, partition);
169 // Lock order: `assigned`, then `next_offsets`, then `positions`, as
170 // `drain_fetch_buffer` in poll.rs.
171 let assigned = self.assigned.lock().await;
172 if !assigned.contains(&key) {
173 let (topic, partition) = key;
174 return Err(ConsumerError::NoCurrentAssignment { topic, partition });
175 }
176 let mut offsets = self.next_offsets.lock().await;
177 let mut positions = self.positions.lock().await;
178 offsets.insert(key.clone(), offset);
179 seek_position(positions.entry(key.clone()).or_default(), leader_epoch);
180 drop(positions);
181 drop(offsets);
182 drop(assigned);
183 Ok(())
184 }
185}
186
187/// Set the leader epoch state of a sought position.
188///
189/// Kafka's `SubscriptionState.TopicPartitionState.validatePosition` waits for
190/// validation when the position has an offset epoch and the current leader
191/// epoch is known. Without an offset epoch there is nothing to validate.
192fn seek_position(position: &mut position::PartitionPosition, leader_epoch: Option<i32>) {
193 position.offset_epoch = LeaderEpoch(leader_epoch.unwrap_or(-1));
194 position.awaiting_validation =
195 position.offset_epoch.is_known() && position.leader_epoch.is_known();
196}
197
198#[cfg(test)]
199mod tests {
200 use krabka_ids::LeaderEpoch;
201
202 use super::{position::PartitionPosition, seek_position};
203
204 /// Kafka's `validatePosition` waits for validation only when both the
205 /// offset epoch and the current leader epoch are known.
206 #[test]
207 fn seek_position_follows_kafkas_validate_position() {
208 for (name, current_leader_epoch, awaiting, seek_epoch, expected) in [
209 ("no epoch", 5, true, None, (-1, false)),
210 ("epoch, leader known", 5, false, Some(4), (4, true)),
211 ("epoch, leader unknown", -1, false, Some(4), (4, false)),
212 ] {
213 let mut position = PartitionPosition {
214 offset_epoch: LeaderEpoch(3),
215 leader_id: 1,
216 leader_epoch: LeaderEpoch(current_leader_epoch),
217 awaiting_validation: awaiting,
218 };
219 seek_position(&mut position, seek_epoch);
220 assert2::assert!(
221 position
222 == PartitionPosition {
223 offset_epoch: LeaderEpoch(expected.0),
224 leader_id: 1,
225 leader_epoch: LeaderEpoch(current_leader_epoch),
226 awaiting_validation: expected.1,
227 },
228 "{name}"
229 );
230 }
231 }
232}