diff options
Diffstat (limited to 'src')
4 files changed, 58 insertions, 66 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 5bd1d3b0..fa34f812 100644 --- a/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp +++ b/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp @@ -33,6 +33,7 @@ #include "src/buildtool/common/remote/retry.hpp" #include "src/buildtool/common/remote/retry_config.hpp" #include "src/buildtool/execution_api/common/bytestream_utils.hpp" +#include "src/buildtool/execution_api/common/content_blob_container.hpp" #include "src/buildtool/execution_api/common/execution_common.hpp" #include "src/buildtool/execution_api/common/message_limits.hpp" #include "src/buildtool/file_system/object_type.hpp" @@ -481,9 +482,9 @@ auto BazelCasClient::FindMissingBlobs(std::string const& instance_name, auto BazelCasClient::BatchUpdateBlobs( std::string const& instance_name, - std::vector<gsl::not_null<BazelBlob const*>>::const_iterator const& begin, - std::vector<gsl::not_null<BazelBlob const*>>::const_iterator const& end) - const noexcept -> std::size_t { + std::unordered_set<BazelBlob>::const_iterator const& begin, + std::unordered_set<BazelBlob>::const_iterator const& end) const noexcept + -> std::size_t { if (begin == end) { return 0; } @@ -496,9 +497,9 @@ auto BazelCasClient::BatchUpdateBlobs( end, "BatchUpdateBlobs", [](bazel_re::BatchUpdateBlobsRequest* request, - BazelBlob const* x) { + BazelBlob const& x) { *(request->add_requests()) = - BazelCasClient::CreateUpdateBlobsSingleRequest(*x); + BazelCasClient::CreateUpdateBlobsSingleRequest(x); }); result.reserve(std::distance(begin, end)); auto batch_update_blobs = @@ -550,8 +551,8 @@ auto BazelCasClient::BatchUpdateBlobs( logger_.Emit(LogLevel::Trace, [begin, end, &result]() { std::ostringstream oss{}; oss << "upload blobs" << std::endl; - std::for_each(begin, end, [&oss](BazelBlob const* blob) { - oss << fmt::format(" - {}", blob->digest.hash()) << std::endl; + std::for_each(begin, end, [&oss](BazelBlob const& blob) { + oss << fmt::format(" - {}", blob.digest.hash()) << std::endl; }); oss << "received blobs" << std::endl; std::for_each( @@ -571,12 +572,12 @@ auto BazelCasClient::BatchUpdateBlobs( logger_.Emit(LogLevel::Trace, "Retrying with missing blobs"); auto received_digests = std::unordered_set<bazel_re::Digest>{result.begin(), result.end()}; - auto missing_blobs = std::vector<gsl::not_null<BazelBlob const*>>{}; + auto missing_blobs = std::unordered_set<BazelBlob>{}; missing_blobs.reserve(missing); std::for_each( begin, end, [&received_digests, &missing_blobs](auto const& blob) { - if (not received_digests.contains(blob->digest)) { - missing_blobs.emplace_back(blob); + if (not received_digests.contains(blob.digest)) { + missing_blobs.emplace(blob); } }); return result.size() + BatchUpdateBlobs(instance_name, @@ -591,7 +592,7 @@ auto BazelCasClient::BatchUpdateBlobs( std::size_t count = 0; std::for_each( begin, end, [this, &count, &instance_name](auto const& blob) { - if (UpdateSingleBlob(instance_name, *blob)) { + if (UpdateSingleBlob(instance_name, blob)) { count += 1; } }); 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 c08adc92..b01e502f 100644 --- a/src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp +++ b/src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp @@ -64,10 +64,9 @@ class BazelCasClient { /// \returns The digests of blobs successfully updated [[nodiscard]] auto BatchUpdateBlobs( std::string const& instance_name, - std::vector<gsl::not_null<BazelBlob const*>>::const_iterator const& - begin, - std::vector<gsl::not_null<BazelBlob const*>>::const_iterator const& end) - const noexcept -> std::size_t; + std::unordered_set<BazelBlob>::const_iterator const& begin, + std::unordered_set<BazelBlob>::const_iterator const& end) const noexcept + -> std::size_t; /// \brief Read 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 c8eb76df..93413c1e 100644 --- a/src/buildtool/execution_api/remote/bazel/bazel_network.cpp +++ b/src/buildtool/execution_api/remote/bazel/bazel_network.cpp @@ -14,9 +14,7 @@ #include "src/buildtool/execution_api/remote/bazel/bazel_network.hpp" -#include <algorithm> #include <functional> -#include <iterator> #include <utility> #include "src/buildtool/execution_api/common/content_blob_container.hpp" @@ -75,34 +73,38 @@ auto BazelNetwork::BlobSpliceSupport() const noexcept -> bool { return cas_->BlobSpliceSupport(hash_function_, instance_name_); } -template <class TIter> -auto BazelNetwork::DoUploadBlobs(TIter const& first, - TIter const& last) noexcept -> bool { +auto BazelNetwork::DoUploadBlobs(std::unordered_set<BazelBlob> blobs) noexcept + -> bool { + if (blobs.empty()) { + return true; + } + try { - // Partition the blobs according to their size. The first group collects - // all the blobs that can be uploaded in batch, the second group gathers - // blobs whose size exceeds the kMaxBatchTransferSize threshold. - // - // The blobs belonging to the second group are uploaded via the - // bytestream api. - std::vector<gsl::not_null<BazelBlob const*>> sorted; - sorted.reserve(std::distance(first, last)); - std::transform( - first, last, std::back_inserter(sorted), [](BazelBlob const& b) { - return &b; - }); - - auto it = std::stable_partition( - sorted.begin(), sorted.end(), [](BazelBlob const* x) { - return x->data->size() <= kMaxBatchTransferSize; - }); - auto digests_count = - cas_->BatchUpdateBlobs(instance_name_, sorted.begin(), it); - - return digests_count == std::distance(sorted.begin(), it) && - std::all_of(it, sorted.end(), [this](BazelBlob const* x) { - return cas_->UpdateSingleBlob(instance_name_, *x); - }); + // Partition the blobs according to their size. + // The first group collects all the blobs that must use bytestream api + // because of their size: + using IteratorType = decltype(blobs)::iterator; + std::vector<IteratorType> to_stream; + to_stream.reserve(blobs.size()); + for (auto it = blobs.begin(); it != blobs.end(); ++it) { + if (it->data->size() > kMaxBatchTransferSize) { + to_stream.push_back(it); + } + } + + for (auto const& it : to_stream) { + if (not cas_->UpdateSingleBlob(instance_name_, *it)) { + return false; + } + blobs.erase(it); + } + to_stream.clear(); + + // After uploading via stream api, only small blobs that may be uploaded + // using batch are in the container: + return cas_->BatchUpdateBlobs( + instance_name_, blobs.begin(), blobs.end()) == blobs.size(); + } catch (...) { Logger::Log(LogLevel::Warning, "Unknown exception"); return false; @@ -111,26 +113,17 @@ auto BazelNetwork::DoUploadBlobs(TIter const& first, auto BazelNetwork::UploadBlobs(std::unordered_set<BazelBlob>&& blobs, bool skip_find_missing) noexcept -> bool { - if (skip_find_missing) { - return DoUploadBlobs(blobs.begin(), blobs.end()); - } - - auto const back_map = BackMap<bazel_re::Digest, BazelBlob>::Make( - &blobs, [](BazelBlob const& blob) { return blob.digest; }); - if (not back_map.has_value()) { - return false; - } - - // find digests of blobs missing in CAS - auto missing_digests = - cas_->FindMissingBlobs(instance_name_, back_map->GetKeys()); - - if (not missing_digests.empty()) { - // update missing blobs - auto missing_blobs = back_map->GetValues(missing_digests); - return DoUploadBlobs(missing_blobs.begin(), missing_blobs.end()); + if (not skip_find_missing) { + auto const back_map = BackMap<bazel_re::Digest, BazelBlob>::Make( + &blobs, [](BazelBlob const& blob) { return blob.digest; }); + if (not back_map.has_value()) { + return false; + } + + // back_map becomes invalid after this call: + blobs = back_map->GetValues(FindMissingBlobs(back_map->GetKeys())); } - return true; + return DoUploadBlobs(std::move(blobs)); } auto BazelNetwork::ExecuteBazelActionSync( diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network.hpp b/src/buildtool/execution_api/remote/bazel/bazel_network.hpp index cd4dd1dc..fc8e459b 100644 --- a/src/buildtool/execution_api/remote/bazel/bazel_network.hpp +++ b/src/buildtool/execution_api/remote/bazel/bazel_network.hpp @@ -98,9 +98,8 @@ class BazelNetwork { ExecutionConfiguration exec_config_{}; HashFunction const& hash_function_; - template <class TIter> - [[nodiscard]] auto DoUploadBlobs(TIter const& first, - TIter const& last) noexcept -> bool; + [[nodiscard]] auto DoUploadBlobs( + std::unordered_set<BazelBlob> blobs) noexcept -> bool; }; #endif // INCLUDED_SRC_BUILDTOOL_EXECUTION_API_REMOTE_BAZEL_BAZEL_NETWORK_HPP |