BazelNetwork: Use ArtifactBlobs in UploadBlobs

8 files changed, 53 insertions, 67 deletions

diff --git a/src/buildtool/execution_api/bazel_msg/bazel_msg_factory.cpp b/src/buildtool/execution_api/bazel_msg/bazel_msg_factory.cpp
index 18e22292..30a1a31a 100644
--- a/src/buildtool/execution_api/bazel_msg/bazel_msg_factory.cpp
+++ b/src/buildtool/execution_api/bazel_msg/bazel_msg_factory.cpp
@@ -913,15 +913,10 @@ auto BazelMsgFactory::CreateActionDigestFromCommandLine(
         return std::nullopt;
     }
 
+    auto result = action->digest;
     if (request.store_blob) {
-        std::invoke(*request.store_blob,
-                    BazelBlob{ArtifactDigestFactory::ToBazel(cmd->digest),
-                              cmd->data,
-                              cmd->is_exec});
-        std::invoke(*request.store_blob,
-                    BazelBlob{ArtifactDigestFactory::ToBazel(action->digest),
-                              action->data,
-                              action->is_exec});
+        std::invoke(*request.store_blob, *std::move(cmd));
+        std::invoke(*request.store_blob, *std::move(action));
     }
-    return action->digest;
+    return result;
 }
diff --git a/src/buildtool/execution_api/bazel_msg/bazel_msg_factory.hpp b/src/buildtool/execution_api/bazel_msg/bazel_msg_factory.hpp
index 6e71ace8..72d9d818 100644
--- a/src/buildtool/execution_api/bazel_msg/bazel_msg_factory.hpp
+++ b/src/buildtool/execution_api/bazel_msg/bazel_msg_factory.hpp
@@ -30,7 +30,6 @@
 #include "src/buildtool/common/artifact_digest.hpp"
 #include "src/buildtool/common/bazel_types.hpp"
 #include "src/buildtool/crypto/hash_function.hpp"
-#include "src/buildtool/execution_api/bazel_msg/bazel_blob_container.hpp"
 #include "src/buildtool/execution_api/bazel_msg/directory_tree.hpp"
 #include "src/buildtool/execution_api/common/artifact_blob_container.hpp"
 #include "src/buildtool/file_system/object_type.hpp"
@@ -203,7 +202,7 @@ class BazelMsgFactory {
 };
 
 struct BazelMsgFactory::ActionDigestRequest final {
-    using BlobStoreFunc = std::function<void(BazelBlob&&)>;
+    using BlobStoreFunc = std::function<void(ArtifactBlob&&)>;
 
     template <typename T>
     using VectorPtr = gsl::not_null<std::vector<T> const*>;
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_action.cpp b/src/buildtool/execution_api/remote/bazel/bazel_action.cpp
index 1b01811e..6fa8eedc 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_action.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_action.cpp
@@ -21,7 +21,6 @@
 
 #include "gsl/gsl"
 #include "src/buildtool/common/artifact_digest_factory.hpp"
-#include "src/buildtool/execution_api/bazel_msg/bazel_blob_container.hpp"
 #include "src/buildtool/execution_api/bazel_msg/bazel_msg_factory.hpp"
 #include "src/buildtool/execution_api/remote/bazel/bazel_execution_client.hpp"
 #include "src/buildtool/execution_api/remote/bazel/bazel_response.hpp"
@@ -53,7 +52,7 @@ BazelAction::BazelAction(
 
 auto BazelAction::Execute(Logger const* logger) noexcept
     -> IExecutionResponse::Ptr {
-    std::unordered_set<BazelBlob> blobs{};
+    std::unordered_set<ArtifactBlob> blobs{};
     auto do_cache = CacheEnabled(cache_flag_);
     auto action = CreateBundlesForAction(&blobs, root_digest_, not do_cache);
     if (not action) {
@@ -139,14 +138,14 @@ auto BazelAction::Execute(Logger const* logger) noexcept
     return nullptr;
 }
 
-auto BazelAction::CreateBundlesForAction(std::unordered_set<BazelBlob>* blobs,
-                                         ArtifactDigest const& exec_dir,
-                                         bool do_not_cache) const noexcept
-    -> std::optional<bazel_re::Digest> {
+auto BazelAction::CreateBundlesForAction(
+    std::unordered_set<ArtifactBlob>* blobs,
+    ArtifactDigest const& exec_dir,
+    bool do_not_cache) const noexcept -> std::optional<bazel_re::Digest> {
     using StoreFunc = BazelMsgFactory::ActionDigestRequest::BlobStoreFunc;
     std::optional<StoreFunc> store_blob = std::nullopt;
     if (blobs != nullptr) {
-        store_blob = [&blobs](BazelBlob&& blob) {
+        store_blob = [&blobs](ArtifactBlob&& blob) {
             blobs->emplace(std::move(blob));
         };
     }
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_action.hpp b/src/buildtool/execution_api/remote/bazel/bazel_action.hpp
index ae32a010..dab0f63c 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_action.hpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_action.hpp
@@ -25,7 +25,7 @@
 
 #include "src/buildtool/common/artifact_digest.hpp"
 #include "src/buildtool/common/bazel_types.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/execution_action.hpp"
 #include "src/buildtool/execution_api/common/execution_response.hpp"
 #include "src/buildtool/execution_api/remote/bazel/bazel_network.hpp"
@@ -67,7 +67,7 @@ class BazelAction final : public IExecutionAction {
         std::map<std::string, std::string> const& properties) noexcept;
 
     [[nodiscard]] auto CreateBundlesForAction(
-        std::unordered_set<BazelBlob>* blobs,
+        std::unordered_set<ArtifactBlob>* blobs,
         ArtifactDigest const& exec_dir,
         bool do_not_cache) const noexcept -> std::optional<bazel_re::Digest>;
 };
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_api.cpp b/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
index bee70424..d46c9e09 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
@@ -31,7 +31,6 @@
 #include "src/buildtool/common/artifact_digest_factory.hpp"
 #include "src/buildtool/common/bazel_types.hpp"
 #include "src/buildtool/common/protocol_traits.hpp"
-#include "src/buildtool/execution_api/bazel_msg/bazel_blob_container.hpp"
 #include "src/buildtool/execution_api/bazel_msg/bazel_common.hpp"
 #include "src/buildtool/execution_api/bazel_msg/directory_tree.hpp"
 #include "src/buildtool/execution_api/common/artifact_blob_container.hpp"
@@ -138,23 +137,6 @@ namespace {
     return true;
 }
 
-[[nodiscard]] auto ConvertToBazelBlobContainer(
-    std::unordered_set<ArtifactBlob>&& container) noexcept
-    -> std::optional<std::unordered_set<BazelBlob>> {
-    std::unordered_set<BazelBlob> blobs;
-    try {
-        blobs.reserve(container.size());
-        for (const auto& blob : container) {
-            blobs.emplace(ArtifactDigestFactory::ToBazel(blob.digest),
-                          blob.data,
-                          blob.is_exec);
-        }
-    } catch (...) {
-        return std::nullopt;
-    }
-    return blobs;
-}
-
 }  // namespace
 
 BazelApi::BazelApi(
@@ -529,10 +511,7 @@ auto BazelApi::CreateAction(
 [[nodiscard]] auto BazelApi::Upload(std::unordered_set<ArtifactBlob>&& blobs,
                                     bool skip_find_missing) const noexcept
     -> bool {
-    auto bazel_blobs = ConvertToBazelBlobContainer(std::move(blobs));
-    return bazel_blobs ? network_->UploadBlobs(*std::move(bazel_blobs),
-                                               skip_find_missing)
-                       : false;
+    return network_->UploadBlobs(std::move(blobs), skip_find_missing);
 }
 
 [[nodiscard]] auto BazelApi::UploadTree(
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network.cpp b/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
index 23da6ad1..dde80242 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
@@ -17,6 +17,8 @@
 #include <functional>
 #include <utility>
 
+#include "src/buildtool/common/artifact_digest_factory.hpp"
+#include "src/buildtool/execution_api/bazel_msg/bazel_blob_container.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"
@@ -73,8 +75,8 @@ auto BazelNetwork::BlobSpliceSupport() const noexcept -> bool {
     return cas_->BlobSpliceSupport(hash_function_, instance_name_);
 }
 
-auto BazelNetwork::DoUploadBlobs(std::unordered_set<BazelBlob> blobs) noexcept
-    -> bool {
+auto BazelNetwork::DoUploadBlobs(
+    std::unordered_set<ArtifactBlob> blobs) noexcept -> bool {
     if (blobs.empty()) {
         return true;
     }
@@ -93,16 +95,28 @@ auto BazelNetwork::DoUploadBlobs(std::unordered_set<BazelBlob> blobs) noexcept
         }
 
         for (auto const& it : to_stream) {
-            if (not cas_->UpdateSingleBlob(instance_name_, *it)) {
+            BazelBlob bazel_blob{ArtifactDigestFactory::ToBazel(it->digest),
+                                 it->data,
+                                 it->is_exec};
+            if (not cas_->UpdateSingleBlob(instance_name_, bazel_blob)) {
                 return false;
             }
             blobs.erase(it);
         }
         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_, blobs) == blobs.size();
+        return cas_->BatchUpdateBlobs(instance_name_, bazel_blobs) ==
+               bazel_blobs.size();
 
     } catch (...) {
         Logger::Log(LogLevel::Warning, "Unknown exception");
@@ -110,11 +124,13 @@ auto BazelNetwork::DoUploadBlobs(std::unordered_set<BazelBlob> blobs) noexcept
     }
 }
 
-auto BazelNetwork::UploadBlobs(std::unordered_set<BazelBlob>&& blobs,
+auto BazelNetwork::UploadBlobs(std::unordered_set<ArtifactBlob>&& blobs,
                                bool skip_find_missing) noexcept -> bool {
     if (not skip_find_missing) {
-        auto const back_map = BackMap<bazel_re::Digest, BazelBlob>::Make(
-            &blobs, [](BazelBlob const& blob) { return blob.digest; });
+        auto const back_map = BackMap<bazel_re::Digest, ArtifactBlob>::Make(
+            &blobs, [](ArtifactBlob const& blob) {
+                return ArtifactDigestFactory::ToBazel(blob.digest);
+            });
         if (not back_map.has_value()) {
             return false;
         }
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network.hpp b/src/buildtool/execution_api/remote/bazel/bazel_network.hpp
index fc8e459b..0d892a55 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_network.hpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_network.hpp
@@ -27,8 +27,8 @@
 #include "src/buildtool/common/remote/port.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/bazel_msg/bazel_common.hpp"
+#include "src/buildtool/execution_api/common/artifact_blob_container.hpp"
 #include "src/buildtool/execution_api/remote/bazel/bazel_ac_client.hpp"
 #include "src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp"
 #include "src/buildtool/execution_api/remote/bazel/bazel_execution_client.hpp"
@@ -72,7 +72,7 @@ class BazelNetwork {
     /// \param blobs              The blobs to upload
     /// \param skip_find_missing  Skip finding missing blobs, just upload all
     /// \returns True if upload was successful, false otherwise
-    [[nodiscard]] auto UploadBlobs(std::unordered_set<BazelBlob>&& blobs,
+    [[nodiscard]] auto UploadBlobs(std::unordered_set<ArtifactBlob>&& blobs,
                                    bool skip_find_missing = false) noexcept
         -> bool;
 
@@ -99,7 +99,7 @@ class BazelNetwork {
     HashFunction const& hash_function_;
 
     [[nodiscard]] auto DoUploadBlobs(
-        std::unordered_set<BazelBlob> blobs) noexcept -> bool;
+        std::unordered_set<ArtifactBlob> blobs) noexcept -> bool;
 };
 
 #endif  // INCLUDED_SRC_BUILDTOOL_EXECUTION_API_REMOTE_BAZEL_BAZEL_NETWORK_HPP
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_response.cpp b/src/buildtool/execution_api/remote/bazel/bazel_response.cpp
index 81f8310d..e4bc9248 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_response.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_response.cpp
@@ -26,11 +26,10 @@
 #include "src/buildtool/common/artifact.hpp"
 #include "src/buildtool/common/artifact_digest.hpp"
 #include "src/buildtool/common/artifact_digest_factory.hpp"
-#include "src/buildtool/common/bazel_digest_factory.hpp"
 #include "src/buildtool/common/protocol_traits.hpp"
 #include "src/buildtool/crypto/hash_function.hpp"
-#include "src/buildtool/execution_api/bazel_msg/bazel_blob_container.hpp"
 #include "src/buildtool/execution_api/bazel_msg/bazel_msg_factory.hpp"
+#include "src/buildtool/execution_api/common/artifact_blob_container.hpp"
 #include "src/buildtool/execution_api/common/common_api.hpp"
 #include "src/buildtool/execution_api/common/content_blob_container.hpp"
 #include "src/buildtool/execution_api/remote/bazel/bazel_network_reader.hpp"
@@ -44,7 +43,7 @@ namespace {
 
 auto ProcessDirectoryMessage(HashFunction hash_function,
                              bazel_re::Directory const& dir) noexcept
-    -> expected<BazelBlob, std::string> {
+    -> expected<ArtifactBlob, std::string> {
     // in compatible mode: check validity of all symlinks
     for (auto const& link : dir.symlinks()) {
         if (not PathIsNonUpwards(link.target())) {
@@ -53,11 +52,11 @@ auto ProcessDirectoryMessage(HashFunction hash_function,
         }
     }
     auto data = dir.SerializeAsString();
-    auto digest =
-        BazelDigestFactory::HashDataAs<ObjectType::File>(hash_function, data);
-    return BazelBlob{std::move(digest),
-                     std::move(data),
-                     /*is_exec=*/false};
+    auto digest = ArtifactDigestFactory::HashDataAs<ObjectType::File>(
+        hash_function, data);
+    return ArtifactBlob{std::move(digest),
+                        std::move(data),
+                        /*is_exec=*/false};
 }
 
 }  // namespace
@@ -282,11 +281,12 @@ auto BazelResponse::Populate() noexcept -> std::optional<std::string> {
 auto BazelResponse::UploadTreeMessageDirectories(
     bazel_re::Tree const& tree) const -> expected<ArtifactDigest, std::string> {
     auto const upload_callback =
-        [&network = *network_](std::unordered_set<BazelBlob>&& blobs) -> bool {
+        [&network =
+             *network_](std::unordered_set<ArtifactBlob>&& blobs) -> bool {
         return network.UploadBlobs(std::move(blobs));
     };
     auto const hash_function = network_->GetHashFunction();
-    std::unordered_set<BazelBlob> dir_blobs{};
+    std::unordered_set<ArtifactBlob> dir_blobs{};
 
     auto rootdir_blob = ProcessDirectoryMessage(hash_function, tree.root());
     if (not rootdir_blob) {
@@ -294,7 +294,7 @@ auto BazelResponse::UploadTreeMessageDirectories(
     }
     auto const root_digest = rootdir_blob->digest;
     // store or upload rootdir blob, taking maximum transfer size into account
-    if (not UpdateContainerAndUpload<bazel_re::Digest>(
+    if (not UpdateContainerAndUpload<ArtifactDigest>(
             &dir_blobs,
             *std::move(rootdir_blob),
             /*exception_is_fatal=*/false,
@@ -310,7 +310,7 @@ auto BazelResponse::UploadTreeMessageDirectories(
             return unexpected{std::move(blob).error()};
         }
         auto const blob_digest = blob->digest;
-        if (not UpdateContainerAndUpload<bazel_re::Digest>(
+        if (not UpdateContainerAndUpload<ArtifactDigest>(
                 &dir_blobs,
                 *std::move(blob),
                 /*exception_is_fatal=*/false,
@@ -327,7 +327,5 @@ auto BazelResponse::UploadTreeMessageDirectories(
             fmt::format("failed to upload blobs for Tree with root digest {}",
                         root_digest.hash())};
     }
-    return ArtifactDigestFactory::FromBazel(hash_function.GetType(),
-                                            root_digest)
-        .value();  // must succeed all the time
+    return root_digest;
 }