BazelCasClient: Use ArtifactBlob in BatchUpdateBlobs

7 files changed, 47 insertions, 45 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 f42b0605..dc920bca 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp
@@ -485,9 +485,9 @@ auto BazelCasClient::FindMissingBlobs(
     return result;
 }
 
-auto BazelCasClient::BatchUpdateBlobs(
-    std::string const& instance_name,
-    std::unordered_set<BazelBlob> const& blobs) const noexcept -> std::size_t {
+auto BazelCasClient::BatchUpdateBlobs(std::string const& instance_name,
+                                      std::unordered_set<ArtifactBlob> const&
+                                          blobs) const noexcept -> std::size_t {
     if (blobs.empty()) {
         return 0;
     }
@@ -501,7 +501,7 @@ auto BazelCasClient::BatchUpdateBlobs(
                 blobs.end(),
                 "BatchUpdateBlobs",
                 [](bazel_re::BatchUpdateBlobsRequest* request,
-                   BazelBlob const& x) {
+                   ArtifactBlob const& x) {
                     *(request->add_requests()) =
                         BazelCasClient::CreateUpdateBlobsSingleRequest(x);
                 });
@@ -572,10 +572,11 @@ auto BazelCasClient::BatchUpdateBlobs(
         // we are going extra defensive here and also consider missing responses
         // to be a failed blob update. Issue a retry for the missing blobs.
         logger_.Emit(LogLevel::Trace, "Retrying with missing blobs");
-        std::unordered_set<BazelBlob> missing_blobs;
+        std::unordered_set<ArtifactBlob> missing_blobs;
         missing_blobs.reserve(missing);
         for (auto const& blob : blobs) {
-            if (not updated.contains(blob.digest)) {
+            auto bazel_digest = ArtifactDigestFactory::ToBazel(blob.digest);
+            if (not updated.contains(bazel_digest)) {
                 missing_blobs.emplace(blob);
             }
         }
@@ -586,11 +587,16 @@ auto BazelCasClient::BatchUpdateBlobs(
         // trying that again; instead, we fall back to uploading each blob
         // sequentially.
         logger_.Emit(LogLevel::Debug, "Falling back to sequential blob upload");
-        return std::count_if(blobs.begin(),
-                             blobs.end(),
-                             [this, &instance_name](BazelBlob const& blob) {
-                                 return UpdateSingleBlob(instance_name, blob);
-                             });
+        return std::count_if(
+            blobs.begin(),
+            blobs.end(),
+            [this, &instance_name](ArtifactBlob const& blob) {
+                BazelBlob bazel_blob{
+                    ArtifactDigestFactory::ToBazel(blob.digest),
+                    blob.data,
+                    blob.is_exec};
+                return UpdateSingleBlob(instance_name, bazel_blob);
+            });
     }
     return updated.size();
 }
@@ -675,10 +681,11 @@ auto BazelCasClient::CreateBatchRequestsMaxSize(
     return result;
 }
 
-auto BazelCasClient::CreateUpdateBlobsSingleRequest(BazelBlob const& b) noexcept
+auto BazelCasClient::CreateUpdateBlobsSingleRequest(
+    ArtifactBlob const& b) noexcept
     -> bazel_re::BatchUpdateBlobsRequest_Request {
     bazel_re::BatchUpdateBlobsRequest_Request r{};
-    (*r.mutable_digest()) = b.digest;
+    (*r.mutable_digest()) = ArtifactDigestFactory::ToBazel(b.digest);
     r.set_data(*b.data);
     return r;
 }
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 7263bcf2..41b323d2 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp
@@ -65,7 +65,7 @@ class BazelCasClient {
     /// \returns The count of successfully updated blobs
     [[nodiscard]] auto BatchUpdateBlobs(
         std::string const& instance_name,
-        std::unordered_set<BazelBlob> const& blobs) const noexcept
+        std::unordered_set<ArtifactBlob> const& blobs) const noexcept
         -> std::size_t;
 
     /// \brief Read multiple blobs in batch transfer
@@ -158,7 +158,7 @@ class BazelCasClient {
             request_builder) const noexcept -> std::vector<TRequest>;
 
     [[nodiscard]] static auto CreateUpdateBlobsSingleRequest(
-        BazelBlob const& b) noexcept
+        ArtifactBlob const& b) noexcept
         -> bazel_re::BatchUpdateBlobsRequest_Request;
 
     [[nodiscard]] static auto CreateGetTreeRequest(
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network.cpp b/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
index a194c2ee..1f65201c 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
@@ -115,18 +115,9 @@ auto BazelNetwork::DoUploadBlobs(
         }
         to_stream.clear();
 
-        std::unordered_set<BazelBlob> bazel_blobs;
-        bazel_blobs.reserve(blobs.size());
-        for (auto const& blob : blobs) {
-            bazel_blobs.emplace(ArtifactDigestFactory::ToBazel(blob.digest),
-                                blob.data,
-                                blob.is_exec);
-        }
-
         // After uploading via stream api, only small blobs that may be uploaded
         // using batch are in the container:
-        return cas_->BatchUpdateBlobs(instance_name_, bazel_blobs) ==
-               bazel_blobs.size();
+        return cas_->BatchUpdateBlobs(instance_name_, blobs) == blobs.size();
 
     } catch (...) {
         Logger::Log(LogLevel::Warning, "Unknown exception");
diff --git a/test/buildtool/execution_api/bazel/TARGETS b/test/buildtool/execution_api/bazel/TARGETS
index cac3c109..88915e7e 100644
--- a/test/buildtool/execution_api/bazel/TARGETS
+++ b/test/buildtool/execution_api/bazel/TARGETS
@@ -12,7 +12,11 @@
     , ["@", "src", "src/buildtool/common/remote", "remote_common"]
     , ["@", "src", "src/buildtool/common/remote", "retry_config"]
     , ["@", "src", "src/buildtool/crypto", "hash_function"]
-    , ["@", "src", "src/buildtool/execution_api/bazel_msg", "bazel_msg"]
+    , [ "@"
+      , "src"
+      , "src/buildtool/execution_api/common"
+      , "artifact_blob_container"
+      ]
     , [ "@"
       , "src"
       , "src/buildtool/execution_api/common"
diff --git a/test/buildtool/execution_api/bazel/bazel_cas_client.test.cpp b/test/buildtool/execution_api/bazel/bazel_cas_client.test.cpp
index ad2e0678..606212a6 100644
--- a/test/buildtool/execution_api/bazel/bazel_cas_client.test.cpp
+++ b/test/buildtool/execution_api/bazel/bazel_cas_client.test.cpp
@@ -29,7 +29,7 @@
 #include "src/buildtool/common/remote/remote_common.hpp"
 #include "src/buildtool/common/remote/retry_config.hpp"
 #include "src/buildtool/crypto/hash_function.hpp"
-#include "src/buildtool/execution_api/bazel_msg/bazel_blob_container.hpp"
+#include "src/buildtool/execution_api/common/artifact_blob_container.hpp"
 #include "src/buildtool/execution_api/common/content_blob_container.hpp"
 #include "src/buildtool/execution_api/remote/config.hpp"
 #include "src/buildtool/file_system/object_type.hpp"
@@ -63,7 +63,7 @@ TEST_CASE("Bazel internals: CAS Client", "[execution_api]") {
             hash_function, content);
 
         // Valid blob
-        BazelBlob blob{bazel_digest, content, /*is_exec=*/false};
+        ArtifactBlob blob{digest, content, /*is_exec=*/false};
 
         // Search blob via digest
         auto digests =
@@ -89,14 +89,8 @@ TEST_CASE("Bazel internals: CAS Client", "[execution_api]") {
             "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef");
         faulty_digest.set_size_bytes(4);
 
-        // Faulty blob
-        BazelBlob faulty_blob{faulty_digest, content, /*is_exec=*/false};
-
         // Search faulty digest
         CHECK(cas_client.FindMissingBlobs(instance_name, {faulty_digest})
                   .size() == 1);
-
-        // Try upload faulty blob
-        CHECK(cas_client.BatchUpdateBlobs(instance_name, {faulty_blob}) == 0U);
     }
 }
diff --git a/test/utils/TARGETS b/test/utils/TARGETS
index 21e4fb9f..029df21f 100644
--- a/test/utils/TARGETS
+++ b/test/utils/TARGETS
@@ -15,12 +15,16 @@
     , "test_remote_config"
     , ["@", "gsl", "", "gsl"]
     , ["@", "protoc", "", "libprotobuf"]
-    , ["@", "src", "src/buildtool/common", "bazel_digest_factory"]
+    , ["@", "src", "src/buildtool/common", "artifact_digest_factory"]
     , ["@", "src", "src/buildtool/common", "bazel_types"]
     , ["@", "src", "src/buildtool/common/remote", "remote_common"]
     , ["@", "src", "src/buildtool/common/remote", "retry_config"]
     , ["@", "src", "src/buildtool/crypto", "hash_function"]
-    , ["@", "src", "src/buildtool/execution_api/bazel_msg", "bazel_msg"]
+    , [ "@"
+      , "src"
+      , "src/buildtool/execution_api/common"
+      , "artifact_blob_container"
+      ]
     , ["@", "src", "src/buildtool/execution_api/remote", "bazel_network"]
     , ["@", "src", "src/buildtool/execution_api/remote", "config"]
     , ["@", "src", "src/buildtool/file_system", "object_type"]
diff --git a/test/utils/remote_execution/bazel_action_creator.hpp b/test/utils/remote_execution/bazel_action_creator.hpp
index b0038de9..524ce82c 100644
--- a/test/utils/remote_execution/bazel_action_creator.hpp
+++ b/test/utils/remote_execution/bazel_action_creator.hpp
@@ -27,12 +27,12 @@
 
 #include "google/protobuf/repeated_ptr_field.h"
 #include "gsl/gsl"
-#include "src/buildtool/common/bazel_digest_factory.hpp"
+#include "src/buildtool/common/artifact_digest_factory.hpp"
 #include "src/buildtool/common/bazel_types.hpp"
 #include "src/buildtool/common/remote/remote_common.hpp"
 #include "src/buildtool/common/remote/retry_config.hpp"
 #include "src/buildtool/crypto/hash_function.hpp"
-#include "src/buildtool/execution_api/bazel_msg/bazel_blob_container.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/execution_api/remote/config.hpp"
 #include "src/buildtool/file_system/object_type.hpp"
@@ -53,7 +53,7 @@
         *(platform->add_properties()) = property;
     }
 
-    std::unordered_set<BazelBlob> blobs;
+    std::unordered_set<ArtifactBlob> blobs;
 
     bazel_re::Command cmd;
     cmd.set_allocated_platform(platform.release());
@@ -71,23 +71,24 @@
                    });
 
     auto cmd_data = cmd.SerializeAsString();
-    auto cmd_id = BazelDigestFactory::HashDataAs<ObjectType::File>(
+    auto cmd_id = ArtifactDigestFactory::HashDataAs<ObjectType::File>(
         hash_function, cmd_data);
     blobs.emplace(cmd_id, cmd_data, /*is_exec=*/false);
 
     bazel_re::Directory empty_dir;
     auto dir_data = empty_dir.SerializeAsString();
-    auto dir_id = BazelDigestFactory::HashDataAs<ObjectType::Tree>(
+    auto dir_id = ArtifactDigestFactory::HashDataAs<ObjectType::Tree>(
         hash_function, dir_data);
     blobs.emplace(dir_id, dir_data, /*is_exec=*/false);
 
     bazel_re::Action action;
-    (*action.mutable_command_digest()) = cmd_id;
+    (*action.mutable_command_digest()) = ArtifactDigestFactory::ToBazel(cmd_id);
     action.set_do_not_cache(false);
-    (*action.mutable_input_root_digest()) = dir_id;
+    (*action.mutable_input_root_digest()) =
+        ArtifactDigestFactory::ToBazel(dir_id);
 
     auto action_data = action.SerializeAsString();
-    auto action_id = BazelDigestFactory::HashDataAs<ObjectType::File>(
+    auto action_id = ArtifactDigestFactory::HashDataAs<ObjectType::File>(
         hash_function, action_data);
     blobs.emplace(action_id, action_data, /*is_exec=*/false);
 
@@ -109,7 +110,8 @@
                               &retry_config);
 
     if (cas_client.BatchUpdateBlobs(instance_name, blobs) == blobs.size()) {
-        return std::make_unique<bazel_re::Digest>(action_id);
+        return std::make_unique<bazel_re::Digest>(
+            ArtifactDigestFactory::ToBazel(action_id));
     }
     return nullptr;
 }