Add CLI option to set write strategy for target-level cache

7 files changed, 56 insertions, 5 deletions

diff --git a/src/buildtool/common/TARGETS b/src/buildtool/common/TARGETS
index 877d740a..978cebc5 100644
--- a/src/buildtool/common/TARGETS
+++ b/src/buildtool/common/TARGETS
@@ -15,6 +15,7 @@
     , ["src/buildtool/compatibility", "compatibility"]
     , ["src/buildtool/logging", "log_level"]
     , ["src/buildtool/logging", "logging"]
+    , ["src/buildtool/main", "build_utils"]
     , ["src/utils/cpp", "path"]
     , ["@", "cli11", "", "cli11"]
     , ["@", "json", "", "json"]
diff --git a/src/buildtool/common/cli.hpp b/src/buildtool/common/cli.hpp
index bdcc5716..9c2bc7e2 100644
--- a/src/buildtool/common/cli.hpp
+++ b/src/buildtool/common/cli.hpp
@@ -30,6 +30,7 @@
 #include "src/buildtool/common/clidefaults.hpp"
 #include "src/buildtool/compatibility/compatibility.hpp"
 #include "src/buildtool/logging/log_level.hpp"
+#include "src/buildtool/main/build_utils.hpp"
 #include "src/utils/cpp/path.hpp"
 
 constexpr auto kDefaultTimeout = std::chrono::milliseconds{300000};
@@ -103,6 +104,12 @@ struct BuildArguments {
     bool show_runfiles{false};
 };
 
+/// \brief Arguments related to target-level caching
+struct TCArguments {
+    TargetCacheWriteStrategy target_cache_write_strategy{
+        TargetCacheWriteStrategy::Sync};
+};
+
 /// \brief Arguments required for staging.
 struct StageArguments {
     std::filesystem::path output_dir{};
@@ -461,6 +468,26 @@ static inline auto SetupBuildArguments(
         ->type_name("LOGICAL_PATH");
 }
 
+static inline auto SetupTCArguments(gsl::not_null<CLI::App*> const& app,
+                                    gsl::not_null<TCArguments*> const& tcargs) {
+    app->add_option_function<std::string>(
+           "--target-cache-write-strategy",
+           [tcargs](auto const& s) {
+               auto strategy = ToTargetCacheWriteStrategy(s);
+               if (strategy) {
+                   tcargs->target_cache_write_strategy = *strategy;
+               }
+               else {
+                   Logger::Log(LogLevel::Warning,
+                               "Ignoring unknown strategy {} to write "
+                               "target-level cache.",
+                               nlohmann::json(s).dump());
+               }
+           },
+           "Strategy for writing target-cache. (Default: sync)")
+        ->type_name("STRATEGY");
+}
+
 static inline auto SetupStageArguments(
     gsl::not_null<CLI::App*> const& app,
     gsl::not_null<StageArguments*> const& clargs) {
diff --git a/src/buildtool/main/build_utils.cpp b/src/buildtool/main/build_utils.cpp
index 46773ca2..50557b71 100644
--- a/src/buildtool/main/build_utils.cpp
+++ b/src/buildtool/main/build_utils.cpp
@@ -48,6 +48,20 @@ auto CollectNonKnownArtifacts(
             std::make_move_iterator(cache_artifacts.end())};
 }
 
+auto ToTargetCacheWriteStrategy(std::string const& strategy)
+    -> std::optional<TargetCacheWriteStrategy> {
+    if (strategy == "disable") {
+        return TargetCacheWriteStrategy::Disable;
+    }
+    if (strategy == "sync") {
+        return TargetCacheWriteStrategy::Sync;
+    }
+    if (strategy == "split") {
+        return TargetCacheWriteStrategy::Split;
+    }
+    return std::nullopt;
+}
+
 #ifndef BOOTSTRAP_BUILD_TOOL
 void WriteTargetCacheEntries(
     std::unordered_map<TargetCacheKey, AnalysedTargetPtr> const& cache_targets,
diff --git a/src/buildtool/main/build_utils.hpp b/src/buildtool/main/build_utils.hpp
index e78d86d0..193d5717 100644
--- a/src/buildtool/main/build_utils.hpp
+++ b/src/buildtool/main/build_utils.hpp
@@ -16,6 +16,8 @@
 #define INCLUDED_SRC_BUILDOOL_MAIN_BUILD_UTILS_HPP
 
 #include <map>
+#include <optional>
+#include <string>
 #include <unordered_map>
 
 #include "src/buildtool/build_engine/analysed_target/analysed_target.hpp"
@@ -45,6 +47,9 @@ enum class TargetCacheWriteStrategy {
            ///< during artifact sync try to use blob splitting, if available
 };
 
+auto ToTargetCacheWriteStrategy(std::string const&)
+    -> std::optional<TargetCacheWriteStrategy>;
+
 #ifndef BOOTSTRAP_BUILD_TOOL
 void WriteTargetCacheEntries(
     std::unordered_map<TargetCacheKey, AnalysedTargetPtr> const& cache_targets,
diff --git a/src/buildtool/main/cli.cpp b/src/buildtool/main/cli.cpp
index 475b7256..494c01d0 100644
--- a/src/buildtool/main/cli.cpp
+++ b/src/buildtool/main/cli.cpp
@@ -57,6 +57,7 @@ auto SetupBuildCommandArguments(
     SetupClientAuthArguments(app, &clargs->cauth);
     SetupCommonBuildArguments(app, &clargs->build);
     SetupBuildArguments(app, &clargs->build);
+    SetupTCArguments(app, &clargs->tc);
     SetupCompatibilityArguments(app);
 }
 
diff --git a/src/buildtool/main/cli.hpp b/src/buildtool/main/cli.hpp
index b2945757..5c8d91ef 100644
--- a/src/buildtool/main/cli.hpp
+++ b/src/buildtool/main/cli.hpp
@@ -41,6 +41,7 @@ struct CommandLineArguments {
     DiagnosticArguments diagnose;
     EndpointArguments endpoint;
     BuildArguments build;
+    TCArguments tc;
     StageArguments stage;
     RebuildArguments rebuild;
     FetchArguments fetch;
diff --git a/src/buildtool/main/main.cpp b/src/buildtool/main/main.cpp
index 4f874da0..ba5eb620 100644
--- a/src/buildtool/main/main.cpp
+++ b/src/buildtool/main/main.cpp
@@ -1027,11 +1027,13 @@ auto main(int argc, char* argv[]) -> int {
                                             trees,
                                             std::move(cache_artifacts));
                 if (build_result) {
-                    WriteTargetCacheEntries(cache_targets,
-                                            build_result->extra_infos,
-                                            jobs,
-                                            traverser.GetLocalApi(),
-                                            traverser.GetRemoteApi());
+                    WriteTargetCacheEntries(
+                        cache_targets,
+                        build_result->extra_infos,
+                        jobs,
+                        traverser.GetLocalApi(),
+                        traverser.GetRemoteApi(),
+                        arguments.tc.target_cache_write_strategy);
 
                     // Repeat taintedness message to make the user aware that
                     // the artifacts are not for production use.