[fix] Fix consumer doesn't acknowledge all chunk message Ids

Author: RobertIndie

lib/AckGroupingTracker.cc

@@ -21,8 +21,10 @@
 
 #include <atomic>
 #include <limits>
+#include <set>
 
 #include "BitSet.h"
+#include "ChunkMessageIdImpl.h"
 #include "ClientConnection.h"
 #include "Commands.h"
 #include "LogUtils.h"
@@ -42,6 +44,14 @@ void AckGroupingTracker::doImmediateAck(const MessageId& msgId, ResultCallback c
         }
         return;
     }
+    if (auto chunkMessageId =
+            std::dynamic_pointer_cast<ChunkMessageIdImpl>(Commands::getMessageIdImpl(msgId))) {
+        auto msgIdList = chunkMessageId->moveChunkedMessageIds();
+        doImmediateAck(std::set<MessageId>(std::make_move_iterator(msgIdList.begin()),
+                                           std::make_move_iterator(msgIdList.end())),
+                       callback);
+        return;
+    }
     const auto& ackSet = Commands::getMessageIdImpl(msgId)->getBitSet();
     if (waitResponse_) {
         const auto requestId = requestIdSupplier_();
@@ -84,29 +94,43 @@ void AckGroupingTracker::doImmediateAck(const std::set<MessageId>& msgIds, Resul
         return;
     }
 
+    std::set<MessageId> ackMsgIds;
+
+    for (const auto& msgId : msgIds) {
+        if (auto chunkMessageId =
+                std::dynamic_pointer_cast<ChunkMessageIdImpl>(Commands::getMessageIdImpl(msgId))) {
+            auto msgIdList = chunkMessageId->moveChunkedMessageIds();
+            doImmediateAck(std::set<MessageId>(std::make_move_iterator(msgIdList.begin()),
+                                               std::make_move_iterator(msgIdList.end())),
+                           callback);
+        } else {
+            ackMsgIds.insert(msgId);
+        }
+    }
+
     if (Commands::peerSupportsMultiMessageAcknowledgement(cnx->getServerProtocolVersion())) {
         if (waitResponse_) {
             const auto requestId = requestIdSupplier_();
-            cnx->sendRequestWithId(Commands::newMultiMessageAck(consumerId_, msgIds, requestId), requestId)
+            cnx->sendRequestWithId(Commands::newMultiMessageAck(consumerId_, ackMsgIds, requestId), requestId)
                 .addListener([callback](Result result, const ResponseData&) {
                     if (callback) {
                         callback(result);
                     }
                 });
         } else {
-            cnx->sendCommand(Commands::newMultiMessageAck(consumerId_, msgIds));
+            cnx->sendCommand(Commands::newMultiMessageAck(consumerId_, ackMsgIds));
             if (callback) {
                 callback(ResultOk);
             }
         }
     } else {
-        auto count = std::make_shared<std::atomic<size_t>>(msgIds.size());
+        auto count = std::make_shared<std::atomic<size_t>>(ackMsgIds.size());
         auto wrappedCallback = [callback, count](Result result) {
             if (--*count == 0 && callback) {
                 callback(result);
             }
         };
-        for (auto&& msgId : msgIds) {
+        for (auto&& msgId : ackMsgIds) {
             doImmediateAck(msgId, wrappedCallback, CommandAck_AckType_Individual);
         }
     }

lib/ChunkMessageIdImpl.h

@@ -38,11 +38,22 @@ class ChunkMessageIdImpl : public MessageIdImpl, public std::enable_shared_from_
         this->partition_ = msgId.partition();
     }
 
+    void setChunkedMessageIds(std::vector<MessageId>&& chunkedMessageIds) {
+        chunkedMessageIds_ = std::move(chunkedMessageIds);
+        setFirstChunkMessageId(chunkedMessageIds_.front());
+        setLastChunkMessageId(chunkedMessageIds_.back());
+    }
+
     std::shared_ptr<const MessageIdImpl> getFirstChunkMessageId() const { return firstChunkMsgId_; }
 
+    std::vector<MessageId> moveChunkedMessageIds() const noexcept {
+        return std::move(chunkedMessageIds_);
+    }
+
     MessageId build() { return MessageId{std::dynamic_pointer_cast<MessageIdImpl>(shared_from_this())}; }
 
    private:
     std::shared_ptr<MessageIdImpl> firstChunkMsgId_;
+    std::vector<MessageId> chunkedMessageIds_;
 };
 }  // namespace pulsar

lib/ConsumerImpl.cc

@@ -479,8 +479,7 @@ boost::optional<SharedBuffer> ConsumerImpl::processMessageChunk(const SharedBuff
     }
 
     ChunkMessageIdImplPtr chunkMsgId = std::make_shared<ChunkMessageIdImpl>();
-    chunkMsgId->setFirstChunkMessageId(chunkedMsgCtx.getChunkedMessageIds().front());
-    chunkMsgId->setLastChunkMessageId(chunkedMsgCtx.getChunkedMessageIds().back());
+    chunkMsgId->setChunkedMessageIds(chunkedMsgCtx.moveChunkedMessageIds());
     messageId = chunkMsgId->build();
 
     LOG_DEBUG("Chunked message completed chunkId: " << chunkId << ", ChunkedMessageCtx: " << chunkedMsgCtx
@@ -1165,6 +1164,9 @@ bool ConsumerImpl::isCumulativeAcknowledgementAllowed(ConsumerType consumerType)
 
 std::pair<MessageId, bool> ConsumerImpl::prepareIndividualAck(const MessageId& messageId) {
     auto messageIdImpl = Commands::getMessageIdImpl(messageId);
+    if (std::dynamic_pointer_cast<ChunkMessageIdImpl>(messageIdImpl)) {
+        return std::make_pair(messageId, true);
+    }
     auto batchedMessageIdImpl = std::dynamic_pointer_cast<BatchedMessageIdImpl>(messageIdImpl);
 
     auto batchSize = messageId.batchSize();

lib/ConsumerImpl.h

@@ -270,6 +270,10 @@ class ConsumerImpl : public ConsumerImplBase {
 
         const std::vector<MessageId>& getChunkedMessageIds() const noexcept { return chunkedMessageIds_; }
 
+        std::vector<MessageId> moveChunkedMessageIds() const noexcept {
+            return std::move(chunkedMessageIds_);
+        }
+
         long getReceivedTimeMs() const noexcept { return receivedTimeMs_; }
 
         friend std::ostream& operator<<(std::ostream& os, const ChunkedMessageCtx& ctx) {

tests/MessageChunkingTest.cc

@@ -81,7 +81,9 @@ class MessageChunkingTest : public ::testing::TestWithParam<CompressionType> {
     }
 
     void createConsumer(const std::string& topic, Consumer& consumer) {
-        ASSERT_EQ(ResultOk, client_.subscribe(topic, "my-sub", consumer));
+        ConsumerConfiguration conf;
+        conf.setBrokerConsumerStatsCacheTimeInMs(1000);
+        ASSERT_EQ(ResultOk, client_.subscribe(topic, "my-sub", conf, consumer));
     }
 
     void createConsumer(const std::string& topic, Consumer& consumer, ConsumerConfiguration& conf) {
@@ -138,6 +140,7 @@ TEST_P(MessageChunkingTest, testEndToEnd) {
             std::dynamic_pointer_cast<ChunkMessageIdImpl>(PulsarFriend::getMessageIdImpl(messageId));
         ASSERT_TRUE(chunkMsgId);
         receivedMessageIds.emplace_back(messageId);
+        consumer.acknowledge(messageId);
     }
     ASSERT_EQ(receivedMessageIds, sendMessageIds);
     ASSERT_EQ(receivedMessageIds.front().ledgerId(), receivedMessageIds.front().ledgerId());
@@ -147,6 +150,18 @@ TEST_P(MessageChunkingTest, testEndToEnd) {
     auto& chunkedMessageCache = PulsarFriend::getChunkedMessageCache(consumer);
     ASSERT_EQ(chunkedMessageCache.size(), 0);
 
+    BrokerConsumerStats consumerStats;
+    waitUntil(
+        std::chrono::seconds(10),
+        [&] {
+            if (consumer.getBrokerConsumerStats(consumerStats) != ResultOk) {
+                return false;
+            }
+            return consumerStats.getMsgBacklog() == 0;
+        },
+        1000);
+    ASSERT_EQ(consumerStats.getMsgBacklog(), 0);
+
     producer.close();
     consumer.close();
 }