BazelNetwork: Use bytestream for reading unknown size blobs

... otherwise actual blob size might exceed the maximum
transfer size of the CAS client. Therefore, we always have
to use the bytestream client if the size is unknown.

2 files changed, 45 insertions, 2 deletions

diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network.cpp b/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
index a5574c77..bdb4f3be 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
@@ -191,8 +191,11 @@ auto BazelNetwork::BlobReader::Next() noexcept -> std::vector<BazelBlob> {
     std::vector<BazelBlob> blobs{};
 
     while (current_ != ids_.end()) {
-        size += gsl::narrow<std::size_t>(current_->size_bytes());
-        if (size > kMaxBatchTransferSize) {
+        auto blob_size = gsl::narrow<std::size_t>(current_->size_bytes());
+        size += blob_size;
+        // read if size is 0 (unknown) or exceeds transfer size
+        if (blob_size == 0 or size > kMaxBatchTransferSize) {
+            // perform read of range [begin_, current_)
             if (begin_ == current_) {
                 auto blob = cas_->ReadSingleBlob(instance_name_, *begin_);
                 if (blob) {
diff --git a/test/buildtool/execution_api/bazel/bazel_network.test.cpp b/test/buildtool/execution_api/bazel/bazel_network.test.cpp
index b67013a0..28bfcdc2 100644
--- a/test/buildtool/execution_api/bazel/bazel_network.test.cpp
+++ b/test/buildtool/execution_api/bazel/bazel_network.test.cpp
@@ -43,3 +43,43 @@ TEST_CASE("Bazel network: write/read blobs", "[execution_api]") {
     CHECK(blobs[3].data == content_bar);
     CHECK(blobs[4].data == content_foo);
 }
+
+TEST_CASE("Bazel network: read blobs with unknown size", "[execution_api]") {
+    if (Compatibility::IsCompatible()) {
+        // only supported in native mode
+        return;
+    }
+
+    auto const& info = RemoteExecutionConfig::RemoteAddress();
+    std::string instance_name{"remote-execution"};
+    auto network = BazelNetwork{instance_name, info->host, info->port, {}};
+
+    std::string content_foo("foo");
+    std::string content_bar(kLargeSize, 'x');  // single larger blob
+
+    BazelBlob foo{ArtifactDigest::Create(content_foo), content_foo};
+    BazelBlob bar{ArtifactDigest::Create(content_bar), content_bar};
+
+    // Upload blobs
+    REQUIRE(network.UploadBlobs(BlobContainer{{foo, bar}}));
+
+    // Set size to unknown
+    foo.digest.set_size_bytes(0);
+    bar.digest.set_size_bytes(0);
+
+    // Read blobs
+    auto reader = network.ReadBlobs({foo.digest, bar.digest});
+    std::vector<BazelBlob> blobs{};
+    while (true) {
+        auto next = reader.Next();
+        if (next.empty()) {
+            break;
+        }
+        blobs.insert(blobs.end(), next.begin(), next.end());
+    }
+
+    // Check order maintained
+    REQUIRE(blobs.size() == 2);
+    CHECK(blobs[0].data == content_foo);
+    CHECK(blobs[1].data == content_bar);
+}