scripts: chromite: add function to deploy compressed ash

The CL introduces deployment code for an experimental model where
ash-chrome main binary is packaged into squashfs. Unlike Lacros, that
packages the whole /opt/google/chrome, this experiment packages the
binary only.

BUG=b:247397013
TEST=unittest
TEST=cros chrome-sdk ...
TEST=chrome_deploy ...

Change-Id: Ib6a49b2becb86c91996ad942c83b3beaa6ab6c9f
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/3899689
Commit-Queue: Daniil Lunev <dlunev@chromium.org>
Reviewed-by: Mike Frysinger <vapier@chromium.org>
Reviewed-by: Achuith Bhandarkar <achuith@chromium.org>
Tested-by: Daniil Lunev <dlunev@chromium.org>
diff --git a/scripts/deploy_chrome.py b/scripts/deploy_chrome.py
index 64b3684..1c30a97 100644
--- a/scripts/deploy_chrome.py
+++ b/scripts/deploy_chrome.py
@@ -75,6 +75,14 @@
 
 DF_COMMAND = "df -k %s"
 
+# This constants are related to an experiment of running compressed ash chrome
+# to save rootfs space. See b/247397013
+COMPRESSED_ASH_SERVICE = "mount-ash-chrome"
+COMPRESSED_ASH_FILE = "chrome.squashfs"
+RAW_ASH_FILE = "chrome"
+COMPRESSED_ASH_PATH = os.path.join(_CHROME_DIR, COMPRESSED_ASH_FILE)
+RAW_ASH_PATH = os.path.join(_CHROME_DIR, RAW_ASH_FILE)
+
 LACROS_DIR = "/usr/local/lacros-chrome"
 _CONF_FILE = "/etc/chrome_dev.conf"
 _KILL_LACROS_CHROME_CMD = "pkill -f %(lacros_dir)s/chrome"
@@ -147,6 +155,9 @@
                 private_key=options.private_key,
                 include_dev_paths=False,
             )
+            if self._ShouldUseCompressedAsh():
+                self.options.compressed_ash = True
+
         self._root_dir_is_still_readonly = multiprocessing.Event()
 
         self._deployment_name = "lacros" if options.lacros else "chrome"
@@ -157,6 +168,13 @@
         # Whether UI was stopped during setup.
         self._stopped_ui = False
 
+    def _ShouldUseCompressedAsh(self):
+        """Detects if the DUT uses compressed-ash setup."""
+        if self.options.lacros:
+            return False
+
+        return self.device.IfFileExists(COMPRESSED_ASH_PATH)
+
     def _GetRemoteMountFree(self, remote_dir):
         result = self.device.run(DF_COMMAND % remote_dir)
         line = result.stdout.splitlines()[1]
@@ -291,6 +309,21 @@
                     # Wait for processes to actually terminate
                     time.sleep(POST_KILL_WAIT)
                     logging.info("Rechecking the chrome binary...")
+                if self.options.compressed_ash:
+                    result = self.device.run(
+                        ["umount", RAW_ASH_PATH],
+                        check=False,
+                        capture_output=True,
+                    )
+                    if result.returncode and not (
+                        result.returncode == 32
+                        and "not mounted" in result.stderr
+                    ):
+                        raise DeployFailure(
+                            "Could not unmount compressed ash. "
+                            f"Error Code: {result.returncode}, "
+                            f"Error Message: {result.stderr}"
+                        )
         except timeout_util.TimeoutError:
             msg = (
                 "Could not kill processes after %s seconds.  Please exit any "
@@ -416,6 +449,9 @@
                 ["chown", "-R", "chronos:chronos", self.options.target_dir]
             )
 
+        if self.options.compressed_ash:
+            self.device.run(["start", COMPRESSED_ASH_SERVICE])
+
         # Send SIGHUP to dbus-daemon to tell it to reload its configs. This won't
         # pick up major changes (bus type, logging, etc.), but all we care about is
         # getting the latest policy from /opt/google/chrome/dbus so that Chrome will
@@ -923,10 +959,17 @@
         action="store",
         default="auto",
         choices=("always", "never", "auto"),
-        help="Choose the data compression behavior. Default "
+        help="Choose the data transfer compression behavior. Default "
         'is set to "auto", that disables compression if '
         "the target device has a gigabit ethernet port.",
     )
+    parser.add_argument(
+        "--compressed-ash",
+        action="store_true",
+        default=False,
+        help="Use compressed-ash deployment scheme. With the flag, ash-chrome "
+        "binary is stored on DUT in squashfs, mounted upon boot.",
+    )
     return parser
 
 
@@ -953,6 +996,8 @@
             parser.error("--lacros does not support --deploy-test-binaries")
         if options.local_pkg_path:
             parser.error("--lacros does not support --local-pkg-path")
+        if options.compressed_ash:
+            parser.error("--lacros does not support --compressed-ash")
     else:
         if not options.board and options.build_dir:
             match = re.search(r"out_([^/]+)/Release$", options.build_dir)
@@ -1194,6 +1239,33 @@
                 cwd=staging_dir,
             )
 
+    if options.compressed_ash:
+        # Setup SDK here so mksquashfs is still found in no-shell + nostrip
+        # configuration.
+        sdk = cros_chrome_sdk.SDKFetcher(
+            options.cache_dir,
+            options.board,
+            use_external_config=options.use_external_config,
+        )
+        with sdk.Prepare(
+            components=[],
+            target_tc=options.target_tc,
+            toolchain_url=options.toolchain_url,
+        ):
+            cros_build_lib.dbg_run(
+                [
+                    "mksquashfs",
+                    RAW_ASH_FILE,
+                    COMPRESSED_ASH_FILE,
+                    "-all-root",
+                    "-no-progress",
+                    "-comp",
+                    "zstd",
+                ],
+                cwd=staging_dir,
+            )
+        os.truncate(os.path.join(staging_dir, RAW_ASH_FILE), 0)
+
     if options.staging_upload:
         _UploadStagingDir(options, tempdir, staging_dir)