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