[feat] PIP-307 added assigned broker urls for CloseProudcer and CloseConsumer commands and handler logic (#389)

Master Issue: #[pip-307](https://github.com/apache/pulsar/blob/master/pip/pip-307.md)

### Motivation

As part of PIP-307, I am adding the client logic to this c++ client.

### Modifications

- Added assigned broker urls to CloseProudcer and CloseConsumer commands.
- Updated the client reconnection logic to directly connect to the assigned broker urls

### Verifying this change

- Added ExtensibleLoadManagerTest to cover this logic

Author: heesung-sohn

lib/BinaryProtoLookupService.cc

@@ -87,9 +87,9 @@ auto BinaryProtoLookupService::findBroker(const std::string& address, bool autho
                                                 << ", from " << cnx->cnxString());
                 if (data->shouldProxyThroughServiceUrl()) {
                     // logicalAddress is the proxy's address, we should still connect through proxy
-                    promise->setValue({responseBrokerAddress, address});
+                    promise->setValue({responseBrokerAddress, address, true});
                 } else {
-                    promise->setValue({responseBrokerAddress, responseBrokerAddress});
+                    promise->setValue({responseBrokerAddress, responseBrokerAddress, false});
                 }
             }
         });

lib/ClientConnection.cc

@@ -1762,6 +1762,30 @@ void ClientConnection::handleError(const proto::CommandError& error) {
     }
 }
 
+boost::optional<std::string> ClientConnection::getAssignedBrokerServiceUrl(
+    const proto::CommandCloseProducer& closeProducer) {
+    if (tlsSocket_) {
+        if (closeProducer.has_assignedbrokerserviceurltls()) {
+            return closeProducer.assignedbrokerserviceurltls();
+        }
+    } else if (closeProducer.has_assignedbrokerserviceurl()) {
+        return closeProducer.assignedbrokerserviceurl();
+    }
+    return boost::none;
+}
+
+boost::optional<std::string> ClientConnection::getAssignedBrokerServiceUrl(
+    const proto::CommandCloseConsumer& closeConsumer) {
+    if (tlsSocket_) {
+        if (closeConsumer.has_assignedbrokerserviceurltls()) {
+            return closeConsumer.assignedbrokerserviceurltls();
+        }
+    } else if (closeConsumer.has_assignedbrokerserviceurl()) {
+        return closeConsumer.assignedbrokerserviceurl();
+    }
+    return boost::none;
+}
+
 void ClientConnection::handleCloseProducer(const proto::CommandCloseProducer& closeProducer) {
     int producerId = closeProducer.producer_id();
 
@@ -1775,7 +1799,8 @@ void ClientConnection::handleCloseProducer(const proto::CommandCloseProducer& cl
         lock.unlock();
 
         if (producer) {
-            producer->disconnectProducer();
+            auto assignedBrokerServiceUrl = getAssignedBrokerServiceUrl(closeProducer);
+            producer->disconnectProducer(assignedBrokerServiceUrl);
         }
     } else {
         LOG_ERROR(cnxString_ << "Got invalid producer Id in closeProducer command: " << producerId);
@@ -1795,7 +1820,8 @@ void ClientConnection::handleCloseConsumer(const proto::CommandCloseConsumer& cl
         lock.unlock();
 
         if (consumer) {
-            consumer->disconnectConsumer();
+            auto assignedBrokerServiceUrl = getAssignedBrokerServiceUrl(closeconsumer);
+            consumer->disconnectConsumer(assignedBrokerServiceUrl);
         }
     } else {
         LOG_ERROR(cnxString_ << "Got invalid consumer Id in closeConsumer command: " << consumerId);

lib/ClientConnection.h

@@ -421,6 +421,10 @@ class PULSAR_PUBLIC ClientConnection : public std::enable_shared_from_this<Clien
     void handleGetTopicOfNamespaceResponse(const proto::CommandGetTopicsOfNamespaceResponse&);
     void handleGetSchemaResponse(const proto::CommandGetSchemaResponse&);
     void handleAckResponse(const proto::CommandAckResponse&);
+    boost::optional<std::string> getAssignedBrokerServiceUrl(
+        const proto::CommandCloseProducer& closeProducer);
+    boost::optional<std::string> getAssignedBrokerServiceUrl(
+        const proto::CommandCloseConsumer& closeConsumer);
 };
 }  // namespace pulsar
 

lib/ClientImpl.cc

@@ -91,7 +91,9 @@ ClientImpl::ClientImpl(const std::string& serviceUrl, const ClientConfiguration&
             ClientImpl::getClientVersion(clientConfiguration)),
       producerIdGenerator_(0),
       consumerIdGenerator_(0),
-      closingError(ResultOk) {
+      closingError(ResultOk),
+      useProxy_(false),
+      lookupCount_(0L) {
     std::unique_ptr<LoggerFactory> loggerFactory = clientConfiguration_.impl_->takeLogger();
     if (loggerFactory) {
         LogUtils::setLoggerFactory(std::move(loggerFactory));
@@ -532,6 +534,8 @@ Future<Result, ClientConnectionPtr> ClientImpl::getConnection(const std::string&
                 promise.setFailed(result);
                 return;
             }
+            useProxy_ = data.proxyThroughServiceUrl;
+            lookupCount_++;
             pool_.getConnectionAsync(data.logicalAddress, data.physicalAddress, key)
                 .addListener([promise](Result result, const ClientConnectionWeakPtr& weakCnx) {
                     if (result == ResultOk) {
@@ -550,6 +554,33 @@ Future<Result, ClientConnectionPtr> ClientImpl::getConnection(const std::string&
     return promise.getFuture();
 }
 
+const std::string& ClientImpl::getPhysicalAddress(const std::string& logicalAddress) {
+    if (useProxy_) {
+        return serviceNameResolver_.resolveHost();
+    } else {
+        return logicalAddress;
+    }
+}
+
+Future<Result, ClientConnectionPtr> ClientImpl::connect(const std::string& logicalAddress, size_t key) {
+    const auto& physicalAddress = getPhysicalAddress(logicalAddress);
+    Promise<Result, ClientConnectionPtr> promise;
+    pool_.getConnectionAsync(logicalAddress, physicalAddress, key)
+        .addListener([promise](Result result, const ClientConnectionWeakPtr& weakCnx) {
+            if (result == ResultOk) {
+                auto cnx = weakCnx.lock();
+                if (cnx) {
+                    promise.setValue(cnx);
+                } else {
+                    promise.setFailed(ResultConnectError);
+                }
+            } else {
+                promise.setFailed(result);
+            }
+        });
+    return promise.getFuture();
+}
+
 void ClientImpl::handleGetPartitions(const Result result, const LookupDataResultPtr partitionMetadata,
                                      TopicNamePtr topicName, GetPartitionsCallback callback) {
     if (result != ResultOk) {
@@ -635,6 +666,7 @@ void ClientImpl::closeAsync(CloseCallback callback) {
     if (*numberOfOpenHandlers == 0 && callback) {
         handleClose(ResultOk, numberOfOpenHandlers, callback);
     }
+    lookupCount_ = 0;
 }
 
 void ClientImpl::handleClose(Result result, SharedInt numberOfOpenHandlers, ResultCallback callback) {
@@ -722,6 +754,7 @@ void ClientImpl::shutdown() {
     partitionListenerExecutorProvider_->close(timeoutProcessor.getLeftTimeout());
     timeoutProcessor.tok();
     LOG_DEBUG("partitionListenerExecutorProvider_ is closed");
+    lookupCount_ = 0;
 }
 
 uint64_t ClientImpl::newProducerId() {

lib/ClientImpl.h

@@ -97,6 +97,8 @@ class ClientImpl : public std::enable_shared_from_this<ClientImpl> {
 
     Future<Result, ClientConnectionPtr> getConnection(const std::string& topic, size_t key);
 
+    Future<Result, ClientConnectionPtr> connect(const std::string& logicalAddress, size_t key);
+
     void closeAsync(CloseCallback callback);
     void shutdown();
 
@@ -124,6 +126,7 @@ class ClientImpl : public std::enable_shared_from_this<ClientImpl> {
     std::shared_ptr<std::atomic<uint64_t>> getRequestIdGenerator() const { return requestIdGenerator_; }
 
     ConnectionPool& getConnectionPool() noexcept { return pool_; }
+    uint64_t getLookupCount() { return lookupCount_; }
 
     static std::chrono::nanoseconds getOperationTimeout(const ClientConfiguration& clientConfiguration);
 
@@ -160,6 +163,8 @@ class ClientImpl : public std::enable_shared_from_this<ClientImpl> {
                                           const std::string& consumerName, const ConsumerConfiguration& conf,
                                           SubscribeCallback callback);
 
+    const std::string& getPhysicalAddress(const std::string& logicalAddress);
+
     static std::string getClientVersion(const ClientConfiguration& clientConfiguration);
 
     enum State
@@ -191,6 +196,8 @@ class ClientImpl : public std::enable_shared_from_this<ClientImpl> {
     SynchronizedHashMap<ConsumerImplBase*, ConsumerImplBaseWeakPtr> consumers_;
 
     std::atomic<Result> closingError;
+    std::atomic<bool> useProxy_;
+    std::atomic<uint64_t> lookupCount_;
 
     friend class Client;
 };

lib/ConsumerImpl.cc

@@ -1234,10 +1234,13 @@ void ConsumerImpl::negativeAcknowledge(const MessageId& messageId) {
     negativeAcksTracker_->add(messageId);
 }
 
-void ConsumerImpl::disconnectConsumer() {
-    LOG_INFO("Broker notification of Closed consumer: " << consumerId_);
+void ConsumerImpl::disconnectConsumer() { disconnectConsumer(boost::none); }
+
+void ConsumerImpl::disconnectConsumer(const boost::optional<std::string>& assignedBrokerUrl) {
+    LOG_INFO("Broker notification of Closed consumer: "
+             << consumerId_ << (assignedBrokerUrl ? (" assignedBrokerUrl: " + assignedBrokerUrl.get()) : ""));
     resetCnx();
-    scheduleReconnection();
+    scheduleReconnection(assignedBrokerUrl);
 }
 
 void ConsumerImpl::closeAsync(ResultCallback originalCallback) {

lib/ConsumerImpl.h

@@ -131,6 +131,7 @@ class ConsumerImpl : public ConsumerImplBase {
     void hasMessageAvailableAsync(HasMessageAvailableCallback callback) override;
 
     virtual void disconnectConsumer();
+    virtual void disconnectConsumer(const boost::optional<std::string>& assignedBrokerUrl);
     Result fetchSingleMessageFromBroker(Message& msg);
 
     virtual bool isCumulativeAcknowledgementAllowed(ConsumerType consumerType);

lib/HandlerBase.cc

@@ -68,7 +68,18 @@ void HandlerBase::setCnx(const ClientConnectionPtr& cnx) {
     connection_ = cnx;
 }
 
-void HandlerBase::grabCnx() {
+void HandlerBase::grabCnx() { grabCnx(boost::none); }
+
+Future<Result, ClientConnectionPtr> HandlerBase::getConnection(
+    const ClientImplPtr& client, const boost::optional<std::string>& assignedBrokerUrl) {
+    if (assignedBrokerUrl && client->getLookupCount() > 0) {
+        return client->connect(assignedBrokerUrl.get(), connectionKeySuffix_);
+    } else {
+        return client->getConnection(topic(), connectionKeySuffix_);
+    }
+}
+
+void HandlerBase::grabCnx(const boost::optional<std::string>& assignedBrokerUrl) {
     bool expectedState = false;
     if (!reconnectionPending_.compare_exchange_strong(expectedState, true)) {
         LOG_INFO(getName() << "Ignoring reconnection attempt since there's already a pending reconnection");
@@ -90,7 +101,7 @@ void HandlerBase::grabCnx() {
         return;
     }
     auto self = shared_from_this();
-    auto cnxFuture = client->getConnection(topic(), connectionKeySuffix_);
+    auto cnxFuture = getConnection(client, assignedBrokerUrl);
     cnxFuture.addListener([this, self](Result result, const ClientConnectionPtr& cnx) {
         if (result == ResultOk) {
             LOG_DEBUG(getName() << "Connected to broker: " << cnx->cnxString());
@@ -141,37 +152,37 @@ void HandlerBase::handleDisconnection(Result result, const ClientConnectionPtr&
             break;
     }
 }
-
-void HandlerBase::scheduleReconnection() {
+void HandlerBase::scheduleReconnection() { scheduleReconnection(boost::none); }
+void HandlerBase::scheduleReconnection(const boost::optional<std::string>& assignedBrokerUrl) {
     const auto state = state_.load();
 
     if (state == Pending || state == Ready) {
-        TimeDuration delay = backoff_.next();
+        TimeDuration delay = assignedBrokerUrl ? std::chrono::milliseconds(0) : backoff_.next();
 
         LOG_INFO(getName() << "Schedule reconnection in " << (toMillis(delay) / 1000.0) << " s");
         timer_->expires_from_now(delay);
         // passing shared_ptr here since time_ will get destroyed, so tasks will be cancelled
         // so we will not run into the case where grabCnx is invoked on out of scope handler
         auto name = getName();
         std::weak_ptr<HandlerBase> weakSelf{shared_from_this()};
-        timer_->async_wait([name, weakSelf](const ASIO_ERROR& ec) {
+        timer_->async_wait([name, weakSelf, assignedBrokerUrl](const ASIO_ERROR& ec) {
             auto self = weakSelf.lock();
             if (self) {
-                self->handleTimeout(ec);
+                self->handleTimeout(ec, assignedBrokerUrl);
             } else {
                 LOG_WARN(name << "Cancel the reconnection since the handler is destroyed");
             }
         });
     }
 }
 
-void HandlerBase::handleTimeout(const ASIO_ERROR& ec) {
+void HandlerBase::handleTimeout(const ASIO_ERROR& ec, const boost::optional<std::string>& assignedBrokerUrl) {
     if (ec) {
         LOG_DEBUG(getName() << "Ignoring timer cancelled event, code[" << ec << "]");
         return;
     } else {
         epoch_++;
-        grabCnx();
+        grabCnx(assignedBrokerUrl);
     }
 }
 

lib/HandlerBase.h

@@ -20,6 +20,7 @@
 #define _PULSAR_HANDLER_BASE_HEADER_
 #include <pulsar/Result.h>
 
+#include <boost/optional.hpp>
 #include <memory>
 #include <mutex>
 #include <string>
@@ -53,11 +54,22 @@ class HandlerBase : public std::enable_shared_from_this<HandlerBase> {
     void resetCnx() { setCnx(nullptr); }
 
    protected:
+    /*
+     * tries reconnection and sets connection_ to valid object
+     * @param assignedBrokerUrl assigned broker url to directly connect to without lookup
+     */
+    void grabCnx(const boost::optional<std::string>& assignedBrokerUrl);
+
     /*
      * tries reconnection and sets connection_ to valid object
      */
     void grabCnx();
 
+    /*
+     * Schedule reconnection after backoff time
+     * @param assignedBrokerUrl assigned broker url to directly connect to without lookup
+     */
+    void scheduleReconnection(const boost::optional<std::string>& assignedBrokerUrl);
     /*
      * Schedule reconnection after backoff time
      */
@@ -89,9 +101,12 @@ class HandlerBase : public std::enable_shared_from_this<HandlerBase> {
    private:
     const std::shared_ptr<std::string> topic_;
 
+    Future<Result, ClientConnectionPtr> getConnection(const ClientImplPtr& client,
+                                                      const boost::optional<std::string>& assignedBrokerUrl);
+
     void handleDisconnection(Result result, const ClientConnectionPtr& cnx);
 
-    void handleTimeout(const ASIO_ERROR& ec);
+    void handleTimeout(const ASIO_ERROR& ec, const boost::optional<std::string>& assignedBrokerUrl);
 
    protected:
     ClientImplWeakPtr client_;

lib/LookupService.h

@@ -42,6 +42,7 @@ class LookupService {
     struct LookupResult {
         std::string logicalAddress;
         std::string physicalAddress;
+        bool proxyThroughServiceUrl;
 
         friend std::ostream& operator<<(std::ostream& os, const LookupResult& lookupResult) {
             return os << "logical address: " << lookupResult.logicalAddress