YT: Use ipv6_address from service discovery response by default

commit_hash:52b5c4b6f88961d2c310f0bd6164d7f248258782

7 files changed, 38 insertions, 4 deletions

diff --git a/yt/yt/core/net/address.cpp b/yt/yt/core/net/address.cpp
index c08fb94226..6f7a44100f 100644
--- a/yt/yt/core/net/address.cpp
+++ b/yt/yt/core/net/address.cpp
@@ -110,6 +110,13 @@ TStringBuf GetServiceHostName(TStringBuf address)
 
 ////////////////////////////////////////////////////////////////////////////////
 
+TString FormatNetworkAddress(TStringBuf address, int port)
+{
+    return Format("[%v]:%v", address, port);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
 const TNetworkAddress NullNetworkAddress;
 
 TNetworkAddress::TNetworkAddress()
diff --git a/yt/yt/core/net/address.h b/yt/yt/core/net/address.h
index 6f507b2900..99e7869f9e 100644
--- a/yt/yt/core/net/address.h
+++ b/yt/yt/core/net/address.h
@@ -43,6 +43,9 @@ TStringBuf GetServiceHostName(TStringBuf address);
 
 ////////////////////////////////////////////////////////////////////////////////
 
+//! Constructs an address of the form |[address]:port|.
+TString FormatNetworkAddress(TStringBuf address, int port);
+
 class TIP6Address;
 
 //! An opaque wrapper for |sockaddr| type.
diff --git a/yt/yt/core/rpc/balancing_channel.cpp b/yt/yt/core/rpc/balancing_channel.cpp
index 7884bba0d2..7e9cd188d0 100644
--- a/yt/yt/core/rpc/balancing_channel.cpp
+++ b/yt/yt/core/rpc/balancing_channel.cpp
@@ -151,7 +151,10 @@ private:
                 continue;
             }
 
-            auto addresses = AddressesFromEndpointSet(endpointSetOrError.Value());
+            auto addresses = AddressesFromEndpointSet(
+                endpointSetOrError.Value(),
+                Config_->Endpoints->UseIPv4,
+                Config_->Endpoints->UseIPv6);
             allAddresses.insert(allAddresses.end(), addresses.begin(), addresses.end());
         }
 
diff --git a/yt/yt/core/rpc/config.cpp b/yt/yt/core/rpc/config.cpp
index 50030c0c3b..b9cb445371 100644
--- a/yt/yt/core/rpc/config.cpp
+++ b/yt/yt/core/rpc/config.cpp
@@ -220,6 +220,10 @@ void TServiceDiscoveryEndpointsConfig::Register(TRegistrar registrar)
     registrar.Parameter("endpoint_set_id", &TThis::EndpointSetId);
     registrar.Parameter("update_period", &TThis::UpdatePeriod)
         .Default(TDuration::Seconds(60));
+    registrar.Parameter("use_ipv4", &TThis::UseIPv4)
+        .Default(false);
+    registrar.Parameter("use_ipv6", &TThis::UseIPv6)
+        .Default(false);
 
     registrar.Postprocessor([] (TThis* config) {
         if (config->Cluster.has_value() == !config->Clusters.empty()) {
diff --git a/yt/yt/core/rpc/config.h b/yt/yt/core/rpc/config.h
index aceeadc725..203d62b8bc 100644
--- a/yt/yt/core/rpc/config.h
+++ b/yt/yt/core/rpc/config.h
@@ -302,6 +302,11 @@ public:
     TString EndpointSetId;
     TDuration UpdatePeriod;
 
+    //! Use IPv4 address of endpoint.
+    bool UseIPv4;
+    //! Use IPv6 address of endpoint.
+    bool UseIPv6;
+
     REGISTER_YSON_STRUCT(TServiceDiscoveryEndpointsConfig);
 
     static void Register(TRegistrar registrar);
diff --git a/yt/yt/core/rpc/helpers.cpp b/yt/yt/core/rpc/helpers.cpp
index d9d2a93a89..ca7b6221e5 100644
--- a/yt/yt/core/rpc/helpers.cpp
+++ b/yt/yt/core/rpc/helpers.cpp
@@ -514,12 +514,21 @@ void SetCurrentAuthenticationIdentity(const IClientRequestPtr& request)
     SetAuthenticationIdentity(request, GetCurrentAuthenticationIdentity());
 }
 
-std::vector<std::string> AddressesFromEndpointSet(const NServiceDiscovery::TEndpointSet& endpointSet)
+std::vector<std::string> AddressesFromEndpointSet(
+    const NServiceDiscovery::TEndpointSet& endpointSet,
+    bool useIPv4,
+    bool useIPv6)
 {
     std::vector<std::string> addresses;
     addresses.reserve(endpointSet.Endpoints.size());
     for (const auto& endpoint : endpointSet.Endpoints) {
-        addresses.push_back(NNet::BuildServiceAddress(endpoint.Fqdn, endpoint.Port));
+        if (useIPv6 && !endpoint.IP6Address.empty()) {
+            addresses.push_back(NNet::FormatNetworkAddress(endpoint.IP6Address, endpoint.Port));
+        } else if (useIPv4 && !endpoint.IP4Address.empty()) {
+            addresses.push_back(NNet::FormatNetworkAddress(endpoint.IP4Address, endpoint.Port));
+        } else {
+            addresses.push_back(NNet::BuildServiceAddress(endpoint.Fqdn, endpoint.Port));
+        }
     }
     return addresses;
 }
diff --git a/yt/yt/core/rpc/helpers.h b/yt/yt/core/rpc/helpers.h
index c631684555..0624017613 100644
--- a/yt/yt/core/rpc/helpers.h
+++ b/yt/yt/core/rpc/helpers.h
@@ -93,7 +93,10 @@ void WriteAuthenticationIdentityToProto(T* proto, const TAuthenticationIdentity&
 template <class T>
 TAuthenticationIdentity ParseAuthenticationIdentityFromProto(const T& proto);
 
-std::vector<std::string> AddressesFromEndpointSet(const NServiceDiscovery::TEndpointSet& endpointSet);
+std::vector<std::string> AddressesFromEndpointSet(
+    const NServiceDiscovery::TEndpointSet& endpointSet,
+    bool useIPv4 = false,
+    bool useIPv6 = false);
 
 ////////////////////////////////////////////////////////////////////////////////