Add Consumer::seek method

Author: benesch

changelog.md

@@ -8,6 +8,7 @@
 * Remove build-time dependency on bindgen, clang, and libclang.
 * Move zstd compression support behind the `zstd` feature flag.
 * Ensure all features are honored in the CMake build system.
+* Add [`Consumer::seek`] method.
 
 
 <a name="0.21.0"></a>

src/client.rs

@@ -273,7 +273,7 @@ impl<C: ClientContext> Client<C> {
 
     /// Returns a NativeTopic from the current client. The NativeTopic shouldn't outlive the client
     /// it was generated from.
-    fn native_topic(&self, topic: &str) -> KafkaResult<NativeTopic> {
+    pub(crate) fn native_topic(&self, topic: &str) -> KafkaResult<NativeTopic> {
         let topic_c = CString::new(topic.to_string())?;
         Ok(unsafe {
             NativeTopic::from_ptr(rdsys::rd_kafka_topic_new(
@@ -291,7 +291,7 @@ impl<C: ClientContext> Client<C> {
     }
 }
 
-struct NativeTopic {
+pub(crate) struct NativeTopic {
     ptr: *mut RDKafkaTopic,
 }
 
@@ -305,7 +305,7 @@ impl NativeTopic {
     }
 
     /// Returns the pointer to the librdkafka RDKafkaTopic structure.
-    fn ptr(&self) -> *mut RDKafkaTopic {
+    pub fn ptr(&self) -> *mut RDKafkaTopic {
         self.ptr
     }
 

src/consumer/base_consumer.rs

@@ -9,8 +9,7 @@ use crate::error::{IsError, KafkaError, KafkaResult};
 use crate::groups::GroupList;
 use crate::message::{BorrowedMessage, Message};
 use crate::metadata::Metadata;
-use crate::topic_partition_list::Offset::Offset;
-use crate::topic_partition_list::TopicPartitionList;
+use crate::topic_partition_list::{Offset, TopicPartitionList};
 use crate::util::{cstr_to_owned, timeout_to_ms};
 
 use std::mem;
@@ -208,6 +207,26 @@ impl<C: ConsumerContext> Consumer<C> for BaseConsumer<C> {
         Ok(())
     }
 
+    fn seek<T>(&self, topic: &str, partition: i32, offset: Offset, timeout: T) -> KafkaResult<()>
+    where
+        T: Into<Option<Duration>>,
+    {
+        let topic = self.client.native_topic(topic)?;
+        let ret_code = unsafe {
+            rdsys::rd_kafka_seek(
+                topic.ptr(),
+                partition,
+                offset.to_raw(),
+                timeout_to_ms(timeout),
+            )
+        };
+        if ret_code.is_error() {
+            let error = unsafe { cstr_to_owned(rdsys::rd_kafka_err2str(ret_code)) };
+            return Err(KafkaError::Seek(error));
+        };
+        Ok(())
+    }
+
     fn commit(
         &self,
         topic_partition_list: &TopicPartitionList,
@@ -332,7 +351,7 @@ impl<C: ConsumerContext> Consumer<C> for BaseConsumer<C> {
         let mut tpl = unsafe { TopicPartitionList::from_ptr(tpl_ptr) };
 
         // Set the timestamp we want in the offset field for every partition as librdkafka expects.
-        tpl.set_all_offsets(Offset(timestamp));
+        tpl.set_all_offsets(Offset::Offset(timestamp));
 
         self.offsets_for_times(tpl, timeout)
     }

src/consumer/mod.rs

@@ -19,7 +19,7 @@ use crate::util::cstr_to_owned;
 use std::ptr;
 use std::time::Duration;
 
-use crate::topic_partition_list::TopicPartitionList;
+use crate::topic_partition_list::{Offset, TopicPartitionList};
 
 /// Rebalance information.
 #[derive(Clone, Debug)]
@@ -139,6 +139,17 @@ pub trait Consumer<C: ConsumerContext = DefaultConsumerContext> {
         self.get_base_consumer().assign(assignment)
     }
 
+    /// Seek to `offset` for the specified `topic` and `partition`. After a
+    /// successful call to `seek`, the next poll of the consumer will return the
+    /// message with `offset`.
+    fn seek<T>(&self, topic: &str, partition: i32, offset: Offset, timeout: T) -> KafkaResult<()>
+    where
+        T: Into<Option<Duration>>,
+    {
+        self.get_base_consumer()
+            .seek(topic, partition, offset, timeout)
+    }
+
     /// Commits the offset of the specified message. The commit can be sync (blocking), or async.
     /// Notice that when a specific offset is committed, all the previous offsets are considered
     /// committed as well. Use this method only if you are processing messages in order.

src/error.rs

@@ -64,6 +64,8 @@ pub enum KafkaError {
     PartitionEOF(i32),
     /// Pause/Resume failed.
     PauseResume(String),
+    /// Seeking a partition failed.
+    Seek(String),
     /// Setting partition offset failed.
     SetPartitionOffset(RDKafkaError),
     /// Offset store failed.
@@ -113,6 +115,9 @@ impl fmt::Debug for KafkaError {
             KafkaError::PauseResume(ref err) => {
                 write!(f, "KafkaError (Pause/resume error: {})", err)
             }
+            KafkaError::Seek(ref err) => {
+                write!(f, "KafkaError (Seek error: {})", err)
+            }
             KafkaError::SetPartitionOffset(err) => {
                 write!(f, "KafkaError (Set partition offset error: {})", err)
             }
@@ -149,6 +154,7 @@ impl fmt::Display for KafkaError {
             KafkaError::OffsetFetch(err) => write!(f, "Offset fetch error: {}", err),
             KafkaError::PartitionEOF(part_n) => write!(f, "Partition EOF: {}", part_n),
             KafkaError::PauseResume(ref err) => write!(f, "Pause/resume error: {}", err),
+            KafkaError::Seek(ref err) => write!(f, "Seek error: {}", err),
             KafkaError::SetPartitionOffset(err) => write!(f, "Set partition offset error: {}", err),
             KafkaError::StoreOffset(err) => write!(f, "Store offset error: {}", err),
             KafkaError::Subscription(ref err) => write!(f, "Subscription error: {}", err),
@@ -175,6 +181,7 @@ impl error::Error for KafkaError {
             KafkaError::OffsetFetch(_) => "Offset fetch error",
             KafkaError::PartitionEOF(_) => "Partition EOF error",
             KafkaError::PauseResume(_) => "Pause/resume error",
+            KafkaError::Seek(_) => "Seek error",
             KafkaError::SetPartitionOffset(_) => "Set partition offset error",
             KafkaError::StoreOffset(_) => "Store offset error",
             KafkaError::Subscription(_) => "Subscription error",
@@ -200,6 +207,7 @@ impl error::Error for KafkaError {
             KafkaError::OffsetFetch(ref err) => Some(err),
             KafkaError::PartitionEOF(_) => None,
             KafkaError::PauseResume(_) => None,
+            KafkaError::Seek(_) => None,
             KafkaError::SetPartitionOffset(ref err) => Some(err),
             KafkaError::StoreOffset(ref err) => Some(err),
             KafkaError::Subscription(_) => None,

tests/test_consumers.rs

@@ -153,6 +153,37 @@ fn test_produce_consume_base() {
         .wait();
 }
 
+// Seeking should allow replaying messages and skipping messages.
+#[test]
+fn test_produce_consume_seek() {
+    let _r = env_logger::try_init();
+
+    let topic_name = rand_test_topic();
+    populate_topic(&topic_name, 5, &value_fn, &key_fn, Some(0), None);
+    let consumer = create_base_consumer(&rand_test_group(), None);
+    consumer.subscribe(&[topic_name.as_str()]).unwrap();
+
+    for (i, message) in consumer.iter().take(3).enumerate() {
+        match message {
+            Ok(message) => assert_eq!(dbg!(message.offset()), i as i64),
+            Err(e) => panic!("Error receiving message: {:?}", e),
+        }
+    }
+
+    consumer.seek(&topic_name, 0, Offset::Offset(1), None).unwrap();
+
+    for (i, message) in consumer.iter().take(3).enumerate() {
+        match message {
+            Ok(message) => assert_eq!(message.offset(), i as i64 + 1),
+            Err(e) => panic!("Error receiving message: {:?}", e),
+        }
+    }
+
+    consumer.seek(&topic_name, 0, Offset::End, None).unwrap();
+
+    ensure_empty(&consumer, "There should be no messages left");
+}
+
 // All produced messages should be consumed.
 #[test]
 fn test_produce_consume_base_assign() {