From ed59ecf1491a0b6f3b933d2c4b2cd97a506f0090 Mon Sep 17 00:00:00 2001
From: Maksim Denisov <denisov.maksim@huawei.com>
Date: Tue, 21 Jan 2025 17:39:49 +0100
Subject: ServeApi: Implement DownloadTree

---
 src/buildtool/serve_api/remote/serve_api.cpp | 50 ++++++++++++++++++++++++++++
 1 file changed, 50 insertions(+)

(limited to 'src/buildtool/serve_api/remote/serve_api.cpp')

diff --git a/src/buildtool/serve_api/remote/serve_api.cpp b/src/buildtool/serve_api/remote/serve_api.cpp
index 37911e1d..cb0ea7ac 100644
--- a/src/buildtool/serve_api/remote/serve_api.cpp
+++ b/src/buildtool/serve_api/remote/serve_api.cpp
@@ -25,6 +25,7 @@
 #include "src/buildtool/execution_api/serve/mr_git_api.hpp"
 #include "src/buildtool/execution_api/utils/rehash_utils.hpp"
 #include "src/buildtool/file_system/object_type.hpp"
+#include "src/buildtool/storage/storage.hpp"
 
 auto ServeApi::UploadTree(ArtifactDigest const& tree,
                           std::filesystem::path const& git_repo) const noexcept
@@ -108,4 +109,53 @@ auto ServeApi::UploadTree(ArtifactDigest const& tree,
     return std::monostate{};
 }
 
+auto ServeApi::DownloadTree(ArtifactDigest const& tree) const noexcept
+    -> expected<std::monostate, std::string> {
+    if (not tree.IsTree() or not ProtocolTraits::IsNative(tree.GetHashType())) {
+        return unexpected{fmt::format("Not a git tree: {}", tree.hash())};
+    }
+
+    // Check the tree is already in native CAS:
+    auto native_config = StorageConfig::Builder::Rebuild(storage_config_)
+                             .SetHashType(HashFunction::Type::GitSHA1)
+                             .Build();
+    if (not native_config.has_value()) {
+        return unexpected{fmt::format("Failed to create native storage: {}",
+                                      std::move(native_config).error())};
+    }
+    if (Storage::Create(&*native_config).CAS().TreePath(tree).has_value()) {
+        return std::monostate{};
+    }
+
+    // Make tree available on the remote end point:
+    auto const on_remote = TreeInRemoteCAS(tree.hash());
+    if (not on_remote.has_value()) {
+        return unexpected{
+            fmt::format("Failed to upload {} from serve to the remote end "
+                        "point.",
+                        tree.hash())};
+    }
+
+    // Download tree from the remote end point:
+    Artifact::ObjectInfo const info{*on_remote, ObjectType::Tree};
+    if (not apis_.remote->RetrieveToCas({info}, *apis_.local)) {
+        return unexpected{fmt::format(
+            "Failed to download {} from the remote end point.", tree.hash())};
+    }
+
+    // The remote end point may operate in the compatible mode. In such a case,
+    // an extra rehashing is needed:
+    if (not ProtocolTraits::IsNative(storage_config_.hash_function.GetType())) {
+        auto rehashed = RehashUtils::RehashDigest(
+            {info}, storage_config_, *native_config, &apis_);
+        if (not rehashed.has_value() or rehashed->size() != 1 or
+            rehashed->front().digest != tree) {
+            return unexpected{fmt::format("Failed to rehash downloaded {}:\n{}",
+                                          on_remote->hash(),
+                                          std::move(rehashed).error())};
+        }
+    }
+    return std::monostate{};
+}
+
 #endif
-- 
cgit v1.2.3