ArtifactBlob: Support incremental reading

4 files changed, 26 insertions, 4 deletions

diff --git a/src/buildtool/common/TARGETS b/src/buildtool/common/TARGETS
index b7c65fce..292fb597 100644
--- a/src/buildtool/common/TARGETS
+++ b/src/buildtool/common/TARGETS
@@ -84,7 +84,11 @@
   , "name": ["artifact_blob"]
   , "hdrs": ["artifact_blob.hpp"]
   , "srcs": ["artifact_blob.cpp"]
-  , "deps": ["common"]
+  , "deps":
+    [ "common"
+    , ["src/utils/cpp", "expected"]
+    , ["src/utils/cpp", "incremental_reader"]
+    ]
   , "private-deps": [["src/utils/cpp", "hash_combine"]]
   , "stage": ["src", "buildtool", "common"]
   }
diff --git a/src/buildtool/common/artifact_blob.cpp b/src/buildtool/common/artifact_blob.cpp
index 894d6d91..3b199e5e 100644
--- a/src/buildtool/common/artifact_blob.cpp
+++ b/src/buildtool/common/artifact_blob.cpp
@@ -21,6 +21,15 @@ auto ArtifactBlob::ReadContent() const noexcept
     return content_;
 }
 
+auto ArtifactBlob::ReadIncrementally(std::size_t chunk_size) const& noexcept
+    -> expected<IncrementalReader, std::string> {
+    if (content_ == nullptr) {
+        return unexpected<std::string>{
+            "ArtifactBlob::ReadIncrementally: missing memory source"};
+    }
+    return IncrementalReader::FromMemory(chunk_size, content_.get());
+}
+
 namespace std {
 auto hash<ArtifactBlob>::operator()(ArtifactBlob const& blob) const noexcept
     -> std::size_t {
diff --git a/src/buildtool/common/artifact_blob.hpp b/src/buildtool/common/artifact_blob.hpp
index d75a6b01..612f851f 100644
--- a/src/buildtool/common/artifact_blob.hpp
+++ b/src/buildtool/common/artifact_blob.hpp
@@ -22,6 +22,8 @@
 #include <utility>
 
 #include "src/buildtool/common/artifact_digest.hpp"
+#include "src/utils/cpp/expected.hpp"
+#include "src/utils/cpp/incremental_reader.hpp"
 
 class ArtifactBlob final {
   public:
@@ -52,6 +54,14 @@ class ArtifactBlob final {
     [[nodiscard]] auto ReadContent() const noexcept
         -> std::shared_ptr<std::string const>;
 
+    /// \brief Create an IncrementalReader that uses this ArtifactBlob's content
+    /// source.
+    /// \param chunk_size   Size of chunk, must be greater than 0.
+    /// \return Valid IncrementalReader on success or an error message on
+    /// failure.
+    [[nodiscard]] auto ReadIncrementally(std::size_t chunk_size) const& noexcept
+        -> expected<IncrementalReader, std::string>;
+
     /// \brief Set executable permission.
     void SetExecutable(bool is_executable) noexcept {
         is_executable_ = is_executable;
diff --git a/src/buildtool/execution_api/remote/bazel/bytestream_client.hpp b/src/buildtool/execution_api/remote/bazel/bytestream_client.hpp
index 99ad4e81..23a296c9 100644
--- a/src/buildtool/execution_api/remote/bazel/bytestream_client.hpp
+++ b/src/buildtool/execution_api/remote/bazel/bytestream_client.hpp
@@ -143,9 +143,8 @@ class ByteStreamClient {
             request.set_resource_name(resource_name);
             request.mutable_data()->reserve(ByteStreamUtils::kChunkSize);
 
-            auto const data_to_read = blob.ReadContent();
-            auto const to_read = ::IncrementalReader::FromMemory(
-                ByteStreamUtils::kChunkSize, data_to_read.get());
+            auto const to_read =
+                blob.ReadIncrementally(ByteStreamUtils::kChunkSize);
             if (not to_read.has_value()) {
                 logger_.Emit(
                     LogLevel::Error,