testing/legion/client_lib.py

   1 # Copyright 2015 The Chromium 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 """Defines the client library."""
   6
   7 import argparse
   8 import datetime
   9 import logging
  10 import os
  11 import socket
  12 import subprocess
  13 import sys
  14 import tempfile
  15 import threading
  16 import xmlrpclib
  17
  18 #pylint: disable=relative-import
  19 import common_lib
  20
  21 THIS_DIR = os.path.dirname(os.path.abspath(__file__))
  22 SWARMING_DIR = os.path.join(THIS_DIR, '..', '..', 'tools/swarming_client')
  23 ISOLATE_PY = os.path.join(SWARMING_DIR, 'isolate.py')
  24 SWARMING_PY = os.path.join(SWARMING_DIR, 'swarming.py')
  25
  26
  27 class Error(Exception):
  28   pass
  29
  30
  31 class ConnectionTimeoutError(Error):
  32   pass
  33
  34
  35 class ClientController(object):
  36   """Creates, configures, and controls a client machine."""
  37
  38   _client_count = 0
  39   _controllers = []
  40
  41   def __init__(self, isolate_file, config_vars, dimensions, priority=100,
  42                idle_timeout_secs=common_lib.DEFAULT_TIMEOUT_SECS,
  43                connection_timeout_secs=common_lib.DEFAULT_TIMEOUT_SECS,
  44                verbosity='ERROR', name=None):
  45     assert isinstance(config_vars, dict)
  46     assert isinstance(dimensions, dict)
  47     type(self)._controllers.append(self)
  48     type(self)._client_count += 1
  49     self.verbosity = verbosity
  50     self._name = name or 'Client%d' % type(self)._client_count
  51     self._priority = priority
  52     self._isolate_file = isolate_file
  53     self._isolated_file = isolate_file + 'd'
  54     self._idle_timeout_secs = idle_timeout_secs
  55     self._config_vars = config_vars
  56     self._dimensions = dimensions
  57     self._connect_event = threading.Event()
  58     self._connected = False
  59     self._ip_address = None
  60     self._otp = self._CreateOTP()
  61     self._rpc = None
  62
  63     parser = argparse.ArgumentParser()
  64     parser.add_argument('--isolate-server')
  65     parser.add_argument('--swarming-server')
  66     parser.add_argument('--client-connection-timeout-secs',
  67                         default=common_lib.DEFAULT_TIMEOUT_SECS)
  68     args, _ = parser.parse_known_args()
  69
  70     self._isolate_server = args.isolate_server
  71     self._swarming_server = args.swarming_server
  72     self._connection_timeout_secs = (connection_timeout_secs or
  73                                     args.client_connection_timeout_secs)
  74
  75   @property
  76   def name(self):
  77     return self._name
  78
  79   @property
  80   def otp(self):
  81     return self._otp
  82
  83   @property
  84   def connected(self):
  85     return self._connected
  86
  87   @property
  88   def connect_event(self):
  89     return self._connect_event
  90
  91   @property
  92   def rpc(self):
  93     return self._rpc
  94
  95   @property
  96   def verbosity(self):
  97     return self._verbosity
  98
  99   @verbosity.setter
 100   def verbosity(self, level):
 101     """Sets the verbosity level as a string.
 102
 103     Either a string ('INFO', 'DEBUG', etc) or a logging level (logging.INFO,
 104     logging.DEBUG, etc) is allowed.
 105     """
 106     assert isinstance(level, (str, int))
 107     if isinstance(level, int):
 108       level = logging.getLevelName(level)
 109     self._verbosity = level  #pylint: disable=attribute-defined-outside-init
 110
 111   @classmethod
 112   def ReleaseAllControllers(cls):
 113     for controller in cls._controllers:
 114       controller.Release()
 115
 116   def _CreateOTP(self):
 117     """Creates the OTP."""
 118     host_name = socket.gethostname()
 119     test_name = os.path.basename(sys.argv[0])
 120     creation_time = datetime.datetime.utcnow()
 121     otp = 'client:%s-host:%s-test:%s-creation:%s' % (
 122         self._name, host_name, test_name, creation_time)
 123     return otp
 124
 125   def Create(self):
 126     """Creates the client machine."""
 127     logging.info('Creating %s', self.name)
 128     self._connect_event.clear()
 129     self._ExecuteIsolate()
 130     self._ExecuteSwarming()
 131
 132   def WaitForConnection(self):
 133     """Waits for the client machine to connect.
 134
 135     Raises:
 136       ConnectionTimeoutError if the client doesn't connect in time.
 137     """
 138     logging.info('Waiting for %s to connect with a timeout of %d seconds',
 139                  self._name, self._connection_timeout_secs)
 140     self._connect_event.wait(self._connection_timeout_secs)
 141     if not self._connect_event.is_set():
 142       raise ConnectionTimeoutError('%s failed to connect' % self.name)
 143
 144   def Release(self):
 145     """Quits the client's RPC server so it can release the machine."""
 146     if self._rpc is not None and self._connected:
 147       logging.info('Releasing %s', self._name)
 148       try:
 149         self._rpc.Quit()
 150       except (socket.error, xmlrpclib.Fault):
 151         logging.error('Unable to connect to %s to call Quit', self.name)
 152       self._rpc = None
 153       self._connected = False
 154
 155   def _ExecuteIsolate(self):
 156     """Executes isolate.py."""
 157     cmd = [
 158         'python',
 159         ISOLATE_PY,
 160         'archive',
 161         '--isolate', self._isolate_file,
 162         '--isolated', self._isolated_file,
 163         ]
 164
 165     if self._isolate_server:
 166       cmd.extend(['--isolate-server', self._isolate_server])
 167     for key, value in self._config_vars.iteritems():
 168       cmd.extend(['--config-var', key, value])
 169
 170     self._ExecuteProcess(cmd)
 171
 172   def _ExecuteSwarming(self):
 173     """Executes swarming.py."""
 174     cmd = [
 175         'python',
 176         SWARMING_PY,
 177         'trigger',
 178         self._isolated_file,
 179         '--priority', str(self._priority),
 180         ]
 181
 182     if self._isolate_server:
 183       cmd.extend(['--isolate-server', self._isolate_server])
 184     if self._swarming_server:
 185       cmd.extend(['--swarming', self._swarming_server])
 186     for key, value in self._dimensions.iteritems():
 187       cmd.extend(['--dimension', key, value])
 188
 189     cmd.extend([
 190         '--',
 191         '--host', common_lib.MY_IP,
 192         '--otp', self._otp,
 193         '--verbosity', self._verbosity,
 194         '--idle-timeout', str(self._idle_timeout_secs),
 195         ])
 196
 197     self._ExecuteProcess(cmd)
 198
 199   def _ExecuteProcess(self, cmd):
 200     """Executes a process, waits for it to complete, and checks for success."""
 201     logging.debug('Running %s', ' '.join(cmd))
 202     p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
 203     _, stderr = p.communicate()
 204     if p.returncode != 0:
 205       stderr.seek(0)
 206       raise Error(stderr)
 207
 208   def OnConnect(self, ip_address):
 209     """Receives client ip address on connection."""
 210     self._ip_address = ip_address
 211     self._connected = True
 212     self._rpc = common_lib.ConnectToServer(self._ip_address)
 213     logging.info('%s connected from %s', self._name, ip_address)
 214     self._connect_event.set()