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