Support batch index acknowledgment (#151)

Fixes #87

### Modifications

- Add an consumer configuration `setBatchIndexAckEnabled` to enable the
  batch index ACK. When it's enabled, passing the original `MessageId`
  instead of the `MessageIdImpl` that trunks the batch index to the ACK
  grouping tracker.
- Since now a `BatchedMessageIdImpl` could be accepted in the ACK
  grouping tracker, fix the compare logic.
- Support passing a `BitSet` in `Commands::newAck` and get the internal
  `BitSet` from `MessageId` in `Commands::newMultiMessageAck`.
- Skip the acknowledged batch indexes when receiving batched messages in
  `ConsumerImpl::receiveIndividualMessagesFromBatch`.

### Verifications

Modify `BitSetTest.testSet` to verify the `BitSet::get` method added in
this commit.

Add `AcknowledgeTest.testBatchIndexAck` to test batch index ACK for all
types of acknowledgment:
- Individual ACK for a single message
- Individual ACK for a list of messages
- Cumulative ACK

Add `AcknowledgeTest.testMixedCumulativeAck` to test the new compare
logic between `BatchedMessageIdImpl` and `MessageIdImpl` works for
cumulative ACK.

Author: BewareMyPower

include/pulsar/ConsumerConfiguration.h

@@ -563,6 +563,23 @@ class PULSAR_PUBLIC ConsumerConfiguration {
      */
     bool isStartMessageIdInclusive() const;
 
+    /**
+     * Enable the batch index acknowledgment.
+     *
+     * It should be noted that this option can only work when the broker side also enables the batch index
+     * acknowledgment. See the `acknowledgmentAtBatchIndexLevelEnabled` config in `broker.conf`.
+     *
+     * Default: false
+     *
+     * @param enabled whether to enable the batch index acknowledgment
+     */
+    ConsumerConfiguration& setBatchIndexAckEnabled(bool enabled);
+
+    /**
+     * The associated getter of setBatchingEnabled
+     */
+    bool isBatchIndexAckEnabled() const;
+
     friend class PulsarWrapper;
 
    private:

lib/AckGroupingTracker.cc

@@ -19,17 +19,20 @@
 
 #include "AckGroupingTracker.h"
 
+#include "BitSet.h"
 #include "ClientConnection.h"
 #include "Commands.h"
 #include "LogUtils.h"
+#include "MessageIdImpl.h"
 
 namespace pulsar {
 
 DECLARE_LOG_OBJECT();
 
 inline void sendAck(ClientConnectionPtr cnx, uint64_t consumerId, const MessageId& msgId,
                     CommandAck_AckType ackType) {
-    auto cmd = Commands::newAck(consumerId, msgId.ledgerId(), msgId.entryId(), ackType, -1);
+    const auto& bitSet = Commands::getMessageIdImpl(msgId)->getBitSet();
+    auto cmd = Commands::newAck(consumerId, msgId.ledgerId(), msgId.entryId(), bitSet, ackType, -1);
     cnx->sendCommand(cmd);
     LOG_DEBUG("ACK request is sent for message - [" << msgId.ledgerId() << ", " << msgId.entryId() << "]");
 }

lib/AckGroupingTrackerEnabled.cc

@@ -19,6 +19,7 @@
 
 #include "AckGroupingTrackerEnabled.h"
 
+#include <climits>
 #include <mutex>
 
 #include "ClientConnection.h"
@@ -27,11 +28,25 @@
 #include "ExecutorService.h"
 #include "HandlerBase.h"
 #include "LogUtils.h"
+#include "MessageIdUtil.h"
 
 namespace pulsar {
 
 DECLARE_LOG_OBJECT();
 
+// Define a customized compare logic whose difference with the default compare logic of MessageId is:
+// When two MessageId objects are in the same entry, if only one of them is a message in the batch, treat
+// it as a smaller one.
+static int compare(const MessageId& lhs, const MessageId& rhs) {
+    int result = compareLedgerAndEntryId(lhs, rhs);
+    if (result != 0) {
+        return result;
+    } else {
+        return internal::compare(lhs.batchIndex() < 0 ? INT_MAX : lhs.batchIndex(),
+                                 rhs.batchIndex() < 0 ? INT_MAX : rhs.batchIndex());
+    }
+}
+
 AckGroupingTrackerEnabled::AckGroupingTrackerEnabled(ClientImplPtr clientPtr,
                                                      const HandlerBasePtr& handlerPtr, uint64_t consumerId,
                                                      long ackGroupingTimeMs, long ackGroupingMaxSize)
@@ -58,7 +73,7 @@ bool AckGroupingTrackerEnabled::isDuplicate(const MessageId& msgId) {
     {
         // Check if the message ID is already ACKed by a previous (or pending) cumulative request.
         std::lock_guard<std::mutex> lock(this->mutexCumulativeAckMsgId_);
-        if (msgId <= this->nextCumulativeAckMsgId_) {
+        if (compare(msgId, this->nextCumulativeAckMsgId_) <= 0) {
             return true;
         }
     }
@@ -88,7 +103,7 @@ void AckGroupingTrackerEnabled::addAcknowledgeList(const MessageIdList& msgIds)
 
 void AckGroupingTrackerEnabled::addAcknowledgeCumulative(const MessageId& msgId) {
     std::lock_guard<std::mutex> lock(this->mutexCumulativeAckMsgId_);
-    if (msgId > this->nextCumulativeAckMsgId_) {
+    if (compare(msgId, this->nextCumulativeAckMsgId_) > 0) {
         this->nextCumulativeAckMsgId_ = msgId;
         this->requireCumulativeAck_ = true;
     }

lib/BatchMessageAcker.h

@@ -43,6 +43,8 @@ class BatchMessageAcker {
         return prevBatchCumulativelyAcked_.compare_exchange_strong(expectedValue, true);
     }
 
+    const BitSet& getBitSet() const noexcept { return bitSet_; }
+
    private:
     // When a batched message is acknowledged cumulatively, the previous message id will be acknowledged
     // without batch index ACK enabled. However, it should be acknowledged only once. Use this flag to

lib/BatchedMessageIdImpl.h

@@ -41,6 +41,8 @@ class BatchedMessageIdImpl : public MessageIdImpl {
 
     bool shouldAckPreviousMessageId() const { return acker_->shouldAckPreviousMessageId(); }
 
+    const BitSet& getBitSet() const noexcept override { return acker_->getBitSet(); }
+
     MessageId getPreviousMessageId() {
         return MessageIdBuilder().ledgerId(ledgerId_).entryId(entryId_ - 1).partition(partition_).build();
     }

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/Commands.cc

@@ -28,6 +28,7 @@
 
 #include "BatchMessageAcker.h"
 #include "BatchedMessageIdImpl.h"
+#include "BitSet.h"
 #include "ChunkMessageIdImpl.h"
 #include "LogUtils.h"
 #include "MessageImpl.h"
@@ -420,7 +421,7 @@ SharedBuffer Commands::newProducer(const std::string& topic, uint64_t producerId
     return writeMessageWithSize(cmd);
 }
 
-SharedBuffer Commands::newAck(uint64_t consumerId, int64_t ledgerId, int64_t entryId,
+SharedBuffer Commands::newAck(uint64_t consumerId, int64_t ledgerId, int64_t entryId, const BitSet& ackSet,
                               CommandAck_AckType ackType, CommandAck_ValidationError validationError) {
     BaseCommand cmd;
     cmd.set_type(BaseCommand::ACK);
@@ -433,6 +434,9 @@ SharedBuffer Commands::newAck(uint64_t consumerId, int64_t ledgerId, int64_t ent
     auto* msgId = ack->add_message_id();
     msgId->set_ledgerid(ledgerId);
     msgId->set_entryid(entryId);
+    for (auto x : ackSet) {
+        msgId->add_ack_set(x);
+    }
     return writeMessageWithSize(cmd);
 }
 
@@ -446,6 +450,9 @@ SharedBuffer Commands::newMultiMessageAck(uint64_t consumerId, const std::set<Me
         auto newMsgId = ack->add_message_id();
         newMsgId->set_ledgerid(msgId.ledgerId());
         newMsgId->set_entryid(msgId.entryId());
+        for (auto x : getMessageIdImpl(msgId)->getBitSet()) {
+            newMsgId->add_ack_set(x);
+        }
     }
     return writeMessageWithSize(cmd);
 }

lib/Commands.h

@@ -39,6 +39,7 @@ class BatchMessageAcker;
 using BatchMessageAckerPtr = std::shared_ptr<BatchMessageAcker>;
 class MessageIdImpl;
 using MessageIdImplPtr = std::shared_ptr<MessageIdImpl>;
+class BitSet;
 
 namespace proto {
 class BaseCommand;
@@ -112,7 +113,7 @@ class Commands {
                                     bool userProvidedProducerName, bool encrypted,
                                     ProducerAccessMode accessMode, boost::optional<uint64_t> topicEpoch);
 
-    static SharedBuffer newAck(uint64_t consumerId, int64_t ledgerId, int64_t entryId,
+    static SharedBuffer newAck(uint64_t consumerId, int64_t ledgerId, int64_t entryId, const BitSet& ackSet,
                                CommandAck_AckType ackType, CommandAck_ValidationError validationError);
     static SharedBuffer newMultiMessageAck(uint64_t consumerId, const std::set<MessageId>& msgIds);
 

lib/ConsumerConfiguration.cc

@@ -287,4 +287,11 @@ const BatchReceivePolicy& ConsumerConfiguration::getBatchReceivePolicy() const {
     return impl_->batchReceivePolicy;
 }
 
+ConsumerConfiguration& ConsumerConfiguration::setBatchIndexAckEnabled(bool enabled) {
+    impl_->batchIndexAckEnabled = enabled;
+    return *this;
+}
+
+bool ConsumerConfiguration::isBatchIndexAckEnabled() const { return impl_->batchIndexAckEnabled; }
+
 }  // namespace pulsar

lib/ConsumerConfigurationImpl.h

@@ -56,6 +56,7 @@ struct ConsumerConfigurationImpl {
     bool autoAckOldestChunkedMessageOnQueueFull{false};
     bool startMessageIdInclusive{false};
     long expireTimeOfIncompleteChunkedMessageMs{60000};
+    bool batchIndexAckEnabled{false};
 };
 }  // namespace pulsar
 #endif /* LIB_CONSUMERCONFIGURATIONIMPL_H_ */