cherry pick vesoft-inc#431 from 2.0

critical27 · Apr 1, 2021 · 4c61100 · 4c61100
1 parent e89a428
commit 4c61100
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 130 deletions.
diff --git a/src/common/thrift/CMakeLists.txt b/src/common/thrift/CMakeLists.txt
@@ -1,5 +1,4 @@
 nebula_add_library(
     thrift_obj OBJECT
-    ReconnectingRequestChannel.cpp
     ThriftClientManager.cpp
 )
diff --git a/src/common/thrift/ReconnectingRequestChannel.cpp b/src/common/thrift/ReconnectingRequestChannel.cpp
diff --git a/src/common/thrift/ThriftClientManager.inl b/src/common/thrift/ThriftClientManager.inl
@@ -4,9 +4,9 @@
  * attached with Common Clause Condition 1.0, found in the LICENSES directory.
  */
 
-#include <thrift/lib/cpp2/async/ReconnectingRequestChannel.h>
 #include <thrift/lib/cpp2/async/HeaderClientChannel.h>
 #include <thrift/lib/cpp/async/TAsyncSocket.h>
+#include <folly/io/async/AsyncSocket.h>
 #include <folly/system/ThreadName.h>
 #include "network/NetworkUtils.h"
 
@@ -18,51 +18,57 @@ namespace thrift {
 template<class ClientType>
 std::shared_ptr<ClientType> ThriftClientManager<ClientType>::client(
         const HostAddr& host, folly::EventBase* evb, bool compatibility, uint32_t timeout) {
-    VLOG(2) << "Getting a client to "
-            << network::NetworkUtils::intToIPv4(host.first)
-            << ":" << host.second;
-
     if (evb == nullptr) {
         evb = folly::EventBaseManager::get()->getEventBase();
     }
 
-    auto it = clientMap_->find(std::make_pair(host, evb));
-    if (it != clientMap_->end()) {
-        return it->second;
-    }
-
-    // Need to create a new client
     auto ipAddr = network::NetworkUtils::intToIPv4(host.first);
     auto port = host.second;
-    VLOG(2) << "There is no existing client to "
-            << ipAddr << ":" << port
-            << ", trying to create one";
-    auto channel = apache::thrift::ReconnectingRequestChannel::newChannel(
-        *evb, [compatibility, ipAddr, port, timeout] (folly::EventBase& eb) mutable {
-            static thread_local int connectionCount = 0;
-            VLOG(2) << "Connecting to " << ipAddr << ":" << port
-                    << " for " << ++connectionCount << " times";
-            std::shared_ptr<apache::thrift::async::TAsyncSocket> socket;
-            eb.runImmediatelyOrRunInEventBaseThreadAndWait(
-                [&socket, &eb, ipAddr, port]() {
-                    socket = apache::thrift::async::TAsyncSocket::newSocket(
-                        &eb, ipAddr, port, FLAGS_conn_timeout_ms);
-                });
-            auto headerClientChannel =  apache::thrift::HeaderClientChannel::newChannel(socket);
-            if (timeout > 0) {
-                headerClientChannel->setTimeout(timeout);
+    // Get client from client manager if it is ok.
+    auto it = clientMap_->find(std::make_pair(host, evb));
+    if (it != clientMap_->end()) {
+        do {
+            auto channel =
+                dynamic_cast<apache::thrift::HeaderClientChannel*>(it->second->getChannel());
+            if (channel == nullptr || !channel->good()) {
+                // Remove bad connection to create a new one.
+                clientMap_->erase(it);
+                VLOG(2) << "Invalid Channel: " << channel << " for host " << ipAddr << ":" << port;
+                break;
             }
-            if (compatibility) {
-                headerClientChannel->setProtocolId(apache::thrift::protocol::T_BINARY_PROTOCOL);
-                headerClientChannel->setClientType(THRIFT_UNFRAMED_DEPRECATED);
+            auto transport = dynamic_cast<folly::AsyncSocket*>(channel->getTransport());
+            if (transport == nullptr || transport->hangup()) {
+                clientMap_->erase(it);
+                VLOG(2) << "Transport is closed by peers " << transport << " for host "
+                        << ipAddr << ":" << port;
+                break;
             }
-            return headerClientChannel;
-        });
-    std::shared_ptr<ClientType> client(new ClientType(std::move(channel)), [evb](auto* p) {
-        evb->runImmediatelyOrRunInEventBaseThreadAndWait([p] {
-            delete p;
-        });
+            VLOG(2) << "Getting a client to " << ipAddr << ":" << port;
+            return it->second;
+        } while (false);
+    }
+
+    // Need to create a new client and insert it to client map.
+    VLOG(2) << "There is no existing client to " << host << ", trying to create one";
+    static thread_local int connectionCount = 0;
+
+    VLOG(2) << "Connecting to " << host << " for " << ++connectionCount << " times";
+    std::shared_ptr<apache::thrift::async::TAsyncSocket> socket;
+    evb->runImmediatelyOrRunInEventBaseThreadAndWait([&socket, evb, ipAddr, port]() {
+        socket = apache::thrift::async::TAsyncSocket::newSocket(
+            evb, ipAddr, port, FLAGS_conn_timeout_ms);
     });
+    auto headerClientChannel = apache::thrift::HeaderClientChannel::newChannel(socket);
+    if (timeout > 0) {
+        headerClientChannel->setTimeout(timeout);
+    }
+    if (compatibility) {
+        headerClientChannel->setProtocolId(apache::thrift::protocol::T_BINARY_PROTOCOL);
+        headerClientChannel->setClientType(THRIFT_UNFRAMED_DEPRECATED);
+    }
+    std::shared_ptr<ClientType> client(
+        new ClientType(std::move(headerClientChannel)),
+        [evb](auto* p) { evb->runImmediatelyOrRunInEventBaseThreadAndWait([p] { delete p; }); });
     clientMap_->emplace(std::make_pair(host, evb), client);
     return client;
 }