diff options
4 files changed, 33 insertions, 16 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 6f07d322..93f27ff2 100644 --- a/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp +++ b/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp @@ -415,7 +415,7 @@ auto BazelCasClient::IncrementalReadSingleBlob(std::string const& instance_name, auto BazelCasClient::ReadSingleBlob(std::string const& instance_name, ArtifactDigest const& digest) const noexcept -> std::optional<ArtifactBlob> { - return stream_->Read(instance_name, digest); + return stream_->Read(instance_name, digest, temp_space_); } auto BazelCasClient::SplitBlob(HashFunction hash_function, diff --git a/src/buildtool/execution_api/remote/bazel/bytestream_client.hpp b/src/buildtool/execution_api/remote/bazel/bytestream_client.hpp index cf85b769..89d61d1f 100644 --- a/src/buildtool/execution_api/remote/bazel/bytestream_client.hpp +++ b/src/buildtool/execution_api/remote/bazel/bytestream_client.hpp @@ -17,6 +17,7 @@ #include <cstddef> #include <cstdint> +#include <fstream> #include <memory> #include <optional> #include <string> @@ -41,6 +42,7 @@ #include "src/buildtool/logging/logger.hpp" #include "src/utils/cpp/expected.hpp" #include "src/utils/cpp/incremental_reader.hpp" +#include "src/utils/cpp/tmp_dir.hpp" /// Implements client side for google.bytestream.ByteStream service. class ByteStreamClient { @@ -105,23 +107,34 @@ class ByteStreamClient { } [[nodiscard]] auto Read(std::string const& instance_name, - ArtifactDigest const& digest) const noexcept + ArtifactDigest const& digest, + TmpDir::Ptr const& temp_space) const noexcept -> std::optional<ArtifactBlob> { - auto reader = IncrementalRead(instance_name, digest); - std::string output{}; - auto data = reader.Next(); - while (data and not data->empty()) { - output.append(data->begin(), data->end()); - data = reader.Next(); + auto temp_file = TmpDir::CreateFile(temp_space, digest.hash()); + if (temp_file == nullptr) { + return std::nullopt; } - if (not data) { + + auto reader = IncrementalRead(instance_name, digest); + try { + std::ofstream stream{temp_file->GetPath(), std::ios_base::binary}; + + auto data = reader.Next(); + while (data.has_value() and not data->empty() and stream.good()) { + stream << *data; + data = reader.Next(); + } + if (not stream.good() or not data.has_value()) { + return std::nullopt; + } + } catch (...) { return std::nullopt; } - auto blob = ArtifactBlob::FromMemory( + auto blob = ArtifactBlob::FromTempFile( HashFunction{digest.GetHashType()}, digest.IsTree() ? ObjectType::Tree : ObjectType::File, - std::move(output)); + std::move(temp_file)); if (not blob.has_value() or blob->GetDigest() != digest) { return std::nullopt; } diff --git a/test/buildtool/execution_api/bazel/TARGETS b/test/buildtool/execution_api/bazel/TARGETS index acbee81f..d5ad72c9 100644 --- a/test/buildtool/execution_api/bazel/TARGETS +++ b/test/buildtool/execution_api/bazel/TARGETS @@ -62,12 +62,13 @@ , ["@", "src", "src/buildtool/execution_api/remote", "bazel_network"] , ["@", "src", "src/buildtool/execution_api/remote", "config"] , ["@", "src", "src/buildtool/file_system", "object_type"] + , ["@", "src", "src/buildtool/storage", "config"] , ["@", "src", "src/utils/cpp", "expected"] , ["utils", "catch-main-remote-execution"] , ["utils", "execution_bazel"] , ["utils", "test_auth_config"] - , ["utils", "test_hash_function_type"] , ["utils", "test_remote_config"] + , ["utils", "test_storage_config"] ] , "stage": ["test", "buildtool", "execution_api", "bazel"] } diff --git a/test/buildtool/execution_api/bazel/bytestream_client.test.cpp b/test/buildtool/execution_api/bazel/bytestream_client.test.cpp index 49a88109..2de1f578 100644 --- a/test/buildtool/execution_api/bazel/bytestream_client.test.cpp +++ b/test/buildtool/execution_api/bazel/bytestream_client.test.cpp @@ -28,12 +28,14 @@ #include "src/buildtool/crypto/hash_function.hpp" #include "src/buildtool/execution_api/remote/config.hpp" #include "src/buildtool/file_system/object_type.hpp" +#include "src/buildtool/storage/config.hpp" #include "src/utils/cpp/expected.hpp" -#include "test/utils/hermeticity/test_hash_function_type.hpp" +#include "test/utils/hermeticity/test_storage_config.hpp" #include "test/utils/remote_execution/test_auth_config.hpp" #include "test/utils/remote_execution/test_remote_config.hpp" TEST_CASE("ByteStream Client: Transfer single blob", "[execution_api]") { + auto storage_config = TestStorageConfig::Create(); auto auth_config = TestAuthConfig::ReadFromEnvironment(); REQUIRE(auth_config); @@ -45,7 +47,8 @@ TEST_CASE("ByteStream Client: Transfer single blob", "[execution_api]") { remote_config->remote_address->port, &*auth_config}; - HashFunction const hash_function{TestHashType::ReadFromEnvironment()}; + HashFunction const hash_function = storage_config.Get().hash_function; + auto const temp_space = storage_config.Get().CreateTypedTmpDir("space"); SECTION("Upload small blob") { std::string instance_name{"remote-execution"}; @@ -59,7 +62,7 @@ TEST_CASE("ByteStream Client: Transfer single blob", "[execution_api]") { CHECK(stream.Write(instance_name, *blob)); auto const downloaded_blob = - stream.Read(instance_name, blob->GetDigest()); + stream.Read(instance_name, blob->GetDigest(), temp_space); REQUIRE(downloaded_blob.has_value()); auto const downloaded_content = downloaded_blob->ReadContent(); @@ -86,7 +89,7 @@ TEST_CASE("ByteStream Client: Transfer single blob", "[execution_api]") { SECTION("Download large blob") { auto const downloaded_blob = - stream.Read(instance_name, blob->GetDigest()); + stream.Read(instance_name, blob->GetDigest(), temp_space); REQUIRE(downloaded_blob.has_value()); auto const downloaded_content = downloaded_blob->ReadContent(); |