blob containers: Store and upload taking into account content size

Update logic populating containers to use the new method which
is aware of the maximum transfer limit.

2 files changed, 42 insertions, 13 deletions

diff --git a/src/buildtool/execution_api/remote/bazel/bazel_api.cpp b/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
index 0a48a748..2cba30be 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
@@ -29,6 +29,7 @@
 #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/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/stream_dumper.hpp"
 #include "src/buildtool/execution_api/common/tree_reader.hpp"
@@ -68,15 +69,19 @@ namespace {
             return false;
         }
         for (auto const& blob : blobs) {
-            try {
-                auto digest = ArtifactDigest{blob.digest};
-                auto exec = info_map.contains(digest)
-                                ? IsExecutableObject(info_map.at(digest).type)
-                                : false;
-                container.Emplace(ArtifactBlob{digest, blob.data, exec});
-            } catch (std::exception const& ex) {
-                Logger::Log(
-                    LogLevel::Warning, "failed to emplace blob: ", ex.what());
+            auto digest = ArtifactDigest{blob.digest};
+            auto exec = info_map.contains(digest)
+                            ? IsExecutableObject(info_map.at(digest).type)
+                            : false;
+            // Collect blob and upload to other CAS if transfer size reached.
+            if (not UpdateContainerAndUpload<ArtifactDigest>(
+                    &container,
+                    ArtifactBlob{std::move(digest), blob.data, exec},
+                    /*exception_is_fatal=*/true,
+                    [&api](ArtifactBlobContainer&& blobs) {
+                        return api->Upload(std::move(blobs),
+                                           /*skip_find_missing=*/true);
+                    })) {
                 return false;
             }
         }
@@ -89,7 +94,7 @@ namespace {
         return false;
     }
 
-    // Upload blobs to other CAS.
+    // Upload remaining blobs to other CAS.
     return api->Upload(std::move(container), /*skip_find_missing=*/true);
 }
 
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_response.cpp b/src/buildtool/execution_api/remote/bazel/bazel_response.cpp
index aea90ba8..df1d4736 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_response.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_response.cpp
@@ -18,7 +18,8 @@
 
 #include "gsl/gsl"
 #include "src/buildtool/compatibility/native_support.hpp"
-#include "src/buildtool/execution_api/bazel_msg/bazel_msg_factory.hpp"
+#include "src/buildtool/execution_api/bazel_msg/bazel_blob_container.hpp"
+#include "src/buildtool/execution_api/common/common_api.hpp"
 #include "src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp"
 #include "src/buildtool/logging/log_level.hpp"
 #include "src/buildtool/logging/logger.hpp"
@@ -219,16 +220,39 @@ auto BazelResponse::UploadTreeMessageDirectories(
         return std::nullopt;
     }
     auto root_digest = rootdir_blob->digest;
-    dir_blobs.Emplace(std::move(*rootdir_blob));
+    // store or upload rootdir blob, taking maximum transfer size into account
+    if (not UpdateContainerAndUpload<bazel_re::Digest>(
+            &dir_blobs,
+            std::move(*rootdir_blob),
+            /*exception_is_fatal=*/false,
+            [&network = network_](BazelBlobContainer&& blobs) {
+                return network->UploadBlobs(std::move(blobs));
+            })) {
+        Logger::Log(LogLevel::Error,
+                    "uploading Tree's Directory messages failed");
+        return std::nullopt;
+    }
 
     for (auto const& subdir : tree.children()) {
         auto subdir_blob = ProcessDirectoryMessage(subdir);
         if (not subdir_blob) {
             return std::nullopt;
         }
-        dir_blobs.Emplace(std::move(*subdir_blob));
+        // store or upload blob, taking maximum transfer size into account
+        if (not UpdateContainerAndUpload<bazel_re::Digest>(
+                &dir_blobs,
+                std::move(*subdir_blob),
+                /*exception_is_fatal=*/false,
+                [&network = network_](BazelBlobContainer&& blobs) {
+                    return network->UploadBlobs(std::move(blobs));
+                })) {
+            Logger::Log(LogLevel::Error,
+                        "uploading Tree's Directory messages failed");
+            return std::nullopt;
+        }
     }
 
+    // upload any remaining blob
     if (not network_->UploadBlobs(std::move(dir_blobs))) {
         Logger::Log(LogLevel::Error,
                     "uploading Tree's Directory messages failed");