zephyr: zmake: Add a --goma flag for hyperspeed compilation

Add a new global zmake flag, --goma, which enables distributed
compilation using go/ma.  This service is only available for
Googlers, at the moment.  Additionally, it's only available for chroot
users at the moment.

On my machine, using the goma flag speeds up a testall greater than
4x, from 8m20s without --goma, and 1m55s with --goma, and additional
improvements are possible in the future.

Markdown documentation for setting up goma and using the flag follows
in the follow-up CL.

BUG=b:214323409
BRANCH=none
TEST=zmake --goma testall

Signed-off-by: Jack Rosenthal <jrosenth@chromium.org>
Change-Id: I9659d62f2c440f1430a30b784d3bdbb502c3952b
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/ec/+/3385775
Reviewed-by: Jeremy Bettis <jbettis@chromium.org>
Reviewed-by: Aaron Massey <aaronmassey@google.com>

3 files changed, 69 insertions, 30 deletions

diff --git a/zephyr/zmake/README.md b/zephyr/zmake/README.md
index d3f99fe753..6fdb8a43a4 100644
--- a/zephyr/zmake/README.md
+++ b/zephyr/zmake/README.md
@@ -6,7 +6,7 @@
 
 ## Usage
 
-**Usage:** `zmake [-h] [--checkout CHECKOUT] [-D] [-j JOBS] [-l {DEBUG,INFO,WARNING,ERROR,CRITICAL}] [-L] [--log-label] [--modules-dir MODULES_DIR] [--zephyr-base ZEPHYR_BASE] subcommand ...`
+**Usage:** `zmake [-h] [--checkout CHECKOUT] [-D] [-j JOBS] [--goma] [-l {DEBUG,INFO,WARNING,ERROR,CRITICAL}] [-L] [--log-label] [--modules-dir MODULES_DIR] [--zephyr-base ZEPHYR_BASE] subcommand ...`
 
 Chromium OS's meta-build tool for Zephyr
 
@@ -24,6 +24,7 @@ Chromium OS's meta-build tool for Zephyr
 | `--checkout CHECKOUT` | Path to ChromiumOS checkout |
 | `-D`, `--debug` | Turn on debug features (e.g., stack trace, verbose logging) |
 | `-j JOBS`, `--jobs JOBS` | Degree of multiprogramming to use |
+| `--goma` | Enable hyperspeed compilation with Goma! (Googlers only) |
 | `-l LOG_LEVEL`, `--log-level LOG_LEVEL` | Set the logging level (default=INFO) |
 | `-L`, `--no-log-label` | Turn off logging labels |
 | `--log-label` | Turn on logging labels |
diff --git a/zephyr/zmake/zmake/__main__.py b/zephyr/zmake/zmake/__main__.py
index d8f33a97d9..a364ccdd97 100644
--- a/zephyr/zmake/zmake/__main__.py
+++ b/zephyr/zmake/zmake/__main__.py
@@ -124,6 +124,12 @@ def get_argparser():
         help="Degree of multiprogramming to use",
     )
     parser.add_argument(
+        "--goma",
+        action="store_true",
+        dest="goma",
+        help="Enable hyperspeed compilation with Goma! (Googlers only)",
+    )
+    parser.add_argument(
         "-l",
         "--log-level",
         choices=list(log_level_map.keys()),
@@ -348,7 +354,8 @@ def main(argv=None):
         zmake = call_with_namespace(zm.Zmake, opts)
         subcommand_method = getattr(zmake, opts.subcommand.replace("-", "_"))
         result = call_with_namespace(subcommand_method, opts)
-        return result
+        wait_rv = zmake.executor.wait()
+        return result or wait_rv
     finally:
         multiproc.wait_for_log_end()
 
diff --git a/zephyr/zmake/zmake/zmake.py b/zephyr/zmake/zmake/zmake.py
index 47a79f179b..b3506d56ad 100644
--- a/zephyr/zmake/zmake/zmake.py
+++ b/zephyr/zmake/zmake/zmake.py
@@ -4,6 +4,7 @@
 
 """Module encapsulating Zmake wrapper object."""
 import difflib
+import functools
 import logging
 import os
 import pathlib
@@ -153,11 +154,15 @@ class Zmake:
         checkout=None,
         jobserver=None,
         jobs=0,
+        goma=False,
+        gomacc="/mnt/host/depot_tools/.cipd_bin/gomacc",
         modules_dir=None,
         zephyr_base=None,
     ):
         zmake.multiproc.reset()
         self._checkout = checkout
+        self.goma = goma
+        self.gomacc = gomacc
         if zephyr_base:
             self.zephyr_base = zephyr_base
         else:
@@ -178,7 +183,7 @@ class Zmake:
 
         self.logger = logging.getLogger(self.__class__.__name__)
         self.executor = zmake.multiproc.Executor()
-        self._sequential = jobs == 1
+        self._sequential = jobs == 1 and not goma
 
     @property
     def checkout(self):
@@ -292,6 +297,13 @@ class Zmake:
             base_config |= zmake.build_config.BuildConfig(
                 cmake_defs={"ALLOW_WARNINGS": "ON"}
             )
+        if self.goma:
+            base_config |= zmake.build_config.BuildConfig(
+                cmake_defs={
+                    "CMAKE_C_COMPILER_LAUNCHER": self.gomacc,
+                    "CMAKE_CXX_COMPILER_LAUNCHER": self.gomacc,
+                },
+            )
 
         if not build_dir.exists():
             build_dir = build_dir.mkdir()
@@ -464,6 +476,9 @@ class Zmake:
             with self.jobserver.get_job():
                 dirs[build_name] = build_dir / "build-{}".format(build_name)
                 cmd = ["/usr/bin/ninja", "-C", dirs[build_name].as_posix()]
+                if self.goma:
+                    # Go nuts ninja, goma does the heavy lifting!
+                    cmd.append("-j1024")
                 self.logger.info(
                     "Building %s:%s: %s",
                     build_dir,
@@ -520,9 +535,50 @@ class Zmake:
 
         return 0
 
+    def _run_test(self, elf_file):
+        """Run a single test, with goma if enabled.
+
+        Args:
+            elf_file: The path to the ELF to run.
+        """
+        cmd = []
+        if self.goma:
+            cmd.append(self.gomacc)
+        cmd.append(elf_file)
+
+        def _run():
+            self.logger.info("Running tests in %s.", elf_file)
+            proc = self.jobserver.popen(
+                cmd,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                encoding="utf-8",
+                errors="replace",
+            )
+            job_id = "test {}".format(elf_file)
+            zmake.multiproc.log_output(
+                self.logger,
+                logging.DEBUG,
+                proc.stdout,
+                job_id=job_id,
+            )
+            zmake.multiproc.log_output(
+                self.logger,
+                logging.ERROR,
+                proc.stderr,
+                job_id=job_id,
+            )
+            if proc.wait():
+                raise OSError(get_process_failure_msg(proc))
+
+        if self.goma:
+            _run()
+        else:
+            with self.jobserver.get_job():
+                _run()
+
     def test(self, build_dir):
         """Test a build directory."""
-        procs = []
         output_files = []
         self.build(build_dir, output_files_out=output_files)
 
@@ -533,33 +589,8 @@ class Zmake:
             return 0
 
         for output_file in output_files:
-            self.logger.info("Running tests in %s.", output_file)
-            with self.jobserver.get_job():
-                proc = self.jobserver.popen(
-                    [output_file],
-                    stdout=subprocess.PIPE,
-                    stderr=subprocess.PIPE,
-                    encoding="utf-8",
-                    errors="replace",
-                )
-                job_id = "test {}".format(output_file)
-                zmake.multiproc.log_output(
-                    self.logger,
-                    logging.DEBUG,
-                    proc.stdout,
-                    job_id=job_id,
-                )
-                zmake.multiproc.log_output(
-                    self.logger,
-                    logging.ERROR,
-                    proc.stderr,
-                    job_id=job_id,
-                )
-                procs.append(proc)
+            self.executor.append(func=functools.partial(self._run_test, output_file))
 
-        for idx, proc in enumerate(procs):
-            if proc.wait():
-                raise OSError(get_process_failure_msg(proc))
         return 0
 
     def testall(self, clobber=False):