From e0833fcedfdb7233047615650900ef2f285a80eb Mon Sep 17 00:00:00 2001
From: Maksim Denisov <denisov.maksim@huawei.com>
Date: Mon, 10 Feb 2025 10:56:56 +0100
Subject: Store BazelCapabilitiesClient in BazelNetwork

...and pass it to BazelCasClient
---
 src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp |  6 ++++--
 src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp |  6 +++++-
 src/buildtool/execution_api/remote/bazel/bazel_network.cpp    | 10 +++++++++-
 src/buildtool/execution_api/remote/bazel/bazel_network.hpp    |  2 ++
 4 files changed, 20 insertions(+), 4 deletions(-)

(limited to 'src')

diff --git a/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp b/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp
index 67361036..347fa204 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp
@@ -215,9 +215,11 @@ BazelCasClient::BazelCasClient(
     std::string const& server,
     Port port,
     gsl::not_null<Auth const*> const& auth,
-    gsl::not_null<RetryConfig const*> const& retry_config) noexcept
+    gsl::not_null<RetryConfig const*> const& retry_config,
+    gsl::not_null<BazelCapabilitiesClient const*> const& capabilities) noexcept
     : stream_{std::make_unique<ByteStreamClient>(server, port, auth)},
-      retry_config_{*retry_config} {
+      retry_config_{*retry_config},
+      capabilities_{*capabilities} {
     stub_ = bazel_re::ContentAddressableStorage::NewStub(
         CreateChannelWithCredentials(server, port, auth));
 }
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp b/src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp
index 6e2f928b..0e719441 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp
@@ -36,6 +36,7 @@
 #include "src/buildtool/common/remote/retry_config.hpp"
 #include "src/buildtool/crypto/hash_function.hpp"
 #include "src/buildtool/execution_api/common/artifact_blob.hpp"
+#include "src/buildtool/execution_api/remote/bazel/bazel_capabilities_client.hpp"
 #include "src/buildtool/execution_api/remote/bazel/bytestream_client.hpp"
 #include "src/buildtool/logging/logger.hpp"
 
@@ -47,7 +48,9 @@ class BazelCasClient {
         std::string const& server,
         Port port,
         gsl::not_null<Auth const*> const& auth,
-        gsl::not_null<RetryConfig const*> const& retry_config) noexcept;
+        gsl::not_null<RetryConfig const*> const& retry_config,
+        gsl::not_null<BazelCapabilitiesClient const*> const&
+            capabilities) noexcept;
 
     /// \brief Find missing blobs
     /// \param[in] instance_name Name of the CAS instance
@@ -143,6 +146,7 @@ class BazelCasClient {
   private:
     std::unique_ptr<ByteStreamClient> stream_;
     RetryConfig const& retry_config_;
+    [[maybe_unused]] BazelCapabilitiesClient const& capabilities_;
     std::unique_ptr<bazel_re::ContentAddressableStorage::Stub> stub_;
     Logger logger_{"RemoteCasClient"};
 
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network.cpp b/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
index 823148f0..f7aa3dcc 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
@@ -31,7 +31,15 @@ BazelNetwork::BazelNetwork(
     ExecutionConfiguration const& exec_config,
     gsl::not_null<HashFunction const*> const& hash_function) noexcept
     : instance_name_{std::move(instance_name)},
-      cas_{std::make_unique<BazelCasClient>(host, port, auth, retry_config)},
+      capabilities_{std::make_unique<BazelCapabilitiesClient>(host,
+                                                              port,
+                                                              auth,
+                                                              retry_config)},
+      cas_{std::make_unique<BazelCasClient>(host,
+                                            port,
+                                            auth,
+                                            retry_config,
+                                            capabilities_.get())},
       ac_{std::make_unique<BazelAcClient>(host, port, auth, retry_config)},
       exec_{std::make_unique<BazelExecutionClient>(host,
                                                    port,
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network.hpp b/src/buildtool/execution_api/remote/bazel/bazel_network.hpp
index 3b359700..1c65cbd7 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_network.hpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_network.hpp
@@ -31,6 +31,7 @@
 #include "src/buildtool/execution_api/bazel_msg/bazel_common.hpp"
 #include "src/buildtool/execution_api/common/artifact_blob.hpp"
 #include "src/buildtool/execution_api/remote/bazel/bazel_ac_client.hpp"
+#include "src/buildtool/execution_api/remote/bazel/bazel_capabilities_client.hpp"
 #include "src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp"
 #include "src/buildtool/execution_api/remote/bazel/bazel_execution_client.hpp"
 #include "src/buildtool/execution_api/remote/bazel/bazel_network_reader.hpp"
@@ -93,6 +94,7 @@ class BazelNetwork {
 
   private:
     std::string const instance_name_;
+    std::unique_ptr<BazelCapabilitiesClient> capabilities_;
     std::unique_ptr<BazelCasClient> cas_;
     std::unique_ptr<BazelAcClient> ac_;
     std::unique_ptr<BazelExecutionClient> exec_;
-- 
cgit v1.2.3