YT-26639: reuse connection object for the same endpoint

commit_hash:aa962f5a7cdef69708c3502aaead051042eec0e5

1 files changed, 98 insertions, 13 deletions

diff --git a/yt/cpp/mapreduce/rpc_client/client_impl.cpp b/yt/cpp/mapreduce/rpc_client/client_impl.cpp
index de01ff0569e..84fb5282173 100644
--- a/yt/cpp/mapreduce/rpc_client/client_impl.cpp
+++ b/yt/cpp/mapreduce/rpc_client/client_impl.cpp
@@ -6,11 +6,65 @@
 #include <yt/cpp/mapreduce/common/retry_lib.h>
 
 #include <yt/cpp/mapreduce/interface/client_method_options.h>
+#include <yt/cpp/mapreduce/interface/logging/yt_log.h>
 #include <yt/cpp/mapreduce/interface/tvm.h>
 
 #include <yt/yt/client/api/rpc_proxy/config.h>
 #include <yt/yt/client/api/rpc_proxy/connection.h>
 
+namespace NYT::NDetail {
+
+////////////////////////////////////////////////////////////////////////////////
+
+struct TConnectionCacheKey
+{
+public:
+    TMaybe<TString> JobProxySocketPath;
+    TString ClusterUrl;
+    TMaybe<TString> RpcProxyRole;
+    TMaybe<TString> ProxyAddress;
+
+    // N.B. we want to compute hash of this struct, so we use sorted map.
+    TMap<TString, TString> ProxyUrlAliasingRules;
+
+public:
+    TConnectionCacheKey() = default;
+
+    TConnectionCacheKey(const TClientContext& context)
+        : JobProxySocketPath(context.JobProxySocketPath)
+        , ClusterUrl(context.ServerName)
+        , RpcProxyRole(context.RpcProxyRole)
+        , ProxyAddress(context.ProxyAddress)
+        , ProxyUrlAliasingRules(context.Config->ProxyUrlAliasingRules.begin(), context.Config->ProxyUrlAliasingRules.end())
+    { }
+
+    bool operator==(const TConnectionCacheKey& other) const = default;
+};
+
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT::NDetail
+
+template <>
+struct THash<NYT::NDetail::TConnectionCacheKey>
+{
+    size_t operator() (const NYT::NDetail::TConnectionCacheKey& key) const
+    {
+        using NYT::HashCombine;
+        size_t result = 0;
+        HashCombine(result, key.JobProxySocketPath);
+        HashCombine(result, key.ClusterUrl);
+        HashCombine(result, key.RpcProxyRole);
+        HashCombine(result, key.ProxyAddress);
+        for (const auto& [k, v] : key.ProxyUrlAliasingRules) {
+            HashCombine(result, k);
+            HashCombine(result, v);
+        }
+        return result;
+    }
+};
+
 namespace NYT {
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -19,29 +73,61 @@ namespace NDetail {
 
 ////////////////////////////////////////////////////////////////////////////////
 
-NYT::NApi::IClientPtr CreateApiClient(const TClientContext& context)
+NApi::IConnectionPtr GetOrCreateConnection(const TConnectionCacheKey& key)
 {
+    static TMutex lock;
+    static THashMap<TConnectionCacheKey, TWeakPtr<NApi::IConnection>> cache;
+    static size_t maxCacheSize = 32;
+
+    auto g = Guard(lock);
+
+    auto it = cache.find(key);
+    if (it != cache.end()) {
+        auto connection = it->second.Lock();
+        if (connection) {
+            return connection;
+        }
+    }
+
     auto connectionConfig = New<NApi::NRpcProxy::TConnectionConfig>();
     connectionConfig->SetDefaults();
-    if (context.JobProxySocketPath) {
-        connectionConfig->ProxyUnixDomainSocket = *context.JobProxySocketPath;
+    if (key.JobProxySocketPath) {
+        connectionConfig->ProxyUnixDomainSocket = *key.JobProxySocketPath;
         connectionConfig->EnableProxyDiscovery = false;
     } else {
-        connectionConfig->ClusterUrl = context.ServerName;
+        connectionConfig->ClusterUrl = key.ClusterUrl;
     }
-    if (context.RpcProxyRole) {
-        connectionConfig->ProxyRole = *context.RpcProxyRole;
+    if (key.RpcProxyRole) {
+        connectionConfig->ProxyRole = *key.RpcProxyRole;
     }
-    if (context.ProxyAddress) {
-        connectionConfig->ProxyAddresses = {*context.ProxyAddress};
+    if (key.ProxyAddress) {
+        connectionConfig->ProxyAddresses = {*key.ProxyAddress};
+    }
+
+    for (const auto& [clusterName, url] : key.ProxyUrlAliasingRules) {
+        connectionConfig->ProxyUrlAliasingRules.emplace(clusterName, url);
     }
 
-    THashMap<std::string, std::string> proxyUrlAliasingRules;
-    for (const auto& [clusterName, url] : context.Config->ProxyUrlAliasingRules) {
-        proxyUrlAliasingRules.emplace(clusterName, url);
+    auto connection = NApi::NRpcProxy::CreateConnection(connectionConfig);
+
+    if (it != cache.end()) {
+        it->second = connection;
+    } else if (cache.size() < maxCacheSize) {
+        // N.B. once cache size exceeds limit we disable caching.
+        // We believe such cases are rare but if it is the case we don't want to leak memory
+        // (memory is leaked since we never clear the cache).
+        cache.emplace(key, connection);
+    } else {
+        YT_LOG_WARNING("Cannot cache IConnection since connection cache reached maximum size");
     }
 
-    connectionConfig->ProxyUrlAliasingRules = std::move(proxyUrlAliasingRules);
+    return connection;
+}
+
+NYT::NApi::IClientPtr CreateApiClient(const TClientContext& context)
+{
+    auto key = TConnectionCacheKey(context);
+    auto connection = GetOrCreateConnection(key);
 
     NApi::TClientOptions clientOptions;
     clientOptions.Token = context.Token;
@@ -55,7 +141,6 @@ NYT::NApi::IClientPtr CreateApiClient(const TClientContext& context)
         clientOptions.MultiproxyTargetCluster = context.MultiproxyTargetCluster;
     }
 
-    auto connection = NApi::NRpcProxy::CreateConnection(connectionConfig);
     return connection->CreateClient(clientOptions);
 }