Execution server: Fix OutputDirectory generation in compatible mode

1 files changed, 151 insertions, 4 deletions

diff --git a/src/buildtool/execution_api/execution_service/execution_server.cpp b/src/buildtool/execution_api/execution_service/execution_server.cpp
index d203e733..ab0ebfc9 100644
--- a/src/buildtool/execution_api/execution_service/execution_server.cpp
+++ b/src/buildtool/execution_api/execution_service/execution_server.cpp
@@ -18,11 +18,14 @@
 #include <fstream>
 #include <iostream>
 #include <string>
+#include <unordered_map>
 #include <utility>
 
 #include "fmt/format.h"
+#include "gsl-lite/gsl-lite.hpp"
 #include "src/buildtool/compatibility/native_support.hpp"
 #include "src/buildtool/execution_api/local/garbage_collector.hpp"
+#include "src/buildtool/file_system/file_system_manager.hpp"
 
 auto ExecutionServiceImpl::GetAction(::bazel_re::ExecuteRequest const* request)
     const noexcept -> std::pair<std::optional<::bazel_re::Action>,
@@ -134,8 +137,133 @@ auto ExecutionServiceImpl::GetIExecutionAction(
     return {std::move(i_execution_action), std::nullopt};
 }
 
-static void AddOutputPaths(::bazel_re::ExecuteResponse* response,
-                           IExecutionResponse::Ptr const& execution) noexcept {
+static auto GetDirectoryFromDigest(::bazel_re::Digest const& digest,
+                                   LocalStorage const& storage) noexcept
+    -> std::optional<::bazel_re::Directory> {
+    // determine directory path from digest
+    auto const& path = storage.BlobPath(digest, /*is_executable=*/false);
+    if (not path) {
+        return std::nullopt;
+    }
+
+    // read directory content from path
+    auto const& content = FileSystemManager::ReadFile(*path);
+    if (not content) {
+        return std::nullopt;
+    }
+
+    // parse directory content
+    ::bazel_re::Directory dir{};
+    if (not dir.ParseFromString(*content)) {
+        return std::nullopt;
+    }
+    return dir;
+}
+
+// NOLINTNEXTLINE(misc-no-recursion)
+static auto CollectChildDirectoriesRecursively(
+    ::bazel_re::Directory const& root,
+    LocalStorage const& storage,
+    gsl::not_null<std::unordered_map<::bazel_re::Digest,
+                                     ::bazel_re::Directory>*> map) noexcept
+    -> bool {
+    return std::all_of(root.directories().begin(),
+                       root.directories().end(),
+                       // NOLINTNEXTLINE(misc-no-recursion)
+                       [&map, &storage](auto const& node) {
+                           if (map->find(node.digest()) != map->end()) {
+                               return true;
+                           }
+                           auto tmp_root =
+                               GetDirectoryFromDigest(node.digest(), storage);
+                           if (not tmp_root) {
+                               return false;
+                           }
+                           if (not CollectChildDirectoriesRecursively(
+                                   *tmp_root, storage, map)) {
+                               return false;
+                           }
+                           try {
+                               map->emplace(node.digest(), *tmp_root);
+                           } catch (...) {
+                               return false;
+                           }
+                           return true;
+                       });
+}
+
+static auto GetChildrenFromDirectory(::bazel_re::Directory const& root,
+                                     LocalStorage const& storage) noexcept
+    -> std::optional<std::vector<::bazel_re::Directory>> {
+    // determine child directories
+    std::unordered_map<::bazel_re::Digest, ::bazel_re::Directory> map{};
+    if (not CollectChildDirectoriesRecursively(root, storage, &map)) {
+        return std::nullopt;
+    }
+
+    // extract digests from child directories
+    std::vector<::bazel_re::Digest> digests{};
+    digests.reserve(map.size());
+    std::transform(map.begin(),
+                   map.end(),
+                   std::back_inserter(digests),
+                   [](auto const& pair) { return pair.first; });
+
+    // sort digests
+    std::sort(digests.begin(),
+              digests.end(),
+              [](auto const& left, auto const& right) {
+                  return left.hash() < right.hash();
+              });
+
+    // extract directory messages
+    std::vector<::bazel_re::Directory> children{};
+    children.reserve(digests.size());
+    std::transform(digests.begin(),
+                   digests.end(),
+                   std::back_inserter(children),
+                   [&map](auto const& digest) { return map[digest]; });
+
+    return children;
+}
+
+static auto CreateTreeDigestFromDirectoryDigest(
+    ::bazel_re::Digest const& dir_digest,
+    LocalStorage const& storage) noexcept -> std::optional<::bazel_re::Digest> {
+    // determine root directory message
+    auto root = GetDirectoryFromDigest(dir_digest, storage);
+    if (not root) {
+        return std::nullopt;
+    }
+
+    // determine child directory messages
+    auto children = GetChildrenFromDirectory(*root, storage);
+    if (not children) {
+        return std::nullopt;
+    }
+
+    // create tree message
+    ::bazel_re::Tree tree{};
+    tree.set_allocated_root(
+        gsl::owner<::bazel_re::Directory*>{new ::bazel_re::Directory{*root}});
+    tree.mutable_children()->Reserve(gsl::narrow<int>((*children).size()));
+    std::copy((*children).begin(),
+              (*children).end(),
+              ::pb::back_inserter(tree.mutable_children()));
+
+    // serialize and store tree message
+    auto content = tree.SerializeAsString();
+    auto tree_digest = storage.StoreBlob(content, /*is_executable=*/false);
+    if (not tree_digest) {
+        return std::nullopt;
+    }
+
+    return tree_digest;
+}
+
+static auto AddOutputPaths(::bazel_re::ExecuteResponse* response,
+                           IExecutionResponse::Ptr const& execution,
+                           LocalStorage const& storage) noexcept -> bool {
     auto const& size = static_cast<int>(execution->Artifacts().size());
     response->mutable_result()->mutable_output_files()->Reserve(size);
     response->mutable_result()->mutable_output_directories()->Reserve(size);
@@ -146,7 +274,20 @@ static void AddOutputPaths(::bazel_re::ExecuteResponse* response,
         if (info.type == ObjectType::Tree) {
             ::bazel_re::OutputDirectory out_dir;
             *(out_dir.mutable_path()) = path;
-            *(out_dir.mutable_tree_digest()) = std::move(dgst);
+            if (not Compatibility::IsCompatible()) {
+                // In native mode: Set the directory digest directly.
+                *(out_dir.mutable_tree_digest()) = std::move(dgst);
+            }
+            else {
+                // In compatible mode: Create a tree digest from directory
+                // digest on the fly and set tree digest.
+                auto digest =
+                    CreateTreeDigestFromDirectoryDigest(dgst, storage);
+                if (not digest) {
+                    return false;
+                }
+                *(out_dir.mutable_tree_digest()) = std::move(*digest);
+            }
             response->mutable_result()->mutable_output_directories()->Add(
                 std::move(out_dir));
         }
@@ -159,6 +300,7 @@ static void AddOutputPaths(::bazel_re::ExecuteResponse* response,
                 std::move(out_file));
         }
     }
+    return true;
 }
 
 auto ExecutionServiceImpl::AddResult(
@@ -166,7 +308,12 @@ auto ExecutionServiceImpl::AddResult(
     IExecutionResponse::Ptr const& i_execution_response,
     std::string const& action_hash) const noexcept
     -> std::optional<std::string> {
-    AddOutputPaths(response, i_execution_response);
+    if (not AddOutputPaths(response, i_execution_response, storage_)) {
+        auto str = fmt::format("Error in creating output paths of action {}",
+                               action_hash);
+        logger_.Emit(LogLevel::Error, str);
+        return std::nullopt;
+    }
     auto* result = response->mutable_result();
     result->set_exit_code(i_execution_response->ExitCode());
     if (i_execution_response->HasStdErr()) {