Use common interface for reading tree entries and leafs

...in LocalApi and BazelApi.

5 files changed, 189 insertions, 180 deletions

diff --git a/src/buildtool/execution_api/remote/bazel/bazel_api.cpp b/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
index 0e3a1ca2..47d024f9 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
@@ -31,11 +31,13 @@
 #include "src/buildtool/execution_api/bazel_msg/bazel_common.hpp"
 #include "src/buildtool/execution_api/bazel_msg/bazel_msg_factory.hpp"
 #include "src/buildtool/execution_api/common/common_api.hpp"
+#include "src/buildtool/execution_api/common/tree_reader.hpp"
 #include "src/buildtool/execution_api/remote/bazel/bazel_ac_client.hpp"
 #include "src/buildtool/execution_api/remote/bazel/bazel_action.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.hpp"
+#include "src/buildtool/execution_api/remote/bazel/bazel_network_reader.hpp"
 #include "src/buildtool/execution_api/remote/bazel/bazel_response.hpp"
 #include "src/buildtool/file_system/file_system_manager.hpp"
 #include "src/buildtool/file_system/object_type.hpp"
@@ -223,10 +225,15 @@ auto BazelApi::CreateAction(
         else {
             if (IsTreeObject(info.type)) {
                 // read object infos from sub tree and call retrieve recursively
-                auto const infos = network_->RecursivelyReadTreeLeafs(
-                    info.digest, output_paths[i], alternative.has_value());
-                if (not infos or
-                    not RetrieveToPaths(infos->second, infos->first)) {
+                auto request_remote_tree = alternative.has_value()
+                                               ? std::make_optional(info.digest)
+                                               : std::nullopt;
+                auto reader = TreeReader<BazelNetworkReader>{
+                    *network_, std::move(request_remote_tree)};
+                auto const result = reader.RecursivelyReadTreeLeafs(
+                    info.digest, output_paths[i]);
+                if (not result or
+                    not RetrieveToPaths(result->infos, result->paths)) {
                     return false;
                 }
             }
@@ -315,9 +322,10 @@ 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 const infos = network_->ReadDirectTreeEntries(
+            auto reader = TreeReader<BazelNetworkReader>{*network_};
+            auto const result = reader.ReadDirectTreeEntries(
                 info.digest, std::filesystem::path{});
-            if (not infos or not RetrieveToCas(infos->second, api)) {
+            if (not result or not RetrieveToCas(result->infos, api)) {
                 return false;
             }
         }
@@ -363,11 +371,12 @@ 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 const infos = network_->ReadDirectTreeEntries(
+                auto reader = TreeReader<BazelNetworkReader>{*network_};
+                auto const result = reader.ReadDirectTreeEntries(
                     info.digest, std::filesystem::path{});
-                if (not infos or
+                if (not result or
                     not ParallelRetrieveToCas(
-                        infos->second, api, jobs, use_blob_splitting)) {
+                        result->infos, api, jobs, use_blob_splitting)) {
                     return false;
                 }
             }
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network.cpp b/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
index 325abd02..485e65dc 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
@@ -307,141 +307,12 @@ auto BazelNetwork::GetCachedActionResult(
         instance_name_, action, false, false, output_files);
 }
 
-auto BazelNetwork::RecursivelyReadTreeLeafs(
-    bazel_re::Digest const& tree_digest,
-    std::filesystem::path const& parent,
-    bool request_remote_tree) const noexcept
-    -> std::optional<std::pair<std::vector<std::filesystem::path>,
-                               std::vector<Artifact::ObjectInfo>>> {
-    std::optional<DirectoryMap> dir_map{std::nullopt};
-    if (Compatibility::IsCompatible() and request_remote_tree) {
-        // Query full tree from remote CAS. Note that this is currently not
-        // supported by Buildbarn revision c3c06bbe2a.
-        auto dirs =
-            cas_->GetTree(instance_name_, tree_digest, kMaxBatchTransferSize);
-
-        // Convert to Directory map
-        dir_map = DirectoryMap{};
-        dir_map->reserve(dirs.size());
-        for (auto& dir : dirs) {
-            try {
-                dir_map->emplace(ArtifactDigest::Create<ObjectType::File>(
-                                     dir.SerializeAsString()),
-                                 std::move(dir));
-            } catch (...) {
-                return std::nullopt;
-            }
-        }
-    }
-
-    std::vector<std::filesystem::path> paths{};
-    std::vector<Artifact::ObjectInfo> infos{};
-
-    auto store_info = [&paths, &infos](auto path, auto info) {
-        paths.emplace_back(path);
-        infos.emplace_back(info);
-        return true;
-    };
-
-    if (ReadObjectInfosRecursively(dir_map, store_info, parent, tree_digest)) {
-        return std::make_pair(std::move(paths), std::move(infos));
-    }
-
-    return std::nullopt;
-}
-
-auto BazelNetwork::ReadDirectTreeEntries(
-    bazel_re::Digest const& tree_digest,
-    std::filesystem::path const& parent) const noexcept
-    -> std::optional<std::pair<std::vector<std::filesystem::path>,
-                               std::vector<Artifact::ObjectInfo>>> {
-    std::vector<std::filesystem::path> paths{};
-    std::vector<Artifact::ObjectInfo> infos{};
-
-    auto store_info = [&paths, &infos](auto path, auto info) {
-        paths.emplace_back(path);
-        infos.emplace_back(info);
-        return true;
-    };
-
-    if (Compatibility::IsCompatible()) {
-        // read from CAS
-        if (auto dir = ReadDirectory(this, tree_digest)) {
-            if (not BazelMsgFactory::ReadObjectInfosFromDirectory(
-                    *dir, [&store_info, &parent](auto path, auto info) {
-                        return store_info(parent / path, info);
-                    })) {
-                return std::nullopt;
-            }
-        }
-    }
-    else {
-        if (auto entries = ReadGitTree(this, tree_digest)) {
-            if (not BazelMsgFactory::ReadObjectInfosFromGitTree(
-                    *entries, [&store_info, &parent](auto path, auto info) {
-                        return store_info(parent / path, info);
-                    })) {
-                return std::nullopt;
-            }
-        }
-    }
-
-    return std::make_pair(std::move(paths), std::move(infos));
-}
-
-// NOLINTNEXTLINE(misc-no-recursion)
-auto BazelNetwork::ReadObjectInfosRecursively(
-    std::optional<DirectoryMap> const& dir_map,
-    BazelMsgFactory::InfoStoreFunc const& store_info,
-    std::filesystem::path const& parent,
-    bazel_re::Digest const& digest) const noexcept -> bool {
-    if (Compatibility::IsCompatible()) {
-        // read from in-memory Directory map
-        if (dir_map) {
-            if (dir_map->contains(digest)) {
-                return BazelMsgFactory::ReadObjectInfosFromDirectory(
-                    dir_map->at(digest),
-                    [this, &dir_map, &store_info, &parent](auto path,
-                                                           auto info) {
-                        return IsTreeObject(info.type)
-                                   ? ReadObjectInfosRecursively(dir_map,
-                                                                store_info,
-                                                                parent / path,
-                                                                info.digest)
-                                   : store_info(parent / path, info);
-                    });
-            }
-        }
-
-        // fallback read from CAS
-        if (auto dir = ReadDirectory(this, digest)) {
-            return BazelMsgFactory::ReadObjectInfosFromDirectory(
-                *dir,
-                [this, &dir_map, &store_info, &parent](auto path, auto info) {
-                    return IsTreeObject(info.type)
-                               ? ReadObjectInfosRecursively(dir_map,
-                                                            store_info,
-                                                            parent / path,
-                                                            info.digest)
-                               : store_info(parent / path, info);
-                });
-        }
-    }
-    else {
-        if (auto entries = ReadGitTree(this, digest)) {
-            return BazelMsgFactory::ReadObjectInfosFromGitTree(
-                *entries,
-                [this, &dir_map, &store_info, &parent](auto path, auto info) {
-                    return IsTreeObject(info.type)
-                               ? ReadObjectInfosRecursively(dir_map,
-                                                            store_info,
-                                                            parent / path,
-                                                            info.digest)
-                               : store_info(parent / path, info);
-                });
-        }
+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 false;
+    return cas_->GetTree(instance_name_, digest, kMaxBatchTransferSize);
 }
 
 auto BazelNetwork::DumpToStream(Artifact::ObjectInfo const& info,
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network.hpp b/src/buildtool/execution_api/remote/bazel/bazel_network.hpp
index 7101c926..6ee44921 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_network.hpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_network.hpp
@@ -19,7 +19,6 @@
 #include <memory>
 #include <optional>
 #include <string>
-#include <unordered_map>
 #include <utility>
 #include <vector>
 
@@ -111,41 +110,16 @@ class BazelNetwork {
         std::vector<std::string> const& output_files) const noexcept
         -> std::optional<bazel_re::ActionResult>;
 
-    /// \brief Traverses a tree recursively and retrieves object infos of all
-    /// found blobs. Tree objects are not added to the result list, but
-    /// converted to a path name.
-    /// \param tree_digest          Digest of the tree.
-    /// \param parent               Local parent path.
-    /// \param request_remote_tree  Query full tree from remote CAS.
-    /// \returns Pair of vectors, first containing filesystem paths, second
-    /// containing object infos.
-    [[nodiscard]] auto RecursivelyReadTreeLeafs(
-        bazel_re::Digest const& tree_digest,
-        std::filesystem::path const& parent,
-        bool request_remote_tree = false) const noexcept
-        -> std::optional<std::pair<std::vector<std::filesystem::path>,
-                                   std::vector<Artifact::ObjectInfo>>>;
-
-    /// \brief Reads the flat content of a tree and returns object infos of all
-    /// its direct entries (trees and blobs).
-    /// \param tree_digest      Digest of the tree.
-    /// \param parent           Local parent path.
-    /// \returns Pair of vectors, first containing filesystem paths, second
-    /// containing object infos.
-    [[nodiscard]] auto ReadDirectTreeEntries(
-        bazel_re::Digest const& tree_digest,
-        std::filesystem::path const& parent) const noexcept
-        -> std::optional<std::pair<std::vector<std::filesystem::path>,
-                                   std::vector<Artifact::ObjectInfo>>>;
-
     [[nodiscard]] auto DumpToStream(Artifact::ObjectInfo const& info,
                                     gsl::not_null<FILE*> const& stream,
                                     bool raw_tree) const noexcept -> bool;
 
-  private:
-    using DirectoryMap =
-        std::unordered_map<bazel_re::Digest, bazel_re::Directory>;
+    /// \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_{};
     std::unique_ptr<BazelCasClient> cas_{};
@@ -155,12 +129,6 @@ class BazelNetwork {
     template <class T_Iter>
     [[nodiscard]] auto DoUploadBlobs(T_Iter const& first,
                                      T_Iter const& last) noexcept -> bool;
-
-    [[nodiscard]] auto ReadObjectInfosRecursively(
-        std::optional<DirectoryMap> const& dir_map,
-        BazelMsgFactory::InfoStoreFunc const& store_info,
-        std::filesystem::path const& parent,
-        bazel_re::Digest const& digest) const noexcept -> bool;
 };
 
 #endif  // INCLUDED_SRC_BUILDTOOL_EXECUTION_API_REMOTE_BAZEL_BAZEL_NETWORK_HPP
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network_reader.cpp b/src/buildtool/execution_api/remote/bazel/bazel_network_reader.cpp
new file mode 100644
index 00000000..ee0fbc78
--- /dev/null
+++ b/src/buildtool/execution_api/remote/bazel/bazel_network_reader.cpp
@@ -0,0 +1,110 @@
+// Copyright 2024 Huawei Cloud Computing Technology Co., Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "src/buildtool/execution_api/remote/bazel/bazel_network_reader.hpp"
+
+#include <algorithm>
+
+#include "src/buildtool/crypto/hash_function.hpp"
+#include "src/buildtool/execution_api/bazel_msg/bazel_blob.hpp"
+#include "src/buildtool/execution_api/bazel_msg/bazel_msg_factory.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(
+    BazelNetwork const& network,
+    std::optional<ArtifactDigest> request_remote_tree) noexcept
+    : network_(network) {
+    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));
+        }
+    }
+}
+
+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);
+        }
+    } catch (...) {
+        // fallthrough
+    }
+
+    auto blobs = network_.ReadBlobs({digest}).Next();
+    if (blobs.size() == 1) {
+        return BazelMsgFactory::MessageFromString<bazel_re::Directory>(
+            blobs.at(0).data);
+    }
+    Logger::Log(
+        LogLevel::Debug, "Directory {} not found in CAS", digest.hash());
+    return std::nullopt;
+}
+
+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);
+    }
+    Logger::Log(LogLevel::Debug, "Tree {} not found in CAS", digest.hash());
+    return std::nullopt;
+}
+
+auto BazelNetworkReader::MakeAuxiliaryMap(
+    std::vector<bazel_re::Directory>&& full_tree) noexcept
+    -> std::optional<DirectoryMap> {
+    DirectoryMap result;
+    result.reserve(full_tree.size());
+    for (auto& dir : full_tree) {
+        try {
+            result.emplace(ArtifactDigest::Create<ObjectType::File>(
+                               dir.SerializeAsString()),
+                           std::move(dir));
+        } catch (...) {
+            return std::nullopt;
+        }
+    }
+    return result;
+}
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network_reader.hpp b/src/buildtool/execution_api/remote/bazel/bazel_network_reader.hpp
new file mode 100644
index 00000000..8d46c9f8
--- /dev/null
+++ b/src/buildtool/execution_api/remote/bazel/bazel_network_reader.hpp
@@ -0,0 +1,51 @@
+// Copyright 2024 Huawei Cloud Computing Technology Co., Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef INCLUDED_SRC_BUILDTOOL_EXECUTION_API_REMOTE_BAZEL_BAZEL_TREE_READER_HPP
+#define INCLUDED_SRC_BUILDTOOL_EXECUTION_API_REMOTE_BAZEL_BAZEL_TREE_READER_HPP
+
+#include <optional>
+#include <unordered_map>
+
+#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/file_system/git_repo.hpp"
+
+class BazelNetworkReader final {
+  public:
+    explicit BazelNetworkReader(
+        BazelNetwork const& network,
+        std::optional<ArtifactDigest> request_remote_tree =
+            std::nullopt) noexcept;
+
+    [[nodiscard]] auto ReadDirectory(ArtifactDigest const& digest)
+        const noexcept -> std::optional<bazel_re::Directory>;
+
+    [[nodiscard]] auto ReadGitTree(ArtifactDigest const& digest) const noexcept
+        -> std::optional<GitRepo::tree_entries_t>;
+
+  private:
+    using DirectoryMap =
+        std::unordered_map<bazel_re::Digest, bazel_re::Directory>;
+
+    BazelNetwork const& network_;
+    std::optional<DirectoryMap> auxiliary_map_;
+
+    [[nodiscard]] static auto MakeAuxiliaryMap(
+        std::vector<bazel_re::Directory>&& full_tree) noexcept
+        -> std::optional<DirectoryMap>;
+};
+
+#endif  // INCLUDED_SRC_BUILDTOOL_EXECUTION_API_REMOTE_BAZEL_BAZEL_TREE_READER_HPP
\ No newline at end of file