Intermediate changes

5 files changed, 213 insertions, 0 deletions

diff --git a/yt/yt/client/federated/cache.cpp b/yt/yt/client/federated/cache.cpp
new file mode 100644
index 00000000000..ab64ca070ac
--- /dev/null
+++ b/yt/yt/client/federated/cache.cpp
@@ -0,0 +1,119 @@
+#include "cache.h"
+#include "client.h"
+
+#include <yt/yt/client/cache/cache_base.h>
+#include <yt/yt/client/cache/options.h>
+#include <yt/yt/client/cache/rpc.h>
+
+#include <util/string/split.h>
+
+namespace NYT::NClient::NFederated {
+
+namespace {
+
+////////////////////////////////////////////////////////////////////////////////
+
+class TClientsCache
+    : public NCache::TClientsCacheBase
+{
+public:
+    TClientsCache(
+        TClustersConfig clustersConfig,
+        NApi::TClientOptions options,
+        TFederationConfigPtr federationConfig,
+        TString clusterSeparator)
+        : ClustersConfig_(std::move(clustersConfig))
+        , Options_(std::move(options))
+        , FederationConfig_(std::move(federationConfig))
+        , ClusterSeparator_(std::move(clusterSeparator))
+        {}
+
+protected:
+    NApi::IClientPtr CreateClient(TStringBuf clusterUrl) override
+    {
+        std::vector<TString> clusters;
+        NYT::NApi::IClientPtr client;
+        StringSplitter(clusterUrl).SplitByString(ClusterSeparator_).SkipEmpty().Collect(&clusters);
+        if (clusters.size() == 1) {
+            return NCache::CreateClient(NCache::MakeClusterConfig(ClustersConfig_, clusterUrl), Options_);
+        } else {
+            std::vector<NYT::NApi::IClientPtr> clients;
+            clients.reserve(clusters.size());
+            for (auto& cluster : clusters) {
+                clients.push_back(GetClient(cluster));
+            }
+            return NFederated::CreateClient(std::move(clients), FederationConfig_);
+        }
+    }
+
+private:
+    const TClustersConfig ClustersConfig_;
+    const NApi::TClientOptions Options_;
+    const TFederationConfigPtr FederationConfig_;
+    const TString ClusterSeparator_;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace
+
+IClientsCachePtr CreateFederatedClientsCache(
+    const TClustersConfig& config,
+    const NYT::NApi::TClientOptions& options,
+    TFederationConfigPtr federatedConfig,
+    TString clusterSeparator)
+{
+    return NYT::New<TClientsCache>(
+        std::move(config),
+        std::move(options),
+        std::move(federatedConfig),
+        std::move(clusterSeparator));
+}
+
+IClientsCachePtr CreateFederatedClientsCache(
+    const TClustersConfig& config,
+    const NYT::NApi::TClientOptions& options,
+    TString chaosBundleName,
+    TString clusterSeparator)
+{
+    auto federationConfig = NYT::New<NYT::NClient::NFederated::TFederationConfig>();
+    if (!chaosBundleName.empty()) {
+        federationConfig->BundleName = std::move(chaosBundleName);
+    }
+    return CreateFederatedClientsCache(
+        std::move(config),
+        std::move(options),
+        std::move(federationConfig),
+        std::move(clusterSeparator));
+}
+
+IClientsCachePtr CreateFederatedClientsCache(
+    const TConfig& config,
+    const NYT::NApi::TClientOptions& options,
+    TString chaosBundleName,
+    TString clusterSeparator)
+{
+    TClustersConfig clustersConfig;
+    *clustersConfig.MutableDefaultConfig() = config;
+
+    return CreateFederatedClientsCache(
+        std::move(clustersConfig),
+        std::move(options),
+        std::move(chaosBundleName),
+        std::move(clusterSeparator));
+}
+
+IClientsCachePtr CreateFederatedClientsCache(
+    TString chaosBundleName,
+    TString clusterSeparator)
+{
+    return CreateFederatedClientsCache(
+        TClustersConfig{},
+        NCache::GetClientOpsFromEnvStatic(),
+        std::move(chaosBundleName),
+        std::move(clusterSeparator));
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT::NClient::NFederated
diff --git a/yt/yt/client/federated/cache.h b/yt/yt/client/federated/cache.h
new file mode 100644
index 00000000000..ebdb6c4b096
--- /dev/null
+++ b/yt/yt/client/federated/cache.h
@@ -0,0 +1,58 @@
+#pragma once
+
+#include "config.h"
+
+#include <yt/yt/client/cache/cache.h>
+
+
+namespace NYT::NClient::NFederated {
+
+using NCache::IClientsCachePtr;
+using NCache::TClustersConfig;
+using NCache::TConfig;
+
+////////////////////////////////////////////////////////////////////////////////
+
+//! Creates clients cache which explicitly given federation and clusters config.
+//! Server name is always overwritten with requested.
+//! Creates FederatedClient when clusters are enumerated with separator(default is '+').
+//! For example: seneca-sas+seneca-vla.
+IClientsCachePtr CreateFederatedClientsCache(
+    TFederationConfigPtr federationConfig,
+    const TClustersConfig& config,
+    const NYT::NApi::TClientOptions& options,
+    TString clusterSeparator = "+");
+
+//! Shortcut to create cache with default federation config.
+IClientsCachePtr CreateFederatedClientsCache(
+    const TClustersConfig& config,
+    const NYT::NApi::TClientOptions& options,
+    TString chaosBundleName,
+    TString clusterSeparator = "+");
+
+//! Creates clients cache which shares same config (except server name).
+//! Shortcut to create cache with default federation config.
+IClientsCachePtr CreateFederatedClientsCache(
+    const TConfig& config,
+    const NYT::NApi::TClientOptions& options,
+    TString chaosBundleName,
+    TString clusterSeparator = "+");
+
+//! Creates clients cache which shares same config (except server name).
+//! Shortcut to create cache with default federation config.
+//! Shortcut to use client options from env.
+IClientsCachePtr CreateFederatedClientsCache(
+    const TConfig& config,
+    TString chaosBundleName,
+    TString clusterSeparator = "+");
+
+//! Shortcut to create cache with default federation config.
+//! Shortcut to use client options from env.
+//! Shortcut to create cache with default clusters config.
+IClientsCachePtr CreateFederatedClientsCache(
+    TString chaosBundleName,
+    TString clusterSeparator = "+");
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT::NClient::NFederated
diff --git a/yt/yt/client/federated/unittests/cache_ut.cpp b/yt/yt/client/federated/unittests/cache_ut.cpp
new file mode 100644
index 00000000000..10be390ecb3
--- /dev/null
+++ b/yt/yt/client/federated/unittests/cache_ut.cpp
@@ -0,0 +1,33 @@
+#include <yt/yt/client/federated/cache.h>
+#include <yt/yt/client/federated/config.h>
+
+#include <yt/yt/client/cache/cache.h>
+
+#include <library/cpp/testing/gtest/gtest.h>
+
+#include <util/system/env.h>
+
+namespace NYT::NClient::NFederated {
+
+using namespace NYT::NApi;
+
+////////////////////////////////////////////////////////////////////////////////
+
+TEST(TFederatedClientsCacheTest, GetSameClient) {
+    SetEnv("YT_TOKEN", "AAAA-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA");
+    auto cache = CreateFederatedClientsCache("my_bundle");
+    auto client1 = cache->GetClient("localhost");
+    auto client2 = cache->GetClient("localhost");
+
+    EXPECT_TRUE(client1 == client2);
+
+    // This is needed for TConnection.OnProxyUpdate to stop
+    // and to remove references to TConnection that it's holding.
+    // It's because we don't actually create YT Server.
+    client1->GetConnection()->Terminate();
+    client2->GetConnection()->Terminate();
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT::NClient::NFederated
diff --git a/yt/yt/client/federated/unittests/ya.make b/yt/yt/client/federated/unittests/ya.make
index 5cb97978d07..1400f640f1e 100644
--- a/yt/yt/client/federated/unittests/ya.make
+++ b/yt/yt/client/federated/unittests/ya.make
@@ -3,6 +3,7 @@ GTEST(unittester-federated-client)
 SRCS(
     client_ut.cpp
     connection_ut.cpp
+    cache_ut.cpp
 )
 
 PEERDIR(
diff --git a/yt/yt/client/federated/ya.make b/yt/yt/client/federated/ya.make
index efea22071ca..19ba6e0f671 100644
--- a/yt/yt/client/federated/ya.make
+++ b/yt/yt/client/federated/ya.make
@@ -4,12 +4,14 @@ SRCS(
     client.cpp
     config.cpp
     connection.cpp
+    cache.cpp
 )
 
 PEERDIR(
     library/cpp/yt/string
     yt/yt/core
     yt/yt/client
+    yt/yt/client/cache
 )
 
 END()