apache · BewareMyPower · Feb 29, 2024 · Feb 26, 2024 · Feb 27, 2024 · Feb 28, 2024
diff --git a/lib/ClientConnection.cc b/lib/ClientConnection.cc
@@ -1800,6 +1800,7 @@ void ClientConnection::handleTopicMigrated(const proto::CommandTopicMigrated& co
         if (it != producers_.end()) {
             auto producer = it->second.lock();
             producer->setRedirectedClusterURI(migratedBrokerServiceUrl);
+            unsafeRemovePendingRequest(producer->firstRequestIdAfterConnect());
             LOG_INFO("Producer id:" << resourceId << " is migrated to " << migratedBrokerServiceUrl);
         } else {
             LOG_WARN("Got invalid producer Id in topicMigrated command: " << resourceId);
@@ -1809,6 +1810,7 @@ void ClientConnection::handleTopicMigrated(const proto::CommandTopicMigrated& co
         if (it != consumers_.end()) {
             auto consumer = it->second.lock();
             consumer->setRedirectedClusterURI(migratedBrokerServiceUrl);
+            unsafeRemovePendingRequest(consumer->firstRequestIdAfterConnect());
             LOG_INFO("Consumer id:" << resourceId << " is migrated to " << migratedBrokerServiceUrl);
         } else {
             LOG_WARN("Got invalid consumer Id in topicMigrated command: " << resourceId);
@@ -2027,4 +2029,14 @@ void ClientConnection::handleAckResponse(const proto::CommandAckResponse& respon
     }
 }
 
+void ClientConnection::unsafeRemovePendingRequest(long requestId) {
+    auto it = pendingRequests_.find(requestId);
+    if (it != pendingRequests_.end()) {
+        it->second.promise.setFailed(ResultDisconnected);
+        ASIO_ERROR ec;
+        it->second.timer->cancel(ec);
+        pendingRequests_.erase(it);
+    }
+}
+
 }  // namespace pulsar
diff --git a/lib/ClientConnection.h b/lib/ClientConnection.h
@@ -426,6 +426,8 @@ class PULSAR_PUBLIC ClientConnection : public std::enable_shared_from_this<Clien
     boost::optional<std::string> getAssignedBrokerServiceUrl(const proto::CommandCloseProducer&);
     boost::optional<std::string> getAssignedBrokerServiceUrl(const proto::CommandCloseConsumer&);
     std::string getMigratedBrokerServiceUrl(const proto::CommandTopicMigrated&);
+    // This method must be called when `mutex_` is held
+    void unsafeRemovePendingRequest(long requestId);
 };
 }  // namespace pulsar
 

diff --git a/lib/ConsumerImpl.cc b/lib/ConsumerImpl.cc
@@ -251,7 +251,7 @@ Future<Result, bool> ConsumerImpl::connectionOpened(const ClientConnectionPtr& c
     unAckedMessageTrackerPtr_->clear();
 
     ClientImplPtr client = client_.lock();
-    uint64_t requestId = client->newRequestId();
+    long requestId = client->newRequestId();
     SharedBuffer cmd = Commands::newSubscribe(
         topic(), subscription_, consumerId_, requestId, getSubType(), getConsumerName(), subscriptionMode_,
         subscribeMessageId, readCompacted_, config_.getProperties(), config_.getSubscriptionProperties(),
@@ -260,6 +260,7 @@ Future<Result, bool> ConsumerImpl::connectionOpened(const ClientConnectionPtr& c
 
     // Keep a reference to ensure object is kept alive.
     auto self = get_shared_this_ptr();
+    setFirstRequestIdAfterConnect(requestId);
     cnx->sendRequestWithId(cmd, requestId)
         .addListener([this, self, cnx, promise](Result result, const ResponseData& responseData) {
             Result handleResult = handleCreateConsumer(cnx, result);
@@ -1719,6 +1720,7 @@ void ConsumerImpl::cancelTimers() noexcept {
     batchReceiveTimer_->cancel(ec);
     checkExpiredChunkedTimer_->cancel(ec);
     unAckedMessageTrackerPtr_->stop();
+    consumerStatsBasePtr_->stop();
 }
 
 void ConsumerImpl::processPossibleToDLQ(const MessageId& messageId, ProcessDLQCallBack cb) {

diff --git a/lib/HandlerBase.h b/lib/HandlerBase.h
@@ -100,6 +100,10 @@ class HandlerBase : public std::enable_shared_from_this<HandlerBase> {
     const std::string& topic() const { return *topic_; }
     const std::shared_ptr<std::string>& getTopicPtr() const { return topic_; }
 
+    long firstRequestIdAfterConnect() const {
+        return firstRequestIdAfterConnect_.load(std::memory_order_acquire);
+    }
+
    private:
     const std::shared_ptr<std::string> topic_;
 
@@ -140,6 +144,10 @@ class HandlerBase : public std::enable_shared_from_this<HandlerBase> {
 
     Result convertToTimeoutIfNecessary(Result result, ptime startTimestamp) const;
 
+    void setFirstRequestIdAfterConnect(long requestId) {
+        firstRequestIdAfterConnect_.store(requestId, std::memory_order_release);
+    }
+
    private:
     DeadlineTimerPtr timer_;
     DeadlineTimerPtr creationTimer_;
@@ -148,6 +156,7 @@ class HandlerBase : public std::enable_shared_from_this<HandlerBase> {
     std::atomic<bool> reconnectionPending_;
     ClientConnectionWeakPtr connection_;
     std::string redirectedClusterURI_;
+    std::atomic<long> firstRequestIdAfterConnect_{-1L};
 
     friend class ClientConnection;
     friend class PulsarFriend;

diff --git a/lib/ProducerImpl.cc b/lib/ProducerImpl.cc
@@ -149,7 +149,7 @@ Future<Result, bool> ProducerImpl::connectionOpened(const ClientConnectionPtr& c
     ClientImplPtr client = client_.lock();
     cnx->registerProducer(producerId_, shared_from_this());
 
-    int requestId = client->newRequestId();
+    long requestId = client->newRequestId();
 
     SharedBuffer cmd = Commands::newProducer(topic(), producerId_, producerName_, requestId,
                                              conf_.getProperties(), conf_.getSchema(), epoch_,
@@ -159,6 +159,7 @@ Future<Result, bool> ProducerImpl::connectionOpened(const ClientConnectionPtr& c
 
     // Keep a reference to ensure object is kept alive.
     auto self = shared_from_this();
+    setFirstRequestIdAfterConnect(requestId);
     cnx->sendRequestWithId(cmd, requestId)
         .addListener([this, self, cnx, promise](Result result, const ResponseData& responseData) {
             Result handleResult = handleCreateProducer(cnx, result, responseData);

diff --git a/lib/stats/ConsumerStatsBase.h b/lib/stats/ConsumerStatsBase.h
@@ -28,6 +28,7 @@ namespace pulsar {
 class ConsumerStatsBase {
    public:
     virtual void start() {}
+    virtual void stop() {}
     virtual void receivedMessage(Message&, Result) = 0;
     virtual void messageAcknowledged(Result, CommandAck_AckType, uint32_t ackNums = 1) = 0;
     virtual ~ConsumerStatsBase() {}

diff --git a/lib/stats/ConsumerStatsImpl.h b/lib/stats/ConsumerStatsImpl.h
@@ -59,6 +59,10 @@ class ConsumerStatsImpl : public std::enable_shared_from_this<ConsumerStatsImpl>
     ConsumerStatsImpl(const ConsumerStatsImpl& stats);
     void flushAndReset(const ASIO_ERROR&);
     void start() override;
+    void stop() override {
+        ASIO_ERROR error;
+        timer_->cancel(error);
+    }
     void receivedMessage(Message&, Result) override;
     void messageAcknowledged(Result, CommandAck_AckType, uint32_t ackNums) override;
     virtual ~ConsumerStatsImpl();

diff --git a/tests/PulsarFriend.h b/tests/PulsarFriend.h
@@ -169,7 +169,7 @@ class PulsarFriend {
     static ClientConnectionWeakPtr getClientConnection(HandlerBase& handler) { return handler.connection_; }
 
     static std::string getConnectionPhysicalAddress(HandlerBase& handler) {
-        auto cnx = handler.connection_.lock();
+        auto cnx = handler.getCnx().lock();
         if (cnx) {
             return cnx->physicalAddress_;
         }

diff --git a/tests/extensibleLM/ExtensibleLoadManagerTest.cc b/tests/extensibleLM/ExtensibleLoadManagerTest.cc
@@ -25,6 +25,7 @@
 #include "include/pulsar/Client.h"
 #include "lib/LogUtils.h"
 #include "lib/Semaphore.h"
+#include "lib/TimeUtils.h"
 #include "tests/HttpHelper.h"
 #include "tests/PulsarFriend.h"
 
@@ -40,6 +41,9 @@ bool checkTime() {
 }
 
 TEST(ExtensibleLoadManagerTest, testPubSubWhileUnloading) {
+    constexpr auto maxWaitTime = std::chrono::seconds(5);
+    constexpr long maxWaitTimeMs = maxWaitTime.count() * 1000L;
+
     const static std::string blueAdminUrl = "http://localhost:8080/";
     const static std::string greenAdminUrl = "http://localhost:8081/";
     const static std::string topicNameSuffix = std::to_string(time(NULL));
@@ -105,12 +109,13 @@ TEST(ExtensibleLoadManagerTest, testPubSubWhileUnloading) {
             std::string content = std::to_string(i);
             const auto msg = MessageBuilder().setContent(content).build();
 
-            ASSERT_TRUE(waitUntil(std::chrono::seconds(60), [&] {
-                Result sendResult = producer.send(msg);
-                return sendResult == ResultOk;
-            }));
+            auto start = TimeUtils::currentTimeMillis();
+            Result sendResult = producer.send(msg);
+            auto elapsed = TimeUtils::currentTimeMillis() - start;
+            LOG_INFO("produce i: " << i << " " << elapsed << " ms");
+            ASSERT_EQ(sendResult, ResultOk);
+            ASSERT_TRUE(elapsed < maxWaitTimeMs);
 
-            LOG_INFO("produced i:" << i);
             producedMsgs.emplace(i, i);
             i++;
         }
@@ -124,18 +129,20 @@ TEST(ExtensibleLoadManagerTest, testPubSubWhileUnloading) {
             if (stopConsumer && producedMsgs.size() == msgCount && consumedMsgs.size() == msgCount) {
                 break;
             }
-            Result receiveResult =
-                consumer.receive(receivedMsg, 1000);  // Assumed that we wait 1000 ms for each message
-            if (receiveResult != ResultOk) {
-                continue;
-            }
+            auto start = TimeUtils::currentTimeMillis();
+            Result receiveResult = consumer.receive(receivedMsg, maxWaitTimeMs);
+            auto elapsed = TimeUtils::currentTimeMillis() - start;
             int i = std::stoi(receivedMsg.getDataAsString());
-            LOG_INFO("received i:" << i);
-            ASSERT_TRUE(waitUntil(std::chrono::seconds(60), [&] {
-                Result ackResult = consumer.acknowledge(receivedMsg);
-                return ackResult == ResultOk;
-            }));
-            LOG_INFO("acked i:" << i);
+            LOG_INFO("receive i: " << i << " " << elapsed << " ms");
+            ASSERT_EQ(receiveResult, ResultOk);
+            ASSERT_TRUE(elapsed < maxWaitTimeMs);
+
+            start = TimeUtils::currentTimeMillis();
+            Result ackResult = consumer.acknowledge(receivedMsg);
+            elapsed = TimeUtils::currentTimeMillis() - start;
+            LOG_INFO("acked i:" << i << " " << elapsed << " ms");
+            ASSERT_TRUE(elapsed < maxWaitTimeMs);
+            ASSERT_EQ(ackResult, ResultOk);
             consumedMsgs.emplace(i, i);
         }
         LOG_INFO("consumer finished");
@@ -153,7 +160,7 @@ TEST(ExtensibleLoadManagerTest, testPubSubWhileUnloading) {
         std::string destinationBroker;
         while (checkTime()) {
             // make sure producers and consumers are ready
-            ASSERT_TRUE(waitUntil(std::chrono::seconds(30),
+            ASSERT_TRUE(waitUntil(maxWaitTime,
                                   [&] { return consumerImpl.isConnected() && producerImpl.isConnected(); }));
 
             std::string url =
@@ -182,7 +189,7 @@ TEST(ExtensibleLoadManagerTest, testPubSubWhileUnloading) {
             }
 
             // make sure producers and consumers are ready
-            ASSERT_TRUE(waitUntil(std::chrono::seconds(30),
+            ASSERT_TRUE(waitUntil(maxWaitTime,
                                   [&] { return consumerImpl.isConnected() && producerImpl.isConnected(); }));
             std::string responseDataAfterUnload;
             ASSERT_TRUE(waitUntil(std::chrono::seconds(60), [&] {
@@ -220,7 +227,7 @@ TEST(ExtensibleLoadManagerTest, testPubSubWhileUnloading) {
         LOG_INFO("res:" << res);
         return res == 200;
     }));
-    ASSERT_TRUE(waitUntil(std::chrono::seconds(130), [&] {
+    ASSERT_TRUE(waitUntil(maxWaitTime, [&] {
         auto &consumerImpl = PulsarFriend::getConsumerImpl(consumer);
         auto &producerImpl = PulsarFriend::getProducerImpl(producer);
         auto consumerConnAddress = PulsarFriend::getConnectionPhysicalAddress(consumerImpl);