Allow passing ProxyUrlAliasingRules to clients directly

commit_hash:5493d9ed32b24050fa5feb08fe13097302ca0a23

7 files changed, 62 insertions, 21 deletions

diff --git a/yt/cpp/mapreduce/client/client.cpp b/yt/cpp/mapreduce/client/client.cpp
index bd5ea379b5c..2d52bbabbca 100644
--- a/yt/cpp/mapreduce/client/client.cpp
+++ b/yt/cpp/mapreduce/client/client.cpp
@@ -65,18 +65,13 @@ namespace {
 
 ////////////////////////////////////////////////////////////////////////////////
 
-THashMap<TString, TString> ParseProxyUrlAliasingRules(TString envConfig)
+void ApplyProxyUrlAliasingRules(
+    TString& url,
+    const THashMap<TString, TString>& proxyUrlAliasingRules)
 {
-    if (envConfig.empty()) {
-        return {};
-    }
-    return NYTree::ConvertTo<THashMap<TString, TString>>(NYson::TYsonString(envConfig));
-}
-
-void ApplyProxyUrlAliasingRules(TString& url)
-{
-    static auto rules = ParseProxyUrlAliasingRules(GetEnv("YT_PROXY_URL_ALIASING_CONFIG"));
-    if (auto ruleIt = rules.find(url); ruleIt != rules.end()) {
+    if (auto ruleIt = proxyUrlAliasingRules.find(url);
+        ruleIt != proxyUrlAliasingRules.end()
+    ) {
         url = ruleIt->second;
     }
 }
@@ -1566,7 +1561,7 @@ void SetupClusterContext(
 {
     context.ServerName = serverName;
     context.MultiproxyTargetCluster = serverName;
-    ApplyProxyUrlAliasingRules(context.ServerName);
+    ApplyProxyUrlAliasingRules(context.ServerName, context.Config->ProxyUrlAliasingRules);
 
     if (context.ServerName.find('.') == TString::npos &&
         context.ServerName.find(':') == TString::npos &&
diff --git a/yt/cpp/mapreduce/interface/config.cpp b/yt/cpp/mapreduce/interface/config.cpp
index cdabeaf3240..901f59fc215 100644
--- a/yt/cpp/mapreduce/interface/config.cpp
+++ b/yt/cpp/mapreduce/interface/config.cpp
@@ -190,6 +190,30 @@ void TConfig::LoadTimings()
     HostListUpdateInterval = TDuration::Seconds(60);
 }
 
+void TConfig::LoadProxyUrlAliasingRules()
+{
+    TString strConfig = GetEnv("YT_PROXY_URL_ALIASING_CONFIG");
+    if (!strConfig) {
+        return;
+    }
+
+    NYT::TNode nodeConfig;
+
+    try {
+        nodeConfig = NodeFromYsonString(strConfig);
+        Y_ENSURE(nodeConfig.IsMap());
+    } catch (const yexception& exc) {
+        ythrow yexception()
+            << "Failed to parse YT_PROXY_URL_ALIASING_CONFIG (it must be yson map): "
+            << exc;
+    }
+
+    for (const auto& [key, value] : nodeConfig.AsMap()) {
+        Y_ENSURE(value.IsString(), "Proxy url is not string");
+        ProxyUrlAliasingRules.emplace(key, value.AsString());
+    }
+}
+
 void TConfig::Reset()
 {
     Hosts = GetEnv("YT_HOSTS", DefaultHosts);
@@ -219,6 +243,7 @@ void TConfig::Reset()
     LoadToken();
     LoadSpec();
     LoadTimings();
+    LoadProxyUrlAliasingRules();
 
     CacheUploadDeduplicationMode = GetUploadingDeduplicationMode("YT_UPLOAD_DEDUPLICATION", EUploadDeduplicationMode::Host);
     CacheUploadDeduplicationThreshold = 10_MB;
diff --git a/yt/cpp/mapreduce/interface/config.h b/yt/cpp/mapreduce/interface/config.h
index 6e6d7b191c2..7aba0ccd60f 100644
--- a/yt/cpp/mapreduce/interface/config.h
+++ b/yt/cpp/mapreduce/interface/config.h
@@ -9,6 +9,7 @@
 
 #include <util/generic/maybe.h>
 #include <util/generic/string.h>
+#include <util/generic/hash.h>
 #include <util/generic/hash_set.h>
 
 #include <util/datetime/base.h>
@@ -111,6 +112,12 @@ struct TConfig
     /// @brief Represents the role involved in RPC proxy configuration.
     TString RpcProxyRole;
 
+    /// @brief Proxy url aliasing rules to be used for connection.
+    ///
+    /// You can pass here "foo" => "fqdn:port" and afterwards use "foo" as handy alias,
+    /// while all connections will be made to "fqdn:port" address.
+    THashMap<TString, TString> ProxyUrlAliasingRules;
+
     ///
     /// For historical reasons mapreduce client uses its own logging system.
     ///
@@ -264,6 +271,7 @@ struct TConfig
     void LoadToken();
     void LoadSpec();
     void LoadTimings();
+    void LoadProxyUrlAliasingRules();
 
     void Reset();
 
diff --git a/yt/yt/client/api/rpc_proxy/config.cpp b/yt/yt/client/api/rpc_proxy/config.cpp
index b3aef0ea36f..5d39bec379b 100644
--- a/yt/yt/client/api/rpc_proxy/config.cpp
+++ b/yt/yt/client/api/rpc_proxy/config.cpp
@@ -50,6 +50,8 @@ void TConnectionConfig::Register(TRegistrar registrar)
         .Optional();
     registrar.Parameter("enable_proxy_discovery", &TThis::EnableProxyDiscovery)
         .Default(true);
+    registrar.Parameter("proxy_url_aliasing_rules", &TThis::ProxyUrlAliasingRules)
+        .Default();
 
     registrar.Parameter("dynamic_channel_pool", &TThis::DynamicChannelPool)
         .DefaultNew();
diff --git a/yt/yt/client/api/rpc_proxy/config.h b/yt/yt/client/api/rpc_proxy/config.h
index 990708b7466..bbac7614512 100644
--- a/yt/yt/client/api/rpc_proxy/config.h
+++ b/yt/yt/client/api/rpc_proxy/config.h
@@ -15,6 +15,8 @@
 #include <yt/yt/client/api/client.h>
 #include <yt/yt/client/api/config.h>
 
+#include <util/generic/hash.h>
+
 namespace NYT::NApi::NRpcProxy {
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -35,6 +37,7 @@ struct TConnectionConfig
     NRpc::TServiceDiscoveryEndpointsConfigPtr ProxyEndpoints;
     std::optional<std::string> ProxyUnixDomainSocket;
     bool EnableProxyDiscovery;
+    THashMap<std::string, std::string> ProxyUrlAliasingRules;
 
     NRpc::TDynamicChannelPoolConfigPtr DynamicChannelPool;
 
diff --git a/yt/yt/client/api/rpc_proxy/connection_impl.cpp b/yt/yt/client/api/rpc_proxy/connection_impl.cpp
index a6275f5d9a5..2403d2e5d14 100644
--- a/yt/yt/client/api/rpc_proxy/connection_impl.cpp
+++ b/yt/yt/client/api/rpc_proxy/connection_impl.cpp
@@ -62,18 +62,24 @@ THashMap<std::string, std::string> ParseProxyUrlAliasingRules(const TString& env
     return ConvertTo<THashMap<std::string, std::string>>(TYsonString(envConfig));
 }
 
-void ApplyProxyUrlAliasingRules(std::string& url, const std::optional<THashMap<std::string, std::string>>& proxyUrlAliasingRules)
+void ApplyProxyUrlAliasingRules(std::string& url, const THashMap<std::string, std::string>& proxyUrlAliasingRules)
 {
     static const auto rulesFromEnv = ParseProxyUrlAliasingRules(GetEnv("YT_PROXY_URL_ALIASING_CONFIG"));
 
-    const auto& rules = proxyUrlAliasingRules.value_or(rulesFromEnv);
+    auto lookup = [&](const auto& rules) {
+        if (auto ruleIt = rules.find(url); ruleIt != rules.end()) {
+            url = ruleIt->second;
+        }
+    };
 
-    if (auto ruleIt = rules.find(url); ruleIt != rules.end()) {
-        url = ruleIt->second;
+    if (proxyUrlAliasingRules) {
+        lookup(proxyUrlAliasingRules);
+    } else if (rulesFromEnv) {
+        lookup(rulesFromEnv);
     }
 }
 
-std::string NormalizeHttpProxyUrl(std::string url, const std::optional<THashMap<std::string, std::string>>& proxyUrlAliasingRules)
+std::string NormalizeHttpProxyUrl(std::string url, const THashMap<std::string, std::string>& proxyUrlAliasingRules)
 {
     ApplyProxyUrlAliasingRules(url, proxyUrlAliasingRules);
 
@@ -395,8 +401,10 @@ std::vector<std::string> TConnection::DiscoverProxiesViaHttp()
                     .OptionalItem("network_name", Config_->ProxyNetworkName)
                 .EndMap().ToString());
 
-        auto url = NormalizeHttpProxyUrl(*Config_->ClusterUrl) + "/api/v4/discover_proxies";
-        auto client = IsProxyUrlSecure(*Config_->ClusterUrl)
+        auto url = NormalizeHttpProxyUrl(
+            *Config_->ClusterUrl,
+            Config_->ProxyUrlAliasingRules) + "/api/v4/discover_proxies";
+        auto client = IsProxyUrlSecure(url)
             ? NHttps::CreateClient(Config_->HttpsClient, std::move(poller))
             : NHttp::CreateClient(Config_->HttpClient, std::move(poller));
         // TODO(babenko): switch to std::string
diff --git a/yt/yt/client/api/rpc_proxy/private.h b/yt/yt/client/api/rpc_proxy/private.h
index 67de1498f9c..ff852633363 100644
--- a/yt/yt/client/api/rpc_proxy/private.h
+++ b/yt/yt/client/api/rpc_proxy/private.h
@@ -29,10 +29,10 @@ YT_DEFINE_GLOBAL(const NLogging::TLogger, RpcProxyClientLogger, "RpcProxyClient"
 THashMap<std::string, std::string> ParseProxyUrlAliasingRules(const TString& envConfig);
 void ApplyProxyUrlAliasingRules(
     std::string& url,
-    const std::optional<THashMap<std::string, std::string>>& proxyUrlAliasingRules = {});
+    const THashMap<std::string, std::string>& proxyUrlAliasingRules = {});
 std::string NormalizeHttpProxyUrl(
     std::string url,
-    const std::optional<THashMap<std::string, std::string>>& proxyUrlAliasingRules = {});
+    const THashMap<std::string, std::string>& proxyUrlAliasingRules = {});
 
 ////////////////////////////////////////////////////////////////////////////////