fix the reader cannot consume after switching to a new cluster

Author: BewareMyPower

lib/ClientConnection.cc

@@ -1292,7 +1292,7 @@ void ClientConnection::handleConsumerStatsTimeout(const ASIO_ERROR& ec,
     startConsumerStatsTimer(consumerStatsRequests);
 }
 
-void ClientConnection::close(Result result, bool detach) {
+void ClientConnection::close(Result result, bool detach, bool switchCluster) {
     Lock lock(mutex_);
     if (isClosed()) {
         return;
@@ -1368,6 +1368,9 @@ void ClientConnection::close(Result result, bool detach) {
     for (ConsumersMap::iterator it = consumers.begin(); it != consumers.end(); ++it) {
         auto consumer = it->second.lock();
         if (consumer) {
+            if (switchCluster) {
+                consumer->onClusterSwitching();
+            }
             consumer->handleDisconnection(result, self);
         }
     }

lib/ClientConnection.h

@@ -157,10 +157,11 @@ class PULSAR_PUBLIC ClientConnection : public std::enable_shared_from_this<Clien
      *
      * @param result all pending futures will complete with this result
      * @param detach remove it from the pool if it's true
+     * @param switchCluster whether the close is triggered by cluster switching
      *
      * `detach` should only be false when the connection pool is closed.
      */
-    void close(Result result = ResultConnectError, bool detach = true);
+    void close(Result result = ResultConnectError, bool detach = true, bool switchCluster = false);
 
     bool isClosed() const;
 

lib/ClientImpl.cc

@@ -896,7 +896,7 @@ void ClientImpl::updateServiceInfo(ServiceInfo&& serviceInfo) {
                         ? std::nullopt
                         : std::make_optional(clientConfiguration_.getTlsTrustCertsFilePath())};
 
-    pool_.resetConnections(clientConfiguration_.getAuthPtr(), clientConfiguration_);
+    pool_.resetForClusterSwitching(clientConfiguration_.getAuthPtr(), clientConfiguration_);
 
     lookupServicePtr_->close();
     for (auto&& it : redirectedClusterLookupServicePtrs_) {

lib/ConnectionPool.cc

@@ -67,16 +67,16 @@ bool ConnectionPool::close() {
     return true;
 }
 
-void ConnectionPool::resetConnections(const AuthenticationPtr& authentication,
-                                      const ClientConfiguration& conf) {
+void ConnectionPool::resetForClusterSwitching(const AuthenticationPtr& authentication,
+                                              const ClientConfiguration& conf) {
     std::unique_lock<std::recursive_mutex> lock(mutex_);
     authentication_ = authentication;
     clientConfiguration_ = conf;
 
     for (auto cnxIt = pool_.begin(); cnxIt != pool_.end(); cnxIt++) {
         auto& cnx = cnxIt->second;
         if (cnx) {
-            cnx->close(ResultDisconnected, false);
+            cnx->close(ResultDisconnected, false, true);
         }
     }
     pool_.clear();

lib/ConnectionPool.h

@@ -55,7 +55,7 @@ class PULSAR_PUBLIC ConnectionPool {
      * Close all existing connections and update the authentication and configuration.
      * Unlike close(), the pool remains open for new connections.
      */
-    void resetConnections(const AuthenticationPtr& authentication, const ClientConfiguration& conf);
+    void resetForClusterSwitching(const AuthenticationPtr& authentication, const ClientConfiguration& conf);
 
     void remove(const std::string& logicalAddress, const std::string& physicalAddress, size_t keySuffix,
                 ClientConnection* value);

lib/ConsumerImpl.cc

@@ -125,7 +125,8 @@ ConsumerImpl::ConsumerImpl(const ClientImplPtr& client, const std::string& topic
       negativeAcksTracker_(std::make_shared<NegativeAcksTracker>(client, *this, conf)),
       ackGroupingTrackerPtr_(newAckGroupingTracker(topic, conf, client)),
       readCompacted_(conf.isReadCompacted()),
-      startMessageId_(pulsar::getStartMessageId(startMessageId, conf.isStartMessageIdInclusive())),
+      startMessageIdFromConfig_(pulsar::getStartMessageId(startMessageId, conf.isStartMessageIdInclusive())),
+      startMessageId_(startMessageIdFromConfig_),
       maxPendingChunkedMessage_(conf.getMaxPendingChunkedMessage()),
       autoAckOldestChunkedMessageOnQueueFull_(conf.isAutoAckOldestChunkedMessageOnQueueFull()),
       expireTimeOfIncompleteChunkedMessageMs_(conf.getExpireTimeOfIncompleteChunkedMessageMs()),
@@ -1134,6 +1135,16 @@ void ConsumerImpl::messageProcessed(Message& msg, bool track) {
     }
 }
 
+void ConsumerImpl::onClusterSwitching() {
+    {
+        LockGuard lock{mutex_};
+        incomingMessages_.clear();
+        startMessageId_ = startMessageIdFromConfig_;
+        lastDequedMessageId_ = MessageId::earliest();
+    }
+    ackGroupingTrackerPtr_->flushAndClean();
+}
+
 /**
  * Clear the internal receiver queue and returns the message id of what was the 1st message in the queue that
  * was

lib/ConsumerImpl.h

@@ -162,6 +162,8 @@ class ConsumerImpl : public ConsumerImplBase {
     void doImmediateAck(const MessageId& msgId, const ResultCallback& callback, CommandAck_AckType ackType);
     void doImmediateAck(const std::set<MessageId>& msgIds, const ResultCallback& callback);
 
+    void onClusterSwitching();
+
    protected:
     // overrided methods from HandlerBase
     Future<Result, bool> connectionOpened(const ClientConnectionPtr& cnx) override;
@@ -266,6 +268,11 @@ class ConsumerImpl : public ConsumerImplBase {
 
     MessageId lastDequedMessageId_{MessageId::earliest()};
     MessageId lastMessageIdInBroker_{MessageId::earliest()};
+
+    // When the consumer switches to a new cluster, we should reset `startMessageId_` to the original value,
+    // otherwise, the message id of the old cluster might be passed in the Subscribe request on the new
+    // cluster.
+    const optional<MessageId> startMessageIdFromConfig_;
     optional<MessageId> startMessageId_;
 
     SeekStatus seekStatus_{SeekStatus::NOT_STARTED};

tests/ClientTest.cc

@@ -522,11 +522,26 @@ TEST(ClientTest, testUpdateServiceInfo) {
     ServiceInfo info3{"pulsar://localhost:6650", std::nullopt, std::nullopt};
 
     Client client{info1.serviceUrl, ClientConfiguration().setAuth(*info1.authentication)};
+
     const auto topicRequiredAuth = "private/auth/testUpdateConnectionInfo-" + std::to_string(time(nullptr));
     Producer producer;
     ASSERT_EQ(ResultOk, client.createProducer(topicRequiredAuth, producer));
-    MessageId msgId;
-    ASSERT_EQ(ResultOk, producer.send(MessageBuilder().setContent("msg-0").build(), msgId));
+
+    Reader reader;
+    ASSERT_EQ(ResultOk, client.createReader(topicRequiredAuth, MessageId::earliest(), {}, reader));
+
+    auto sendAndReceive = [&](const std::string &value) {
+        MessageId msgId;
+        ASSERT_EQ(ResultOk, producer.send(MessageBuilder().setContent(value).build(), msgId));
+        LOG_INFO("Sent " << value << " to " << msgId);
+
+        Message msg;
+        ASSERT_EQ(ResultOk, reader.readNext(msg, 3000));
+        LOG_INFO("Read " << msg.getDataAsString() << " from " << msgId);
+        ASSERT_EQ(value, msg.getDataAsString());
+    };
+
+    sendAndReceive("msg-0");
 
     // Switch to cluster 2 (started by ./build-support/start-mim-test-service-inside-container.sh)
     ASSERT_FALSE(PulsarFriend::getConnections(client).empty());
@@ -537,8 +552,7 @@ TEST(ClientTest, testUpdateServiceInfo) {
     ASSERT_EQ(info2.tlsTrustCertsFilePath, client.getServiceInfo().tlsTrustCertsFilePath);
 
     // Now the same will access the same topic in cluster 2
-    ASSERT_EQ(ResultOk, producer.send(MessageBuilder().setContent("msg-1").build(), msgId));
-    ASSERT_EQ(ResultOk, producer.close());
+    sendAndReceive("msg-1");
 
     // Switch back to cluster 1 without any authentication, the previous authentication info configured for
     // cluster 2 will be cleared.
@@ -548,8 +562,9 @@ TEST(ClientTest, testUpdateServiceInfo) {
     ASSERT_EQ(info3.tlsTrustCertsFilePath, client.getServiceInfo().tlsTrustCertsFilePath);
 
     const auto topicNoAuth = "testUpdateConnectionInfo-" + std::to_string(time(nullptr));
+    producer.close();
     ASSERT_EQ(ResultOk, client.createProducer(topicNoAuth, producer));
-    ASSERT_EQ(ResultOk, producer.send(MessageBuilder().setContent("msg-2").build(), msgId));
+    ASSERT_EQ(ResultOk, producer.send(MessageBuilder().setContent("msg-2").build()));
 
     client.close();