BazelCasClient: Use ArtifactDigest in FindMissingBlobs

3 files changed, 21 insertions, 24 deletions

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 179ac0db..2a9fc182 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp
@@ -427,19 +427,26 @@ auto BazelCasClient::BlobSpliceSupport(
 
 auto BazelCasClient::FindMissingBlobs(
     std::string const& instance_name,
-    std::unordered_set<bazel_re::Digest> const& digests) const noexcept
-    -> std::unordered_set<bazel_re::Digest> {
-    std::unordered_set<bazel_re::Digest> result;
+    std::unordered_set<ArtifactDigest> const& digests) const noexcept
+    -> std::unordered_set<ArtifactDigest> {
+    std::unordered_set<ArtifactDigest> result;
     if (digests.empty()) {
         return result;
     }
+
+    auto const back_map = BackMap<bazel_re::Digest, ArtifactDigest>::Make(
+        &digests, ArtifactDigestFactory::ToBazel);
+    if (not back_map.has_value()) {
+        return digests;
+    }
+
     try {
         result.reserve(digests.size());
         auto requests =
             CreateBatchRequestsMaxSize<bazel_re::FindMissingBlobsRequest>(
                 instance_name,
-                digests.begin(),
-                digests.end(),
+                back_map->GetKeys().begin(),
+                back_map->GetKeys().end(),
                 "FindMissingBlobs",
                 [](bazel_re::FindMissingBlobsRequest* request,
                    bazel_re::Digest const& x) {
@@ -456,10 +463,11 @@ auto BazelCasClient::FindMissingBlobs(
                 retry_config_,
                 logger_);
             if (ok) {
-                auto batch =
-                    ProcessResponseContents<bazel_re::Digest>(response);
-                for (auto&& x : batch) {
-                    result.emplace(std::move(x));
+                for (auto& batch :
+                     ProcessResponseContents<bazel_re::Digest>(response)) {
+                    if (auto value = back_map->GetReference(batch)) {
+                        result.emplace(*value.value());
+                    }
                 }
             }
             else {
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 370bb7ca..279164bf 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp
@@ -55,8 +55,8 @@ class BazelCasClient {
     /// \returns The digests of blobs not found in CAS
     [[nodiscard]] auto FindMissingBlobs(
         std::string const& instance_name,
-        std::unordered_set<bazel_re::Digest> const& digests) const noexcept
-        -> std::unordered_set<bazel_re::Digest>;
+        std::unordered_set<ArtifactDigest> const& digests) const noexcept
+        -> std::unordered_set<ArtifactDigest>;
 
     /// \brief Upload multiple blobs in batch transfer
     /// \param[in] instance_name Name of the CAS instance
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network.cpp b/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
index b435fb8e..0be8b539 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 <functional>
 #include <utility>
 
-#include "src/buildtool/common/artifact_digest_factory.hpp"
 #include "src/buildtool/execution_api/common/content_blob_container.hpp"
 #include "src/buildtool/execution_api/common/message_limits.hpp"
 #include "src/buildtool/logging/log_level.hpp"
@@ -44,23 +43,13 @@ BazelNetwork::BazelNetwork(
 
 auto BazelNetwork::IsAvailable(ArtifactDigest const& digest) const noexcept
     -> bool {
-    return cas_
-        ->FindMissingBlobs(instance_name_,
-                           {ArtifactDigestFactory::ToBazel(digest)})
-        .empty();
+    return cas_->FindMissingBlobs(instance_name_, {digest}).empty();
 }
 
 auto BazelNetwork::FindMissingBlobs(
     std::unordered_set<ArtifactDigest> const& digests) const noexcept
     -> std::unordered_set<ArtifactDigest> {
-    auto const back_map = BackMap<bazel_re::Digest, ArtifactDigest>::Make(
-        &digests, ArtifactDigestFactory::ToBazel);
-    if (not back_map.has_value()) {
-        return digests;
-    }
-    auto missing_digests =
-        cas_->FindMissingBlobs(instance_name_, back_map->GetKeys());
-    return back_map->GetValues(missing_digests);
+    return cas_->FindMissingBlobs(instance_name_, digests);
 }
 
 auto BazelNetwork::SplitBlob(bazel_re::Digest const& blob_digest) const noexcept