Fix root directory upload to honor tree invariant.

2 files changed, 101 insertions, 16 deletions

diff --git a/src/buildtool/execution_api/remote/bazel/bazel_api.cpp b/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
index 93e1be2a..f1442111 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_api.cpp
@@ -17,10 +17,14 @@
 #include <algorithm>
 #include <map>
 #include <memory>
+#include <optional>
 #include <string>
+#include <unordered_map>
 #include <vector>
 
+#include "fmt/core.h"
 #include "src/buildtool/common/bazel_types.hpp"
+#include "src/buildtool/compatibility/compatibility.hpp"
 #include "src/buildtool/execution_api/bazel_msg/bazel_blob.hpp"
 #include "src/buildtool/execution_api/bazel_msg/bazel_blob_container.hpp"
 #include "src/buildtool/execution_api/bazel_msg/bazel_common.hpp"
@@ -235,31 +239,108 @@ auto BazelApi::CreateAction(
     return network_->UploadBlobs(blobs, skip_find_missing);
 }
 
+/// NOLINTNEXTLINE(misc-no-recursion)
+[[nodiscard]] auto BazelApi::UploadBlobTree(
+    BlobTreePtr const& blob_tree) noexcept -> bool {
+
+    // Create digest list from blobs for batch availability check.
+    std::vector<bazel_re::Digest> digests;
+    digests.reserve(blob_tree->size());
+    std::unordered_map<bazel_re::Digest, BlobTreePtr> tree_map;
+    for (auto const& node : *blob_tree) {
+        auto digest = node->Blob().digest;
+        digests.emplace_back(digest);
+        try {
+            tree_map.emplace(std::move(digest), node);
+        } catch (...) {
+            return false;
+        }
+    }
+
+    // Find missing digests.
+    auto missing_digests = network_->IsAvailable(digests);
+
+    // Process missing blobs.
+    BlobContainer container;
+    for (auto const& digest : missing_digests) {
+        if (auto it = tree_map.find(digest); it != tree_map.end()) {
+            auto const& node = it->second;
+            // Process trees.
+            if (node->IsTree()) {
+                if (not UploadBlobTree(node)) {
+                    return false;
+                }
+            }
+            // Store blob.
+            try {
+                container.Emplace(node->Blob());
+            } catch (...) {
+                return false;
+            }
+        }
+    }
+
+    return network_->UploadBlobs(container, /*skip_find_missing=*/true);
+}
+
 [[nodiscard]] auto BazelApi::UploadTree(
     std::vector<DependencyGraph::NamedArtifactNodePtr> const&
         artifacts) noexcept -> std::optional<ArtifactDigest> {
-    BlobContainer blobs{};
-    auto digest = BazelMsgFactory::CreateDirectoryDigestFromTree(
-        artifacts,
-        [&blobs](BazelBlob&& blob) { blobs.Emplace(std::move(blob)); });
-    if (not digest) {
-        Logger::Log(LogLevel::Debug, "failed to create digest for tree.");
+    auto build_root = DirectoryTree::FromNamedArtifacts(artifacts);
+    if (not build_root) {
+        Logger::Log(LogLevel::Debug,
+                    "failed to create build root from artifacts.");
         return std::nullopt;
     }
 
-    Logger::Log(LogLevel::Trace, [&digest]() {
-        std::ostringstream oss{};
-        oss << "upload root directory" << std::endl;
-        oss << fmt::format(" - root digest: {}", digest->hash()) << std::endl;
-        return oss.str();
-    });
+    if (Compatibility::IsCompatible()) {
+        BlobContainer blobs{};
+        auto digest = BazelMsgFactory::CreateDirectoryDigestFromTree(
+            *build_root,
+            [&blobs](BazelBlob&& blob) { blobs.Emplace(std::move(blob)); });
+        if (not digest) {
+            Logger::Log(LogLevel::Debug,
+                        "failed to create digest for build root.");
+            return std::nullopt;
+        }
+        Logger::Log(LogLevel::Trace, [&digest]() {
+            std::ostringstream oss{};
+            oss << "upload root directory" << std::endl;
+            oss << fmt::format(" - root digest: {}", digest->hash())
+                << std::endl;
+            return oss.str();
+        });
+        if (not Upload(blobs, /*skip_find_missing=*/false)) {
+            Logger::Log(LogLevel::Debug,
+                        "failed to upload blobs for build root.");
+            return std::nullopt;
+        }
+        return ArtifactDigest{*digest};
+    }
 
-    if (not Upload(blobs, /*skip_find_missing=*/false)) {
-        Logger::Log(LogLevel::Debug, "failed to upload blobs for tree.");
+    auto blob_tree = BlobTree::FromDirectoryTree(*build_root);
+    if (not blob_tree) {
+        Logger::Log(LogLevel::Debug,
+                    "failed to create blob tree for build root.");
         return std::nullopt;
     }
-
-    return ArtifactDigest{*digest};
+    auto tree_blob = (*blob_tree)->Blob();
+    // Upload blob tree if tree is not available at the remote side (content
+    // first).
+    if (not network_->IsAvailable(tree_blob.digest)) {
+        if (not UploadBlobTree(*blob_tree)) {
+            Logger::Log(LogLevel::Debug,
+                        "failed to upload blob tree for build root.");
+            return std::nullopt;
+        }
+        if (not Upload(BlobContainer{{tree_blob}},
+                       /*skip_find_missing=*/true)) {
+            Logger::Log(LogLevel::Debug,
+                        "failed to upload tree blob for build root.");
+            return std::nullopt;
+        }
+    }
+    return ArtifactDigest{tree_blob.digest};
 }
 
 [[nodiscard]] auto BazelApi::IsAvailable(
diff --git a/src/buildtool/execution_api/remote/bazel/bazel_api.hpp b/src/buildtool/execution_api/remote/bazel/bazel_api.hpp
index 7f7a0045..a435824e 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_api.hpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_api.hpp
@@ -23,6 +23,7 @@
 #include "gsl-lite/gsl-lite.hpp"
 #include "src/buildtool/common/artifact_digest.hpp"
 #include "src/buildtool/execution_api/bazel_msg/bazel_common.hpp"
+#include "src/buildtool/execution_api/bazel_msg/blob_tree.hpp"
 #include "src/buildtool/execution_api/common/execution_api.hpp"
 #include "src/buildtool/execution_api/remote/config.hpp"
 
@@ -81,6 +82,9 @@ class BazelApi final : public IExecutionApi {
 
   private:
     std::shared_ptr<BazelNetwork> network_;
+
+    [[nodiscard]] auto UploadBlobTree(BlobTreePtr const& blob_tree) noexcept
+        -> bool;
 };
 
 #endif  // INCLUDED_SRC_BUILDTOOL_EXECUTION_API_REMOTE_BAZEL_BAZEL_API_HPP