diff options
| -rw-r--r-- | src/buildtool/execution_api/remote/bazel/bazel_network.cpp | 65 |
1 files changed, 20 insertions, 45 deletions
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network.cpp b/src/buildtool/execution_api/remote/bazel/bazel_network.cpp index de88c79b..d6a5f78d 100644 --- a/src/buildtool/execution_api/remote/bazel/bazel_network.cpp +++ b/src/buildtool/execution_api/remote/bazel/bazel_network.cpp @@ -163,53 +163,28 @@ template <class T_Iter> auto BazelNetwork::DoUploadBlobs(T_Iter const& first, T_Iter const& last) noexcept -> bool { try { - auto num_blobs = gsl::narrow<std::size_t>(std::distance(first, last)); - - std::vector<bazel_re::Digest> digests{}; - digests.reserve(num_blobs); - - auto begin = first; - auto current = first; - std::size_t transfer_size{}; - while (current != last) { - auto const& blob = *current; - transfer_size += blob.data.size(); - if (transfer_size > kMaxBatchTransferSize) { - if (begin == current) { - if (cas_->UpdateSingleBlob(instance_name_, blob)) { - digests.emplace_back(blob.digest); - } - ++current; - } - else { - for (auto& digest : cas_->BatchUpdateBlobs( - instance_name_, begin, current)) { - digests.emplace_back(std::move(digest)); - } - } - begin = current; - transfer_size = 0; - } - else { - ++current; - } - } - if (begin != current) { - for (auto& digest : - cas_->BatchUpdateBlobs(instance_name_, begin, current)) { - digests.emplace_back(std::move(digest)); - } - } - - if (digests.size() == num_blobs) { - return true; - } + // 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<typename T_Iter::value_type> sorted(first, last); + auto it = std::stable_partition( + sorted.begin(), sorted.end(), [](auto const& x) { + return x.data.size() <= kMaxBatchTransferSize; + }); + auto digests = + cas_->BatchUpdateBlobs(instance_name_, sorted.begin(), it); + + return digests.size() == std::distance(sorted.begin(), it) && + std::all_of(it, sorted.end(), [this](auto const& x) { + return cas_->UpdateSingleBlob(instance_name_, x); + }); } catch (...) { - Logger::Log(LogLevel::Warning, "unknonwn exception"); + Logger::Log(LogLevel::Error, "Unknown exception"); + return false; } - - Logger::Log(LogLevel::Warning, "Failed to update all blobs"); - return false; } auto BazelNetwork::UploadBlobs(BlobContainer const& blobs, |