Support seek for chunk messages

Author: RobertIndie

lib/Commands.cc

@@ -28,6 +28,7 @@
 
 #include "BatchMessageAcker.h"
 #include "BatchedMessageIdImpl.h"
+#include "ChunkMessageIdImpl.h"
 #include "LogUtils.h"
 #include "MessageImpl.h"
 #include "PulsarApi.pb.h"
@@ -512,8 +513,17 @@ SharedBuffer Commands::newSeek(uint64_t consumerId, uint64_t requestId, const Me
     commandSeek->set_request_id(requestId);
 
     MessageIdData& messageIdData = *commandSeek->mutable_message_id();
-    messageIdData.set_ledgerid(messageId.ledgerId());
-    messageIdData.set_entryid(messageId.entryId());
+
+    auto chunkMsgId = std::dynamic_pointer_cast<ChunkMessageIdImpl>(messageId.impl_);
+    if (chunkMsgId) {
+        auto firstId = chunkMsgId->getFirstChunkMessageId();
+        messageIdData.set_ledgerid(firstId->ledgerId_);
+        messageIdData.set_entryid(firstId->entryId_);
+    } else {
+        messageIdData.set_ledgerid(messageId.ledgerId());
+        messageIdData.set_entryid(messageId.entryId());
+    }
+
     return writeMessageWithSize(cmd);
 }
 

lib/ConsumerImpl.h

@@ -308,8 +308,7 @@ class ConsumerImpl : public ConsumerImplBase {
     boost::optional<SharedBuffer> processMessageChunk(const SharedBuffer& payload,
                                                       const proto::MessageMetadata& metadata,
                                                       const proto::MessageIdData& messageIdData,
-                                                      const ClientConnectionPtr& cnx,
-                                                      MessageId& messageId);
+                                                      const ClientConnectionPtr& cnx, MessageId& messageId);
 
     friend class PulsarFriend;
 

lib/MessageId.cc

@@ -70,15 +70,14 @@ void MessageId::serialize(std::string& result) const {
     }
 
     auto chunkMsgId = std::dynamic_pointer_cast<ChunkMessageIdImpl>(impl_);
-    if(chunkMsgId) {
-        auto* firstChunkIdData = new proto::MessageIdData();
+    if (chunkMsgId) {
+        proto::MessageIdData& firstChunkIdData = *idData.mutable_first_chunk_message_id();
         auto firstChunkId = chunkMsgId->getFirstChunkMessageId();
-        firstChunkIdData->set_ledgerid(firstChunkId->ledgerId_);
-        firstChunkIdData->set_entryid(firstChunkId->entryId_);
-        if(chunkMsgId->partition_!=-1){
-            firstChunkIdData->set_partition(firstChunkId->partition_);
+        firstChunkIdData.set_ledgerid(firstChunkId->ledgerId_);
+        firstChunkIdData.set_entryid(firstChunkId->entryId_);
+        if (chunkMsgId->partition_ != -1) {
+            firstChunkIdData.set_partition(firstChunkId->partition_);
         }
-        idData.set_allocated_first_chunk_message_id(firstChunkIdData);
     }
 
     idData.SerializeToString(&result);
@@ -95,7 +94,7 @@ MessageId MessageId::deserialize(const std::string& serializedMessageId) {
 
     MessageId msgId = MessageIdBuilder::from(idData).build();
 
-    if(idData.has_first_chunk_message_id()) {
+    if (idData.has_first_chunk_message_id()) {
         ChunkMessageIdImplPtr chunkMsgId = std::make_shared<ChunkMessageIdImpl>();
         chunkMsgId->setFirstChunkMessageId(MessageIdBuilder::from(idData.first_chunk_message_id()).build());
         chunkMsgId->setLastChunkMessageId(msgId);
@@ -117,10 +116,10 @@ int32_t MessageId::batchSize() const { return impl_->batchSize_; }
 
 PULSAR_PUBLIC std::ostream& operator<<(std::ostream& s, const pulsar::MessageId& messageId) {
     auto chunkMsgId = std::dynamic_pointer_cast<ChunkMessageIdImpl>(messageId.impl_);
-    if(chunkMsgId) {
+    if (chunkMsgId) {
         auto firstId = chunkMsgId->getFirstChunkMessageId();
-        s << '(' << firstId->ledgerId_ << ',' << firstId->entryId_ << ','
-          << firstId->partition_ << ',' << firstId->batchIndex_ << ");";
+        s << '(' << firstId->ledgerId_ << ',' << firstId->entryId_ << ',' << firstId->partition_ << ','
+          << firstId->batchIndex_ << ");";
     }
     s << '(' << messageId.impl_->ledgerId_ << ',' << messageId.impl_->entryId_ << ','
       << messageId.impl_->partition_ << ',' << messageId.impl_->batchIndex_ << ')';

tests/MessageChunkingTest.cc

@@ -23,6 +23,7 @@
 #include <ctime>
 #include <random>
 
+#include "ChunkMessageIdImpl.h"
 #include "PulsarFriend.h"
 #include "WaitUtils.h"
 #include "lib/LogUtils.h"
@@ -117,7 +118,8 @@ TEST_P(MessageChunkingTest, testEndToEnd) {
     for (int i = 0; i < numMessages; i++) {
         MessageId messageId;
         ASSERT_EQ(ResultOk, producer.send(MessageBuilder().setContent(largeMessage).build(), messageId));
-        auto chunkMsgId = std::dynamic_pointer_cast<ChunkMessageIdImpl>(PulsarFriend::getMessageIdImpl(messageId));
+        auto chunkMsgId =
+            std::dynamic_pointer_cast<ChunkMessageIdImpl>(PulsarFriend::getMessageIdImpl(messageId));
         ASSERT_TRUE(chunkMsgId);
         LOG_INFO("Send " << i << " to " << messageId);
         sendMessageIds.emplace_back(messageId);
@@ -132,7 +134,8 @@ TEST_P(MessageChunkingTest, testEndToEnd) {
         ASSERT_EQ(msg.getMessageId().batchIndex(), -1);
         ASSERT_EQ(msg.getMessageId().batchSize(), 0);
         auto messageId = msg.getMessageId();
-        auto chunkMsgId = std::dynamic_pointer_cast<ChunkMessageIdImpl>(PulsarFriend::getMessageIdImpl(messageId));
+        auto chunkMsgId =
+            std::dynamic_pointer_cast<ChunkMessageIdImpl>(PulsarFriend::getMessageIdImpl(messageId));
         ASSERT_TRUE(chunkMsgId);
         receivedMessageIds.emplace_back(messageId);
     }
@@ -261,6 +264,56 @@ TEST_P(MessageChunkingTest, testMaxPendingChunkMessages) {
     consumer.close();
 }
 
+TEST_P(MessageChunkingTest, testSeekChunkMessages) {
+    const std::string topic =
+        "MessageChunkingTest-testSeekChunkMessages-" + toString(GetParam()) + std::to_string(time(nullptr));
+
+    constexpr int numMessages = 10;
+
+    Consumer consumer1;
+    ConsumerConfiguration consumer1Conf;
+    consumer1Conf.setStartMessageIdInclusive(true);
+    createConsumer(topic, consumer1, consumer1Conf);
+
+    Producer producer;
+    createProducer(topic, producer);
+
+    for (int i = 0; i < numMessages; i++) {
+        MessageId messageId;
+        ASSERT_EQ(ResultOk, producer.send(MessageBuilder().setContent(largeMessage).build(), messageId));
+        LOG_INFO("Send " << i << " to " << messageId);
+    }
+
+    Message msg;
+    std::vector<MessageId> receivedMessageIds;
+    for (int i = 0; i < numMessages; i++) {
+        ASSERT_EQ(ResultOk, consumer1.receive(msg, 3000));
+        LOG_INFO("Receive " << msg.getLength() << " bytes from " << msg.getMessageId());
+        receivedMessageIds.emplace_back(msg.getMessageId());
+    }
+
+    consumer1.seek(receivedMessageIds[1]);
+    for (int i = 1; i < numMessages; i++) {
+        Message msgAfterSeek;
+        ASSERT_EQ(ResultOk, consumer1.receive(msgAfterSeek, 3000));
+        ASSERT_EQ(msgAfterSeek.getMessageId(), receivedMessageIds[i]);
+    }
+
+    consumer1.close();
+    Consumer consumer2;
+    createConsumer(topic, consumer2);
+
+    consumer2.seek(receivedMessageIds[1]);
+    for (int i = 2; i < numMessages; i++) {
+        Message msgAfterSeek;
+        ASSERT_EQ(ResultOk, consumer2.receive(msgAfterSeek, 3000));
+        ASSERT_EQ(msgAfterSeek.getMessageId(), receivedMessageIds[i]);
+    }
+
+    consumer2.close();
+    producer.close();
+}
+
 TEST(ChunkMessageIdTest, testSetChunkMessageId) {
     MessageId msgId;
     {

tests/PulsarFriend.h

@@ -180,9 +180,7 @@ class PulsarFriend {
 
     static proto::MessageMetadata& getMessageMetadata(Message& message) { return message.impl_->metadata; }
 
-    static std::shared_ptr<MessageIdImpl> getMessageIdImpl(MessageId& msgId) {
-        return msgId.impl_;
-    }
+    static std::shared_ptr<MessageIdImpl> getMessageIdImpl(MessageId& msgId) { return msgId.impl_; }
 };
 }  // namespace pulsar