1#pylint: disable-msg=C0111 2 3import cPickle 4import logging 5import os 6import time 7 8import common 9from autotest_lib.scheduler import drone_utility, email_manager 10from autotest_lib.client.bin import local_host 11from autotest_lib.client.common_lib import error, global_config, utils 12from autotest_lib.client.common_lib.cros.graphite import autotest_stats 13 14 15AUTOTEST_INSTALL_DIR = global_config.global_config.get_config_value('SCHEDULER', 16 'drone_installation_directory') 17DEFAULT_CONTAINER_PATH = global_config.global_config.get_config_value( 18 'AUTOSERV', 'container_path') 19 20class DroneUnreachable(Exception): 21 """The drone is non-sshable.""" 22 pass 23 24 25class _BaseAbstractDrone(object): 26 """ 27 Attributes: 28 * allowed_users: set of usernames allowed to use this drone. if None, 29 any user can use this drone. 30 """ 31 def __init__(self, timestamp_remote_calls=True): 32 """Instantiate an abstract drone. 33 34 @param timestamp_remote_calls: If true, drone_utility is invoked with 35 the --call_time option and the current time. Currently this is only 36 used for testing. 37 """ 38 self._calls = [] 39 self.hostname = None 40 self.enabled = True 41 self.max_processes = 0 42 self.active_processes = 0 43 self.allowed_users = None 44 self._autotest_install_dir = AUTOTEST_INSTALL_DIR 45 self._host = None 46 self.timestamp_remote_calls = timestamp_remote_calls 47 # If drone supports server-side packaging. The property support_ssp will 48 # init self._support_ssp later. 49 self._support_ssp = None 50 51 52 def shutdown(self): 53 pass 54 55 56 @property 57 def _drone_utility_path(self): 58 return os.path.join(self._autotest_install_dir, 59 'scheduler', 'drone_utility.py') 60 61 62 def used_capacity(self): 63 """Gets the capacity used by this drone 64 65 Returns a tuple of (percentage_full, -max_capacity). This is to aid 66 direct comparisons, so that a 0/10 drone is considered less heavily 67 loaded than a 0/2 drone. 68 69 This value should never be used directly. It should only be used in 70 direct comparisons using the basic comparison operators, or using the 71 cmp() function. 72 """ 73 if self.max_processes == 0: 74 return (1.0, 0) 75 return (float(self.active_processes) / self.max_processes, 76 -self.max_processes) 77 78 79 def usable_by(self, user): 80 if self.allowed_users is None: 81 return True 82 return user in self.allowed_users 83 84 85 def _execute_calls_impl(self, calls): 86 if not self._host: 87 raise ValueError('Drone cannot execute calls without a host.') 88 drone_utility_cmd = self._drone_utility_path 89 if self.timestamp_remote_calls: 90 drone_utility_cmd = '%s --call_time %s' % ( 91 drone_utility_cmd, time.time()) 92 logging.info("Running drone_utility on %s", self.hostname) 93 result = self._host.run('python %s' % drone_utility_cmd, 94 stdin=cPickle.dumps(calls), stdout_tee=None, 95 connect_timeout=300) 96 try: 97 return cPickle.loads(result.stdout) 98 except Exception: # cPickle.loads can throw all kinds of exceptions 99 logging.critical('Invalid response:\n---\n%s\n---', result.stdout) 100 raise 101 102 103 def _execute_calls(self, calls): 104 autotest_stats.Gauge('drone_execute_call_count').send( 105 self.hostname.replace('.', '_'), len(calls)) 106 return_message = self._execute_calls_impl(calls) 107 for warning in return_message['warnings']: 108 subject = 'Warning from drone %s' % self.hostname 109 logging.warning(subject + '\n' + warning) 110 email_manager.manager.enqueue_notify_email(subject, warning) 111 return return_message['results'] 112 113 114 def get_calls(self): 115 """Returns the calls queued against this drone. 116 117 @return: A list of calls queued against the drone. 118 """ 119 return self._calls 120 121 122 def call(self, method, *args, **kwargs): 123 return self._execute_calls( 124 [drone_utility.call(method, *args, **kwargs)]) 125 126 127 def queue_call(self, method, *args, **kwargs): 128 self._calls.append(drone_utility.call(method, *args, **kwargs)) 129 130 131 def clear_call_queue(self): 132 self._calls = [] 133 134 135 def execute_queued_calls(self): 136 if not self._calls: 137 return 138 results = self._execute_calls(self._calls) 139 self.clear_call_queue() 140 return results 141 142 143 def set_autotest_install_dir(self, path): 144 pass 145 146 147 @property 148 def support_ssp(self): 149 """Check if the drone supports server-side packaging with container. 150 151 @return: True if the drone supports server-side packaging with container 152 """ 153 if not self._host: 154 raise ValueError('Can not determine if drone supports server-side ' 155 'packaging before host is set.') 156 if self._support_ssp is None: 157 try: 158 # TODO(crbug.com/471316): We need a better way to check if drone 159 # supports container, and install/upgrade base container. The 160 # check of base container folder is not reliable and shall be 161 # obsoleted once that bug is fixed. 162 self._host.run('which lxc-start') 163 # Test if base container is setup. 164 base_container = os.path.join(DEFAULT_CONTAINER_PATH, 'base') 165 # SSP uses privileged containers, sudo access is required. If 166 # the process can't run sudo command without password, SSP can't 167 # work properly. sudo command option -n will avoid user input. 168 # If password is required, the command will fail and raise 169 # AutoservRunError exception. 170 self._host.run('sudo -n ls "%s"' % base_container) 171 self._support_ssp = True 172 except (error.AutoservRunError, error.AutotestHostRunError): 173 # Local drone raises AutotestHostRunError, while remote drone 174 # raises AutoservRunError. 175 self._support_ssp = False 176 return self._support_ssp 177 178 179SiteDrone = utils.import_site_class( 180 __file__, 'autotest_lib.scheduler.site_drones', 181 '_SiteAbstractDrone', _BaseAbstractDrone) 182 183 184class _AbstractDrone(SiteDrone): 185 pass 186 187 188class _LocalDrone(_AbstractDrone): 189 def __init__(self, timestamp_remote_calls=True): 190 super(_LocalDrone, self).__init__( 191 timestamp_remote_calls=timestamp_remote_calls) 192 self.hostname = 'localhost' 193 self._host = local_host.LocalHost() 194 self._drone_utility = drone_utility.DroneUtility() 195 196 197 def send_file_to(self, drone, source_path, destination_path, 198 can_fail=False): 199 if drone.hostname == self.hostname: 200 self.queue_call('copy_file_or_directory', source_path, 201 destination_path) 202 else: 203 self.queue_call('send_file_to', drone.hostname, source_path, 204 destination_path, can_fail) 205 206 207class _RemoteDrone(_AbstractDrone): 208 def __init__(self, hostname, timestamp_remote_calls=True): 209 super(_RemoteDrone, self).__init__( 210 timestamp_remote_calls=timestamp_remote_calls) 211 self.hostname = hostname 212 self._host = drone_utility.create_host(hostname) 213 if not self._host.is_up(): 214 logging.error('Drone %s is unpingable, kicking out', hostname) 215 raise DroneUnreachable 216 217 218 def set_autotest_install_dir(self, path): 219 self._autotest_install_dir = path 220 221 222 def shutdown(self): 223 super(_RemoteDrone, self).shutdown() 224 self._host.close() 225 226 227 def send_file_to(self, drone, source_path, destination_path, 228 can_fail=False): 229 if drone.hostname == self.hostname: 230 self.queue_call('copy_file_or_directory', source_path, 231 destination_path) 232 elif isinstance(drone, _LocalDrone): 233 drone.queue_call('get_file_from', self.hostname, source_path, 234 destination_path) 235 else: 236 self.queue_call('send_file_to', drone.hostname, source_path, 237 destination_path, can_fail) 238 239 240def get_drone(hostname): 241 """ 242 Use this factory method to get drone objects. 243 """ 244 if hostname == 'localhost': 245 return _LocalDrone() 246 try: 247 return _RemoteDrone(hostname) 248 except DroneUnreachable: 249 return None 250