lib/remoteexec_util.py

# Copyright 2021 The Chromium OS Authors. All rights reserved.
# 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(object):
  """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 = Path(
          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)