lib/remoteexec_util.py

# Copyright 2021 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Module to use remoteexec from builders."""

import os
from pathlib import Path
from typing import Union

from chromite.lib import constants
from chromite.lib import cros_build_lib


class Remoteexec:
    """Interface to use remoteexec on bots."""

    def __init__(
        self,
        reclient_dir: Union[str, os.PathLike],
        reproxy_cfg_file: Union[str, os.PathLike],
    ):
        """Initializes a Remoteexec instance.

        Args:
            reclient_dir: Path to the re-client directory that contains the
                reproxy, bootstrap, rewrapper binaries.
            reproxy_cfg_file: Path to the config file for starting reproxy.

        Raises:
            ValueError
        """

        reclient_dir = Path(reclient_dir)
        reproxy_cfg_file = Path(reproxy_cfg_file)

        if not reclient_dir.is_dir():
            raise ValueError(
                f"reclient_dir does not point to a directory: {reclient_dir}"
            )

        self.reclient_dir = reclient_dir

        # TODO(crbug.com/1256966): `reproxy_cfg_file` is either the full path
        # of a config generated by recipes or the name of a committed file
        # in `chromite/sdk/reclient_cfgs. When other builders, other than the
        # initial informational builder, are added, they should all be generated
        # by recipes rather than committed.
        if not reproxy_cfg_file.exists():
            reproxy_cfg_file = (
                constants.CHROMITE_DIR
                / "sdk"
                / "reclient_cfgs"
                / reproxy_cfg_file
            )
            if not reproxy_cfg_file.exists():
                raise ValueError(
                    f"reproxy_cfg_file does not exist: {reproxy_cfg_file}"
                )
        self.reproxy_cfg_file = reproxy_cfg_file

    def __eq__(self, other):
        if self.__class__ is other.__class__:
            return (
                self.reclient_dir == other.reclient_dir
                and self.reproxy_cfg_file == other.reproxy_cfg_file
            )
        return NotImplemented

    def GetChrootExtraEnv(self):
        """Extra env vars set to do remoteexec inside chroot."""
        # These paths should match the paths in chroot that the
        # reclient directory and reproxy config file get mapped to
        # in sdk_lib/enter_chroot.sh
        reclient_dir = Path.home() / "reclient"
        reproxy_cfg_file = Path.home() / "reclient_cfgs" / "reproxy_chroot.cfg"
        return {
            "RECLIENT_DIR": str(reclient_dir),
            "REPROXY_CFG": str(reproxy_cfg_file),
        }

    def _RunRemoteExec(self, shutdown: bool = False) -> None:
        """Run RemoteExec command.

        Run the start or shutdown command for the remoteexec.

        Args:
            shutdown: If true, add shutdown command.

        Raises:
            cros_build_lib.RunCommandError
        """
        bootstrap_path = self.reclient_dir / "bootstrap"
        reproxy_path = self.reclient_dir / "reproxy"
        remoteexec_cmd = [
            bootstrap_path,
            "--cfg",
            self.reproxy_cfg_file,
            "--re_proxy",
            reproxy_path,
        ]
        if shutdown:
            remoteexec_cmd.append("--shutdown")
        cros_build_lib.run(remoteexec_cmd)

    def Start(self):
        """Start RemoteExec.

        Run the command to start the remoteexec.

        Raises:
            cros_build_lib.RunCommandError
        """
        self._RunRemoteExec()

    def Stop(self):
        """Shutdown RemoteExec.

        Run the command to shutdown the remoteexec.

        Raises:
            cros_build_lib.RunCommandError
        """
        self._RunRemoteExec(shutdown=True)