[fix] Remove closed producers and consumers from client

Fixes apache#55

### Motivation

1. When a producer or consumer is closed, the reference is still stored in `ClientImpl`. If a client kept creating producers or consumers, the memory usage would not reduce.
2. When the `HandlerBase::connection_` field is modified, the `removeProducer` or `removeConsumer` method is not called. Then these producers and consumers will be cached in the connection until the connection is closed.

### Modifications

In `ClientImpl`, use `SynchronizedHashMap` to store references of producers and consumers. The key is the pointer, the value is the `weak_ptr` objects. Then, add `cleanupProducer` and `cleanupConsumer` methods to remove a specific reference for a given pointer.

To call these two methods, this PR does a refactoring that all `ClientImpl` references in producer or consumer are changed to a trivial reference `ClientImpl&` instead of a `weak_ptr`. It's because the lifetime of `ClientImpl` should be longer than them. Using `weak_ptr` is redundant and will increase many repeat codes like:

```c++
auto client = client_.lock();
if (client) {
    client->cleanupProducer(this);
}
```

After that, unify the `shutdown` implementations for producers and consumers, which resets some states, cancels the timers, and unregister itself from both `ClientImpl` and `ClientConnection`. Add adding customized non-null callback in `closeAsync` and `unsubscribeAsync`.

For `ClientConnection`, add a `beforeConnectionChange` virtual method to `HandlerBase`, which is called before `connection_` field is modified. And disallow the direct access to `connection_` from derived classes.

### Verifications

`ClientCloseTest` is added to verify the following cases:
- a single topic
- a partitioned topic (multiple topics)
- a partitioned topic with regex subscription

`ClientCloseTest` is to verify these handlers are removed from `ClientImpl` and `ClientConnection`. `testShutdown` is to verify the handlers will be closed after the owned `Client` instance destroyed.

Author: BewareMyPower

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

@@ -68,7 +68,7 @@ jobs:
         run: ./pulsar-test-service-start.sh
 
       - name: Run unit tests
-        run: ./run-unit-tests.sh
+        run: RETRY_FAILED=3 ./run-unit-tests.sh
 
       - name: Stop Pulsar service
         run: ./pulsar-test-service-stop.sh

lib/AckGroupingTrackerDisabled.cc

@@ -22,6 +22,7 @@
 #include "HandlerBase.h"
 #include "PulsarApi.pb.h"
 #include <pulsar/MessageId.h>
+#include "LogUtils.h"
 
 namespace pulsar {
 

lib/AckGroupingTrackerEnabled.cc

@@ -33,9 +33,9 @@ namespace pulsar {
 
 DECLARE_LOG_OBJECT();
 
-AckGroupingTrackerEnabled::AckGroupingTrackerEnabled(ClientImplPtr clientPtr,
-                                                     const HandlerBasePtr& handlerPtr, uint64_t consumerId,
-                                                     long ackGroupingTimeMs, long ackGroupingMaxSize)
+AckGroupingTrackerEnabled::AckGroupingTrackerEnabled(ClientImpl& client, const HandlerBasePtr& handlerPtr,
+                                                     uint64_t consumerId, long ackGroupingTimeMs,
+                                                     long ackGroupingMaxSize)
     : AckGroupingTracker(),
       handlerWeakPtr_(handlerPtr),
       consumerId_(consumerId),
@@ -46,7 +46,7 @@ AckGroupingTrackerEnabled::AckGroupingTrackerEnabled(ClientImplPtr clientPtr,
       rmutexPendingIndAcks_(),
       ackGroupingTimeMs_(ackGroupingTimeMs),
       ackGroupingMaxSize_(ackGroupingMaxSize),
-      executor_(clientPtr->getIOExecutorProvider()->get()),
+      executor_(client.getIOExecutorProvider()->get()),
       timer_(),
       mutexTimer_() {
     LOG_DEBUG("ACK grouping is enabled, grouping time " << ackGroupingTimeMs << "ms, grouping max size "

lib/AckGroupingTrackerEnabled.h

@@ -41,13 +41,13 @@ class AckGroupingTrackerEnabled : public AckGroupingTracker {
 
     /**
      * Constructing ACK grouping tracker for peresistent topics.
-     * @param[in] clientPtr pointer to client object.
+     * @param[in] client the Client object.
      * @param[in] handlerPtr the shared pointer to connection handler.
      * @param[in] consumerId consumer ID that this tracker belongs to.
      * @param[in] ackGroupingTimeMs ACK grouping time window in milliseconds.
      * @param[in] ackGroupingMaxSize max. number of ACK requests can be grouped.
      */
-    AckGroupingTrackerEnabled(ClientImplPtr clientPtr, const HandlerBasePtr& handlerPtr, uint64_t consumerId,
+    AckGroupingTrackerEnabled(ClientImpl& client, const HandlerBasePtr& handlerPtr, uint64_t consumerId,
                               long ackGroupingTimeMs, long ackGroupingMaxSize);
 
     void start() override;

lib/ClientConnection.h

@@ -314,7 +314,7 @@ class PULSAR_PUBLIC ClientConnection : public std::enable_shared_from_this<Clien
     typedef std::map<long, Promise<Result, NamespaceTopicsPtr>> PendingGetNamespaceTopicsMap;
     PendingGetNamespaceTopicsMap pendingGetNamespaceTopicsRequests_;
 
-    std::mutex mutex_;
+    mutable std::mutex mutex_;
     typedef std::unique_lock<std::mutex> Lock;
 
     // Pending buffers to write on the socket

lib/ClientImpl.cc

@@ -170,10 +170,10 @@ void ClientImpl::handleCreateProducer(const Result result, const LookupDataResul
     if (!result) {
         ProducerImplBasePtr producer;
         if (partitionMetadata->getPartitions() > 0) {
-            producer = std::make_shared<PartitionedProducerImpl>(shared_from_this(), topicName,
+            producer = std::make_shared<PartitionedProducerImpl>(*this, topicName,
                                                                  partitionMetadata->getPartitions(), conf);
         } else {
-            producer = std::make_shared<ProducerImpl>(shared_from_this(), *topicName, conf);
+            producer = std::make_shared<ProducerImpl>(*this, *topicName, conf);
         }
         producer->getProducerCreatedFuture().addListener(
             std::bind(&ClientImpl::handleProducerCreated, shared_from_this(), std::placeholders::_1,
@@ -189,9 +189,15 @@ void ClientImpl::handleCreateProducer(const Result result, const LookupDataResul
 void ClientImpl::handleProducerCreated(Result result, ProducerImplBaseWeakPtr producerBaseWeakPtr,
                                        CreateProducerCallback callback, ProducerImplBasePtr producer) {
     if (result == ResultOk) {
-        Lock lock(mutex_);
-        producers_.push_back(producer);
-        lock.unlock();
+        auto pair = producers_.emplace(producer.get(), producer);
+        if (!pair.second) {
+            auto existingProducer = pair.first->second.lock();
+            LOG_ERROR("Unexpected existing producer at the same address: "
+                      << pair.first->first << ", producer: "
+                      << (existingProducer ? existingProducer->getProducerName() : "(null)"));
+            callback(ResultUnknownError, {});
+            return;
+        }
         callback(result, Producer(producer));
     } else {
         callback(result, {});
@@ -236,14 +242,23 @@ void ClientImpl::handleReaderMetadataLookup(const Result result, const LookupDat
         return;
     }
 
-    ReaderImplPtr reader = std::make_shared<ReaderImpl>(shared_from_this(), topicName->toString(), conf,
+    ReaderImplPtr reader = std::make_shared<ReaderImpl>(*this, topicName->toString(), conf,
                                                         getListenerExecutorProvider()->get(), callback);
     ConsumerImplBasePtr consumer = reader->getConsumer().lock();
     auto self = shared_from_this();
     reader->start(startMessageId, [this, self](const ConsumerImplBaseWeakPtr& weakConsumerPtr) {
-        Lock lock(mutex_);
-        consumers_.push_back(weakConsumerPtr);
-        lock.unlock();
+        auto consumer = weakConsumerPtr.lock();
+        if (consumer) {
+            auto pair = consumers_.emplace(consumer.get(), consumer);
+            if (!pair.second) {
+                auto existingConsumer = pair.first->second.lock();
+                LOG_ERROR("Unexpected existing consumer at the same address: "
+                          << pair.first->first
+                          << ", consumer: " << (existingConsumer ? existingConsumer->getName() : "(null)"));
+            }
+        } else {
+            LOG_ERROR("Unexpected case: the consumer is somehow expired");
+        }
     });
 }
 
@@ -286,7 +301,7 @@ void ClientImpl::createPatternMultiTopicsConsumer(const Result result, const Nam
             PatternMultiTopicsConsumerImpl::topicsPatternFilter(*topics, pattern);
 
         consumer = std::make_shared<PatternMultiTopicsConsumerImpl>(
-            shared_from_this(), regexPattern, *matchTopics, subscriptionName, conf, lookupServicePtr_);
+            *this, regexPattern, *matchTopics, subscriptionName, conf, lookupServicePtr_);
 
         consumer->getConsumerCreatedFuture().addListener(
             std::bind(&ClientImpl::handleConsumerCreated, shared_from_this(), std::placeholders::_1,
@@ -324,7 +339,7 @@ void ClientImpl::subscribeAsync(const std::vector<std::string>& topics, const st
     }
 
     ConsumerImplBasePtr consumer = std::make_shared<MultiTopicsConsumerImpl>(
-        shared_from_this(), topics, subscriptionName, topicNamePtr, conf, lookupServicePtr_);
+        *this, topics, subscriptionName, topicNamePtr, conf, lookupServicePtr_);
 
     consumer->getConsumerCreatedFuture().addListener(std::bind(&ClientImpl::handleConsumerCreated,
                                                                shared_from_this(), std::placeholders::_1,
@@ -374,12 +389,12 @@ void ClientImpl::handleSubscribe(const Result result, const LookupDataResultPtr
                 callback(ResultInvalidConfiguration, Consumer());
                 return;
             }
-            consumer = std::make_shared<MultiTopicsConsumerImpl>(shared_from_this(), topicName,
+            consumer = std::make_shared<MultiTopicsConsumerImpl>(*this, topicName,
                                                                  partitionMetadata->getPartitions(),
                                                                  subscriptionName, conf, lookupServicePtr_);
         } else {
-            auto consumerImpl = std::make_shared<ConsumerImpl>(
-                shared_from_this(), topicName->toString(), subscriptionName, conf, topicName->isPersistent());
+            auto consumerImpl = std::make_shared<ConsumerImpl>(*this, topicName->toString(), subscriptionName,
+                                                               conf, topicName->isPersistent());
             consumerImpl->setPartitionIndex(topicName->getPartitionIndex());
             consumer = consumerImpl;
         }
@@ -397,9 +412,15 @@ void ClientImpl::handleSubscribe(const Result result, const LookupDataResultPtr
 void ClientImpl::handleConsumerCreated(Result result, ConsumerImplBaseWeakPtr consumerImplBaseWeakPtr,
                                        SubscribeCallback callback, ConsumerImplBasePtr consumer) {
     if (result == ResultOk) {
-        Lock lock(mutex_);
-        consumers_.push_back(consumer);
-        lock.unlock();
+        auto pair = consumers_.emplace(consumer.get(), consumer);
+        if (!pair.second) {
+            auto existingConsumer = pair.first->second.lock();
+            LOG_ERROR("Unexpected existing consumer at the same address: "
+                      << pair.first->first
+                      << ", consumer: " << (existingConsumer ? existingConsumer->getName() : "(null)"));
+            callback(ResultUnknownError, {});
+            return;
+        }
         callback(result, Consumer(consumer));
     } else {
         callback(result, {});
@@ -478,12 +499,11 @@ void ClientImpl::getPartitionsForTopicAsync(const std::string& topic, GetPartiti
 
 void ClientImpl::closeAsync(CloseCallback callback) {
     Lock lock(mutex_);
-    ProducersList producers(producers_);
-    ConsumersList consumers(consumers_);
-
-    if (state_ != Open && callback) {
+    if (state_ != Open) {
         lock.unlock();
-        callback(ResultAlreadyClosed);
+        if (callback) {
+            callback(ResultAlreadyClosed);
+        }
         return;
     }
     // Set the state to Closing so that no producers could get added
@@ -492,12 +512,15 @@ void ClientImpl::closeAsync(CloseCallback callback) {
 
     memoryLimitController_.close();
 
+    auto producers = producers_.move();
+    auto consumers = consumers_.move();
+
     SharedInt numberOfOpenHandlers = std::make_shared<int>(producers.size() + consumers.size());
     LOG_INFO("Closing Pulsar client with " << producers.size() << " producers and " << consumers.size()
                                            << " consumers");
 
-    for (ProducersList::iterator it = producers.begin(); it != producers.end(); ++it) {
-        ProducerImplBasePtr producer = it->lock();
+    for (auto&& kv : producers) {
+        ProducerImplBasePtr producer = kv.second.lock();
         if (producer && !producer->isClosed()) {
             producer->closeAsync(std::bind(&ClientImpl::handleClose, shared_from_this(),
                                            std::placeholders::_1, numberOfOpenHandlers, callback));
@@ -507,8 +530,8 @@ void ClientImpl::closeAsync(CloseCallback callback) {
         }
     }
 
-    for (ConsumersList::iterator it = consumers.begin(); it != consumers.end(); ++it) {
-        ConsumerImplBasePtr consumer = it->lock();
+    for (auto&& kv : consumers) {
+        ConsumerImplBasePtr consumer = kv.second.lock();
         if (consumer && !consumer->isClosed()) {
             consumer->closeAsync(std::bind(&ClientImpl::handleClose, shared_from_this(),
                                            std::placeholders::_1, numberOfOpenHandlers, callback));
@@ -562,23 +585,18 @@ void ClientImpl::handleClose(Result result, SharedInt numberOfOpenHandlers, Resu
 }
 
 void ClientImpl::shutdown() {
-    Lock lock(mutex_);
-    ProducersList producers;
-    ConsumersList consumers;
-
-    producers.swap(producers_);
-    consumers.swap(consumers_);
-    lock.unlock();
+    auto producers = producers_.move();
+    auto consumers = consumers_.move();
 
-    for (ProducersList::iterator it = producers.begin(); it != producers.end(); ++it) {
-        ProducerImplBasePtr producer = it->lock();
+    for (auto&& kv : producers) {
+        ProducerImplBasePtr producer = kv.second.lock();
         if (producer) {
             producer->shutdown();
         }
     }
 
-    for (ConsumersList::iterator it = consumers.begin(); it != consumers.end(); ++it) {
-        ConsumerImplBasePtr consumer = it->lock();
+    for (auto&& kv : consumers) {
+        ConsumerImplBasePtr consumer = kv.second.lock();
         if (consumer) {
             consumer->shutdown();
         }
@@ -631,26 +649,24 @@ uint64_t ClientImpl::newRequestId() {
 }
 
 uint64_t ClientImpl::getNumberOfProducers() {
-    Lock lock(mutex_);
     uint64_t numberOfAliveProducers = 0;
-    for (const auto& producer : producers_) {
+    producers_.forEachValue([&numberOfAliveProducers](const ProducerImplBaseWeakPtr& producer) {
         const auto& producerImpl = producer.lock();
         if (producerImpl) {
             numberOfAliveProducers += producerImpl->getNumberOfConnectedProducer();
         }
-    }
+    });
     return numberOfAliveProducers;
 }
 
 uint64_t ClientImpl::getNumberOfConsumers() {
-    Lock lock(mutex_);
     uint64_t numberOfAliveConsumers = 0;
-    for (const auto& consumer : consumers_) {
+    consumers_.forEachValue([&numberOfAliveConsumers](const ConsumerImplBaseWeakPtr& consumer) {
         const auto consumerImpl = consumer.lock();
         if (consumerImpl) {
             numberOfAliveConsumers += consumerImpl->getNumberOfConnectedConsumer();
         }
-    }
+    });
     return numberOfAliveConsumers;
 }
 

lib/ClientImpl.h

@@ -31,6 +31,7 @@
 #include <atomic>
 #include <vector>
 #include "ServiceNameResolver.h"
+#include "SynchronizedHashMap.h"
 
 namespace pulsar {
 
@@ -91,6 +92,11 @@ class ClientImpl : public std::enable_shared_from_this<ClientImpl> {
     ExecutorServiceProviderPtr getListenerExecutorProvider();
     ExecutorServiceProviderPtr getPartitionListenerExecutorProvider();
     LookupServicePtr getLookup();
+
+    void cleanupProducer(ProducerImplBase* address) { producers_.remove(address); }
+
+    void cleanupConsumer(ConsumerImplBase* address) { consumers_.remove(address); }
+
     friend class PulsarFriend;
 
    private:
@@ -147,11 +153,8 @@ class ClientImpl : public std::enable_shared_from_this<ClientImpl> {
     uint64_t consumerIdGenerator_;
     uint64_t requestIdGenerator_;
 
-    typedef std::vector<ProducerImplBaseWeakPtr> ProducersList;
-    ProducersList producers_;
-
-    typedef std::vector<ConsumerImplBaseWeakPtr> ConsumersList;
-    ConsumersList consumers_;
+    SynchronizedHashMap<ProducerImplBase*, ProducerImplBaseWeakPtr> producers_;
+    SynchronizedHashMap<ConsumerImplBase*, ConsumerImplBaseWeakPtr> consumers_;
 
     std::atomic<Result> closingError;
 

lib/ConnectionPool.h

@@ -77,7 +77,7 @@ class PULSAR_PUBLIC ConnectionPool {
     std::mutex mutex_;
     std::atomic_bool closed_{false};
 
-    friend class ConnectionPoolTest;
+    friend class PulsarFriend;
 };
 }  // namespace pulsar
 #endif  //_PULSAR_CONNECTION_POOL_HEADER_