diff options
author | Maksim Denisov <denisov.maksim@huawei.com> | 2025-02-20 12:15:26 +0100 |
---|---|---|
committer | Maksim Denisov <denisov.maksim@huawei.com> | 2025-02-21 14:46:30 +0100 |
commit | e8c157aed3fed0472d14b205c6b1dc78c393a430 (patch) | |
tree | 6ff7466196947ffe6d29ad773289c4b1f3551b9c | |
parent | 986b6e25526e38f21a6b3f11beefbb9679a200ab (diff) | |
download | justbuild-e8c157aed3fed0472d14b205c6b1dc78c393a430.tar.gz |
ByteStreamClient: Use ArtifactBlob in Write
5 files changed, 42 insertions, 39 deletions
diff --git a/src/buildtool/execution_api/remote/TARGETS b/src/buildtool/execution_api/remote/TARGETS index 43f3f381..01b0bf2b 100644 --- a/src/buildtool/execution_api/remote/TARGETS +++ b/src/buildtool/execution_api/remote/TARGETS @@ -30,6 +30,7 @@ , ["src/buildtool/execution_api/bazel_msg", "execution_config"] , ["src/buildtool/execution_api/common", "artifact_blob"] , ["src/buildtool/execution_api/common", "bytestream_utils"] + , ["src/buildtool/execution_api/common", "ids"] , ["src/buildtool/execution_api/common", "message_limits"] , ["src/buildtool/file_system", "git_repo"] , ["src/buildtool/logging", "log_level"] @@ -53,7 +54,6 @@ , ["src/buildtool/common", "protocol_traits"] , ["src/buildtool/common/remote", "retry"] , ["src/buildtool/execution_api/bazel_msg", "bazel_msg_factory"] - , ["src/buildtool/execution_api/common", "ids"] , ["src/buildtool/file_system", "object_type"] , ["src/utils/cpp", "back_map"] , ["src/utils/cpp", "gsl"] 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 d6814b08..817a0f05 100644 --- a/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp +++ b/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp @@ -33,7 +33,6 @@ #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/ids.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" @@ -383,25 +382,14 @@ auto BazelCasClient::UpdateSingleBlob(std::string const& instance_name, return oss.str(); }); - thread_local static std::string uuid{}; - if (uuid.empty()) { - auto id = CreateProcessUniqueId(); - if (not id) { - logger_.Emit(LogLevel::Debug, "Failed creating process unique id."); - return false; - } - uuid = CreateUUIDVersion4(*id); - } - auto ok = stream_->Write( - ByteStreamUtils::WriteRequest{instance_name, uuid, blob.digest}, - *blob.data); - if (not ok) { + if (not stream_->Write(instance_name, blob)) { logger_.Emit(LogLevel::Error, "Failed to write {}:{}", blob.digest.hash(), blob.digest.size()); + return false; } - return ok; + return true; } auto BazelCasClient::IncrementalReadSingleBlob(std::string const& instance_name, diff --git a/src/buildtool/execution_api/remote/bazel/bytestream_client.hpp b/src/buildtool/execution_api/remote/bazel/bytestream_client.hpp index 4343393c..452c0abe 100644 --- a/src/buildtool/execution_api/remote/bazel/bytestream_client.hpp +++ b/src/buildtool/execution_api/remote/bazel/bytestream_client.hpp @@ -31,7 +31,9 @@ #include "src/buildtool/auth/authentication.hpp" #include "src/buildtool/common/remote/client_common.hpp" #include "src/buildtool/common/remote/port.hpp" +#include "src/buildtool/execution_api/common/artifact_blob.hpp" #include "src/buildtool/execution_api/common/bytestream_utils.hpp" +#include "src/buildtool/execution_api/common/ids.hpp" #include "src/buildtool/logging/log_level.hpp" #include "src/buildtool/logging/logger.hpp" #include "src/utils/cpp/expected.hpp" @@ -111,19 +113,34 @@ class ByteStreamClient { return output; } - [[nodiscard]] auto Write(ByteStreamUtils::WriteRequest&& write_request, - std::string const& data) const noexcept -> bool { + [[nodiscard]] auto Write(std::string const& instance_name, + ArtifactBlob const& blob) const noexcept -> bool { + thread_local static std::string uuid{}; + if (uuid.empty()) { + auto id = CreateProcessUniqueId(); + if (not id) { + logger_.Emit(LogLevel::Debug, + "Failed creating process unique id."); + return false; + } + uuid = CreateUUIDVersion4(*id); + } + try { grpc::ClientContext ctx; google::bytestream::WriteResponse response{}; auto writer = stub_->Write(&ctx, &response); + auto const resource_name = + ByteStreamUtils::WriteRequest{instance_name, uuid, blob.digest} + .ToString(); + google::bytestream::WriteRequest request{}; - request.set_resource_name(std::move(write_request).ToString()); + request.set_resource_name(resource_name); request.mutable_data()->reserve(ByteStreamUtils::kChunkSize); auto const to_read = ::IncrementalReader::FromMemory( - ByteStreamUtils::kChunkSize, &data); + ByteStreamUtils::kChunkSize, &*blob.data); if (not to_read.has_value()) { logger_.Emit( LogLevel::Error, @@ -147,7 +164,8 @@ class ByteStreamClient { *request.mutable_data() = *chunk; request.set_write_offset(static_cast<int>(pos)); - request.set_finish_write(pos + chunk->size() >= data.size()); + request.set_finish_write(pos + chunk->size() >= + blob.data->size()); if (writer->Write(request)) { pos += chunk->size(); ++it; @@ -184,12 +202,12 @@ class ByteStreamClient { return false; } if (gsl::narrow<std::size_t>(response.committed_size()) != - data.size()) { + blob.data->size()) { logger_.Emit( LogLevel::Warning, "Commited size {} is different from the original one {}.", response.committed_size(), - data.size()); + blob.data->size()); return false; } return true; diff --git a/test/buildtool/execution_api/bazel/TARGETS b/test/buildtool/execution_api/bazel/TARGETS index 7b1a139e..edbaf311 100644 --- a/test/buildtool/execution_api/bazel/TARGETS +++ b/test/buildtool/execution_api/bazel/TARGETS @@ -54,8 +54,10 @@ , ["@", "grpc", "", "grpc"] , ["@", "gsl", "", "gsl"] , ["@", "src", "src/buildtool/common", "artifact_digest_factory"] + , ["@", "src", "src/buildtool/common", "common"] , ["@", "src", "src/buildtool/common/remote", "remote_common"] , ["@", "src", "src/buildtool/crypto", "hash_function"] + , ["@", "src", "src/buildtool/execution_api/common", "artifact_blob"] , ["@", "src", "src/buildtool/execution_api/common", "bytestream_utils"] , ["@", "src", "src/buildtool/execution_api/common", "common"] , ["@", "src", "src/buildtool/execution_api/common", "ids"] diff --git a/test/buildtool/execution_api/bazel/bytestream_client.test.cpp b/test/buildtool/execution_api/bazel/bytestream_client.test.cpp index 07d85671..96a95701 100644 --- a/test/buildtool/execution_api/bazel/bytestream_client.test.cpp +++ b/test/buildtool/execution_api/bazel/bytestream_client.test.cpp @@ -22,11 +22,12 @@ #include "catch2/catch_test_macros.hpp" #include "gsl/gsl" +#include "src/buildtool/common/artifact_digest.hpp" #include "src/buildtool/common/artifact_digest_factory.hpp" #include "src/buildtool/common/remote/remote_common.hpp" #include "src/buildtool/crypto/hash_function.hpp" +#include "src/buildtool/execution_api/common/artifact_blob.hpp" #include "src/buildtool/execution_api/common/bytestream_utils.hpp" -#include "src/buildtool/execution_api/common/ids.hpp" #include "src/buildtool/execution_api/remote/config.hpp" #include "src/buildtool/file_system/object_type.hpp" #include "test/utils/hermeticity/test_hash_function_type.hpp" @@ -44,7 +45,6 @@ TEST_CASE("ByteStream Client: Transfer single blob", "[execution_api]") { auto stream = ByteStreamClient{remote_config->remote_address->host, remote_config->remote_address->port, &*auth_config}; - auto uuid = CreateUUIDVersion4(*CreateProcessUniqueId()); HashFunction const hash_function{TestHashType::ReadFromEnvironment()}; @@ -56,16 +56,13 @@ TEST_CASE("ByteStream Client: Transfer single blob", "[execution_api]") { auto digest = ArtifactDigestFactory::HashDataAs<ObjectType::File>( hash_function, content); - CHECK(stream.Write( - ByteStreamUtils::WriteRequest{instance_name, uuid, digest}, - content)); + ArtifactBlob const blob{digest, content, /*is_exec=*/false}; + CHECK(stream.Write(instance_name, blob)); - SECTION("Download small blob") { - auto const data = stream.Read( - ByteStreamUtils::ReadRequest{instance_name, digest}); + auto const data = + stream.Read(ByteStreamUtils::ReadRequest{instance_name, digest}); - CHECK(data == content); - } + CHECK(data == content); } SECTION("Small blob with wrong digest") { @@ -76,10 +73,9 @@ TEST_CASE("ByteStream Client: Transfer single blob", "[execution_api]") { // Valid digest, but for a different string auto digest = ArtifactDigestFactory::HashDataAs<ObjectType::File>( hash_function, other_content); + ArtifactBlob const blob{digest, content, /*is_exec=*/false}; - CHECK(not stream.Write( - ByteStreamUtils::WriteRequest{instance_name, uuid, digest}, - content)); + CHECK(not stream.Write(instance_name, blob)); } SECTION("Upload large blob") { @@ -95,10 +91,9 @@ TEST_CASE("ByteStream Client: Transfer single blob", "[execution_api]") { // digest of "instance_nameinstance_nameinstance_..." auto digest = ArtifactDigestFactory::HashDataAs<ObjectType::File>( hash_function, content); + ArtifactBlob const blob{digest, content, /*is_exec=*/false}; - CHECK(stream.Write( - ByteStreamUtils::WriteRequest{instance_name, uuid, digest}, - content)); + CHECK(stream.Write(instance_name, blob)); SECTION("Download large blob") { auto const data = stream.Read( |