[improve] modify the negativeACK structure to reduce memory overhead (#497)

Author: syslucas

.github/workflows/ci-pr-validation.yaml

@@ -167,7 +167,7 @@ jobs:
           sudo apt-get install -y libcurl4-openssl-dev libssl-dev \
             protobuf-compiler libprotobuf-dev libboost-dev \
             libboost-dev libboost-program-options-dev \
-            libzstd-dev libsnappy-dev libgmock-dev libgtest-dev
+            libzstd-dev libsnappy-dev libgmock-dev libgtest-dev libroaring-dev
       - name: CMake
         run: cmake -B build -DBUILD_PERF_TOOLS=ON -DCMAKE_CXX_STANDARD=20
       - name: Build
@@ -188,16 +188,16 @@ jobs:
       matrix:
         include:
           - name: 'Windows x64'
-            os: windows-2019
+            os: windows-2022
             triplet: x64-windows-static
             suffix: 'windows-win64'
-            generator: 'Visual Studio 16 2019'
+            generator: 'Visual Studio 17 2022'
             arch: '-A x64'
           - name: 'Windows x86'
-            os: windows-2019
+            os: windows-2022
             triplet: x86-windows-static
             suffix: 'windows-win32'
-            generator: 'Visual Studio 16 2019'
+            generator: 'Visual Studio 17 2022'
             arch: '-A Win32'
 
     steps:

.github/workflows/codeql-analysis.yml

@@ -60,7 +60,7 @@ jobs:
         sudo apt-get install -y libcurl4-openssl-dev libssl-dev \
           protobuf-compiler libprotobuf-dev libboost-dev \
           libboost-dev libboost-program-options-dev \
-          libzstd-dev libsnappy-dev
+          libzstd-dev libsnappy-dev libroaring-dev
 
     - name: Build
       run: |

CMakeLists.txt

@@ -124,13 +124,15 @@ if (INTEGRATE_VCPKG)
     find_package(protobuf CONFIG REQUIRED)
     find_package(zstd CONFIG REQUIRED)
     find_package(Snappy CONFIG REQUIRED)
+    find_package(roaring CONFIG REQUIRED)
     set(COMMON_LIBS CURL::libcurl
         ZLIB::ZLIB
         OpenSSL::SSL
         OpenSSL::Crypto
         protobuf::libprotobuf
         $<IF:$<TARGET_EXISTS:zstd::libzstd_shared>,zstd::libzstd_shared,zstd::libzstd_static>
         Snappy::snappy
+        roaring::roaring
         )
     if (USE_ASIO)
         find_package(asio CONFIG REQUIRED)

LegacyFindPackages.cmake

@@ -117,6 +117,25 @@ if (NOT ZLIB_INCLUDE_DIRS OR NOT ZLIB_LIBRARIES)
     message(FATAL_ERROR "Could not find zlib")
 endif ()
 
+find_package(roaring QUIET)
+if (NOT ROARING_FOUND)
+    find_path(ROARING_INCLUDE_DIRS NAMES roaring/roaring.hh)
+    find_library(ROARING_LIBRARIES NAMES roaring libroaring)
+endif ()
+message("ROARING_INCLUDE_DIRS: " ${ROARING_INCLUDE_DIRS})
+message("ROARING_LIBRARIES: " ${ROARING_LIBRARIES})
+if (NOT ROARING_INCLUDE_DIRS OR NOT ROARING_LIBRARIES)
+    message(FATAL_ERROR "Could not find libroaring")
+endif ()
+file(READ "${ROARING_INCLUDE_DIRS}/roaring/roaring.hh" ROARING_HEADER_CONTENTS)
+string(REGEX MATCH "namespace roaring" ROARING_HAS_NAMESPACE "${ROARING_HEADER_CONTENTS}")
+if (ROARING_HAS_NAMESPACE)
+    message(STATUS "Roaring64Map is in namespace roaring")
+else ()
+    message(STATUS "Roaring64Map is in global namespace")
+    add_definitions(-DROARING_NAMESPACE_GLOBAL)
+endif ()
+
 if (LINK_STATIC AND NOT VCPKG_TRIPLET)
     find_library(LIB_ZSTD NAMES libzstd.a)
     message(STATUS "ZStd: ${LIB_ZSTD}")
@@ -129,6 +148,7 @@ if (LINK_STATIC AND NOT VCPKG_TRIPLET)
 elseif (LINK_STATIC AND VCPKG_TRIPLET)
     find_package(Protobuf REQUIRED)
     message(STATUS "Found protobuf static library: " ${Protobuf_LIBRARIES})
+    find_package(roaring REQUIRED)
     if (MSVC AND (${CMAKE_BUILD_TYPE} STREQUAL Debug))
         find_library(ZLIB_LIBRARIES NAMES zlibd)
     else ()
@@ -231,6 +251,7 @@ include_directories(
   ${Boost_INCLUDE_DIRS}
   ${OPENSSL_INCLUDE_DIR}
   ${ZLIB_INCLUDE_DIRS}
+  ${ROARING_INCLUDE_DIRS}
   ${CURL_INCLUDE_DIRS}
   ${Protobuf_INCLUDE_DIRS}
   ${GTEST_INCLUDE_PATH}
@@ -246,6 +267,7 @@ set(COMMON_LIBS
   ${CURL_LIBRARIES}
   ${OPENSSL_LIBRARIES}
   ${ZLIB_LIBRARIES}
+  ${ROARING_LIBRARIES}
   ${ADDITIONAL_LIBRARIES}
   ${CMAKE_DL_LIBS}
 )

dependencies.yaml

@@ -26,3 +26,4 @@ zstd: 1.5.5
 snappy: 1.1.10
 openssl: 1.1.1w
 curl: 8.6.0
+roaring: 4.3.1

include/pulsar/ConsumerConfiguration.h

@@ -277,6 +277,21 @@ class PULSAR_PUBLIC ConsumerConfiguration {
      */
     long getNegativeAckRedeliveryDelayMs() const;
 
+    /**
+     * Set the precision bit count for negative ack redelivery delay.
+     * The lower bits of the redelivery time will be trimmed to reduce the memory occupation.
+     * @param negativeAckPrecisionBitCnt
+     *            negative ack precision bit count
+     */
+    void setNegativeAckPrecisionBitCnt(int negativeAckPrecisionBitCnt);
+
+    /**
+     * Get the configured precision bit count for negative ack redelivery delay.
+     *
+     * @return redelivery time precision bit count
+     */
+    int getNegativeAckPrecisionBitCnt() const;
+
     /**
      * Set time window in milliseconds for grouping message ACK requests. An ACK request is not sent
      * to broker until the time window reaches its end, or the number of grouped messages reaches

lib/ConsumerConfiguration.cc

@@ -134,6 +134,16 @@ void ConsumerConfiguration::setAckGroupingTimeMs(long ackGroupingMillis) {
     impl_->ackGroupingTimeMs = ackGroupingMillis;
 }
 
+int ConsumerConfiguration::getNegativeAckPrecisionBitCnt() const { return impl_->negativeAckPrecisionBitCnt; }
+
+void ConsumerConfiguration::setNegativeAckPrecisionBitCnt(int negativeAckPrecisionBitCnt) {
+    if (negativeAckPrecisionBitCnt < 0) {
+        throw std::invalid_argument(
+            "Consumer Config Exception: NegativeAckPrecisionBitCnt should be nonnegative number.");
+    }
+    impl_->negativeAckPrecisionBitCnt = negativeAckPrecisionBitCnt;
+}
+
 long ConsumerConfiguration::getAckGroupingTimeMs() const { return impl_->ackGroupingTimeMs; }
 
 void ConsumerConfiguration::setAckGroupingMaxSize(long maxGroupingSize) {

lib/ConsumerConfigurationImpl.h

@@ -26,6 +26,7 @@ struct ConsumerConfigurationImpl {
     long unAckedMessagesTimeoutMs{0};
     long tickDurationInMs{1000};
     long negativeAckRedeliveryDelayMs{60000};
+    int negativeAckPrecisionBitCnt{8};
     long ackGroupingTimeMs{100};
     long ackGroupingMaxSize{1000};
     long brokerConsumerStatsCacheTimeInMs{30 * 1000L};  // 30 seconds

lib/NegativeAcksTracker.cc

@@ -19,14 +19,19 @@
 
 #include "NegativeAcksTracker.h"
 
+#include <cstdint>
 #include <functional>
 #include <set>
+#include <utility>
 
 #include "ClientImpl.h"
 #include "ConsumerImpl.h"
 #include "ExecutorService.h"
 #include "LogUtils.h"
 #include "MessageIdUtil.h"
+#include "pulsar/MessageBuilder.h"
+#include "pulsar/MessageId.h"
+#include "pulsar/MessageIdBuilder.h"
 DECLARE_LOG_OBJECT()
 
 namespace pulsar {
@@ -41,6 +46,7 @@ NegativeAcksTracker::NegativeAcksTracker(const ClientImplPtr &client, ConsumerIm
     nackDelay_ =
         std::chrono::milliseconds(std::max(conf.getNegativeAckRedeliveryDelayMs(), MIN_NACK_DELAY_MILLIS));
     timerInterval_ = std::chrono::milliseconds((long)(nackDelay_.count() / 3));
+    nackPrecisionBit_ = conf.getNegativeAckPrecisionBitCnt();
     LOG_DEBUG("Created negative ack tracker with delay: " << nackDelay_.count() << " ms - Timer interval: "
                                                           << timerInterval_.count());
 }
@@ -75,13 +81,22 @@ void NegativeAcksTracker::handleTimer(const ASIO_ERROR &ec) {
 
     auto now = Clock::now();
 
+    // The map is sorted by time, so we can exit immediately when we traverse to a time that does not match
     for (auto it = nackedMessages_.begin(); it != nackedMessages_.end();) {
-        if (it->second < now) {
-            messagesToRedeliver.insert(it->first);
-            it = nackedMessages_.erase(it);
-        } else {
-            ++it;
+        if (it->first > now) {
+            // We are done with all the messages that need to be redelivered
+            break;
         }
+
+        auto ledgerMap = it->second;
+        for (auto ledgerIt = ledgerMap.begin(); ledgerIt != ledgerMap.end(); ++ledgerIt) {
+            auto entrySet = ledgerIt->second;
+            for (auto setIt = entrySet.begin(); setIt != entrySet.end(); ++setIt) {
+                messagesToRedeliver.insert(
+                    MessageIdBuilder().ledgerId(ledgerIt->first).entryId(*setIt).build());
+            }
+        }
+        it = nackedMessages_.erase(it);
     }
     lock.unlock();
 
@@ -92,14 +107,27 @@ void NegativeAcksTracker::handleTimer(const ASIO_ERROR &ec) {
     scheduleTimer();
 }
 
+std::chrono::steady_clock::time_point trimLowerBit(const std::chrono::steady_clock::time_point &tp,
+                                                   int bits) {
+    // get origin timestamp in nanoseconds
+    auto timestamp = std::chrono::duration_cast<std::chrono::nanoseconds>(tp.time_since_epoch()).count();
+
+    // trim lower bits
+    auto trimmedTimestamp = timestamp & (~((1LL << bits) - 1));
+
+    return std::chrono::steady_clock::time_point(std::chrono::nanoseconds(trimmedTimestamp));
+}
+
 void NegativeAcksTracker::add(const MessageId &m) {
     auto msgId = discardBatch(m);
     auto now = Clock::now();
 
     {
         std::lock_guard<std::mutex> lock{mutex_};
+        auto trimmedTimestamp = trimLowerBit(now + nackDelay_, nackPrecisionBit_);
+        // If the timestamp is already in the map, we can just add the message to the existing entry
         // Erase batch id to group all nacks from same batch
-        nackedMessages_[msgId] = now + nackDelay_;
+        nackedMessages_[trimmedTimestamp][msgId.ledgerId()].add((uint64_t)msgId.entryId());
     }
 
     scheduleTimer();

lib/NegativeAcksTracker.h

@@ -27,6 +27,8 @@
 #include <map>
 #include <memory>
 #include <mutex>
+#include <roaring/roaring64map.hh>
+#include <unordered_map>
 
 #include "AsioDefines.h"
 #include "AsioTimer.h"
@@ -39,6 +41,12 @@ class ClientImpl;
 using ClientImplPtr = std::shared_ptr<ClientImpl>;
 class ExecutorService;
 using ExecutorServicePtr = std::shared_ptr<ExecutorService>;
+using LedgerId = int64_t;
+#ifdef ROARING_NAMESPACE_GLOBAL
+using ConditionalRoaringMap = Roaring64Map;
+#else
+using ConditionalRoaringMap = roaring::Roaring64Map;
+#endif
 
 class NegativeAcksTracker : public std::enable_shared_from_this<NegativeAcksTracker> {
    public:
@@ -64,8 +72,9 @@ class NegativeAcksTracker : public std::enable_shared_from_this<NegativeAcksTrac
 
     std::chrono::milliseconds nackDelay_;
     std::chrono::milliseconds timerInterval_;
+    int nackPrecisionBit_;
     typedef typename std::chrono::steady_clock Clock;
-    std::map<MessageId, Clock::time_point> nackedMessages_;
+    std::map<Clock::time_point, std::unordered_map<LedgerId, ConditionalRoaringMap>> nackedMessages_;
 
     const DeadlineTimerPtr timer_;
     std::atomic_bool closed_{false};