Add tests

Author: BewareMyPower

lib/BitSet.h

@@ -39,6 +39,8 @@ class BitSet {
 
     BitSet(int32_t numBits) : words_((numBits / 64) + ((numBits % 64 == 0) ? 0 : 1)) { assert(numBits > 0); }
 
+    BitSet(Data&& words) : words_(std::move(words)), wordsInUse_(words_.size()) {}
+
     // Support range loop like:
     // ```c++
     // BitSet bitSet(129);
@@ -55,6 +57,15 @@ class BitSet {
      */
     bool isEmpty() const noexcept { return wordsInUse_ == 0; }
 
+    /**
+     * Returns the value of the bit with the specific index. The value is {@code true} if the bit with the
+     * index {@code bitIndex} is currently set in this {@code BitSet}; otherwise, the result is {@code false}.
+     *
+     * @param  bitIndex   the bit index
+     * @return the value of the bit with the specified index
+     */
+    bool get(int32_t bitIndex) const;
+
     /**
      * Sets the bits from the specified {@code fromIndex} (inclusive) to the
      * specified {@code toIndex} (exclusive) to {@code true}.
@@ -164,6 +175,12 @@ class BitSet {
     }
 };
 
+inline bool BitSet::get(int32_t bitIndex) const {
+    assert(bitIndex >= 0);
+    auto wordIndex_ = wordIndex(bitIndex);
+    return (wordIndex_ < wordsInUse_) && ((words_[wordIndex_] & (1L << bitIndex)) != 0);
+}
+
 inline void BitSet::set(int32_t fromIndex, int32_t toIndex) {
     assert(fromIndex < toIndex && fromIndex >= 0 && toIndex >= 0);
     if (fromIndex == toIndex) {

lib/ConsumerImpl.cc

@@ -506,8 +506,13 @@ void ConsumerImpl::messageReceived(const ClientConnectionPtr& cnx, const proto::
     }
 
     if (metadata.has_num_messages_in_batch()) {
+        BitSet::Data words(msg.ack_set_size());
+        for (int i = 0; i < words.size(); i++) {
+            words[i] = msg.ack_set(i);
+        }
+        BitSet ackSet{std::move(words)};
         Lock lock(mutex_);
-        numOfMessageReceived = receiveIndividualMessagesFromBatch(cnx, m, msg.redelivery_count());
+        numOfMessageReceived = receiveIndividualMessagesFromBatch(cnx, m, ackSet, msg.redelivery_count());
     } else {
         // try convery key value data.
         m.impl_->convertPayloadToKeyValue(config_.getSchema());
@@ -622,7 +627,8 @@ void ConsumerImpl::notifyPendingReceivedCallback(Result result, Message& msg,
 
 // Zero Queue size is not supported with Batch Messages
 uint32_t ConsumerImpl::receiveIndividualMessagesFromBatch(const ClientConnectionPtr& cnx,
-                                                          Message& batchedMessage, int redeliveryCount) {
+                                                          Message& batchedMessage, const BitSet& ackSet,
+                                                          int redeliveryCount) {
     auto batchSize = batchedMessage.impl_->metadata.num_messages_in_batch();
     LOG_DEBUG("Received Batch messages of size - " << batchSize
                                                    << " -- msgId: " << batchedMessage.getMessageId());
@@ -653,6 +659,13 @@ uint32_t ConsumerImpl::receiveIndividualMessagesFromBatch(const ClientConnection
             }
         }
 
+        if (!ackSet.isEmpty() && !ackSet.get(i)) {
+            LOG_DEBUG(getName() << "Ignoring message from " << i
+                                << "th message, which has been acknowledged");
+            ++skippedMessages;
+            continue;
+        }
+
         executeNotifyCallback(msg);
     }
 

lib/ConsumerImpl.h

@@ -48,6 +48,7 @@ typedef std::shared_ptr<Backoff> BackoffPtr;
 
 class AckGroupingTracker;
 using AckGroupingTrackerPtr = std::shared_ptr<AckGroupingTracker>;
+class BitSet;
 class ConsumerStatsBase;
 using ConsumerStatsBasePtr = std::shared_ptr<ConsumerStatsBase>;
 class UnAckedMessageTracker;
@@ -158,7 +159,7 @@ class ConsumerImpl : public ConsumerImplBase {
     void increaseAvailablePermits(const ClientConnectionPtr& currentCnx, int delta = 1);
     void drainIncomingMessageQueue(size_t count);
     uint32_t receiveIndividualMessagesFromBatch(const ClientConnectionPtr& cnx, Message& batchedMessage,
-                                                int redeliveryCount);
+                                                const BitSet& ackSet, int redeliveryCount);
     bool isPriorBatchIndex(int32_t idx);
     bool isPriorEntryIndex(int64_t idx);
     void brokerConsumerStatsListener(Result, BrokerConsumerStatsImpl, BrokerConsumerStatsCallback);

tests/AcknowledgeTest.cc

@@ -210,4 +210,45 @@ TEST_F(AcknowledgeTest, testBatchedMessageId) {
     ASSERT_EQ(consumers[3].getNumAcked(CommandAck_AckType_Cumulative), 1);
 }
 
+TEST_F(AcknowledgeTest, testBatchIndexAck) {
+    Client client(lookupUrl);
+    const std::string topic = "test-batch-index-ack-" + unique_str();
+    Producer producer;
+    ASSERT_EQ(ResultOk, client.createProducer(
+                            topic,
+                            ProducerConfiguration().setBatchingMaxMessages(100).setBatchingMaxPublishDelayMs(
+                                3600 * 1000 /* 1h */),
+                            producer));
+    std::vector<ConsumerWrapper> consumers{3};
+    for (size_t i = 0; i < consumers.size(); i++) {
+        consumers[i].initialize(client, topic, "sub-" + std::to_string(i), true /* enable batch index ack */);
+    }
+    constexpr int numMessages = 5;
+    for (int i = 0; i < numMessages; i++) {
+        producer.sendAsync(MessageBuilder().setContent("msg-" + std::to_string(i)).build(), nullptr);
+    }
+    producer.flush();
+    for (int i = 0; i < consumers.size(); i++) {
+        consumers[i].receiveAtMost(numMessages);
+        if (i >= 0) {
+            ASSERT_EQ(consumers[0].messageIdList(), consumers[i].messageIdList());
+        }
+    }
+    auto msgIds = consumers[0].messageIdList();
+
+    consumers[0].acknowledgeAndRestart({0, 2, 4}, AckType::INDIVIDUAL);
+    consumers[0].receiveAtMost(2);
+    ASSERT_EQ(subMessageIdList(msgIds, {1, 3}), consumers[0].messageIdList());
+    Message msg;
+    ASSERT_EQ(ResultTimeout, consumers[0].receive(msg));
+
+    consumers[1].acknowledgeAndRestart({0, 3}, AckType::INDIVIDUAL_LIST);
+    consumers[1].receiveAtMost(3);
+    ASSERT_EQ(subMessageIdList(msgIds, {1, 2, 4}), consumers[1].messageIdList());
+
+    consumers[2].acknowledgeAndRestart({3}, AckType::CUMULATIVE);
+    consumers[2].receiveAtMost(1);
+    ASSERT_EQ(subMessageIdList(msgIds, {4}), consumers[2].messageIdList());
+}
+
 INSTANTIATE_TEST_SUITE_P(BasicEndToEndTest, AcknowledgeTest, testing::Values(100, 0));

tests/BitSetTest.cc

@@ -62,6 +62,13 @@ TEST(BitSetTest, testSet) {
     // range contains one word
     bitSet.set(3, 29);
     ASSERT_EQ(toLongVector(bitSet), std::vector<uint64_t>{0x1ffffff8});
+    for (int i = 0; i < 64 * 5 + 1; i++) {
+        if (i >= 3 && i < 29) {
+            ASSERT_TRUE(bitSet.get(i));
+        } else {
+            ASSERT_FALSE(bitSet.get(i));
+        }
+    }
 
     // range contains multiple words
     bitSet.set(64 * 2 + 11, 64 * 4 + 19);

tests/ConsumerWrapper.h

@@ -35,13 +35,26 @@ enum class AckType
     CUMULATIVE
 };
 
+inline MessageIdList subMessageIdList(const MessageIdList& messageIdList,
+                                      const std::vector<size_t>& indexes) {
+    std::vector<MessageId> subMessageIdList;
+    for (size_t index : indexes) {
+        subMessageIdList.emplace_back(messageIdList.at(index));
+    }
+    return subMessageIdList;
+}
+
 class ConsumerWrapper {
    public:
-    void initialize(Client& client, const std::string& topic, const std::string& subscription) {
+    void initialize(Client& client, const std::string& topic, const std::string& subscription,
+                    bool enableBatchIndexAck = false) {
+        client_ = &client;
+        topic_ = topic;
+        subscription_ = subscription;
         // Enable the stats for cumulative ack
-        ConsumerConfiguration conf;
-        conf.setUnAckedMessagesTimeoutMs(10000);
-        ASSERT_EQ(ResultOk, client.subscribe(topic, subscription, conf, consumer_));
+        conf_.setUnAckedMessagesTimeoutMs(10000);
+        conf_.setBatchIndexAckEnabled(enableBatchIndexAck);
+        ASSERT_EQ(ResultOk, client_->subscribe(topic_, subscription_, conf_, consumer_));
     }
 
     const std::vector<MessageId>& messageIdList() const noexcept { return messageIdList_; }
@@ -58,11 +71,8 @@ class ConsumerWrapper {
 
     unsigned long getNumAcked(CommandAck_AckType ackType) const;
 
-    void acknowledgeAndRedeliver(const std::vector<size_t>& indexes, AckType ackType) {
-        std::vector<MessageId> msgIds;
-        for (size_t index : indexes) {
-            msgIds.emplace_back(messageIdList_.at(index));
-        }
+    void acknowledge(const std::vector<size_t>& indexes, AckType ackType) {
+        auto msgIds = subMessageIdList(messageIdList_, indexes);
         if (ackType == AckType::INDIVIDUAL_LIST) {
             consumer_.acknowledge(msgIds);
         } else {
@@ -76,10 +86,27 @@ class ConsumerWrapper {
         }
         // Wait until the acknowledge command is sent
         std::this_thread::sleep_for(std::chrono::milliseconds(100));
+    }
+
+    void acknowledgeAndRedeliver(const std::vector<size_t>& indexes, AckType ackType) {
+        acknowledge(indexes, ackType);
         consumer_.redeliverUnacknowledgedMessages();
     }
 
+    // NOTE: Currently Pulsar broker doesn't support redelivery with batch index ACK well, so here we verify
+    // the acknowledgment by restarting the consumer.
+    void acknowledgeAndRestart(const std::vector<size_t>& indexes, AckType ackType) {
+        acknowledge(indexes, ackType);
+        messageIdList_.clear();
+        consumer_.close();
+        ASSERT_EQ(ResultOk, client_->subscribe(topic_, subscription_, conf_, consumer_));
+    }
+
    private:
+    Client* client_;
+    std::string topic_;
+    std::string subscription_;
+    ConsumerConfiguration conf_;
     Consumer consumer_;
     std::vector<MessageId> messageIdList_;
 };