Fix a null ACK grouping tracker can be accessed after consumer is closed

BewareMyPower · BewareMyPower · commit 51b901629227 · 2025-10-21T22:11:43.000+08:00
Fixes #516
diff --git a/lib/AckGroupingTracker.cc b/lib/AckGroupingTracker.cc
@@ -36,7 +36,7 @@ DECLARE_LOG_OBJECT();
 
 void AckGroupingTracker::doImmediateAck(const MessageId& msgId, const ResultCallback& callback,
                                         CommandAck_AckType ackType) const {
-    const auto cnx = connectionSupplier_();
+    const auto cnx = getConnection();
     if (!cnx) {
         LOG_DEBUG("Connection is not ready, ACK failed for " << msgId);
         if (callback) {
@@ -89,7 +89,7 @@ static std::ostream& operator<<(std::ostream& os, const std::set<MessageId>& msg
 
 void AckGroupingTracker::doImmediateAck(const std::set<MessageId>& msgIds,
                                         const ResultCallback& callback) const {
-    const auto cnx = connectionSupplier_();
+    const auto cnx = getConnection();
     if (!cnx) {
         LOG_DEBUG("Connection is not ready, ACK failed for " << msgIds);
         if (callback) {
diff --git a/lib/AckGroupingTracker.h b/lib/AckGroupingTracker.h
@@ -27,13 +27,15 @@
 #include <set>
 
 #include "ProtoApiEnums.h"
+#include "lib/HandlerBase.h"
 
 namespace pulsar {
 
 class ClientConnection;
 using ClientConnectionPtr = std::shared_ptr<ClientConnection>;
 using ClientConnectionWeakPtr = std::weak_ptr<ClientConnection>;
 using ResultCallback = std::function<void(Result)>;
+using HandlerBaseWeakPtr = std::weak_ptr<HandlerBase>;
 
 /**
  * @class AckGroupingTracker
@@ -42,19 +44,19 @@ using ResultCallback = std::function<void(Result)>;
  */
 class AckGroupingTracker : public std::enable_shared_from_this<AckGroupingTracker> {
    public:
-    AckGroupingTracker(std::function<ClientConnectionPtr()> connectionSupplier,
-                       std::function<uint64_t()> requestIdSupplier, uint64_t consumerId, bool waitResponse)
-        : connectionSupplier_(std::move(connectionSupplier)),
-          requestIdSupplier_(std::move(requestIdSupplier)),
+    AckGroupingTracker(std::function<uint64_t()> requestIdSupplier, uint64_t consumerId, bool waitResponse)
+        : requestIdSupplier_(std::move(requestIdSupplier)),
           consumerId_(consumerId),
           waitResponse_(waitResponse) {}
 
     virtual ~AckGroupingTracker() = default;
 
     /**
      * Start tracking the ACK requests.
+     *
+     * @param[in] handler the handler to get a ClientConnection for sending the ACK requests.
      */
-    virtual void start() {}
+    virtual void start(const HandlerBaseWeakPtr& handler) { handler_ = handler; }
 
     /**
      * Since ACK requests are grouped and delayed, we need to do some best-effort duplicate check to
@@ -99,15 +101,39 @@ class AckGroupingTracker : public std::enable_shared_from_this<AckGroupingTracke
      */
     virtual void flushAndClean() {}
 
+    /**
+     * Close the ACK grouping tracker, which will prevent further ACK requests being sent.
+     */
+    virtual void close() { isClosed_.store(true, std::memory_order_relaxed); }
+
    protected:
     void doImmediateAck(const MessageId& msgId, const ResultCallback& callback,
                         CommandAck_AckType ackType) const;
     void doImmediateAck(const std::set<MessageId>& msgIds, const ResultCallback& callback) const;
+    bool isClosed() const noexcept { return isClosed_.load(std::memory_order_relaxed); }
+    bool validateClosed(const ResultCallback& callback) const {
+        if (isClosed()) {
+            if (callback) {
+                callback(ResultAlreadyClosed);
+            }
+            return true;
+        }
+        return false;
+    }
 
    private:
-    const std::function<ClientConnectionPtr()> connectionSupplier_;
+    std::weak_ptr<HandlerBase> handler_;
     const std::function<uint64_t()> requestIdSupplier_;
     const uint64_t consumerId_;
+    std::atomic_bool isClosed_{false};
+
+    ClientConnectionPtr getConnection() const {
+        auto handler = handler_.lock();
+        if (!handler) {
+            return nullptr;
+        }
+        return handler->getCnx().lock();
+    }
 
    protected:
     const bool waitResponse_;
diff --git a/lib/AckGroupingTrackerEnabled.cc b/lib/AckGroupingTrackerEnabled.cc
@@ -45,7 +45,10 @@ static int compare(const MessageId& lhs, const MessageId& rhs) {
     }
 }
 
-void AckGroupingTrackerEnabled::start() { this->scheduleTimer(); }
+void AckGroupingTrackerEnabled::start(const HandlerBaseWeakPtr& handler) {
+    AckGroupingTracker::start(handler);
+    this->scheduleTimer();
+}
 
 bool AckGroupingTrackerEnabled::isDuplicate(const MessageId& msgId) {
     {
@@ -62,6 +65,9 @@ bool AckGroupingTrackerEnabled::isDuplicate(const MessageId& msgId) {
 }
 
 void AckGroupingTrackerEnabled::addAcknowledge(const MessageId& msgId, const ResultCallback& callback) {
+    if (validateClosed(callback)) {
+        return;
+    }
     std::lock_guard<std::recursive_mutex> lock(this->rmutexPendingIndAcks_);
     this->pendingIndividualAcks_.insert(msgId);
     if (waitResponse_) {
@@ -76,6 +82,9 @@ void AckGroupingTrackerEnabled::addAcknowledge(const MessageId& msgId, const Res
 
 void AckGroupingTrackerEnabled::addAcknowledgeList(const MessageIdList& msgIds,
                                                    const ResultCallback& callback) {
+    if (validateClosed(callback)) {
+        return;
+    }
     std::lock_guard<std::recursive_mutex> lock(this->rmutexPendingIndAcks_);
     for (const auto& msgId : msgIds) {
         this->pendingIndividualAcks_.emplace(msgId);
@@ -92,6 +101,9 @@ void AckGroupingTrackerEnabled::addAcknowledgeList(const MessageIdList& msgIds,
 
 void AckGroupingTrackerEnabled::addAcknowledgeCumulative(const MessageId& msgId,
                                                          const ResultCallback& callback) {
+    if (validateClosed(callback)) {
+        return;
+    }
     std::unique_lock<std::mutex> lock(this->mutexCumulativeAckMsgId_);
     bool completeCallback = true;
     if (compare(msgId, this->nextCumulativeAckMsgId_) > 0) {
@@ -115,10 +127,15 @@ void AckGroupingTrackerEnabled::addAcknowledgeCumulative(const MessageId& msgId,
         callback(ResultOk);
     }
 }
-
 AckGroupingTrackerEnabled::~AckGroupingTrackerEnabled() {
-    isClosed_ = true;
-    this->flush();
+    std::lock_guard<std::mutex> lock(this->mutexTimer_);
+    if (this->timer_) {
+        cancelTimer(*this->timer_);
+    }
+}
+
+void AckGroupingTrackerEnabled::close() {
+    AckGroupingTracker::close();
     std::lock_guard<std::mutex> lock(this->mutexTimer_);
     if (this->timer_) {
         cancelTimer(*this->timer_);
@@ -165,7 +182,7 @@ void AckGroupingTrackerEnabled::flushAndClean() {
 }
 
 void AckGroupingTrackerEnabled::scheduleTimer() {
-    if (isClosed_) {
+    if (isClosed()) {
         return;
     }
 
diff --git a/lib/AckGroupingTrackerEnabled.h b/lib/AckGroupingTrackerEnabled.h
@@ -45,11 +45,10 @@ using HandlerBaseWeakPtr = std::weak_ptr<HandlerBase>;
  */
 class AckGroupingTrackerEnabled : public AckGroupingTracker {
    public:
-    AckGroupingTrackerEnabled(const std::function<ClientConnectionPtr()>& connectionSupplier,
-                              const std::function<uint64_t()>& requestIdSupplier, uint64_t consumerId,
+    AckGroupingTrackerEnabled(const std::function<uint64_t()>& requestIdSupplier, uint64_t consumerId,
                               bool waitResponse, long ackGroupingTimeMs, long ackGroupingMaxSize,
                               const ExecutorServicePtr& executor)
-        : AckGroupingTracker(connectionSupplier, requestIdSupplier, consumerId, waitResponse),
+        : AckGroupingTracker(requestIdSupplier, consumerId, waitResponse),
           ackGroupingTimeMs_(ackGroupingTimeMs),
           ackGroupingMaxSize_(ackGroupingMaxSize),
           executor_(executor) {
@@ -58,21 +57,20 @@ class AckGroupingTrackerEnabled : public AckGroupingTracker {
 
     ~AckGroupingTrackerEnabled();
 
-    void start() override;
+    void start(const HandlerBaseWeakPtr& handler) override;
     bool isDuplicate(const MessageId& msgId) override;
     void addAcknowledge(const MessageId& msgId, const ResultCallback& callback) override;
     void addAcknowledgeList(const MessageIdList& msgIds, const ResultCallback& callback) override;
     void addAcknowledgeCumulative(const MessageId& msgId, const ResultCallback& callback) override;
-    void flush();
     void flushAndClean() override;
+    void close() override;
+
+   private:
+    void flush();
 
    protected:
-    //! Method for scheduling grouping timer.
     void scheduleTimer();
 
-    //! State
-    std::atomic_bool isClosed_{false};
-
     //! Next message ID to be cumulatively cumulatively.
     MessageId nextCumulativeAckMsgId_{MessageId::earliest()};
     bool requireCumulativeAck_{false};
diff --git a/lib/ConsumerImpl.cc b/lib/ConsumerImpl.cc
@@ -74,6 +74,25 @@ static boost::optional<MessageId> getStartMessageId(const boost::optional<Messag
     return startMessageId;
 }
 
+static AckGroupingTracker* newAckGroupingTracker(const ClientImplPtr& client, const std::string& topic,
+                                                 uint64_t consumerId, const ConsumerConfiguration& config) {
+    const auto requestIdGenerator = client->getRequestIdGenerator();
+    const auto requestIdSupplier = [requestIdGenerator] { return (*requestIdGenerator)++; };
+
+    if (TopicName::get(topic)->isPersistent()) {
+        if (config.getAckGroupingTimeMs() > 0) {
+            return new AckGroupingTrackerEnabled(
+                requestIdSupplier, consumerId, config.isAckReceiptEnabled(), config.getAckGroupingTimeMs(),
+                config.getAckGroupingMaxSize(), client->getIOExecutorProvider()->get());
+        } else {
+            return new AckGroupingTrackerDisabled(requestIdSupplier, consumerId,
+                                                  config.isAckReceiptEnabled());
+        }
+    } else {
+        return new AckGroupingTracker(requestIdSupplier, consumerId, config.isAckReceiptEnabled());
+    }
+}
+
 ConsumerImpl::ConsumerImpl(const ClientImplPtr& client, const std::string& topic,
                            const std::string& subscriptionName, const ConsumerConfiguration& conf,
                            bool isPersistent, const ConsumerInterceptorsPtr& interceptors,
@@ -105,6 +124,7 @@ ConsumerImpl::ConsumerImpl(const ClientImplPtr& client, const std::string& topic
       consumerStr_("[" + topic + ", " + subscriptionName + ", " + std::to_string(consumerId_) + "] "),
       messageListenerRunning_(!conf.isStartPaused()),
       negativeAcksTracker_(std::make_shared<NegativeAcksTracker>(client, *this, conf)),
+      ackGroupingTrackerPtr_(newAckGroupingTracker(client, topic, consumerId_, conf)),
       readCompacted_(conf.isReadCompacted()),
       startMessageId_(getStartMessageId(startMessageId, conf.isStartMessageIdInclusive())),
       maxPendingChunkedMessage_(conf.getMaxPendingChunkedMessage()),
@@ -198,38 +218,7 @@ const std::string& ConsumerImpl::getTopic() const { return topic(); }
 
 void ConsumerImpl::start() {
     HandlerBase::start();
-
-    std::weak_ptr<ConsumerImpl> weakSelf{get_shared_this_ptr()};
-    auto connectionSupplier = [weakSelf]() -> ClientConnectionPtr {
-        auto self = weakSelf.lock();
-        if (!self) {
-            return nullptr;
-        }
-        return self->getCnx().lock();
-    };
-
-    // NOTE: start() is always called in `ClientImpl`'s method, so lock() returns not null
-    const auto requestIdGenerator = client_.lock()->getRequestIdGenerator();
-    const auto requestIdSupplier = [requestIdGenerator] { return (*requestIdGenerator)++; };
-
-    // Initialize ackGroupingTrackerPtr_ here because the get_shared_this_ptr() was not initialized until the
-    // constructor completed.
-    if (TopicName::get(topic())->isPersistent()) {
-        if (config_.getAckGroupingTimeMs() > 0) {
-            ackGroupingTrackerPtr_.reset(new AckGroupingTrackerEnabled(
-                connectionSupplier, requestIdSupplier, consumerId_, config_.isAckReceiptEnabled(),
-                config_.getAckGroupingTimeMs(), config_.getAckGroupingMaxSize(),
-                client_.lock()->getIOExecutorProvider()->get()));
-        } else {
-            ackGroupingTrackerPtr_.reset(new AckGroupingTrackerDisabled(
-                connectionSupplier, requestIdSupplier, consumerId_, config_.isAckReceiptEnabled()));
-        }
-    } else {
-        LOG_INFO(getName() << "ACK will NOT be sent to broker for this non-persistent topic.");
-        ackGroupingTrackerPtr_.reset(new AckGroupingTracker(connectionSupplier, requestIdSupplier,
-                                                            consumerId_, config_.isAckReceiptEnabled()));
-    }
-    ackGroupingTrackerPtr_->start();
+    ackGroupingTrackerPtr_->start(std::static_pointer_cast<HandlerBase>(shared_from_this()));
 }
 
 void ConsumerImpl::beforeConnectionChange(ClientConnection& cnx) { cnx.removeConsumer(consumerId_); }
@@ -591,17 +580,16 @@ void ConsumerImpl::messageReceived(const ClientConnectionPtr& cnx, const proto::
     LOG_DEBUG(getName() << " metadata.has_num_messages_in_batch() = "
                         << metadata.has_num_messages_in_batch());
 
-    uint32_t numOfMessageReceived = m.impl_->metadata.num_messages_in_batch();
-    auto ackGroupingTrackerPtr = ackGroupingTrackerPtr_;
-    if (ackGroupingTrackerPtr == nullptr) {  // The consumer is closing
+    const auto state = state_.load(std::memory_order_relaxed);
+    if (state == Closing || state == Closed) {
         return;
     }
-    if (ackGroupingTrackerPtr->isDuplicate(m.getMessageId())) {
+    uint32_t numOfMessageReceived = m.impl_->metadata.num_messages_in_batch();
+    if (ackGroupingTrackerPtr_->isDuplicate(m.getMessageId())) {
         LOG_DEBUG(getName() << " Ignoring message as it was ACKed earlier by same consumer.");
         increaseAvailablePermits(cnx, numOfMessageReceived);
         return;
     }
-    ackGroupingTrackerPtr.reset();
 
     if (metadata.has_num_messages_in_batch()) {
         BitSet::Data words(msg.ack_set_size());
@@ -1340,12 +1328,8 @@ void ConsumerImpl::closeAsync(const ResultCallback& originalCallback) {
     incomingMessages_.close();
 
     // Flush pending grouped ACK requests.
-    if (ackGroupingTrackerPtr_.use_count() != 1) {
-        LOG_ERROR("AckGroupingTracker is shared by other "
-                  << (ackGroupingTrackerPtr_.use_count() - 1)
-                  << " threads, which will prevent flushing the ACKs");
-    }
-    ackGroupingTrackerPtr_.reset();
+    ackGroupingTrackerPtr_->flushAndClean();
+    ackGroupingTrackerPtr_->close();
     negativeAcksTracker_->close();
 
     ClientConnectionPtr cnx = getCnx().lock();
@@ -1369,13 +1353,12 @@ void ConsumerImpl::closeAsync(const ResultCallback& originalCallback) {
     cnx->sendRequestWithId(Commands::newCloseConsumer(consumerId_, requestId), requestId)
         .addListener([self, callback](Result result, const ResponseData&) { callback(result); });
 }
-
 const std::string& ConsumerImpl::getName() const { return consumerStr_; }
 
 void ConsumerImpl::shutdown() { internalShutdown(); }
 
 void ConsumerImpl::internalShutdown() {
-    ackGroupingTrackerPtr_.reset();
+    ackGroupingTrackerPtr_->close();
     incomingMessages_.clear();
     possibleSendToDeadLetterTopicMessages_.clear();
     resetCnx();
diff --git a/lib/ConsumerImpl.h b/lib/ConsumerImpl.h
@@ -246,7 +246,7 @@ class ConsumerImpl : public ConsumerImplBase {
     UnAckedMessageTrackerPtr unAckedMessageTrackerPtr_;
     BrokerConsumerStatsImpl brokerConsumerStats_;
     std::shared_ptr<NegativeAcksTracker> negativeAcksTracker_;
-    AckGroupingTrackerPtr ackGroupingTrackerPtr_;
+    const AckGroupingTrackerPtr ackGroupingTrackerPtr_;
 
     MessageCryptoPtr msgCrypto_;
     const bool readCompacted_;
diff --git a/lib/OnceUniquePtr.h b/lib/OnceUniquePtr.h
@@ -0,0 +1,46 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+#pragma once
+
+#include <memory>
+
+namespace pulsar {
+
+// A wrapper around std::unique_ptr that allows initialization only once.
+// This wrapper is used when the initialization needs to be deferred.
+template <typename T>
+class OnceUniquePtr {
+   public:
+    void init(T* ptr) {
+        if (ptr && !ptr_) {
+            ptr_.reset(ptr);
+        }
+    }
+
+    const T& operator*() const noexcept { return *ptr_; }
+    T& operator*() noexcept { return *ptr_; }
+
+    const T* operator->() const noexcept { return ptr_.get(); }
+    T* operator->() noexcept { return ptr_.get(); }
+
+   private:
+    std::unique_ptr<T> ptr_;
+};
+
+}  // namespace pulsar
diff --git a/tests/BasicEndToEndTest.cc b/tests/BasicEndToEndTest.cc
diff --git a/tests/ConsumerTest.cc b/tests/ConsumerTest.cc
