[C++] UnAckedMessageTracker handles batch message incorrectly

pulsar-client-cpp/lib/ConsumerImpl.h

@@ -223,6 +223,7 @@ class ConsumerImpl : public ConsumerImplBase,
 
     FRIEND_TEST(ConsumerTest, testPartitionedConsumerUnAckedMessageRedelivery);
     FRIEND_TEST(ConsumerTest, testMultiTopicsConsumerUnAckedMessageRedelivery);
+    FRIEND_TEST(ConsumerTest, testBatchUnAckedMessageTracker);
 };
 
 } /* namespace pulsar */

pulsar-client-cpp/lib/UnAckedMessageTrackerEnabled.cc

@@ -90,12 +90,13 @@ UnAckedMessageTrackerEnabled::UnAckedMessageTrackerEnabled(long timeoutMs, long
     timeoutHandler();
 }
 
-bool UnAckedMessageTrackerEnabled::add(const MessageId& m) {
+bool UnAckedMessageTrackerEnabled::add(const MessageId& msgId) {
     std::lock_guard<std::mutex> acquire(lock_);
-    if (messageIdPartitionMap.count(m) == 0) {
+    MessageId id(msgId.partition(), msgId.ledgerId(), msgId.entryId(), -1);
+    if (messageIdPartitionMap.count(id) == 0) {
         std::set<MessageId>& partition = timePartitions.back();
-        bool emplace = messageIdPartitionMap.emplace(m, partition).second;
-        bool insert = partition.insert(m).second;
+        bool emplace = messageIdPartitionMap.emplace(id, partition).second;
+        bool insert = partition.insert(id).second;
         return emplace && insert;
     }
     return false;
@@ -106,13 +107,14 @@ bool UnAckedMessageTrackerEnabled::isEmpty() {
     return messageIdPartitionMap.empty();
 }
 
-bool UnAckedMessageTrackerEnabled::remove(const MessageId& m) {
+bool UnAckedMessageTrackerEnabled::remove(const MessageId& msgId) {
     std::lock_guard<std::mutex> acquire(lock_);
+    MessageId id(msgId.partition(), msgId.ledgerId(), msgId.entryId(), -1);
     bool removed = false;
 
-    std::map<MessageId, std::set<MessageId>&>::iterator exist = messageIdPartitionMap.find(m);
+    std::map<MessageId, std::set<MessageId>&>::iterator exist = messageIdPartitionMap.find(id);
     if (exist != messageIdPartitionMap.end()) {
-        removed = exist->second.erase(m);
+        removed = exist->second.erase(id);
         messageIdPartitionMap.erase(exist);
     }
     return removed;

pulsar-client-cpp/lib/UnAckedMessageTrackerEnabled.h

@@ -29,8 +29,8 @@ class UnAckedMessageTrackerEnabled : public UnAckedMessageTrackerInterface {
     ~UnAckedMessageTrackerEnabled();
     UnAckedMessageTrackerEnabled(long timeoutMs, const ClientImplPtr, ConsumerImplBase&);
     UnAckedMessageTrackerEnabled(long timeoutMs, long tickDuration, const ClientImplPtr, ConsumerImplBase&);
-    bool add(const MessageId& m);
-    bool remove(const MessageId& m);
+    bool add(const MessageId& msgId);
+    bool remove(const MessageId& msgId);
     void removeMessagesTill(const MessageId& msgId);
     void removeTopicMessage(const std::string& topic);
     void timeoutHandler();
@@ -52,6 +52,7 @@ class UnAckedMessageTrackerEnabled : public UnAckedMessageTrackerInterface {
 
     FRIEND_TEST(ConsumerTest, testPartitionedConsumerUnAckedMessageRedelivery);
     FRIEND_TEST(ConsumerTest, testMultiTopicsConsumerUnAckedMessageRedelivery);
+    FRIEND_TEST(ConsumerTest, testBatchUnAckedMessageTracker);
 };
 }  // namespace pulsar
 

pulsar-client-cpp/tests/ConsumerTest.cc

@@ -20,6 +20,8 @@
 #include <thread>
 #include <time.h>
 #include <set>
+#include <map>
+#include <vector>
 
 #include "gtest/gtest.h"
 
@@ -345,4 +347,70 @@ TEST(ConsumerTest, testMultiTopicsConsumerUnAckedMessageRedelivery) {
     client.close();
 }
 
+TEST(ConsumerTest, testBatchUnAckedMessageTracker) {
+    Client client(lookupUrl);
+    const std::string topic = "testBatchUnAckedMessageTracker" + std::to_string(time(nullptr));
+    std::string subName = "sub-batch-un-acked-msg-tracker";
+    constexpr int numOfMessages = 50;
+    constexpr int batchSize = 5;
+    constexpr int batchCount = numOfMessages / batchSize;
+    constexpr int unAckedMessagesTimeoutMs = 10000;
+    constexpr int tickDurationInMs = 1000;
+
+    Consumer consumer;
+    ConsumerConfiguration consumerConfig;
+    consumerConfig.setUnAckedMessagesTimeoutMs(unAckedMessagesTimeoutMs);
+    consumerConfig.setTickDurationInMs(tickDurationInMs);
+    ASSERT_EQ(ResultOk, client.subscribe(topic, subName, consumerConfig, consumer));
+    auto consumerImplPtr = PulsarFriend::getConsumerImplPtr(consumer);
+    auto tracker =
+        static_cast<UnAckedMessageTrackerEnabled*>(consumerImplPtr->unAckedMessageTrackerPtr_.get());
+
+    // send messages
+    ProducerConfiguration producerConfig;
+    producerConfig.setBatchingEnabled(true);
+    producerConfig.setBlockIfQueueFull(true);
+    producerConfig.setBatchingMaxMessages(batchSize);
+    Producer producer;
+    ASSERT_EQ(ResultOk, client.createProducer(topic, producerConfig, producer));
+    std::string prefix = "message-";
+    for (int i = 0; i < numOfMessages; i++) {
+        std::string messageContent = prefix + std::to_string(i);
+        Message msg = MessageBuilder().setContent(messageContent).build();
+        producer.sendAsync(msg, NULL);
+    }
+    producer.close();
+
+    std::map<MessageId, std::vector<MessageId>> msgIdInBatchMap;
+    std::vector<MessageId> messageIds;
+    for (auto i = 0; i < numOfMessages; ++i) {
+        Message msg;
+        ASSERT_EQ(ResultOk, consumer.receive(msg, 1000));
+        MessageId msgId = msg.getMessageId();
+        MessageId id(msgId.partition(), msgId.ledgerId(), msgId.entryId(), -1);
+        msgIdInBatchMap[id].emplace_back(msgId);
+    }
+
+    ASSERT_EQ(batchCount, msgIdInBatchMap.size());
+    ASSERT_EQ(batchCount, tracker->size());
+    for (const auto& iter : msgIdInBatchMap) {
+        ASSERT_EQ(iter.second.size(), batchSize);
+    }
+
+    int ackedBatchCount = 0;
+    for (auto iter = msgIdInBatchMap.begin(); iter != msgIdInBatchMap.end(); ++iter) {
+        ASSERT_EQ(batchSize, iter->second.size());
+        for (auto i = 0; i < iter->second.size(); ++i) {
+            ASSERT_EQ(ResultOk, consumer.acknowledge(iter->second[i]));
+        }
+        ackedBatchCount++;
+        ASSERT_EQ(batchCount - ackedBatchCount, tracker->size());
+    }
+    ASSERT_EQ(0, tracker->size());
+    ASSERT_TRUE(tracker->isEmpty());
+
+    consumer.close();
+    client.close();
+}
+
 }  // namespace pulsar