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
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

1 files changed, 24 insertions, 0 deletions

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):