diff options
Diffstat (limited to 'src')
4 files changed, 32 insertions, 9 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 347fa204..566b9a29 100644 --- a/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp +++ b/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp @@ -233,6 +233,8 @@ auto BazelCasClient::BatchReadBlobs( return result; } + auto const max_content_size = GetMaxBatchTransferSize(instance_name); + auto const back_map = BackMap<bazel_re::Digest, ArtifactDigest>::Make( &blobs, ArtifactDigestFactory::ToBazel); if (not back_map.has_value()) { @@ -257,7 +259,7 @@ auto BazelCasClient::BatchReadBlobs( it, back_map->GetKeys().end(), MessageLimits::kMaxGrpcLength, - MessageLimits::kMaxGrpcLength); + max_content_size); logger_.Emit(LogLevel::Trace, "BatchReadBlobs - Request size: {} bytes\n", request.ByteSizeLong()); @@ -570,6 +572,8 @@ auto BazelCasClient::BatchUpdateBlobs(std::string const& instance_name, return 0; } + auto const max_content_size = GetMaxBatchTransferSize(instance_name); + auto request_creator = [&instance_name](ArtifactBlob const& blob) { bazel_re::BatchUpdateBlobsRequest request; request.set_instance_name(instance_name); @@ -591,7 +595,7 @@ auto BazelCasClient::BatchUpdateBlobs(std::string const& instance_name, it, blobs.end(), MessageLimits::kMaxGrpcLength, - MessageLimits::kMaxGrpcLength); + max_content_size); logger_.Emit(LogLevel::Trace, "BatchUpdateBlobs - Request size: {} bytes\n", request.ByteSizeLong()); @@ -686,6 +690,11 @@ auto BazelCasClient::BatchUpdateBlobs(std::string const& instance_name, return updated.size(); } +auto BazelCasClient::GetMaxBatchTransferSize( + std::string const& instance_name) const noexcept -> std::size_t { + return capabilities_.GetCapabilities(instance_name)->MaxBatchTransferSize; +} + auto BazelCasClient::CreateGetTreeRequest( std::string const& instance_name, bazel_re::Digest const& root_digest, 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 0e719441..67428364 100644 --- a/src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp +++ b/src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp @@ -143,10 +143,13 @@ class BazelCasClient { HashFunction hash_function, std::string const& instance_name) const noexcept -> bool; + [[nodiscard]] auto GetMaxBatchTransferSize( + std::string const& instance_name) const noexcept -> std::size_t; + private: std::unique_ptr<ByteStreamClient> stream_; RetryConfig const& retry_config_; - [[maybe_unused]] BazelCapabilitiesClient const& capabilities_; + BazelCapabilitiesClient const& capabilities_; std::unique_ptr<bazel_re::ContentAddressableStorage::Stub> stub_; Logger logger_{"RemoteCasClient"}; diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network_reader.cpp b/src/buildtool/execution_api/remote/bazel/bazel_network_reader.cpp index 239595a2..1464c134 100644 --- a/src/buildtool/execution_api/remote/bazel/bazel_network_reader.cpp +++ b/src/buildtool/execution_api/remote/bazel/bazel_network_reader.cpp @@ -246,14 +246,21 @@ auto BazelNetworkReader::Validate(ArtifactBlob const& blob) const noexcept return rehashed == blob.digest; } +auto BazelNetworkReader::GetMaxBatchTransferSize() const noexcept + -> std::size_t { + return cas_.GetMaxBatchTransferSize(instance_name_); +} + namespace { [[nodiscard]] auto FindCurrentIterator( + std::size_t content_limit, std::vector<ArtifactDigest>::const_iterator const& begin, std::vector<ArtifactDigest>::const_iterator const& end) noexcept { - auto it = std::find_if(begin, end, [](ArtifactDigest const& digest) { - auto const size = digest.size(); - return size == 0 or size > MessageLimits::kMaxGrpcLength; - }); + auto it = + std::find_if(begin, end, [content_limit](ArtifactDigest const& digest) { + auto const size = digest.size(); + return size == 0 or size > content_limit; + }); if (it == begin and begin != end) { ++it; } @@ -266,7 +273,8 @@ BazelNetworkReader::IncrementalReader::Iterator::Iterator( std::vector<ArtifactDigest>::const_iterator begin, std::vector<ArtifactDigest>::const_iterator end) noexcept : owner_{owner}, begin_{begin}, end_{end} { - current_ = FindCurrentIterator(begin_, end_); + current_ = + FindCurrentIterator(owner_.GetMaxBatchTransferSize(), begin_, end_); } auto BazelNetworkReader::IncrementalReader::Iterator::operator*() const noexcept @@ -286,6 +294,7 @@ auto BazelNetworkReader::IncrementalReader::Iterator::operator*() const noexcept auto BazelNetworkReader::IncrementalReader::Iterator::operator++() noexcept -> Iterator& { begin_ = current_; - current_ = FindCurrentIterator(begin_, end_); + current_ = + FindCurrentIterator(owner_.GetMaxBatchTransferSize(), begin_, end_); return *this; } diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network_reader.hpp b/src/buildtool/execution_api/remote/bazel/bazel_network_reader.hpp index c87b2984..e1a75df9 100644 --- a/src/buildtool/execution_api/remote/bazel/bazel_network_reader.hpp +++ b/src/buildtool/execution_api/remote/bazel/bazel_network_reader.hpp @@ -95,6 +95,8 @@ class BazelNetworkReader final { std::vector<ArtifactDigest> const& digests) const noexcept -> std::vector<ArtifactBlob>; + [[nodiscard]] auto GetMaxBatchTransferSize() const noexcept -> std::size_t; + [[nodiscard]] auto Validate(ArtifactBlob const& blob) const noexcept -> bool; }; |