Use BaselCasClient in BazelNetworkReader

...bypassing BazelNetwork layer.

5 files changed, 101 insertions, 79 deletions

diff --git a/src/buildtool/execution_api/remote/bazel/bazel_api.cpp b/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
index 2cba30be..f539f325 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
@@ -249,7 +249,7 @@ auto BazelApi::CreateAction(
                                                ? std::make_optional(info.digest)
                                                : std::nullopt;
                 auto reader = TreeReader<BazelNetworkReader>{
-                    *network_, std::move(request_remote_tree)};
+                    network_->CreateReader(), std::move(request_remote_tree)};
                 auto const result = reader.RecursivelyReadTreeLeafs(
                     info.digest, output_paths[i]);
                 if (not result or
@@ -304,7 +304,7 @@ auto BazelApi::CreateAction(
     std::vector<Artifact::ObjectInfo> const& artifacts_info,
     std::vector<int> const& fds,
     bool raw_tree) noexcept -> bool {
-    auto dumper = StreamDumper<BazelNetworkReader>{*network_};
+    auto dumper = StreamDumper<BazelNetworkReader>{network_->CreateReader()};
     return CommonRetrieveToFds(
         artifacts_info,
         fds,
@@ -343,7 +343,8 @@ auto BazelApi::CreateAction(
     for (auto const& dgst : missing_artifacts_info->digests) {
         auto const& info = missing_artifacts_info->back_map[dgst];
         if (IsTreeObject(info.type)) {
-            auto reader = TreeReader<BazelNetworkReader>{*network_};
+            auto reader =
+                TreeReader<BazelNetworkReader>{network_->CreateReader()};
             auto const result = reader.ReadDirectTreeEntries(
                 info.digest, std::filesystem::path{});
             if (not result or not RetrieveToCas(result->infos, api)) {
@@ -392,7 +393,8 @@ auto BazelApi::CreateAction(
         for (auto const& dgst : missing_artifacts_info->digests) {
             auto const& info = missing_artifacts_info->back_map[dgst];
             if (IsTreeObject(info.type)) {
-                auto reader = TreeReader<BazelNetworkReader>{*network_};
+                auto reader =
+                    TreeReader<BazelNetworkReader>{network_->CreateReader()};
                 auto const result = reader.ReadDirectTreeEntries(
                     info.digest, std::filesystem::path{});
                 if (not result or
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network.cpp b/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
index 0b543d9b..ba30b564 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
@@ -17,7 +17,6 @@
 #include <algorithm>
 #include <cstddef>
 
-#include "src/buildtool/compatibility/compatibility.hpp"
 #include "src/buildtool/execution_api/common/message_limits.hpp"
 #include "src/buildtool/logging/log_level.hpp"
 #include "src/buildtool/logging/logger.hpp"
@@ -181,9 +180,8 @@ auto BazelNetwork::ReadBlobs(std::vector<bazel_re::Digest> ids) const noexcept
     return BlobReader{instance_name_, cas_.get(), std::move(ids)};
 }
 
-auto BazelNetwork::IncrementalReadSingleBlob(bazel_re::Digest const& id)
-    const noexcept -> ByteStreamClient::IncrementalReader {
-    return cas_->IncrementalReadSingleBlob(instance_name_, id);
+auto BazelNetwork::CreateReader() const noexcept -> BazelNetworkReader {
+    return BazelNetworkReader{instance_name_, *cas_};
 }
 
 auto BazelNetwork::GetCachedActionResult(
@@ -193,11 +191,3 @@ auto BazelNetwork::GetCachedActionResult(
     return ac_->GetActionResult(
         instance_name_, action, false, false, output_files);
 }
-
-auto BazelNetwork::QueryFullTree(bazel_re::Digest const& digest) const noexcept
-    -> std::optional<std::vector<bazel_re::Directory>> {
-    if (not Compatibility::IsCompatible()) {
-        return std::nullopt;
-    }
-    return cas_->GetTree(instance_name_, digest, kMaxBatchTransferSize);
-}
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network.hpp b/src/buildtool/execution_api/remote/bazel/bazel_network.hpp
index 1aa5c1f7..422f59dd 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/remote/bazel/bazel_ac_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"
 
 /// \brief Contains all network clients and is responsible for all network IO.
 class BazelNetwork {
@@ -101,19 +102,13 @@ class BazelNetwork {
     [[nodiscard]] auto ReadBlobs(
         std::vector<bazel_re::Digest> ids) const noexcept -> BlobReader;
 
-    [[nodiscard]] auto IncrementalReadSingleBlob(bazel_re::Digest const& id)
-        const noexcept -> ByteStreamClient::IncrementalReader;
+    [[nodiscard]] auto CreateReader() const noexcept -> BazelNetworkReader;
 
     [[nodiscard]] auto GetCachedActionResult(
         bazel_re::Digest const& action,
         std::vector<std::string> const& output_files) const noexcept
         -> std::optional<bazel_re::ActionResult>;
 
-    /// \brief Query full tree from remote CAS. Note that this is currently not
-    // supported by Buildbarn revision c3c06bbe2a.
-    [[nodiscard]] auto QueryFullTree(bazel_re::Digest const& digest)
-        const noexcept -> std::optional<std::vector<bazel_re::Directory>>;
-
   private:
     std::string const instance_name_{};
     ExecutionConfiguration exec_config_{};
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network_reader.cpp b/src/buildtool/execution_api/remote/bazel/bazel_network_reader.cpp
index 9c5b4b4b..f4ffcb87 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_network_reader.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_network_reader.cpp
@@ -18,37 +18,41 @@
 
 #include "src/buildtool/crypto/hash_function.hpp"
 #include "src/buildtool/execution_api/bazel_msg/bazel_msg_factory.hpp"
+#include "src/buildtool/execution_api/common/message_limits.hpp"
 #include "src/buildtool/file_system/file_system_manager.hpp"
 #include "src/buildtool/logging/log_level.hpp"
 #include "src/buildtool/logging/logger.hpp"
 #include "src/utils/cpp/path.hpp"
 
+BazelNetworkReader::BazelNetworkReader(std::string instance_name,
+                                       BazelCasClient const& cas) noexcept
+    : instance_name_{std::move(instance_name)}, cas_(cas) {}
+
 BazelNetworkReader::BazelNetworkReader(
-    BazelNetwork const& network,
+    BazelNetworkReader&& other,
     std::optional<ArtifactDigest> request_remote_tree) noexcept
-    : network_(network) {
+    : instance_name_{other.instance_name_}, cas_(other.cas_) {
     if (Compatibility::IsCompatible() and request_remote_tree) {
-        auto full_tree = network_.QueryFullTree(*request_remote_tree);
-        if (full_tree) {
-            auxiliary_map_ = MakeAuxiliaryMap(std::move(*full_tree));
-        }
+        // Query full tree from remote CAS. Note that this is currently not
+        // supported by Buildbarn revision c3c06bbe2a.
+        auto full_tree = cas_.GetTree(
+            instance_name_, *request_remote_tree, kMaxBatchTransferSize);
+        auxiliary_map_ = MakeAuxiliaryMap(std::move(full_tree));
     }
 }
 
 auto BazelNetworkReader::ReadDirectory(ArtifactDigest const& digest)
     const noexcept -> std::optional<bazel_re::Directory> {
-    try {
-        if (auxiliary_map_ and auxiliary_map_->contains(digest)) {
-            return auxiliary_map_->at(digest);
+    if (auxiliary_map_) {
+        auto it = auxiliary_map_->find(digest);
+        if (it != auxiliary_map_->end()) {
+            return it->second;
         }
-    } catch (...) {
-        // fallthrough
     }
 
-    auto blobs = network_.ReadBlobs({digest}).Next();
-    if (blobs.size() == 1) {
+    if (auto blob = ReadSingleBlob(digest)) {
         return BazelMsgFactory::MessageFromString<bazel_re::Directory>(
-            *blobs.at(0).data);
+            *blob->data);
     }
     Logger::Log(
         LogLevel::Debug, "Directory {} not found in CAS", digest.hash());
@@ -57,52 +61,43 @@ auto BazelNetworkReader::ReadDirectory(ArtifactDigest const& digest)
 
 auto BazelNetworkReader::ReadGitTree(ArtifactDigest const& digest)
     const noexcept -> std::optional<GitRepo::tree_entries_t> {
-    auto blobs = network_.ReadBlobs({digest}).Next();
-    if (blobs.size() == 1) {
-        auto const& content = *blobs.at(0).data;
-        auto check_symlinks = [this](std::vector<bazel_re::Digest> const& ids) {
-            auto size = ids.size();
-            auto reader = network_.ReadBlobs(ids);
-            auto blobs = reader.Next();
-            std::size_t count{};
-            while (not blobs.empty()) {
-                if (count + blobs.size() > size) {
-                    Logger::Log(LogLevel::Debug,
-                                "received more blobs than requested.");
-                    return false;
-                }
-                for (auto const& blob : blobs) {
-                    if (not PathIsNonUpwards(*blob.data)) {
-                        return false;
-                    }
-                }
-                count += blobs.size();
-                blobs = reader.Next();
-            }
-            return true;
-        };
-        return GitRepo::ReadTreeData(
-            content,
-            HashFunction::ComputeTreeHash(content).Bytes(),
-            check_symlinks,
-            /*is_hex_id=*/false);
+    auto read_blob = ReadSingleBlob(digest);
+    if (not read_blob) {
+        Logger::Log(LogLevel::Debug, "Tree {} not found in CAS", digest.hash());
+        return std::nullopt;
     }
-    Logger::Log(LogLevel::Debug, "Tree {} not found in CAS", digest.hash());
-    return std::nullopt;
+    auto check_symlinks = [this](std::vector<bazel_re::Digest> const& ids) {
+        // TODO(denisov) Fix non-incremental read
+        auto blobs = BatchReadBlobs(ids);
+        if (blobs.size() > ids.size()) {
+            Logger::Log(LogLevel::Debug, "received more blobs than requested.");
+            return false;
+        }
+        return std::all_of(
+            blobs.begin(), blobs.end(), [](ArtifactBlob const& blob) {
+                return PathIsNonUpwards(*blob.data);
+            });
+    };
+
+    std::string const& content = *read_blob->data;
+    return GitRepo::ReadTreeData(content,
+                                 HashFunction::ComputeTreeHash(content).Bytes(),
+                                 check_symlinks,
+                                 /*is_hex_id=*/false);
 }
 
 auto BazelNetworkReader::DumpRawTree(Artifact::ObjectInfo const& info,
                                      DumpCallback const& dumper) const noexcept
     -> bool {
-    auto blobs = network_.ReadBlobs({info.digest}).Next();
-    if (blobs.size() != 1) {
+    auto read_blob = ReadSingleBlob(info.digest);
+    if (not read_blob) {
         Logger::Log(
             LogLevel::Debug, "Object {} not found in CAS", info.digest.hash());
         return false;
     }
 
     try {
-        return std::invoke(dumper, *blobs.at(0).data);
+        return std::invoke(dumper, *read_blob->data);
     } catch (...) {
         return false;
     }
@@ -111,7 +106,7 @@ auto BazelNetworkReader::DumpRawTree(Artifact::ObjectInfo const& info,
 auto BazelNetworkReader::DumpBlob(Artifact::ObjectInfo const& info,
                                   DumpCallback const& dumper) const noexcept
     -> bool {
-    auto reader = network_.IncrementalReadSingleBlob(info.digest);
+    auto reader = cas_.IncrementalReadSingleBlob(instance_name_, info.digest);
     auto data = reader.Next();
     while (data and not data->empty()) {
         try {
@@ -142,3 +137,31 @@ auto BazelNetworkReader::MakeAuxiliaryMap(
     }
     return result;
 }
+
+auto BazelNetworkReader::ReadSingleBlob(ArtifactDigest const& digest)
+    const noexcept -> std::optional<ArtifactBlob> {
+    if (auto blob = cas_.ReadSingleBlob(instance_name_, digest)) {
+        return ArtifactBlob{
+            ArtifactDigest{blob->digest}, blob->data, blob->is_exec};
+    }
+    return std::nullopt;
+}
+
+auto BazelNetworkReader::BatchReadBlobs(
+    std::vector<bazel_re::Digest> const& blobs) const noexcept
+    -> std::vector<ArtifactBlob> {
+    std::vector<BazelBlob> result =
+        cas_.BatchReadBlobs(instance_name_, blobs.begin(), blobs.end());
+
+    std::vector<ArtifactBlob> artifacts;
+    artifacts.reserve(result.size());
+    std::transform(result.begin(),
+                   result.end(),
+                   std::back_inserter(artifacts),
+                   [](BazelBlob const& blob) {
+                       return ArtifactBlob{ArtifactDigest{blob.digest},
+                                           blob.data,
+                                           blob.is_exec};
+                   });
+    return artifacts;
+}
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network_reader.hpp b/src/buildtool/execution_api/remote/bazel/bazel_network_reader.hpp
index fb5f5184..fb807a79 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_network_reader.hpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_network_reader.hpp
@@ -20,21 +20,25 @@
 #include <optional>
 #include <string>
 #include <unordered_map>
+#include <vector>
 
 #include "src/buildtool/common/artifact.hpp"
 #include "src/buildtool/common/artifact_digest.hpp"
 #include "src/buildtool/common/bazel_types.hpp"
-#include "src/buildtool/execution_api/remote/bazel/bazel_network.hpp"
+#include "src/buildtool/execution_api/common/artifact_blob_container.hpp"
+#include "src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp"
 #include "src/buildtool/file_system/git_repo.hpp"
 
 class BazelNetworkReader final {
   public:
     using DumpCallback = std::function<bool(std::string const&)>;
 
-    explicit BazelNetworkReader(
-        BazelNetwork const& network,
-        std::optional<ArtifactDigest> request_remote_tree =
-            std::nullopt) noexcept;
+    BazelNetworkReader(std::string instance_name,
+                       BazelCasClient const& cas) noexcept;
+
+    BazelNetworkReader(
+        BazelNetworkReader&& other,
+        std::optional<ArtifactDigest> request_remote_tree) noexcept;
 
     [[nodiscard]] auto ReadDirectory(ArtifactDigest const& digest)
         const noexcept -> std::optional<bazel_re::Directory>;
@@ -50,16 +54,24 @@ class BazelNetworkReader final {
                                 DumpCallback const& dumper) const noexcept
         -> bool;
 
+    [[nodiscard]] auto ReadSingleBlob(ArtifactDigest const& digest)
+        const noexcept -> std::optional<ArtifactBlob>;
+
   private:
     using DirectoryMap =
-        std::unordered_map<bazel_re::Digest, bazel_re::Directory>;
+        std::unordered_map<ArtifactDigest, bazel_re::Directory>;
 
-    BazelNetwork const& network_;
+    std::string const instance_name_;
+    BazelCasClient const& cas_;
     std::optional<DirectoryMap> auxiliary_map_;
 
     [[nodiscard]] static auto MakeAuxiliaryMap(
         std::vector<bazel_re::Directory>&& full_tree) noexcept
         -> std::optional<DirectoryMap>;
+
+    [[nodiscard]] auto BatchReadBlobs(
+        std::vector<bazel_re::Digest> const& blobs) const noexcept
+        -> std::vector<ArtifactBlob>;
 };
 
 #endif  // INCLUDED_SRC_BUILDTOOL_EXECUTION_API_REMOTE_BAZEL_BAZEL_TREE_READER_HPP