src/third_party/chromite/lib/vm.py

   1 # Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4
   5 """VM-related helper functions/classes."""
   6
   7 import logging
   8 import os
   9 import shutil
  10 import time
  11
  12 from chromite.buildbot import constants
  13 from chromite.lib import cros_build_lib
  14 from chromite.lib import osutils
  15 from chromite.lib import remote_access
  16
  17
  18 class VMError(Exception):
  19   """A base exception for VM errors."""
  20
  21
  22 class VMCreationError(VMError):
  23   """Raised when failed to create a VM image."""
  24
  25
  26 def VMIsUpdatable(path):
  27   """Check if the existing VM image is updatable.
  28
  29   Args:
  30     path: Path to the VM image.
  31
  32   Returns:
  33     True if VM is updatable; False otherwise.
  34   """
  35   table = cros_build_lib.GetImageDiskPartitionInfo(path, unit='MB')
  36   # Assume if size of the two root partitions match, the image
  37   # is updatable.
  38   return table['ROOT-B'].size == table['ROOT-A'].size
  39
  40
  41 def CreateVMImage(image=None, board=None, updatable=True, dest_dir=None):
  42   """Returns the path of the image built to run in a VM.
  43
  44   By default, the returned VM is a test image that can run full update
  45   testing on it. If there exists a VM image with the matching
  46   |updatable| setting, this method returns the path to the existing
  47   image. If |dest_dir| is set, it will copy/create the VM image to the
  48   |dest_dir|.
  49
  50   Args:
  51     image: Path to the (non-VM) image. Defaults to None to use the latest
  52       image for the board.
  53     board: Board that the image was built with. If None, attempts to use the
  54       configured default board.
  55     updatable: Create a VM image that supports AU.
  56     dest_dir: If set, create/copy the VM image to |dest|; otherwise,
  57       use the folder where |image| resides.
  58   """
  59   if not image and not board:
  60     raise VMCreationError(
  61         'Cannot create VM when both image and board are None.')
  62
  63   image_dir = os.path.dirname(image)
  64   src_path = dest_path = os.path.join(image_dir, constants.VM_IMAGE_BIN)
  65
  66   if dest_dir:
  67     dest_path = os.path.join(dest_dir, constants.VM_IMAGE_BIN)
  68
  69   exists = False
  70   # Do not create a new VM image if a matching image already exists.
  71   exists = os.path.exists(src_path) and (
  72       not updatable or VMIsUpdatable(src_path))
  73
  74   if exists and dest_dir:
  75     # Copy the existing VM image to dest_dir.
  76     shutil.copyfile(src_path, dest_path)
  77
  78   if not exists:
  79     # No existing VM image that we can reuse. Create a new VM image.
  80     logging.info('Creating %s', dest_path)
  81     cmd = ['./image_to_vm.sh', '--test_image']
  82
  83     if image:
  84       cmd.append('--from=%s' % cros_build_lib.ToChrootPath(image_dir))
  85
  86     if updatable:
  87       cmd.extend(['--disk_layout', '2gb-rootfs-updatable'])
  88
  89     if board:
  90       cmd.extend(['--board', board])
  91
  92     # image_to_vm.sh only runs in chroot, but dest_dir may not be
  93     # reachable from chroot. In that case, we copy it to a temporary
  94     # directory in chroot, and then move it to dest_dir .
  95     tempdir = None
  96     if dest_dir:
  97       # Create a temporary directory in chroot to store the VM
  98       # image. This is to avoid the case where dest_dir is not
  99       # reachable within chroot.
 100       tempdir = cros_build_lib.RunCommand(
 101           ['mktemp', '-d'],
 102           capture_output=True,
 103           enter_chroot=True).output.strip()
 104       cmd.append('--to=%s' % tempdir)
 105
 106     msg = 'Failed to create the VM image'
 107     try:
 108       cros_build_lib.RunCommand(cmd, enter_chroot=True,
 109                                 cwd=constants.SOURCE_ROOT)
 110     except cros_build_lib.RunCommandError as e:
 111       logging.error('%s: %s', msg, e)
 112       if tempdir:
 113         osutils.RmDir(
 114             cros_build_lib.FromChrootPath(tempdir), ignore_missing=True)
 115       raise VMCreationError(msg)
 116
 117     if dest_dir:
 118       # Move VM from tempdir to dest_dir.
 119       shutil.move(
 120           cros_build_lib.FromChrootPath(
 121             os.path.join(tempdir, constants.VM_IMAGE_BIN)), dest_path)
 122       osutils.RmDir(cros_build_lib.FromChrootPath(tempdir), ignore_missing=True)
 123
 124   if not os.path.exists(dest_path):
 125     raise VMCreationError(msg)
 126
 127   return dest_path
 128
 129
 130 class VMStartupError(VMError):
 131   """Raised when failed to start a VM instance."""
 132
 133
 134 class VMStopError(VMError):
 135   """Raised when failed to stop a VM instance."""
 136
 137
 138 class VMInstance(object):
 139   """This is a wrapper of a VM instance."""
 140
 141   MAX_LAUNCH_ATTEMPTS = 5
 142   TIME_BETWEEN_LAUNCH_ATTEMPTS = 30
 143
 144   # VM needs a longer timeout.
 145   SSH_CONNECT_TIMEOUT = 120
 146
 147   def __init__(self, image_path, port=None, tempdir=None,
 148                debug_level=logging.DEBUG):
 149     """Initializes VMWrapper with a VM image path.
 150
 151     Args:
 152       image_path: Path to the VM image.
 153       port: SSH port of the VM.
 154       tempdir: Temporary working directory.
 155       debug_level: Debug level for logging.
 156     """
 157     self.image_path = image_path
 158     self.tempdir = tempdir
 159     self._tempdir_obj = None
 160     if not self.tempdir:
 161       self._tempdir_obj = osutils.TempDir(prefix='vm_wrapper', sudo_rm=True)
 162       self.tempdir = self._tempdir_obj.tempdir
 163     self.kvm_pid_path = os.path.join(self.tempdir, 'kvm.pid')
 164     self.port = (remote_access.GetUnusedPort() if port is None
 165                  else remote_access.NormalizePort(port))
 166     self.debug_level = debug_level
 167     self.ssh_settings = remote_access.CompileSSHConnectSettings(
 168         ConnectTimeout=self.SSH_CONNECT_TIMEOUT)
 169     self.agent = remote_access.RemoteAccess(
 170         remote_access.LOCALHOST, self.tempdir, self.port,
 171         debug_level=self.debug_level, interactive=False)
 172
 173   def _Start(self):
 174     """Run the command to start VM."""
 175     cmd = [os.path.join(constants.CROSUTILS_DIR, 'bin', 'cros_start_vm'),
 176            '--ssh_port', str(self.port),
 177            '--image_path', self.image_path,
 178            '--no_graphics',
 179            '--kvm_pid', self.kvm_pid_path]
 180     try:
 181       self._RunCommand(cmd, capture_output=True)
 182     except cros_build_lib.RunCommandError as e:
 183       msg = 'VM failed to start'
 184       logging.warning('%s: %s', msg, e)
 185       raise VMStartupError(msg)
 186
 187   def Connect(self):
 188     """Returns True if we can connect to VM via SSH."""
 189     try:
 190       self.agent.RemoteSh(['true'], connect_settings=self.ssh_settings)
 191     except Exception:
 192       return False
 193
 194     return True
 195
 196   def Stop(self, ignore_error=False):
 197     """Stops a running VM.
 198
 199     Args:
 200       ignore_error: If set True, do not raise an exception on error.
 201     """
 202     cmd = [os.path.join(constants.CROSUTILS_DIR, 'bin', 'cros_stop_vm'),
 203            '--kvm_pid', self.kvm_pid_path]
 204     result = self._RunCommand(cmd, capture_output=True, error_code_ok=True)
 205     if result.returncode:
 206       msg = 'Failed to stop VM'
 207       if ignore_error:
 208         logging.warning('%s: %s', msg, result.error)
 209       else:
 210         logging.error('%s: %s', msg, result.error)
 211         raise VMStopError(msg)
 212
 213   def Start(self):
 214     """Start VM and wait until we can ssh into it.
 215
 216     This command is more robust than just naively starting the VM as it will
 217     try to start the VM multiple times if the VM fails to start up. This is
 218     inspired by retry_until_ssh in crosutils/lib/cros_vm_lib.sh.
 219     """
 220     for _ in range(self.MAX_LAUNCH_ATTEMPTS):
 221       try:
 222         self._Start()
 223       except VMStartupError:
 224         logging.warning('VM failed to start.')
 225         continue
 226
 227       if self.Connect():
 228         # VM is started up successfully if we can connect to it.
 229         break
 230
 231       logging.warning('Cannot connect to VM...')
 232       self.Stop(ignore_error=True)
 233       time.sleep(self.TIME_BETWEEN_LAUNCH_ATTEMPTS)
 234     else:
 235       raise VMStartupError('Max attempts (%d) to start VM exceeded.'
 236                            % self.MAX_LAUNCH_ATTEMPTS)
 237
 238     logging.info('VM started at port %d', self.port)
 239
 240   def _RunCommand(self, *args, **kwargs):
 241     """Runs a commmand on the host machine."""
 242     kwargs.setdefault('debug_level', self.debug_level)
 243     return cros_build_lib.RunCommand(*args, **kwargs)