From 952c60a8ecc4318fbd20a01fe642b6629ddf42cc Mon Sep 17 00:00:00 2001 From: Maksim Denisov Date: Wed, 29 Jan 2025 16:49:02 +0100 Subject: CommonApi: Remove template parameter from UploadAndUpdateContainer ...since it works with ArtifactBlobs only. --- src/buildtool/execution_api/common/common_api.cpp | 59 +++++++++++++++++++++-- 1 file changed, 56 insertions(+), 3 deletions(-) (limited to 'src/buildtool/execution_api/common/common_api.cpp') diff --git a/src/buildtool/execution_api/common/common_api.cpp b/src/buildtool/execution_api/common/common_api.cpp index 8bd8f0e0..3a713fc4 100644 --- a/src/buildtool/execution_api/common/common_api.cpp +++ b/src/buildtool/execution_api/common/common_api.cpp @@ -28,8 +28,10 @@ #include #include "fmt/core.h" -#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/common/message_limits.hpp" #include "src/buildtool/file_system/object_type.hpp" +#include "src/buildtool/logging/log_level.hpp" auto CommonRetrieveToFds( std::vector const& artifacts_info, @@ -123,7 +125,7 @@ auto CommonUploadBlobTree(BlobTreePtr const& blob_tree, } // Optimize store & upload by taking into account the maximum // transfer size. - if (not UpdateContainerAndUpload( + if (not UpdateContainerAndUpload( &container, node->Blob(), /*exception_is_fatal=*/false, @@ -148,7 +150,7 @@ auto CommonUploadTreeCompatible( // Store and upload blobs, taking into account the maximum transfer size. auto digest = BazelMsgFactory::CreateDirectoryDigestFromTree( build_root, resolve_links, [&blobs, &api](ArtifactBlob&& blob) { - return UpdateContainerAndUpload( + return UpdateContainerAndUpload( &blobs, std::move(blob), /*exception_is_fatal=*/false, @@ -202,3 +204,54 @@ auto CommonUploadTreeNative(IExecutionApi const& api, } return tree_blob.digest; } + +auto UpdateContainerAndUpload( + gsl::not_null*> const& container, + ArtifactBlob&& blob, + bool exception_is_fatal, + std::function&&)> const& uploader, + Logger const* logger) noexcept -> bool { + // Optimize upload of blobs with respect to the maximum transfer limit, such + // that we never store unnecessarily more data in the container than we need + // per remote transfer. + try { + if (blob.data->size() > kMaxBatchTransferSize) { + // large blobs use individual stream upload + if (not uploader( + std::unordered_set{{std::move(blob)}})) { + return false; + } + } + else { + if (not container->contains(blob)) { + std::size_t content_size = 0; + for (auto const& blob : *container) { + content_size += blob.data->size(); + } + + if (content_size + blob.data->size() > kMaxBatchTransferSize) { + // swap away from original container to allow move during + // upload + std::unordered_set tmp_container{}; + std::swap(*container, tmp_container); + // if we would surpass the transfer limit, upload the + // current container and clear it before adding more blobs + if (not uploader(std::move(tmp_container))) { + return false; + } + } + // add current blob to container + container->emplace(std::move(blob)); + } + } + } catch (std::exception const& ex) { + if (exception_is_fatal) { + Logger::Log(logger, + LogLevel::Error, + "failed to emplace blob with\n:{}", + ex.what()); + } + return false; + } + return true; // success! +} -- cgit v1.2.3