Introduced batch availability check for remote execution api

7 files changed, 61 insertions, 1 deletions

diff --git a/src/buildtool/execution_api/common/execution_api.hpp b/src/buildtool/execution_api/common/execution_api.hpp
index 3f1da5b8..816b5031 100644
--- a/src/buildtool/execution_api/common/execution_api.hpp
+++ b/src/buildtool/execution_api/common/execution_api.hpp
@@ -73,6 +73,10 @@ class IExecutionApi {
 
     [[nodiscard]] virtual auto IsAvailable(
         ArtifactDigest const& digest) const noexcept -> bool = 0;
+
+    [[nodiscard]] virtual auto IsAvailable(
+        std::vector<ArtifactDigest> const& digests) const noexcept
+        -> std::vector<ArtifactDigest> = 0;
 };
 
 #endif  // INCLUDED_SRC_BUILDTOOL_EXECUTION_API_COMMON_EXECUTION_APIHPP
diff --git a/src/buildtool/execution_api/local/local_api.hpp b/src/buildtool/execution_api/local/local_api.hpp
index 20886e5a..72b7e28c 100644
--- a/src/buildtool/execution_api/local/local_api.hpp
+++ b/src/buildtool/execution_api/local/local_api.hpp
@@ -150,6 +150,17 @@ class LocalApi final : public IExecutionApi {
         return storage_->BlobPath(digest, false).has_value();
     }
 
+    [[nodiscard]] auto IsAvailable(std::vector<ArtifactDigest> const& digests)
+        const noexcept -> std::vector<ArtifactDigest> final {
+        std::vector<ArtifactDigest> result;
+        for (auto const& digest : digests) {
+            if (not storage_->BlobPath(digest, false).has_value()) {
+                result.push_back(digest);
+            }
+        }
+        return result;
+    }
+
   private:
     std::shared_ptr<LocalTreeMap> tree_map_{std::make_shared<LocalTreeMap>()};
     std::shared_ptr<LocalStorage> storage_{
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_api.cpp b/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
index ca5d4946..51e85176 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
@@ -170,7 +170,7 @@ auto BazelApi::CreateAction(
         return std::nullopt;
     }
     if (tree_map_->AddTree(*digest, std::move(tree))) {
-        return ArtifactDigest{std::move(*digest)};
+        return ArtifactDigest{*digest};
     }
     return std::nullopt;
 }
@@ -179,3 +179,23 @@ auto BazelApi::CreateAction(
     ArtifactDigest const& digest) const noexcept -> bool {
     return network_->IsAvailable(digest);
 }
+
+[[nodiscard]] auto BazelApi::IsAvailable(
+    std::vector<ArtifactDigest> const& digests) const noexcept
+    -> std::vector<ArtifactDigest> {
+    std::vector<bazel_re::Digest> bazel_digests;
+    bazel_digests.reserve(digests.size());
+    std::transform(digests.begin(),
+                   digests.end(),
+                   std::back_inserter(bazel_digests),
+                   [](ArtifactDigest const& digest) { return digest; });
+    auto bazel_result = network_->IsAvailable(bazel_digests);
+    std::vector<ArtifactDigest> result;
+    result.reserve(bazel_result.size());
+    std::transform(
+        bazel_result.begin(),
+        bazel_result.end(),
+        std::back_inserter(result),
+        [](bazel_re::Digest& digest) { return ArtifactDigest{digest}; });
+    return result;
+}
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_api.hpp b/src/buildtool/execution_api/remote/bazel/bazel_api.hpp
index 8582db65..f953618e 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_api.hpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_api.hpp
@@ -57,6 +57,9 @@ class BazelApi final : public IExecutionApi {
     [[nodiscard]] auto IsAvailable(ArtifactDigest const& digest) const noexcept
         -> bool final;
 
+    [[nodiscard]] auto IsAvailable(std::vector<ArtifactDigest> const& digests)
+        const noexcept -> std::vector<ArtifactDigest> final;
+
   private:
     std::shared_ptr<BazelNetwork> network_;
     std::shared_ptr<LocalTreeMap> tree_map_;
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network.cpp b/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
index 669f93dd..b1ef59c8 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
@@ -70,6 +70,11 @@ auto BazelNetwork::IsAvailable(bazel_re::Digest const& digest) const noexcept
         .empty();
 }
 
+auto BazelNetwork::IsAvailable(std::vector<bazel_re::Digest> const& digests)
+    const noexcept -> std::vector<bazel_re::Digest> {
+    return cas_->FindMissingBlobs(instance_name_, digests);
+}
+
 template <class T_Iter>
 auto BazelNetwork::DoUploadBlobs(T_Iter const& first,
                                  T_Iter const& last) noexcept -> bool {
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network.hpp b/src/buildtool/execution_api/remote/bazel/bazel_network.hpp
index 644af2b4..13a5bcf7 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_network.hpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_network.hpp
@@ -54,6 +54,9 @@ class BazelNetwork {
     [[nodiscard]] auto IsAvailable(
         bazel_re::Digest const& digest) const noexcept -> bool;
 
+    [[nodiscard]] auto IsAvailable(std::vector<bazel_re::Digest> const& digests)
+        const noexcept -> std::vector<bazel_re::Digest>;
+
     /// \brief Uploads blobs to CAS
     /// \param blobs              The blobs to upload
     /// \param skip_find_missing  Skip finding missing blobs, just upload all
diff --git a/test/buildtool/execution_engine/executor/executor.test.cpp b/test/buildtool/execution_engine/executor/executor.test.cpp
index e055dfd7..0ffe5e4d 100755
--- a/test/buildtool/execution_engine/executor/executor.test.cpp
+++ b/test/buildtool/execution_engine/executor/executor.test.cpp
@@ -154,6 +154,20 @@ class TestApi : public IExecutionApi {
             return false;
         }
     }
+    [[nodiscard]] auto IsAvailable(std::vector<ArtifactDigest> const& digests)
+        const noexcept -> std::vector<ArtifactDigest> final {
+        std::vector<ArtifactDigest> result;
+        try {
+            for (auto const& digest : digests) {
+                if (not config_.artifacts.at(digest.hash()).available) {
+                    result.push_back(digest);
+                }
+            }
+        } catch (std::exception const& /* unused */) {
+            return result;
+        }
+        return result;
+    }
 
   private:
     TestApiConfig config_{};