Kleaf: fix caching for --config=local

The --config=local mode was broken in that it would not actually cache
the OUT_DIR across executions. This was due to the OUT_DIR being defined
in the sandbox, relative to symlinked sources. It is not clear why the
local mode worked before and what combination of changes broke it. This
change restores is by adding a starlark flag --cache_dir that is used to
define OUT_DIR in KernelEnv. Said flag is also understood by the
bazel.py wrapper and set to a reasonable default within WORKSPACE if not
defined. Yet in order to have $WORKSPACE/out/cache we need to also
execute the KernelEnv rule in config=local mode.

This lets --config=local and --config=fast cache OUT_DIR in
$WORKSPACE/out/cached/<some_generated_name> across invocations. Further,
by specifying --cache_dir=/some/fast/disk consistently across
invocations, one can make use of a fast filesystem for caching; perhaps
a tmpfs.

In an experiment, this saves 50s when comparing --config=local
and non --config=local incremental builds.

Bug: 235632059
Bug: 241603176
Test: manual with the following:
  $ bazel build //common:kernel_dist                && echo >> common/init/main.c && time bazel build //common:kernel_dist
  Takes 114 seconds for the second build; kernel_build() takes 97s.

  $ bazel build --config=local //common:kernel_dist && echo >> common/init/main.c && time bazel build --config=local //common:kernel_dist
  Takes 64  seconds for the second build; kernel_build() takes 40s.

Signed-off-by: Matthias Maennich <maennich@google.com>
Signed-off-by: Yifan Hong <elsk@google.com>
Change-Id: I64fbfd928efd78fb5d3d3c3af6a5582fca22b0a5
(cherry picked from commit c2f2105abc6f729357a56a6c354ac9516a4a0907)

7 files changed, 69 insertions, 7 deletions

diff --git a/kleaf/BUILD.bazel b/kleaf/BUILD.bazel
index 36d1a72..730bc3f 100644
--- a/kleaf/BUILD.bazel
+++ b/kleaf/BUILD.bazel
@@ -30,6 +30,11 @@ string_flag(
     values = LTO_VALUES,
 )
 
+string_flag(
+    name = "cache_dir",
+    build_setting_default = "",
+)
+
 # If true, set `KBUILD_SYMTYPES=1` for kernel_build() with
 # kbuild_symtypes="auto".
 bool_flag(
diff --git a/kleaf/bazel.py b/kleaf/bazel.py
index 046d1c1..ed457f0 100755
--- a/kleaf/bazel.py
+++ b/kleaf/bazel.py
@@ -14,6 +14,8 @@
 
 import argparse
 import os
+import pathlib
+import shutil
 import sys
 from typing import Tuple, Optional
 
@@ -22,6 +24,13 @@ _BAZEL_JDK_REL_PATH = "prebuilts/jdk/jdk11/linux-x86"
 _BAZEL_RC_NAME = "build/kernel/kleaf/common.bazelrc"
 
 
+def _require_absolute_path(p: str) -> pathlib.Path:
+    p = pathlib.Path(p)
+    if not p.is_absolute():
+        raise argparse.ArgumentTypeError("need to specify an absolute path")
+    return p
+
+
 def _partition(lst: list[str], index: Optional[int]) \
         -> Tuple[list[str], Optional[str], list[str]]:
     """Returns the triple split by index.
@@ -90,12 +99,17 @@ class BazelWrapper(object):
         - env: A dictionary containing the new environment variables for the subprocess.
         """
 
+        absolute_cache_dir = f"{self.absolute_out_dir}/cache"
+
         # Arguments known by this bazel wrapper.
         parser = argparse.ArgumentParser(add_help=False)
         parser.add_argument("--use_prebuilt_gki")
         parser.add_argument("--experimental_strip_sandbox_path",
                             action='store_true')
         parser.add_argument("--make_jobs", type=int, default=None)
+        parser.add_argument("--cache_dir",
+                            type=_require_absolute_path,
+                            default=absolute_cache_dir)
 
         # known_args: List of arguments known by this bazel wrapper. These
         #   are stripped from the final bazel invocation.
@@ -111,6 +125,9 @@ class BazelWrapper(object):
         if self.known_args.make_jobs is not None:
             self.env["KLEAF_MAKE_JOBS"] = str(self.known_args.make_jobs)
 
+        self.transformed_command_args.append(
+            f"--//build/kernel/kleaf:cache_dir={self.known_args.cache_dir}")
+
     def _build_final_args(self) -> list[str]:
         """Builds the final arguments for the subprocess."""
         # final_args:
@@ -130,6 +147,13 @@ class BazelWrapper(object):
             final_args.append(self.dash_dash)
         final_args += self.target_patterns
 
+        if self.command == "clean":
+            sys.stderr.write(
+                f"INFO: Removing cache directory for $OUT_DIR: {self.known_args.cache_dir}\n")
+            shutil.rmtree(self.known_args.cache_dir, ignore_errors=True)
+        else:
+            os.makedirs(self.known_args.cache_dir, exist_ok=True)
+
         return final_args
 
     def run(self):
diff --git a/kleaf/bazelrc/local.bazelrc b/kleaf/bazelrc/local.bazelrc
index f3e4f8e..bd5bb21 100644
--- a/kleaf/bazelrc/local.bazelrc
+++ b/kleaf/bazelrc/local.bazelrc
@@ -17,6 +17,6 @@
 # Let rules know what configs we are using now
 build:local --//build/kernel/kleaf:config_local
 
-# Skip sandboxing to cache $OUT_DIR.
+# A list of actions to skip sandboxing.
 build:local --strategy KernelConfig=local
 build:local --strategy KernelBuild=local
diff --git a/kleaf/docs/sandbox.md b/kleaf/docs/sandbox.md
index 689f1a3..a469d86 100644
--- a/kleaf/docs/sandbox.md
+++ b/kleaf/docs/sandbox.md
@@ -62,6 +62,18 @@ you'll get a build error.
 
 See [scmversion.md](scmversion.md).
 
+## Local cache dir
+
+The `--config=local` mode makes use of a persistent `$OUT_DIR`
+across invocations to cache the rule execution state. The default cache
+directory is `$WORKSPACE/out/cache`, but can be overridden by passing
+`--cache_dir=/some/fast/disk` in order to make use of a file system that
+performs well or better for the kernel build workload. Full example:
+
+```shell
+$ tools/bazel run --config=local --cache_dir=/some/fast/disk //common:kernel_aarch64_dist
+```
+
 ## Other flags
 
 The flag `--config=local` is also implied by other flags, e.g.:
diff --git a/kleaf/impl/kernel_build.bzl b/kleaf/impl/kernel_build.bzl
index 87b0e8f..eb7e476 100644
--- a/kleaf/impl/kernel_build.bzl
+++ b/kleaf/impl/kernel_build.bzl
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 load("@bazel_skylib//lib:paths.bzl", "paths")
+load("@bazel_skylib//rules:common_settings.bzl", "BuildSettingInfo")
 load("//build/kernel/kleaf:hermetic_tools.bzl", "HermeticToolsInfo")
 load(
     "//build/kernel/kleaf/artifact_tests:kernel_test.bzl",
@@ -603,6 +604,23 @@ def _kernel_build_impl(ctx):
     command = ""
     command += ctx.attr.config[KernelEnvInfo].setup
 
+    # Use a local cache directory for ${OUT_DIR} so that, even when this _kernel_build
+    # target needs to be rebuilt, we are using $OUT_DIR from previous invocations. This
+    # boosts --config=local builds. See (b/235632059).
+    if ctx.attr._config_is_local[BuildSettingInfo].value:
+        if not ctx.attr._cache_dir[BuildSettingInfo].value:
+            fail("--config=local requires --cache_dir.")
+        command += """
+              KLEAF_CACHED_OUT_DIR={cache_dir}/{name}
+              mkdir -p "${{KLEAF_CACHED_OUT_DIR}}"
+              rsync -aL "${{OUT_DIR}}/" "${{KLEAF_CACHED_OUT_DIR}}/"
+              export OUT_DIR=${{KLEAF_CACHED_OUT_DIR}}
+              unset KLEAF_CACHED_OUT_DIR
+        """.format(
+            cache_dir = ctx.attr._cache_dir[BuildSettingInfo].value,
+            name = utils.sanitize_label_as_filename(ctx.label),
+        )
+
     interceptor_command_prefix = ""
     if interceptor_output:
         interceptor_command_prefix = "interceptor -r -l {interceptor_output} --".format(
@@ -862,6 +880,8 @@ _kernel_build = rule(
         "_compare_to_symbol_list": attr.label(default = "//build/kernel:abi/compare_to_symbol_list", allow_single_file = True),
         "_hermetic_tools": attr.label(default = "//build/kernel:hermetic-tools", providers = [HermeticToolsInfo]),
         "_debug_print_scripts": attr.label(default = "//build/kernel/kleaf:debug_print_scripts"),
+        "_config_is_local": attr.label(default = "//build/kernel/kleaf:config_local"),
+        "_cache_dir": attr.label(default = "//build/kernel/kleaf:cache_dir"),
         # Though these rules are unrelated to the `_kernel_build` rule, they are added as fake
         # dependencies so KernelBuildExtModuleInfo and KernelBuildUapiInfo works.
         # There are no real dependencies. Bazel does not build these targets before building the
diff --git a/kleaf/impl/kernel_env.bzl b/kleaf/impl/kernel_env.bzl
index 17763dd..24b6fe3 100644
--- a/kleaf/impl/kernel_env.bzl
+++ b/kleaf/impl/kernel_env.bzl
@@ -26,11 +26,6 @@ load(":stamp.bzl", "stamp")
 load(":status.bzl", "status")
 load(":utils.bzl", "utils")
 
-def _sanitize_label_as_filename(label):
-    """Sanitize a Bazel label so it is safe to be used as a filename."""
-    label_text = str(label)
-    return "".join([c if c.isalnum() else "_" for c in label_text.elems()])
-
 def _get_kbuild_symtypes(ctx):
     if ctx.attr.kbuild_symtypes == "auto":
         return ctx.attr._kbuild_symtypes_flag[BuildSettingInfo].value
@@ -118,7 +113,7 @@ def _kernel_env_impl(ctx):
     # with --spawn_strategy=local, try to isolate their OUT_DIRs.
     command += """
           export OUT_DIR_SUFFIX={name}
-    """.format(name = _sanitize_label_as_filename(ctx.label).removesuffix("_env"))
+    """.format(name = utils.sanitize_label_as_filename(ctx.label).removesuffix("_env"))
 
     set_source_date_epoch_ret = stamp.set_source_date_epoch(ctx)
     command += set_source_date_epoch_ret.cmd
diff --git a/kleaf/impl/utils.bzl b/kleaf/impl/utils.bzl
index 7fb56bc..38ca1a5 100644
--- a/kleaf/impl/utils.bzl
+++ b/kleaf/impl/utils.bzl
@@ -108,6 +108,11 @@ def _compare_file_names(files, expected_file_names, what):
             expected_file_names,
         ))
 
+def _sanitize_label_as_filename(label):
+    """Sanitize a Bazel label so it is safe to be used as a filename."""
+    label_text = str(label)
+    return "".join([c if c.isalnum() else "_" for c in label_text.elems()])
+
 # Utilities that applies to all Bazel stuff in general. These functions are
 # not Kleaf specific.
 utils = struct(
@@ -117,6 +122,7 @@ utils = struct(
     find_file = find_file,
     find_files = find_files,
     compare_file_names = _compare_file_names,
+    sanitize_label_as_filename = _sanitize_label_as_filename,
 )
 
 def _filter_module_srcs(files):