Adapted TargetCache to download known artifacts from remote CAS to local CAS

4 files changed, 103 insertions, 15 deletions

diff --git a/src/buildtool/build_engine/target_map/target_cache.cpp b/src/buildtool/build_engine/target_map/target_cache.cpp
index 8ba54557..ecb0d787 100644
--- a/src/buildtool/build_engine/target_map/target_cache.cpp
+++ b/src/buildtool/build_engine/target_map/target_cache.cpp
@@ -1,5 +1,7 @@
 #include "src/buildtool/build_engine/target_map/target_cache.hpp"
 
+#include <gsl-lite/gsl-lite.hpp>
+
 #include "src/buildtool/common/repository_config.hpp"
 #include "src/buildtool/execution_api/local/config.hpp"
 #include "src/buildtool/execution_api/local/file_storage.hpp"
@@ -52,6 +54,12 @@ auto TargetCache::Entry::ToResult() const -> std::optional<TargetResult> {
 
 auto TargetCache::Store(Key const& key, Entry const& value) const noexcept
     -> bool {
+    // Before a target-cache entry is stored in local CAS, make sure any created
+    // artifact for this target is downloaded from the remote CAS to the local
+    // CAS.
+    if (not DownloadKnownArtifacts(value)) {
+        return false;
+    }
     if (auto digest = CAS().StoreBlobFromBytes(value.ToJson().dump(2))) {
         auto data =
             Artifact::ObjectInfo{ArtifactDigest{*digest}, ObjectType::File}
@@ -97,6 +105,45 @@ auto TargetCache::Read(Key const& key) const noexcept
     return std::nullopt;
 }
 
+auto TargetCache::DownloadKnownArtifactsFromMap(
+    Expression::map_t const& expr_map) const noexcept -> bool {
+
+    // Get object infos of KNOWN artifacts from map.
+    std::vector<Artifact::ObjectInfo> infos;
+    infos.reserve(expr_map.size());
+    for (auto const& item : expr_map) {
+        try {
+            auto const& desc = item.second->Artifact();
+            // The assumption is that all artifacts mentioned in a target cache
+            // entry are KNOWN to the remote side. So they can be fetched to the
+            // local CAS.
+            gsl_ExpectsAudit(desc.IsKnown());
+            infos.push_back(*desc.ToArtifact().Info());
+        } catch (...) {
+            return false;
+        }
+    }
+
+#ifndef BOOTSTRAP_BUILD_TOOL
+    // Sync KNOWN artifacts from remote to local CAS.
+    return remote_api_->RetrieveToCas(infos, local_api_);
+#else
+    return true;
+#endif
+}
+
+auto TargetCache::DownloadKnownArtifacts(Entry const& value) const noexcept
+    -> bool {
+    auto const& result = value.ToResult();
+    if (not DownloadKnownArtifactsFromMap(result->artifact_stage->Map())) {
+        return false;
+    }
+    if (not DownloadKnownArtifactsFromMap(result->runfiles->Map())) {
+        return false;
+    }
+    return true;
+}
+
 auto TargetCache::ComputeCacheDir() -> std::filesystem::path {
     return LocalExecutionConfig::TargetCacheDir() / ExecutionBackendId();
 }
diff --git a/src/buildtool/build_engine/target_map/target_cache.hpp b/src/buildtool/build_engine/target_map/target_cache.hpp
index eb31b831..71b56321 100644
--- a/src/buildtool/build_engine/target_map/target_cache.hpp
+++ b/src/buildtool/build_engine/target_map/target_cache.hpp
@@ -8,6 +8,9 @@
 #include "src/buildtool/build_engine/analysed_target/analysed_target.hpp"
 #include "src/buildtool/build_engine/base_maps/entity_name.hpp"
 #include "src/buildtool/common/artifact.hpp"
+#ifndef BOOTSTRAP_BUILD_TOOL
+#include "src/buildtool/execution_api/common/execution_api.hpp"
+#endif
 #include "src/buildtool/execution_api/local/file_storage.hpp"
 #include "src/buildtool/execution_api/local/local_cas.hpp"
 #include "src/buildtool/file_system/file_system_manager.hpp"
@@ -85,13 +88,33 @@ class TargetCache {
     [[nodiscard]] auto Read(Key const& key) const noexcept
         -> std::optional<std::pair<Entry, Artifact::ObjectInfo>>;
 
+#ifndef BOOTSTRAP_BUILD_TOOL
+    auto SetLocalApi(gsl::not_null<IExecutionApi*> const& api) noexcept
+        -> void {
+        local_api_ = api;
+    };
+
+    auto SetRemoteApi(gsl::not_null<IExecutionApi*> const& api) noexcept
+        -> void {
+        remote_api_ = api;
+    };
+#endif
+
   private:
     Logger logger_{"TargetCache"};
     FileStorage<ObjectType::File,
                 StoreMode::LastWins,
                 /*kSetEpochTime=*/false>
         file_store_{ComputeCacheDir()};
+#ifndef BOOTSTRAP_BUILD_TOOL
+    IExecutionApi* local_api_{};
+    IExecutionApi* remote_api_{};
+#endif
 
+    [[nodiscard]] auto DownloadKnownArtifacts(Entry const& value) const noexcept
+        -> bool;
+    [[nodiscard]] auto DownloadKnownArtifactsFromMap(
+        Expression::map_t const& expr_map) const noexcept -> bool;
     [[nodiscard]] static auto CAS() noexcept -> LocalCAS<ObjectType::File>& {
         return LocalCAS<ObjectType::File>::Instance();
     }
diff --git a/src/buildtool/graph_traverser/graph_traverser.hpp b/src/buildtool/graph_traverser/graph_traverser.hpp
index 0e81d46e..718df364 100644
--- a/src/buildtool/graph_traverser/graph_traverser.hpp
+++ b/src/buildtool/graph_traverser/graph_traverser.hpp
@@ -58,13 +58,17 @@ class GraphTraverser {
 
     explicit GraphTraverser(CommandLineArguments clargs)
         : clargs_{std::move(clargs)},
-          api_{CreateExecutionApi(RemoteExecutionConfig::RemoteAddress())},
+          local_api_{CreateExecutionApi(std::nullopt)},
+          remote_api_{
+              CreateExecutionApi(RemoteExecutionConfig::RemoteAddress())},
           reporter_{[](auto done, auto cv) {}} {}
 
     explicit GraphTraverser(CommandLineArguments clargs,
                             progress_reporter_t reporter)
         : clargs_{std::move(clargs)},
-          api_{CreateExecutionApi(RemoteExecutionConfig::RemoteAddress())},
+          local_api_{CreateExecutionApi(std::nullopt)},
+          remote_api_{
+              CreateExecutionApi(RemoteExecutionConfig::RemoteAddress())},
           reporter_{std::move(reporter)} {}
 
     /// \brief Parses actions and blobs into graph, traverses it and retrieves
@@ -186,13 +190,18 @@ class GraphTraverser {
             artifact_descriptions, {}, action_descriptions, blobs, trees);
     }
 
-    [[nodiscard]] auto ExecutionApi() const -> gsl::not_null<IExecutionApi*> {
-        return &(*api_);
+    [[nodiscard]] auto GetLocalApi() const -> gsl::not_null<IExecutionApi*> {
+        return &(*local_api_);
+    }
+
+    [[nodiscard]] auto GetRemoteApi() const -> gsl::not_null<IExecutionApi*> {
+        return &(*remote_api_);
     }
 
   private:
     CommandLineArguments const clargs_;
-    gsl::not_null<IExecutionApi::Ptr> const api_;
+    gsl::not_null<IExecutionApi::Ptr> const local_api_;
+    gsl::not_null<IExecutionApi::Ptr> const remote_api_;
     progress_reporter_t reporter_;
 
     /// \brief Reads contents of graph description file as json object. In case
@@ -277,7 +286,7 @@ class GraphTraverser {
                 return false;
             }
         }
-        return api_->Upload(container);
+        return remote_api_->Upload(container);
     }
 
     /// \brief Adds the artifacts to be retrieved to the graph
@@ -320,7 +329,7 @@ class GraphTraverser {
     [[nodiscard]] auto Traverse(
         DependencyGraph const& g,
         std::vector<ArtifactIdentifier> const& artifact_ids) const -> bool {
-        Executor executor{&(*api_),
+        Executor executor{&(*remote_api_),
                           RemoteExecutionConfig::PlatformProperties(),
                           clargs_.build.timeout};
         bool traversing{};
@@ -346,7 +355,7 @@ class GraphTraverser {
         // setup rebuilder with api for cache endpoint
         auto api_cached =
             CreateExecutionApi(RemoteExecutionConfig::CacheAddress());
-        Rebuilder executor{&(*api_),
+        Rebuilder executor{&(*remote_api_),
                            &(*api_cached),
                            RemoteExecutionConfig::PlatformProperties(),
                            clargs_.build.timeout};
@@ -550,7 +559,7 @@ class GraphTraverser {
         auto output_paths = PrepareOutputPaths(rel_paths);
 
         if (not output_paths or
-            not api_->RetrieveToPaths(object_infos, *output_paths)) {
+            not remote_api_->RetrieveToPaths(object_infos, *output_paths)) {
             Logger::Log(LogLevel::Error, "Could not retrieve outputs.");
             return std::nullopt;
         }
@@ -626,9 +635,10 @@ class GraphTraverser {
                 if (paths[i] == *(clargs_.build.print_to_stdout)) {
                     auto info = artifacts[i]->Content().Info();
                     if (info) {
-                        if (not api_->RetrieveToFds({*info},
-                                                    {dup(fileno(stdout))},
-                                                    /*raw_tree=*/false)) {
+                        if (not remote_api_->RetrieveToFds(
+                                {*info},
+                                {dup(fileno(stdout))},
+                                /*raw_tree=*/false)) {
                             Logger::Log(LogLevel::Error,
                                         "Failed to retrieve {}",
                                         *(clargs_.build.print_to_stdout));
diff --git a/src/buildtool/main/main.cpp b/src/buildtool/main/main.cpp
index 41b2a0ef..f066b13e 100644
--- a/src/buildtool/main/main.cpp
+++ b/src/buildtool/main/main.cpp
@@ -994,13 +994,18 @@ auto CollectNonKnownArtifacts(
             std::make_move_iterator(cache_artifacts.end())};
 }
 
+#ifndef BOOTSTRAP_BUILD_TOOL
 void WriteTargetCacheEntries(
     std::unordered_map<TargetCache::Key, AnalysedTargetPtr> const&
         cache_targets,
     std::unordered_map<ArtifactDescription, Artifact::ObjectInfo> const&
         extra_infos,
-    std::size_t jobs) {
+    std::size_t jobs,
+    gsl::not_null<IExecutionApi*> const& local_api,
+    gsl::not_null<IExecutionApi*> const& remote_api) {
     auto ts = TaskSystem{jobs};
+    TargetCache::Instance().SetLocalApi(local_api);
+    TargetCache::Instance().SetRemoteApi(remote_api);
     for (auto const& [key, target] : cache_targets) {
         ts.QueueTask([&key = key, &target = target, &extra_infos]() {
             if (auto entry =
@@ -1019,6 +1024,7 @@ void WriteTargetCacheEntries(
         });
     }
 }
+#endif
 
 }  // namespace
 
@@ -1065,7 +1071,7 @@ auto main(int argc, char* argv[]) -> int {
                                        BaseProgressReporter::Reporter()};
 
         if (arguments.cmd == SubCommand::kInstallCas) {
-            return FetchAndInstallArtifacts(traverser.ExecutionApi(),
+            return FetchAndInstallArtifacts(traverser.GetRemoteApi(),
                                             arguments.fetch)
                        ? kExitSuccess
                        : kExitFailure;
@@ -1188,7 +1194,9 @@ auto main(int argc, char* argv[]) -> int {
                 if (build_result) {
                     WriteTargetCacheEntries(cache_targets,
                                             build_result->extra_infos,
-                                            arguments.common.jobs);
+                                            arguments.common.jobs,
+                                            traverser.GetLocalApi(),
+                                            traverser.GetRemoteApi());
 
                     // Repeat taintedness message to make the user aware that
                     // the artifacts are not for production use.