Refactor and refine chunked message id

Author: RobertIndie

lib/ChunkMessageIdImpl.h

@@ -30,23 +30,17 @@ class ChunkMessageIdImpl : public MessageIdImpl, public std::enable_shared_from_
    public:
     ChunkMessageIdImpl() : firstChunkMsgId_(std::make_shared<MessageIdImpl>()) {}
 
-    void setFirstChunkMessageId(const MessageId& msgId) { *firstChunkMsgId_ = *msgId.impl_; }
-
-    void setLastChunkMessageId(const MessageId& msgId) {
-        this->ledgerId_ = msgId.ledgerId();
-        this->entryId_ = msgId.entryId();
-        this->partition_ = msgId.partition();
-    }
-
     void setChunkedMessageIds(std::vector<MessageId>&& chunkedMessageIds) {
         chunkedMessageIds_ = std::move(chunkedMessageIds);
-        setFirstChunkMessageId(chunkedMessageIds_.front());
-        setLastChunkMessageId(chunkedMessageIds_.back());
+        auto lastChunkMsgId = chunkedMessageIds_.back();
+        this->ledgerId_ = lastChunkMsgId.ledgerId();
+        this->entryId_ = lastChunkMsgId.entryId();
+        this->partition_ = lastChunkMsgId.partition();
     }
 
-    std::shared_ptr<const MessageIdImpl> getFirstChunkMessageId() const { return firstChunkMsgId_; }
+    std::shared_ptr<const MessageIdImpl> getFirstChunkMessageId() const { return chunkedMessageIds_.front().impl_; }
 
-    std::vector<MessageId> moveChunkedMessageIds() const noexcept {
+    std::vector<MessageId> moveChunkedMessageIds() noexcept {
         return std::move(chunkedMessageIds_);
     }
 

lib/ConsumerImpl.h

@@ -270,7 +270,7 @@ class ConsumerImpl : public ConsumerImplBase {
 
         const std::vector<MessageId>& getChunkedMessageIds() const noexcept { return chunkedMessageIds_; }
 
-        std::vector<MessageId> moveChunkedMessageIds() const noexcept {
+        std::vector<MessageId> moveChunkedMessageIds() noexcept {
             return std::move(chunkedMessageIds_);
         }
 

lib/MessageId.cc

@@ -100,8 +100,8 @@ MessageId MessageId::deserialize(const std::string& serializedMessageId) {
 
     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);
+        chunkMsgId->setChunkedMessageIds(
+            {MessageIdBuilder::from(idData.first_chunk_message_id()).build(), msgId});
         return chunkMsgId->build();
     }
 

lib/OpSendMsg.h

@@ -55,6 +55,7 @@ struct OpSendMsg {
     const SendCallback sendCallback;
     std::vector<std::function<void(Result)>> trackerCallbacks;
     ChunkMessageIdImplPtr chunkedMessageId;
+    std::vector<MessageId> chunkMessageIdList;
     // Use shared_ptr here because producer might resend the message with the same arguments
     const std::shared_ptr<SendArguments> sendArgs;
 

lib/ProducerImpl.cc

@@ -886,7 +886,7 @@ bool ProducerImpl::ackReceived(uint64_t sequenceId, MessageId& rawMessageId) {
         return true;
     }
 
-    const auto& op = *pendingMessagesQueue_.front();
+    auto& op = *pendingMessagesQueue_.front();
     if (op.result != ResultOk) {
         LOG_ERROR("Unexpected OpSendMsg whose result is " << op.result << " for " << sequenceId << " and "
                                                           << rawMessageId);
@@ -912,10 +912,9 @@ bool ProducerImpl::ackReceived(uint64_t sequenceId, MessageId& rawMessageId) {
 
     if (op.chunkedMessageId) {
         // Handling the chunk message id.
-        if (op.chunkId == 0) {
-            op.chunkedMessageId->setFirstChunkMessageId(messageId);
-        } else if (op.chunkId == op.numChunks - 1) {
-            op.chunkedMessageId->setLastChunkMessageId(messageId);
+        op.chunkMessageIdList.push_back(messageId);
+        if (op.chunkId == op.numChunks - 1) {
+            op.chunkedMessageId->setChunkedMessageIds(std::move(op.chunkMessageIdList));
             messageId = op.chunkedMessageId->build();
         }
     }

tests/MessageChunkingTest.cc

@@ -154,10 +154,8 @@ TEST_P(MessageChunkingTest, testEndToEnd) {
     waitUntil(
         std::chrono::seconds(10),
         [&] {
-            if (consumer.getBrokerConsumerStats(consumerStats) != ResultOk) {
-                return false;
-            }
-            return consumerStats.getMsgBacklog() == 0;
+            return consumer.getBrokerConsumerStats(consumerStats) == ResultOk &&
+                   consumerStats.getMsgBacklog() == 0;
         },
         1000);
     ASSERT_EQ(consumerStats.getMsgBacklog(), 0);
@@ -333,8 +331,8 @@ TEST(ChunkMessageIdTest, testSetChunkMessageId) {
     MessageId msgId;
     {
         ChunkMessageIdImplPtr chunkMsgId = std::make_shared<ChunkMessageIdImpl>();
-        chunkMsgId->setFirstChunkMessageId(MessageIdBuilder().ledgerId(1).entryId(2).partition(3).build());
-        chunkMsgId->setLastChunkMessageId(MessageIdBuilder().ledgerId(4).entryId(5).partition(6).build());
+        chunkMsgId->setChunkedMessageIds({MessageIdBuilder().ledgerId(1).entryId(2).partition(3).build(),
+                                          MessageIdBuilder().ledgerId(4).entryId(5).partition(6).build()});
         msgId = chunkMsgId->build();
         // Test the destructor of the underlying message id should also work for the generated messageId.
     }