2 files changed, 67 insertions, 18 deletions
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 892cf3d9..a82d5c6e 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_network_reader.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_network_reader.cpp
@@ -91,25 +91,18 @@ auto BazelNetworkReader::ReadGitTree(ArtifactDigest const& digest)
     }
 
     auto check_symlinks = [this](std::vector<ArtifactDigest> const& ids) {
-        size_t const size = ids.size();
-        size_t count = 0;
-        for (auto blobs : ReadIncrementally(&ids)) {
-            if (count + blobs.size() > size) {
-                Logger::Log(LogLevel::Debug,
-                            "received more blobs than requested.");
-                return false;
-            }
-            bool valid = std::all_of(
-                blobs.begin(), blobs.end(), [](ArtifactBlob const& blob) {
-                    auto const content = blob.ReadContent();
-                    return content != nullptr and PathIsNonUpwards(*content);
-                });
-            if (not valid) {
-                return false;
-            }
-            count += blobs.size();
+        auto const blobs = ReadOrdered(ids);
+        if (blobs.size() != ids.size()) {
+            Logger::Log(LogLevel::Debug,
+                        "BazelNetworkReader::ReadGitTree: read wrong number of "
+                        "symlinks.");
+            return false;
         }
-        return true;
+        return std::all_of(
+            blobs.begin(), blobs.end(), [](ArtifactBlob const& blob) {
+                auto const content = blob.ReadContent();
+                return content != nullptr and PathIsNonUpwards(*content);
+            });
     };
 
     return GitRepo::ReadTreeData(*content,
@@ -182,6 +175,55 @@ auto BazelNetworkReader::ReadSingleBlob(ArtifactDigest const& digest)
     return cas_.ReadSingleBlob(instance_name_, digest);
 }
 
+auto BazelNetworkReader::Read(std::unordered_set<ArtifactDigest> const& digests)
+    const noexcept -> std::unordered_set<ArtifactBlob> {
+    std::unordered_set<ArtifactBlob> read_result;
+    read_result.reserve(digests.size());
+
+    std::unordered_set<ArtifactDigest> to_batch;
+    to_batch.reserve(digests.size());
+    // Upload blobs that don't fit for batching: size is larger than limit or
+    // unknown
+    std::size_t const limit = cas_.GetMaxBatchTransferSize(instance_name_);
+    for (auto const& digest : digests) {
+        if (digest.size() == 0 or digest.size() > limit) {
+            auto blob = cas_.ReadSingleBlob(instance_name_, digest);
+            if (blob.has_value()) {
+                read_result.emplace(*std::move(blob));
+            }
+        }
+        else {
+            to_batch.emplace(digest);
+        }
+    }
+
+    // Batch remaining blobs:
+    read_result.merge(cas_.BatchReadBlobs(instance_name_, to_batch));
+    return read_result;
+}
+
+auto BazelNetworkReader::ReadOrdered(std::vector<ArtifactDigest> const& digests)
+    const noexcept -> std::vector<ArtifactBlob> {
+    auto const read_result =
+        Read(std::unordered_set(digests.begin(), digests.end()));
+
+    auto const back_map = BackMap<ArtifactDigest, ArtifactBlob>::Make(
+        &read_result,
+        [](ArtifactBlob const& blob) { return blob.GetDigest(); });
+    if (back_map == nullptr) {
+        return {};
+    }
+
+    std::vector<ArtifactBlob> artifacts;
+    artifacts.reserve(digests.size());
+    for (auto const& digest : digests) {
+        if (auto value = back_map->GetReference(digest)) {
+            artifacts.emplace_back(*value.value());
+        }
+    }
+    return artifacts;
+}
+
 auto BazelNetworkReader::ReadIncrementally(
     gsl::not_null<std::vector<ArtifactDigest> const*> const& digests)
     const noexcept -> IncrementalReader {
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 cf674ca1..c11da691 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_network_reader.hpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_network_reader.hpp
@@ -22,6 +22,7 @@
 #include <optional>
 #include <string>
 #include <unordered_map>
+#include <unordered_set>
 #include <vector>
 
 #include "gsl/gsl"
@@ -73,6 +74,12 @@ class BazelNetworkReader final {
     [[nodiscard]] auto ReadSingleBlob(ArtifactDigest const& digest)
         const noexcept -> std::optional<ArtifactBlob>;
 
+    [[nodiscard]] auto Read(std::unordered_set<ArtifactDigest> const& digests)
+        const noexcept -> std::unordered_set<ArtifactBlob>;
+
+    [[nodiscard]] auto ReadOrdered(std::vector<ArtifactDigest> const& digests)
+        const noexcept -> std::vector<ArtifactBlob>;
+
     [[nodiscard]] auto ReadIncrementally(
         gsl::not_null<std::vector<ArtifactDigest> const*> const& digests)
         const noexcept -> IncrementalReader;