From e62914c686d76d48fad2f4ca019cd95bac8913bf Mon Sep 17 00:00:00 2001
From: Sascha Roloff <sascha.roloff@huawei.com>
Date: Wed, 23 Nov 2022 16:23:57 +0100
Subject: Fix root directory upload to honor tree invariant.

---
 .../execution_api/remote/bazel/bazel_api.cpp       | 113 ++++++++++++++++++---
 1 file changed, 97 insertions(+), 16 deletions(-)

(limited to 'src/buildtool/execution_api/remote/bazel/bazel_api.cpp')

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(
-- 
cgit v1.2.3