Implement blob splicing protocol at just client side

6 files changed, 166 insertions, 0 deletions

diff --git a/src/buildtool/execution_api/remote/bazel/bazel_api.cpp b/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
index 157016f7..afbc9d2e 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
@@ -599,3 +599,27 @@ auto BazelApi::CreateAction(
 [[nodiscard]] auto BazelApi::BlobSplitSupport() const noexcept -> bool {
     return network_->BlobSplitSupport();
 }
+
+[[nodiscard]] auto BazelApi::SpliceBlob(
+    ArtifactDigest const& blob_digest,
+    std::vector<ArtifactDigest> const& chunk_digests) const noexcept
+    -> std::optional<ArtifactDigest> {
+    auto digests = std::vector<bazel_re::Digest>{};
+    digests.reserve(chunk_digests.size());
+    std::transform(chunk_digests.cbegin(),
+                   chunk_digests.cend(),
+                   std::back_inserter(digests),
+                   [](auto const& artifact_digest) {
+                       return static_cast<bazel_re::Digest>(artifact_digest);
+                   });
+    auto digest = network_->SpliceBlob(
+        static_cast<bazel_re::Digest>(blob_digest), digests);
+    if (not digest) {
+        return std::nullopt;
+    }
+    return ArtifactDigest{*digest};
+}
+
+[[nodiscard]] auto BazelApi::BlobSpliceSupport() const noexcept -> bool {
+    return network_->BlobSpliceSupport();
+}
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_api.hpp b/src/buildtool/execution_api/remote/bazel/bazel_api.hpp
index 2d514ecd..dcef334e 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_api.hpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_api.hpp
@@ -101,6 +101,13 @@ class BazelApi final : public IExecutionApi {
 
     [[nodiscard]] auto BlobSplitSupport() const noexcept -> bool final;
 
+    [[nodiscard]] auto SpliceBlob(
+        ArtifactDigest const& blob_digest,
+        std::vector<ArtifactDigest> const& chunk_digests) const noexcept
+        -> std::optional<ArtifactDigest> final;
+
+    [[nodiscard]] auto BlobSpliceSupport() const noexcept -> bool final;
+
   private:
     std::shared_ptr<BazelNetwork> network_;
 
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 e6830261..e034bd14 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp
@@ -14,6 +14,7 @@
 
 #include "src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp"
 
+#include <algorithm>
 #include <mutex>
 #include <shared_mutex>
 #include <sstream>
@@ -103,6 +104,72 @@ namespace {
     return supported;
 }
 
+// In order to determine whether blob splicing is supported at the remote, a
+// trial request to the remote CAS service is issued. This is just a workaround
+// until the blob splice API extension is accepted as part of the official
+// remote execution protocol. Then, the ordinary way to determine server
+// capabilities can be employed by using the capabilities service.
+[[nodiscard]] auto BlobSpliceSupport(
+    std::string const& instance_name,
+    std::unique_ptr<bazel_re::ContentAddressableStorage::Stub> const&
+        stub) noexcept -> bool {
+    // Create empty blob.
+    std::string empty_str{};
+    std::string hash = HashFunction::ComputeBlobHash(empty_str).HexString();
+    bazel_re::Digest digest{};
+    digest.set_hash(NativeSupport::Prefix(hash, false));
+    digest.set_size_bytes(empty_str.size());
+
+    // Upload empty blob.
+    grpc::ClientContext update_context{};
+    bazel_re::BatchUpdateBlobsRequest update_request{};
+    bazel_re::BatchUpdateBlobsResponse update_response{};
+    update_request.set_instance_name(instance_name);
+    auto* request = update_request.add_requests();
+    request->mutable_digest()->CopyFrom(digest);
+    request->set_data(empty_str);
+    grpc::Status update_status = stub->BatchUpdateBlobs(
+        &update_context, update_request, &update_response);
+    if (not update_status.ok()) {
+        return false;
+    }
+
+    // Request splicing empty blob.
+    grpc::ClientContext splice_context{};
+    bazel_re::SpliceBlobRequest splice_request{};
+    bazel_re::SpliceBlobResponse splice_response{};
+    splice_request.set_instance_name(instance_name);
+    splice_request.mutable_blob_digest()->CopyFrom(digest);
+    splice_request.add_chunk_digests()->CopyFrom(digest);
+    grpc::Status splice_status =
+        stub->SpliceBlob(&splice_context, splice_request, &splice_response);
+    return splice_status.ok();
+}
+
+// Cached version of blob-splice support request.
+[[nodiscard]] auto BlobSpliceSupportCached(
+    std::string const& instance_name,
+    std::unique_ptr<bazel_re::ContentAddressableStorage::Stub> const& stub,
+    Logger const* logger) noexcept -> bool {
+    static auto mutex = std::shared_mutex{};
+    static auto blob_splice_support_map =
+        std::unordered_map<std::string, bool>{};
+    {
+        auto lock = std::shared_lock(mutex);
+        if (blob_splice_support_map.contains(instance_name)) {
+            return blob_splice_support_map[instance_name];
+        }
+    }
+    auto supported = ::BlobSpliceSupport(instance_name, stub);
+    logger->Emit(LogLevel::Debug,
+                 "Blob splice support for \"{}\": {}",
+                 instance_name,
+                 supported);
+    auto lock = std::unique_lock(mutex);
+    blob_splice_support_map[instance_name] = supported;
+    return supported;
+}
+
 }  // namespace
 
 BazelCasClient::BazelCasClient(std::string const& server, Port port) noexcept
@@ -336,11 +403,47 @@ auto BazelCasClient::SplitBlob(std::string const& instance_name,
     return ProcessResponseContents<bazel_re::Digest>(response);
 }
 
+auto BazelCasClient::SpliceBlob(
+    std::string const& instance_name,
+    bazel_re::Digest const& blob_digest,
+    std::vector<bazel_re::Digest> const& chunk_digests) const noexcept
+    -> std::optional<bazel_re::Digest> {
+    if (not BlobSpliceSupportCached(instance_name, stub_, &logger_)) {
+        return std::nullopt;
+    }
+    bazel_re::SpliceBlobRequest request{};
+    request.set_instance_name(instance_name);
+    request.mutable_blob_digest()->CopyFrom(blob_digest);
+    std::copy(chunk_digests.cbegin(),
+              chunk_digests.cend(),
+              pb::back_inserter(request.mutable_chunk_digests()));
+    bazel_re::SpliceBlobResponse response{};
+    auto [ok, status] = WithRetry(
+        [this, &response, &request]() {
+            grpc::ClientContext context;
+            return stub_->SpliceBlob(&context, request, &response);
+        },
+        logger_);
+    if (not ok) {
+        LogStatus(&logger_, LogLevel::Error, status, "SpliceBlob");
+        return std::nullopt;
+    }
+    if (not response.has_blob_digest()) {
+        return std::nullopt;
+    }
+    return response.blob_digest();
+}
+
 auto BazelCasClient::BlobSplitSupport(
     std::string const& instance_name) const noexcept -> bool {
     return ::BlobSplitSupportCached(instance_name, stub_, &logger_);
 }
 
+auto BazelCasClient::BlobSpliceSupport(
+    std::string const& instance_name) const noexcept -> bool {
+    return ::BlobSpliceSupportCached(instance_name, stub_, &logger_);
+}
+
 template <class T_ForwardIter>
 auto BazelCasClient::FindMissingBlobs(std::string const& instance_name,
                                       T_ForwardIter const& start,
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 77196022..1509b359 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_cas_client.hpp
@@ -136,9 +136,23 @@ class BazelCasClient {
                                  bazel_re::Digest const& blob_digest)
         const noexcept -> std::optional<std::vector<bazel_re::Digest>>;
 
+    /// @brief Splice blob from chunks at the remote side
+    /// @param[in] instance_name    Name of the CAS instance
+    /// @param[in] blob_digest      Expected digest of the spliced blob
+    /// @param[in] chunk_digests    The chunk digests of the splitted blob
+    /// @return Whether the splice call was successful
+    [[nodiscard]] auto SpliceBlob(
+        std::string const& instance_name,
+        bazel_re::Digest const& blob_digest,
+        std::vector<bazel_re::Digest> const& chunk_digests) const noexcept
+        -> std::optional<bazel_re::Digest>;
+
     [[nodiscard]] auto BlobSplitSupport(
         std::string const& instance_name) const noexcept -> bool;
 
+    [[nodiscard]] auto BlobSpliceSupport(
+        std::string const& instance_name) const noexcept -> bool;
+
   private:
     std::unique_ptr<ByteStreamClient> stream_{};
     std::unique_ptr<bazel_re::ContentAddressableStorage::Stub> stub_;
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network.cpp b/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
index 04c13589..903bb59d 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_network.cpp
@@ -159,10 +159,21 @@ auto BazelNetwork::SplitBlob(bazel_re::Digest const& blob_digest) const noexcept
     return cas_->SplitBlob(instance_name_, blob_digest);
 }
 
+auto BazelNetwork::SpliceBlob(
+    bazel_re::Digest const& blob_digest,
+    std::vector<bazel_re::Digest> const& chunk_digests) const noexcept
+    -> std::optional<bazel_re::Digest> {
+    return cas_->SpliceBlob(instance_name_, blob_digest, chunk_digests);
+}
+
 auto BazelNetwork::BlobSplitSupport() const noexcept -> bool {
     return cas_->BlobSplitSupport(instance_name_);
 }
 
+auto BazelNetwork::BlobSpliceSupport() const noexcept -> bool {
+    return cas_->BlobSpliceSupport(instance_name_);
+}
+
 template <class T_Iter>
 auto BazelNetwork::DoUploadBlobs(T_Iter const& first,
                                  T_Iter const& last) noexcept -> bool {
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_network.hpp b/src/buildtool/execution_api/remote/bazel/bazel_network.hpp
index 0808d2e9..e2625b6c 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_network.hpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_network.hpp
@@ -78,8 +78,15 @@ class BazelNetwork {
     [[nodiscard]] auto SplitBlob(bazel_re::Digest const& blob_digest)
         const noexcept -> std::optional<std::vector<bazel_re::Digest>>;
 
+    [[nodiscard]] auto SpliceBlob(
+        bazel_re::Digest const& blob_digest,
+        std::vector<bazel_re::Digest> const& chunk_digests) const noexcept
+        -> std::optional<bazel_re::Digest>;
+
     [[nodiscard]] auto BlobSplitSupport() const noexcept -> bool;
 
+    [[nodiscard]] auto BlobSpliceSupport() const noexcept -> bool;
+
     /// \brief Uploads blobs to CAS
     /// \param blobs              The blobs to upload
     /// \param skip_find_missing  Skip finding missing blobs, just upload all