Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 1 | # Copyright 2016 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 | |
Mike Frysinger | 666566c | 2016-09-21 00:00:21 -0400 | [diff] [blame] | 5 | """Script for VM Management.""" |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 6 | |
| 7 | from __future__ import print_function |
| 8 | |
| 9 | import os |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 10 | |
| 11 | from chromite.lib import commandline |
| 12 | from chromite.lib import cros_build_lib |
| 13 | from chromite.lib import cros_logging as logging |
Achuith Bhandarkar | f4950ba | 2016-10-11 15:40:07 -0700 | [diff] [blame] | 14 | from chromite.lib import osutils |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 15 | from chromite.lib import remote_access |
Achuith Bhandarkar | 2f8352f | 2017-06-02 12:47:18 -0700 | [diff] [blame] | 16 | from chromite.lib import retry_util |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 17 | |
| 18 | |
| 19 | class VMError(Exception): |
| 20 | """Exception for VM failures.""" |
| 21 | |
| 22 | def __init__(self, message): |
| 23 | super(VMError, self).__init__() |
| 24 | logging.error(message) |
| 25 | |
| 26 | |
| 27 | class VM(object): |
| 28 | """Class for managing a VM.""" |
| 29 | |
| 30 | SSH_PORT = 9222 |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 31 | |
Achuith Bhandarkar | f4950ba | 2016-10-11 15:40:07 -0700 | [diff] [blame] | 32 | def __init__(self, image_path=None, qemu_path=None, enable_kvm=True, |
Achuith Bhandarkar | b891adb | 2016-10-24 18:43:22 -0700 | [diff] [blame] | 33 | display=True, ssh_port=SSH_PORT, dry_run=False): |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 34 | """Initialize VM. |
| 35 | |
| 36 | Args: |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 37 | image_path: path of vm image. |
Achuith Bhandarkar | f4950ba | 2016-10-11 15:40:07 -0700 | [diff] [blame] | 38 | qemu_path: path to qemu binary. |
| 39 | enable_kvm: enable kvm (kernel support for virtualization). |
Achuith Bhandarkar | b891adb | 2016-10-24 18:43:22 -0700 | [diff] [blame] | 40 | display: display video output. |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 41 | ssh_port: ssh port to use. |
| 42 | dry_run: disable VM commands. |
| 43 | """ |
| 44 | |
Achuith Bhandarkar | f4950ba | 2016-10-11 15:40:07 -0700 | [diff] [blame] | 45 | self.qemu_path = qemu_path |
| 46 | self.enable_kvm = enable_kvm |
Achuith Bhandarkar | ee3163d | 2016-10-19 12:58:35 -0700 | [diff] [blame] | 47 | # Software emulation doesn't need sudo access. |
| 48 | self.use_sudo = enable_kvm |
Achuith Bhandarkar | b891adb | 2016-10-24 18:43:22 -0700 | [diff] [blame] | 49 | self.display = display |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 50 | self.image_path = image_path |
| 51 | self.ssh_port = ssh_port |
| 52 | self.dry_run = dry_run |
| 53 | |
Achuith Bhandarkar | ee3163d | 2016-10-19 12:58:35 -0700 | [diff] [blame] | 54 | self.vm_dir = os.path.join(osutils.GetGlobalTempDir(), 'cros_vm') |
| 55 | if os.path.exists(self.vm_dir): |
| 56 | # For security, ensure that vm_dir is not a symlink, and is owned by us or |
| 57 | # by root. |
| 58 | assert not os.path.islink(self.vm_dir), \ |
| 59 | 'VM state dir is misconfigured; please recreate: %s' % self.vm_dir |
| 60 | st_uid = os.stat(self.vm_dir).st_uid |
| 61 | assert st_uid == 0 or st_uid == os.getuid(), \ |
| 62 | 'VM state dir is misconfigured; please recreate: %s' % self.vm_dir |
| 63 | |
| 64 | self.pidfile = os.path.join(self.vm_dir, 'kvm.pid') |
| 65 | self.kvm_monitor = os.path.join(self.vm_dir, 'kvm.monitor') |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 66 | self.kvm_pipe_in = '%s.in' % self.kvm_monitor # to KVM |
| 67 | self.kvm_pipe_out = '%s.out' % self.kvm_monitor # from KVM |
| 68 | self.kvm_serial = '%s.serial' % self.kvm_monitor |
| 69 | |
Achuith Bhandarkar | 65d1a89 | 2017-05-08 14:13:12 -0700 | [diff] [blame] | 70 | self.remote = remote_access.RemoteDevice(remote_access.LOCALHOST, |
| 71 | port=ssh_port) |
| 72 | |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 73 | # TODO(achuith): support nographics, snapshot, mem_path, usb_passthrough, |
| 74 | # moblab, etc. |
| 75 | |
Achuith Bhandarkar | ee3163d | 2016-10-19 12:58:35 -0700 | [diff] [blame] | 76 | |
| 77 | def _RunCommand(self, *args, **kwargs): |
| 78 | """Use SudoRunCommand or RunCommand as necessary.""" |
| 79 | if self.use_sudo: |
| 80 | return cros_build_lib.SudoRunCommand(*args, **kwargs) |
| 81 | else: |
| 82 | return cros_build_lib.RunCommand(*args, **kwargs) |
| 83 | |
| 84 | def _CleanupFiles(self, recreate): |
| 85 | """Cleanup vm_dir. |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 86 | |
| 87 | Args: |
Achuith Bhandarkar | ee3163d | 2016-10-19 12:58:35 -0700 | [diff] [blame] | 88 | recreate: recreate vm_dir. |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 89 | """ |
Achuith Bhandarkar | ee3163d | 2016-10-19 12:58:35 -0700 | [diff] [blame] | 90 | self._RunCommand(['rm', '-rf', self.vm_dir]) |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 91 | if recreate: |
Achuith Bhandarkar | ee3163d | 2016-10-19 12:58:35 -0700 | [diff] [blame] | 92 | self._RunCommand(['mkdir', self.vm_dir]) |
| 93 | self._RunCommand(['chmod', '777', self.vm_dir]) |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 94 | |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 95 | def PerformAction(self, start=False, stop=False, cmd=None): |
| 96 | """Performs an action, one of start, stop, or run a command in the VM. |
| 97 | |
| 98 | Args: |
| 99 | start: start the VM. |
| 100 | stop: stop the VM. |
| 101 | cmd: list or scalar command to run in the VM. |
| 102 | |
| 103 | Returns: |
| 104 | cmd output. |
| 105 | """ |
| 106 | |
| 107 | if not start and not stop and not cmd: |
| 108 | raise VMError('Must specify one of start, stop, or cmd.') |
| 109 | if start: |
| 110 | self.Start() |
| 111 | if stop: |
| 112 | self.Stop() |
| 113 | if cmd: |
| 114 | return self.RemoteCommand(cmd.split()) |
| 115 | |
| 116 | def Start(self): |
| 117 | """Start the VM.""" |
| 118 | |
| 119 | self.Stop() |
| 120 | |
Achuith Bhandarkar | 022d69c | 2016-10-05 14:28:14 -0700 | [diff] [blame] | 121 | logging.debug('Start VM') |
Achuith Bhandarkar | f4950ba | 2016-10-11 15:40:07 -0700 | [diff] [blame] | 122 | if not self.qemu_path: |
| 123 | self.qemu_path = osutils.Which('qemu-system-x86_64') |
| 124 | if not self.qemu_path: |
| 125 | raise VMError('qemu not found.') |
| 126 | logging.debug('qemu path=%s', self.qemu_path) |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 127 | |
| 128 | if not self.image_path: |
| 129 | self.image_path = os.environ.get('VM_IMAGE_PATH', '') |
| 130 | logging.debug('vm image path=%s', self.image_path) |
| 131 | if not self.image_path or not os.path.exists(self.image_path): |
| 132 | raise VMError('VM image path %s does not exist.' % self.image_path) |
| 133 | |
| 134 | self._CleanupFiles(recreate=True) |
| 135 | open(self.kvm_serial, 'w') |
| 136 | for pipe in [self.kvm_pipe_in, self.kvm_pipe_out]: |
| 137 | os.mkfifo(pipe, 0600) |
| 138 | |
Achuith Bhandarkar | f4950ba | 2016-10-11 15:40:07 -0700 | [diff] [blame] | 139 | args = [self.qemu_path, '-m', '2G', '-smp', '4', '-vga', 'cirrus', |
| 140 | '-daemonize', |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 141 | '-pidfile', self.pidfile, |
| 142 | '-chardev', 'pipe,id=control_pipe,path=%s' % self.kvm_monitor, |
| 143 | '-serial', 'file:%s' % self.kvm_serial, |
Achuith Bhandarkar | f4950ba | 2016-10-11 15:40:07 -0700 | [diff] [blame] | 144 | '-mon', 'chardev=control_pipe', |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 145 | '-net', 'nic,model=virtio', |
| 146 | '-net', 'user,hostfwd=tcp::%d-:22' % self.ssh_port, |
| 147 | '-drive', 'file=%s,index=0,media=disk,cache=unsafe' |
| 148 | % self.image_path] |
Achuith Bhandarkar | f4950ba | 2016-10-11 15:40:07 -0700 | [diff] [blame] | 149 | if self.enable_kvm: |
| 150 | args.append('-enable-kvm') |
Achuith Bhandarkar | b891adb | 2016-10-24 18:43:22 -0700 | [diff] [blame] | 151 | if not self.display: |
| 152 | args.extend(['-display', 'none']) |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 153 | logging.info(' '.join(args)) |
| 154 | logging.info('Pid file: %s', self.pidfile) |
| 155 | if not self.dry_run: |
Achuith Bhandarkar | ee3163d | 2016-10-19 12:58:35 -0700 | [diff] [blame] | 156 | self._RunCommand(args) |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 157 | |
| 158 | def _GetVMPid(self): |
| 159 | """Get the pid of the VM. |
| 160 | |
| 161 | Returns: |
| 162 | pid of the VM. |
| 163 | """ |
Achuith Bhandarkar | ee3163d | 2016-10-19 12:58:35 -0700 | [diff] [blame] | 164 | if not os.path.exists(self.vm_dir): |
| 165 | logging.debug('%s not present.', self.vm_dir) |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 166 | return 0 |
| 167 | |
| 168 | if not os.path.exists(self.pidfile): |
Achuith Bhandarkar | f4950ba | 2016-10-11 15:40:07 -0700 | [diff] [blame] | 169 | logging.info('%s does not exist.', self.pidfile) |
| 170 | return 0 |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 171 | |
Achuith Bhandarkar | ee3163d | 2016-10-19 12:58:35 -0700 | [diff] [blame] | 172 | pid = self._RunCommand(['cat', self.pidfile], |
| 173 | redirect_stdout=True).output.rstrip() |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 174 | if not pid.isdigit(): |
Achuith Bhandarkar | f4950ba | 2016-10-11 15:40:07 -0700 | [diff] [blame] | 175 | logging.error('%s in %s is not a pid.', pid, self.pidfile) |
| 176 | return 0 |
| 177 | |
Achuith Bhandarkar | 022d69c | 2016-10-05 14:28:14 -0700 | [diff] [blame] | 178 | return int(pid) |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 179 | |
| 180 | def IsRunning(self): |
| 181 | """Returns True if there's a running VM. |
| 182 | |
| 183 | Returns: |
| 184 | True if there's a running VM. |
| 185 | """ |
Achuith Bhandarkar | f4950ba | 2016-10-11 15:40:07 -0700 | [diff] [blame] | 186 | pid = self._GetVMPid() |
Achuith Bhandarkar | 022d69c | 2016-10-05 14:28:14 -0700 | [diff] [blame] | 187 | if not pid: |
| 188 | return False |
| 189 | |
| 190 | # Make sure the process actually exists. |
Achuith Bhandarkar | ee3163d | 2016-10-19 12:58:35 -0700 | [diff] [blame] | 191 | res = self._RunCommand(['kill', '-0', str(pid)], error_code_ok=True) |
Achuith Bhandarkar | 022d69c | 2016-10-05 14:28:14 -0700 | [diff] [blame] | 192 | return res.returncode == 0 |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 193 | |
| 194 | def Stop(self): |
| 195 | """Stop the VM.""" |
Achuith Bhandarkar | 022d69c | 2016-10-05 14:28:14 -0700 | [diff] [blame] | 196 | logging.debug('Stop VM') |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 197 | |
| 198 | pid = self._GetVMPid() |
Achuith Bhandarkar | f4950ba | 2016-10-11 15:40:07 -0700 | [diff] [blame] | 199 | if pid: |
| 200 | logging.info('Killing %d.', pid) |
| 201 | if not self.dry_run: |
Achuith Bhandarkar | ee3163d | 2016-10-19 12:58:35 -0700 | [diff] [blame] | 202 | self._RunCommand(['kill', '-9', str(pid)], error_code_ok=True) |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 203 | |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 204 | self._CleanupFiles(recreate=False) |
| 205 | |
Achuith Bhandarkar | 2f8352f | 2017-06-02 12:47:18 -0700 | [diff] [blame] | 206 | def _WaitForProcs(self): |
| 207 | """Wait for expected processes to launch.""" |
| 208 | class _TooFewPidsException(Exception): |
| 209 | """Exception for _GetRunningPids to throw.""" |
| 210 | |
| 211 | def _GetRunningPids(exe, numpids): |
| 212 | pids = self.remote.GetRunningPids(exe, full_path=False) |
| 213 | logging.info('%s pids: %s', exe, repr(pids)) |
| 214 | if len(pids) < numpids: |
| 215 | raise _TooFewPidsException() |
| 216 | |
| 217 | def _WaitForProc(exe, numpids): |
| 218 | try: |
| 219 | retry_util.RetryException( |
| 220 | exc_retry=_TooFewPidsException, |
| 221 | max_retry=20, |
| 222 | functor=lambda: _GetRunningPids(exe, numpids), |
| 223 | sleep=2) |
| 224 | except _TooFewPidsException: |
| 225 | raise VMError('_WaitForProcs failed: timed out while waiting for ' |
| 226 | '%d %s processes to start.' % (numpids, exe)) |
| 227 | |
| 228 | # We could also wait for session_manager, nacl_helper, etc, but chrome is |
| 229 | # the long pole. We expect the parent, 2 zygotes, gpu-process, renderer. |
| 230 | # This could potentially break with Mustash. |
| 231 | _WaitForProc('chrome', 5) |
| 232 | |
| 233 | def WaitForBoot(self): |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 234 | """Wait for the VM to boot up. |
| 235 | |
Achuith Bhandarkar | 2f8352f | 2017-06-02 12:47:18 -0700 | [diff] [blame] | 236 | Wait for ssh connection to become active, and wait for all expected chrome |
| 237 | processes to be launched. |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 238 | """ |
Achuith Bhandarkar | ee3163d | 2016-10-19 12:58:35 -0700 | [diff] [blame] | 239 | if not os.path.exists(self.vm_dir): |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 240 | self.Start() |
| 241 | |
Achuith Bhandarkar | 2f8352f | 2017-06-02 12:47:18 -0700 | [diff] [blame] | 242 | try: |
| 243 | result = retry_util.RetryException( |
| 244 | exc_retry=remote_access.SSHConnectionError, |
| 245 | max_retry=10, |
| 246 | functor=lambda: self.RemoteCommand(cmd=['echo']), |
| 247 | sleep=5) |
| 248 | except remote_access.SSHConnectionError: |
| 249 | raise VMError('WaitForBoot timed out trying to connect to VM.') |
| 250 | |
| 251 | if result.returncode != 0: |
| 252 | raise VMError('WaitForBoot failed: %s.' % result.error) |
| 253 | |
| 254 | self._WaitForProcs() |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 255 | |
| 256 | def RemoteCommand(self, cmd): |
| 257 | """Run a remote command in the VM. |
| 258 | |
| 259 | Args: |
| 260 | cmd: command to run, of list type. |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 261 | """ |
| 262 | if not isinstance(cmd, list): |
| 263 | raise VMError('cmd must be a list.') |
| 264 | |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 265 | if not self.dry_run: |
Achuith Bhandarkar | 65d1a89 | 2017-05-08 14:13:12 -0700 | [diff] [blame] | 266 | return self.remote.RunCommand(cmd, debug_level=logging.INFO, |
| 267 | combine_stdout_stderr=True, |
| 268 | log_output=True, |
| 269 | error_code_ok=True) |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 270 | |
| 271 | def ParseCommandLine(argv): |
| 272 | """Parse the command line. |
| 273 | |
| 274 | Args: |
| 275 | argv: Command arguments. |
| 276 | |
| 277 | Returns: |
| 278 | List of parsed args. |
| 279 | """ |
| 280 | parser = commandline.ArgumentParser(description=__doc__) |
Achuith Bhandarkar | f4950ba | 2016-10-11 15:40:07 -0700 | [diff] [blame] | 281 | parser.add_argument('--start', action='store_true', default=False, |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 282 | help='Start the VM.') |
Achuith Bhandarkar | f4950ba | 2016-10-11 15:40:07 -0700 | [diff] [blame] | 283 | parser.add_argument('--stop', action='store_true', default=False, |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 284 | help='Stop the VM.') |
| 285 | parser.add_argument('--cmd', help='Run this command in the VM.') |
| 286 | parser.add_argument('--image-path', type='path', |
| 287 | help='Path to VM image to launch with --start.') |
Achuith Bhandarkar | f4950ba | 2016-10-11 15:40:07 -0700 | [diff] [blame] | 288 | parser.add_argument('--qemu-path', type='path', |
| 289 | help='Path of qemu binary to launch with --start.') |
Achuith Bhandarkar | b891adb | 2016-10-24 18:43:22 -0700 | [diff] [blame] | 290 | parser.add_argument('--disable-kvm', dest='enable_kvm', |
| 291 | action='store_false', default=True, |
Achuith Bhandarkar | f4950ba | 2016-10-11 15:40:07 -0700 | [diff] [blame] | 292 | help='Disable KVM, use software emulation.') |
Achuith Bhandarkar | b891adb | 2016-10-24 18:43:22 -0700 | [diff] [blame] | 293 | parser.add_argument('--no-display', dest='display', |
| 294 | action='store_false', default=True, |
| 295 | help='Do not display video output.') |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 296 | parser.add_argument('--ssh-port', type=int, default=VM.SSH_PORT, |
| 297 | help='ssh port to communicate with VM.') |
Achuith Bhandarkar | f4950ba | 2016-10-11 15:40:07 -0700 | [diff] [blame] | 298 | parser.add_argument('--dry-run', action='store_true', default=False, |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 299 | help='dry run for debugging.') |
| 300 | return parser.parse_args(argv) |
| 301 | |
| 302 | |
| 303 | def main(argv): |
| 304 | args = ParseCommandLine(argv) |
Achuith Bhandarkar | b891adb | 2016-10-24 18:43:22 -0700 | [diff] [blame] | 305 | vm = VM(image_path=args.image_path, qemu_path=args.qemu_path, |
| 306 | enable_kvm=args.enable_kvm, display=args.display, |
Achuith Bhandarkar | d8d1929 | 2016-05-03 14:32:58 -0700 | [diff] [blame] | 307 | ssh_port=args.ssh_port, dry_run=args.dry_run) |
| 308 | vm.PerformAction(start=args.start, stop=args.stop, cmd=args.cmd) |