Merge pull request #596 from fede1024/scanterog/seek-partitions

Change signature for seek_partitions

Author: davidblewett

changelog.md

@@ -4,6 +4,10 @@ See also the [rdkafka-sys changelog](rdkafka-sys/changelog.md).
 
 ## Unreleased
 
+* **Breaking change.** Change signature for `seek_partitions`. Following
+  librdkafka, individual partition errors should be reported in the per-partition
+  `error` field of `TopicPartitionList` elements.
+
 ## 0.33.0 (2023-06-30)
 
 * Add interface to specify custom partitioners by extending `ProducerContext`

src/consumer/base_consumer.rs

@@ -352,9 +352,9 @@ where
 
     fn seek_partitions<T: Into<Timeout>>(
         &self,
-        topic_partition_list: &TopicPartitionList,
+        topic_partition_list: TopicPartitionList,
         timeout: T,
-    ) -> KafkaResult<()> {
+    ) -> KafkaResult<TopicPartitionList> {
         let ret = unsafe {
             RDKafkaError::from_ptr(rdsys::rd_kafka_seek_partitions(
                 self.client.native_ptr(),
@@ -366,7 +366,7 @@ where
             let error = ret.name();
             return Err(KafkaError::Seek(error));
         }
-        Ok(())
+        Ok(topic_partition_list)
     }
 
     fn commit(

src/consumer/mod.rs

@@ -261,11 +261,13 @@ where
     /// in the `offset` field of `TopicPartitionListElem`.
     /// The offset can be either absolute (>= 0) or a logical offset.
     /// Seek should only be performed on already assigned/consumed partitions.
+    /// Individual partition errors are reported in the per-partition `error` field of
+    /// `TopicPartitionListElem`.
     fn seek_partitions<T: Into<Timeout>>(
         &self,
-        topic_partition_list: &TopicPartitionList,
+        topic_partition_list: TopicPartitionList,
         timeout: T,
-    ) -> KafkaResult<()>;
+    ) -> KafkaResult<TopicPartitionList>;
 
     /// Commits the offset of the specified message. The commit can be sync
     /// (blocking), or async. Notice that when a specific offset is committed,

src/consumer/stream_consumer.rs

@@ -412,9 +412,9 @@ where
 
     fn seek_partitions<T: Into<Timeout>>(
         &self,
-        topic_partition_list: &TopicPartitionList,
+        topic_partition_list: TopicPartitionList,
         timeout: T,
-    ) -> KafkaResult<()> {
+    ) -> KafkaResult<TopicPartitionList> {
         self.base.seek_partitions(topic_partition_list, timeout)
     }
 

src/topic_partition_list.rs

@@ -387,11 +387,12 @@ impl fmt::Debug for TopicPartitionList {
             }
             write!(
                 f,
-                "{}/{}: offset={:?} metadata={:?}",
+                "{}/{}: offset={:?} metadata={:?}, error={:?}",
                 elem.topic(),
                 elem.partition(),
                 elem.offset(),
                 elem.metadata(),
+                elem.error(),
             )?;
         }
         write!(f, "}}")

tests/test_low_consumers.rs

@@ -1,6 +1,7 @@
 //! Test data consumption using low level consumers.
 
 use std::collections::HashMap;
+use std::convert::TryInto;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 use std::thread;
@@ -90,6 +91,67 @@ async fn test_produce_consume_seek() {
     }
 }
 
+// Seeking should allow replaying messages and skipping messages.
+#[tokio::test]
+async fn test_produce_consume_seek_partitions() {
+    let _r = env_logger::try_init();
+
+    let topic_name = rand_test_topic();
+    populate_topic(&topic_name, 30, &value_fn, &key_fn, None, None).await;
+
+    let consumer = create_base_consumer(&rand_test_group(), None);
+    consumer.subscribe(&[topic_name.as_str()]).unwrap();
+
+    let mut partition_offset_map = HashMap::new();
+    for message in consumer.iter().take(30) {
+        match message {
+            Ok(m) => {
+                let offset = partition_offset_map.entry(m.partition()).or_insert(0);
+                assert_eq!(m.offset(), *offset);
+                *offset += 1;
+            }
+            Err(e) => panic!("Error receiving message: {:?}", e),
+        }
+    }
+
+    let mut tpl = TopicPartitionList::new();
+    tpl.add_partition_offset(&topic_name, 0, Offset::Beginning)
+        .unwrap();
+    tpl.add_partition_offset(&topic_name, 1, Offset::End)
+        .unwrap();
+    tpl.add_partition_offset(&topic_name, 2, Offset::Offset(2))
+        .unwrap();
+
+    let r_tpl = consumer.seek_partitions(tpl, None).unwrap();
+    assert_eq!(r_tpl.elements().len(), 3);
+    for tpe in r_tpl.elements().iter() {
+        assert!(tpe.error().is_ok());
+    }
+
+    let msg_cnt_p0 = partition_offset_map.get(&0).unwrap();
+    let msg_cnt_p2 = partition_offset_map.get(&2).unwrap();
+    let total_msgs_to_read = msg_cnt_p0 + (msg_cnt_p2 - 2);
+    let mut poffset_map = HashMap::new();
+    for message in consumer.iter().take(total_msgs_to_read.try_into().unwrap()) {
+        match message {
+            Ok(m) => {
+                let offset = poffset_map.entry(m.partition()).or_insert(0);
+                if m.partition() == 0 {
+                    assert_eq!(m.offset(), *offset);
+                } else if m.partition() == 2 {
+                    assert_eq!(m.offset(), *offset + 2);
+                } else if m.partition() == 1 {
+                    panic!("Unexpected message from partition 1")
+                }
+                *offset += 1;
+            }
+            Err(e) => panic!("Error receiving message: {:?}", e),
+        }
+    }
+    assert_eq!(msg_cnt_p0, poffset_map.get(&0).unwrap());
+    assert_eq!(msg_cnt_p2 - 2, *poffset_map.get(&2).unwrap());
+}
+
 // All produced messages should be consumed.
 #[tokio::test]
 async fn test_produce_consume_iter() {