Fix possible zombie consumer when closing after reconnection

Author: BewareMyPower

lib/ConsumerImpl.cc

@@ -255,6 +255,8 @@ Future<Result, bool> ConsumerImpl::connectionOpened(const ClientConnectionPtr& c
     if (duringSeek()) {
         ackGroupingTrackerPtr_->flushAndClean();
     }
+    // Set connection after flushing the ack tracker to avoid sending ACKs on the new connection
+    setCnx(cnx);
 
     Lock lockForMessageId(mutexForMessageId_);
     clearReceiveQueue();
@@ -312,7 +314,6 @@ Result ConsumerImpl::handleCreateConsumer(const ClientConnectionPtr& cnx, Result
         LOG_INFO(getName() << "Created consumer on broker " << cnx->cnxString());
         {
             Lock mutexLock(mutex_);
-            setCnx(cnx);
             incomingMessages_.clear();
             possibleSendToDeadLetterTopicMessages_.clear();
             state_ = Ready;
@@ -338,8 +339,14 @@ Result ConsumerImpl::handleCreateConsumer(const ClientConnectionPtr& cnx, Result
         } else if (messageListener_) {
             sendFlowPermitsToBroker(cnx, 1);
         }
+        {
+            std::lock_guard<std::mutex> lock{subscribeMutex_};
+            subscribePromise_.setValue(true);  // trigger pending requests like seek or getLastMessageId
+        }
         consumerCreatedPromise_.setValue(get_shared_this_ptr());
     } else {
+        // TODO: it's not tested yet because it's hard to mock subscribe failure after reconnection
+        subscribePromise_.setFailed(result);
         if (result == ResultTimeout) {
             // Creating the consumer has timed out. We need to ensure the broker closes the consumer
             // in case it was indeed created, otherwise it might prevent new subscribe operation,
@@ -1311,6 +1318,10 @@ void ConsumerImpl::disconnectConsumer(const boost::optional<std::string>& assign
     LOG_INFO("Broker notification of Closed consumer: "
              << consumerId_ << (assignedBrokerUrl ? (" assignedBrokerUrl: " + assignedBrokerUrl.get()) : ""));
     resetCnx();
+    {
+        std::lock_guard<std::mutex> lock{subscribeMutex_};
+        subscribePromise_ = Promise<Result, bool>{};
+    }
     scheduleReconnection(assignedBrokerUrl);
 }
 
@@ -1542,13 +1553,23 @@ void ConsumerImpl::seekAsync(const MessageId& msgId, const ResultCallback& callb
         return;
     }
 
-    ClientImplPtr client = client_.lock();
-    if (!client) {
-        LOG_ERROR(getName() << "Client is expired when seekAsync " << msgId);
-        return;
-    }
-    const auto requestId = client->newRequestId();
-    seekAsyncInternal(requestId, Commands::newSeek(consumerId_, requestId, msgId), SeekArg{msgId}, callback);
+    auto self = shared_from_this();
+    sendRequestAfterSubscribed([this, self, msgId, callback](Result result) {
+        if (result != ResultOk) {
+            if (callback) {
+                callback(result);
+            }
+            return;
+        }
+        ClientImplPtr client = client_.lock();
+        if (!client) {
+            LOG_ERROR(getName() << "Client is expired when seekAsync " << msgId);
+            return;
+        }
+        const auto requestId = client->newRequestId();
+        seekAsyncInternal(requestId, Commands::newSeek(consumerId_, requestId, msgId), SeekArg{msgId},
+                          callback);
+    });
 }
 
 void ConsumerImpl::seekAsync(uint64_t timestamp, const ResultCallback& callback) {
@@ -1561,14 +1582,23 @@ void ConsumerImpl::seekAsync(uint64_t timestamp, const ResultCallback& callback)
         return;
     }
 
-    ClientImplPtr client = client_.lock();
-    if (!client) {
-        LOG_ERROR(getName() << "Client is expired when seekAsync " << timestamp);
-        return;
-    }
-    const auto requestId = client->newRequestId();
-    seekAsyncInternal(requestId, Commands::newSeek(consumerId_, requestId, timestamp), SeekArg{timestamp},
-                      callback);
+    auto self = shared_from_this();
+    sendRequestAfterSubscribed([this, self, timestamp, callback](Result result) {
+        if (result != ResultOk) {
+            if (callback) {
+                callback(result);
+            }
+            return;
+        }
+        ClientImplPtr client = client_.lock();
+        if (!client) {
+            LOG_ERROR(getName() << "Client is expired when seekAsync " << timestamp);
+            return;
+        }
+        const auto requestId = client->newRequestId();
+        seekAsyncInternal(requestId, Commands::newSeek(consumerId_, requestId, timestamp), SeekArg{timestamp},
+                          callback);
+    });
 }
 
 bool ConsumerImpl::isReadCompacted() { return readCompacted_; }
@@ -1645,11 +1675,20 @@ void ConsumerImpl::getLastMessageIdAsync(const BrokerGetLastMessageIdCallback& c
         return;
     }
 
-    TimeDuration operationTimeout = seconds(client_.lock()->conf().getOperationTimeoutSeconds());
-    BackoffPtr backoff = std::make_shared<Backoff>(milliseconds(100), operationTimeout * 2, milliseconds(0));
-    DeadlineTimerPtr timer = executor_->createDeadlineTimer();
-
-    internalGetLastMessageIdAsync(backoff, operationTimeout, timer, callback);
+    auto self = shared_from_this();
+    sendRequestAfterSubscribed([this, self, callback](Result result) {
+        if (result != ResultOk) {
+            if (callback) {
+                callback(result, {});
+            }
+            return;
+        }
+        TimeDuration operationTimeout = seconds(client_.lock()->conf().getOperationTimeoutSeconds());
+        BackoffPtr backoff =
+            std::make_shared<Backoff>(milliseconds(100), operationTimeout * 2, milliseconds(0));
+        DeadlineTimerPtr timer = executor_->createDeadlineTimer();
+        internalGetLastMessageIdAsync(backoff, operationTimeout, timer, callback);
+    });
 }
 
 void ConsumerImpl::internalGetLastMessageIdAsync(const BackoffPtr& backoff, TimeDuration remainTime,

lib/ConsumerImpl.h

@@ -266,6 +266,21 @@ class ConsumerImpl : public ConsumerImplBase {
     Synchronized<MessageId> seekMessageId_{MessageId::earliest()};
     std::atomic<bool> hasSoughtByTimestamp_{false};
 
+    mutable std::mutex subscribeMutex_;
+    Promise<Result, bool> subscribePromise_;
+
+    template <typename SendRequestFunc>
+    inline void sendRequestAfterSubscribed(SendRequestFunc&& sendRequest) {
+        Lock lock{subscribeMutex_};
+        if (subscribePromise_.isComplete()) {
+            lock.unlock();
+            sendRequest(ResultOk);
+        } else {
+            subscribePromise_.getFuture().addListener(
+                [sendRequest](Result result, bool) { sendRequest(result); });
+        }
+    }
+
     bool hasSoughtByTimestamp() const { return hasSoughtByTimestamp_.load(std::memory_order_acquire); }
     bool duringSeek() const { return seekStatus_ != SeekStatus::NOT_STARTED; }
 

tests/ConsumerTest.cc

@@ -1560,4 +1560,31 @@ TEST(ConsumerTest, testConsumerListenerShouldNotSegfaultAfterClose) {
     ASSERT_EQ(ResultOk, client.close());
 }
 
+TEST(ConsumerTest, testCloseAfterSeek) {
+    const auto topic = "test-close-after-seek-" + std::to_string(time(nullptr));
+    const auto subscription = "sub";
+    Client client(lookupUrl);
+    Consumer consumer;
+    ASSERT_EQ(ResultOk, client.subscribe(topic, subscription, consumer));
+    ASSERT_EQ(ResultOk, consumer.seek(TimeUtils::currentTimeMillis()));
+    ASSERT_EQ(ResultOk, consumer.close());
+
+    // Test the previous consumer will be closed even after seek is done, at the moment the connection might
+    // not be established.
+    ASSERT_EQ(ResultOk, client.subscribe(topic, subscription, consumer));
+    client.close();
+}
+
+TEST(ConsumerTest, testSeekAfterSeek) {
+    const auto topic = "test-seek-after-seek-" + std::to_string(time(nullptr));
+    const auto subscription = "sub";
+    Client client(lookupUrl);
+    Consumer consumer;
+    ASSERT_EQ(ResultOk, client.subscribe(topic, subscription, consumer));
+    ASSERT_EQ(ResultOk, consumer.seek(TimeUtils::currentTimeMillis()));
+    ASSERT_EQ(ResultOk, consumer.seek(MessageId::earliest()));
+    ASSERT_EQ(ResultOk, consumer.seek(TimeUtils::currentTimeMillis()));
+    client.close();
+}
+
 }  // namespace pulsar