[fix] Fix deadlock when closing the partitioned producer

Author: RobertIndie

lib/PartitionedProducerImpl.cc

@@ -267,73 +267,54 @@ void PartitionedProducerImpl::closeAsync(CloseCallback originalCallback) {
         }
     };
 
-    if (state_ == Closed || state_.exchange(Closing) == Closing) {
+    if (state_ == Closed || state_ == Failed || state_.exchange(Closing) == Closing) {
         closeCallback(ResultAlreadyClosed);
         return;
     }
 
-    cancelTimers();
+    auto remainProducersCount = std::make_shared<std::atomic<unsigned long>>(producers_.size());
+
+    auto subProducerClosedCallback = [remainProducersCount, closeCallback](unsigned int partitions,
+                                                                           Result result) {
+        (*remainProducersCount)--;
+        LOG_DEBUG("Closed Producer for partition "
+                  << partitions << " successfully. Remain producers to be closed: " << *remainProducersCount);
+        if (!*remainProducersCount) {
+            closeCallback(result);
+        }
+    };
 
-    unsigned int producerAlreadyClosed = 0;
+    cancelTimers();
 
     // Here we don't need `producersMutex` to protect `producers_`, because `producers_` can only be increased
     // when `state_` is Ready
     for (auto& producer : producers_) {
         if (!producer->isClosed()) {
             auto self = shared_from_this();
             const auto partition = static_cast<unsigned int>(producer->partition());
-            producer->closeAsync([this, self, partition, closeCallback](Result result) {
-                handleSinglePartitionProducerClose(result, partition, closeCallback);
+            producer->closeAsync([this, self, partition, subProducerClosedCallback](Result result) {
+                handleSinglePartitionProducerClose(result, partition,
+                                                   [partition, subProducerClosedCallback](Result result) {
+                                                       subProducerClosedCallback(partition, result);
+                                                   });
             });
         } else {
-            producerAlreadyClosed++;
+            subProducerClosedCallback(producer->partition(), ResultOk);
         }
     }
-    const auto numProducers = producers_.size();
-
-    /*
-     * No need to set state since:-
-     * a. If closeAsync before creation then state == Closed, since producers_.size() = producerAlreadyClosed
-     * = 0
-     * b. If closeAsync called after all sub partitioned producer connected -
-     * handleSinglePartitionProducerClose handles the closing
-     * c. If closeAsync called due to failure in creating just one sub producer then state is set by
-     * handleSinglePartitionProducerCreated
-     */
-    if (producerAlreadyClosed == numProducers) {
-        closeCallback(ResultOk);
-    }
 }
 
-// `callback` is a wrapper of user provided callback, it's not null and will call `shutdown()`
 void PartitionedProducerImpl::handleSinglePartitionProducerClose(Result result,
                                                                  const unsigned int partitionIndex,
                                                                  CloseCallback callback) {
-    if (state_ == Failed) {
-        // we should have already notified the client by callback
-        return;
-    }
     if (result != ResultOk) {
         LOG_ERROR("Closing the producer failed for partition - " << partitionIndex);
-        callback(result);
         state_ = Failed;
-        return;
-    }
-    assert(partitionIndex < getNumPartitionsWithLock());
-    if (numProducersCreated_ > 0) {
-        numProducersCreated_--;
-    }
-    // closed all successfully
-    if (!numProducersCreated_) {
-        // set the producerCreatedPromise to failure, if client called
-        // closeAsync and it's not failure to create producer, the promise
-        // is set second time here, first time it was successful. So check
-        // if there's any adverse effect of setting it again. It should not
-        // be but must check. MUSTCHECK changeme
-        partitionedProducerCreatedPromise_.setFailed(ResultUnknownError);
         callback(result);
         return;
     }
+    assert(partitionIndex < getNumPartitionsWithLock());
+    callback(result);
 }
 
 // override

tests/ProducerTest.cc

@@ -22,6 +22,7 @@
 #include <thread>
 
 #include "HttpHelper.h"
+#include "PulsarFriend.h"
 #include "lib/Future.h"
 #include "lib/Latch.h"
 #include "lib/LogUtils.h"
@@ -467,4 +468,27 @@ TEST(ProducerTest, testCloseProducerBeforeCreated) {
     client.close();
 }
 
+TEST(ProducerTest, testNoDeadlockWhenClosingPartitionedProducerAfterPartitionsUpdate) {
+    const std::string topic = "public/default/partition-test" + std::to_string(time(nullptr));
+    std::string topicOperateUrl = adminUrl + "admin/v2/persistent/" + topic + "/partitions";
+
+    int res = makePutRequest(topicOperateUrl, "2");
+    ASSERT_TRUE(res == 204 || res == 409) << "res: " << res;
+
+    ClientConfiguration clientConf;
+    clientConf.setPartititionsUpdateInterval(1);
+    Client client(serviceUrl, clientConf);
+    ProducerConfiguration conf;
+    Producer producer;
+    client.createProducer(topic, conf, producer);
+    PartitionedProducerImpl& partitionedProducer = PulsarFriend::getPartitionedProducerImpl(producer);
+
+    // TODO: Replace by producer interceptor to reproduce the issue then we can remove
+    // PulsarFriend::updatePartitions
+    PulsarFriend::updatePartitions(partitionedProducer, 3);
+
+    producer.close();
+    client.close();
+}
+
 INSTANTIATE_TEST_CASE_P(Pulsar, ProducerTest, ::testing::Values(true, false));

tests/PulsarFriend.h

@@ -79,6 +79,12 @@ class PulsarFriend {
         return *producerImpl;
     }
 
+    static PartitionedProducerImpl& getPartitionedProducerImpl(Producer producer) {
+        PartitionedProducerImpl* partitionedProducer =
+            static_cast<PartitionedProducerImpl*>(producer.impl_.get());
+        return *partitionedProducer;
+    }
+
     static ProducerImpl& getInternalProducerImpl(Producer producer, int index) {
         PartitionedProducerImpl* producerImpl = static_cast<PartitionedProducerImpl*>(producer.impl_.get());
         return *(producerImpl->producers_[index]);
@@ -181,6 +187,12 @@ class PulsarFriend {
     static proto::MessageMetadata& getMessageMetadata(Message& message) { return message.impl_->metadata; }
 
     static std::shared_ptr<MessageIdImpl> getMessageIdImpl(MessageId& msgId) { return msgId.impl_; }
+
+    static void updatePartitions(PartitionedProducerImpl& partitionedProducer, int newPartitions) {
+        LookupDataResultPtr lookupData = std::make_shared<LookupDataResult>();
+        lookupData->setPartitions(newPartitions);
+        partitionedProducer.handleGetPartitions(ResultOk, lookupData);
+    }
 };
 }  // namespace pulsar