Fix default timeout for beginningOffsets/endOffsets

morganstanley · Nov 8, 2022 · 0ec1511 · 0ec1511
1 parent 3a38e1c
commit 0ec1511
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 6 deletions.
diff --git a/include/kafka/KafkaConsumer.h b/include/kafka/KafkaConsumer.h
@@ -139,15 +139,23 @@ class KafkaConsumer: public KafkaClient
      * Throws KafkaException with errors:
      *   - RD_KAFKA_RESP_ERR__FAIL:  Generic failure
      */
-    std::map<TopicPartition, Offset> beginningOffsets(const TopicPartitions& topicPartitions) const { return getOffsets(topicPartitions, true); }
+    std::map<TopicPartition, Offset> beginningOffsets(const TopicPartitions&    topicPartitions,
+                                                      std::chrono::milliseconds timeout = std::chrono::milliseconds(DEFAULT_QUERY_TIMEOUT_MS)) const
+    {
+        return getOffsets(topicPartitions, true, timeout);
+    }
 
     /**
      * Get the last offset for the given partitions.  The last offset of a partition is the offset of the upcoming message, i.e. the offset of the last available message + 1.
      * This method does not change the current consumer position of the partitions.
      * Throws KafkaException with errors:
      *   - RD_KAFKA_RESP_ERR__FAIL:  Generic failure
      */
-    std::map<TopicPartition, Offset> endOffsets(const TopicPartitions& topicPartitions) const { return getOffsets(topicPartitions, false); }
+    std::map<TopicPartition, Offset> endOffsets(const TopicPartitions&    topicPartitions,
+                                                std::chrono::milliseconds timeout = std::chrono::milliseconds(DEFAULT_QUERY_TIMEOUT_MS)) const
+    {
+        return getOffsets(topicPartitions, false, timeout);
+    }
 
     /**
      * Get the offsets for the given partitions by time-point.
@@ -284,7 +292,9 @@ class KafkaConsumer: public KafkaClient
     void storeOffsetsIfNecessary(const std::vector<consumer::ConsumerRecord>& records);
 
     void seekToBeginningOrEnd(const TopicPartitions& topicPartitions, bool toBeginning, std::chrono::milliseconds timeout);
-    std::map<TopicPartition, Offset> getOffsets(const TopicPartitions& topicPartitions, bool atBeginning) const;
+    std::map<TopicPartition, Offset> getOffsets(const TopicPartitions&    topicPartitions,
+                                                bool                      atBeginning,
+                                                std::chrono::milliseconds timeout) const;
 
     enum class PartitionsRebalanceEvent { Assign, Revoke, IncrementalAssign, IncrementalUnassign };
     void changeAssignment(PartitionsRebalanceEvent event, const TopicPartitions& tps);
@@ -739,14 +749,16 @@ KafkaConsumer::offsetsForTime(const TopicPartitions& topicPartitions,
 }
 
 inline std::map<TopicPartition, Offset>
-KafkaConsumer::getOffsets(const TopicPartitions& topicPartitions, bool atBeginning) const
+KafkaConsumer::getOffsets(const TopicPartitions&    topicPartitions,
+                          bool                      atBeginning,
+                          std::chrono::milliseconds timeout) const
 {
     std::map<TopicPartition, Offset> result;
 
     for (const auto& topicPartition: topicPartitions)
     {
         Offset beginning{}, end{};
-        Error error{ rd_kafka_query_watermark_offsets(getClientHandle(), topicPartition.first.c_str(), topicPartition.second, &beginning, &end, 0) };
+        Error error{ rd_kafka_query_watermark_offsets(getClientHandle(), topicPartition.first.c_str(), topicPartition.second, &beginning, &end, timeout.count()) };
         KAFKA_THROW_IF_WITH_ERROR(error);
 
         result[topicPartition] = (atBeginning ? beginning : end);

diff --git a/tests/integration/TestKafkaConsumer.cc b/tests/integration/TestKafkaConsumer.cc
@@ -213,9 +213,12 @@ TEST(KafkaConsumer, SeekAndPoll)
     auto records = KafkaTestUtility::ConsumeMessagesUntilTimeout(consumer, std::chrono::seconds(1));
     EXPECT_EQ(0, records.size());
 
-    // Try to get the beginning offsets
+    // Try to get the beginning/end offsets
     const kafka::TopicPartition tp{topic, partition};
     std::cout << "[" << kafka::utility::getCurrentTime() << "] Consumer get the beginningOffset[" << consumer.beginningOffsets({tp})[tp] << "]" << std::endl;;
+    std::cout << "[" << kafka::utility::getCurrentTime() << "] Consumer get the endOffset[" << consumer.endOffsets({tp})[tp] << "]" << std::endl;;
+    EXPECT_EQ(0, consumer.beginningOffsets({tp})[tp]);
+    EXPECT_EQ(0, consumer.endOffsets({tp})[tp]);
 
     // Prepare some messages to send
     std::vector<std::tuple<kafka::Headers, std::string, std::string>> messages = {
@@ -227,6 +230,12 @@ TEST(KafkaConsumer, SeekAndPoll)
     // Send the messages
     KafkaTestUtility::ProduceMessages(topic, partition, messages);
 
+    // Try to get the beginning/end offsets
+    std::cout << "[" << kafka::utility::getCurrentTime() << "] Consumer get the beginningOffset[" << consumer.beginningOffsets({tp})[tp] << "]" << std::endl;;
+    std::cout << "[" << kafka::utility::getCurrentTime() << "] Consumer get the endOffset[" << consumer.endOffsets({tp})[tp] << "]" << std::endl;;
+    EXPECT_EQ(0, consumer.beginningOffsets({tp})[tp]);
+    EXPECT_EQ(messages.size(), consumer.endOffsets({tp})[tp]);
+
     // Poll these messages
     records = KafkaTestUtility::ConsumeMessagesUntilTimeout(consumer);
     EXPECT_EQ(messages.size(), records.size());