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 12 13CONFIG = global_config.global_config 14AUTOTEST_INSTALL_DIR = CONFIG.get_config_value('SCHEDULER', 15 'drone_installation_directory') 16DEFAULT_CONTAINER_PATH = CONFIG.get_config_value('AUTOSERV', 'container_path') 17 18SSP_REQUIRED = CONFIG.get_config_value('SCHEDULER', 'exit_on_failed_ssp_setup', 19 default=False) 20 21class DroneUnreachable(Exception): 22 """The drone is non-sshable.""" 23 pass 24 25 26class SiteDrone(object): 27 """ 28 Attributes: 29 * allowed_users: set of usernames allowed to use this drone. if None, 30 any user can use this drone. 31 """ 32 def __init__(self, timestamp_remote_calls=True): 33 """Instantiate an abstract drone. 34 35 @param timestamp_remote_calls: If true, drone_utility is invoked with 36 the --call_time option and the current time. Currently this is only 37 used for testing. 38 """ 39 self._calls = [] 40 self.hostname = None 41 self.enabled = True 42 self.max_processes = 0 43 self.active_processes = 0 44 self.allowed_users = None 45 self._autotest_install_dir = AUTOTEST_INSTALL_DIR 46 self._host = None 47 self.timestamp_remote_calls = timestamp_remote_calls 48 # If drone supports server-side packaging. The property support_ssp will 49 # init self._support_ssp later. 50 self._support_ssp = None 51 self._processes_to_kill = [] 52 53 54 def shutdown(self): 55 pass 56 57 58 @property 59 def _drone_utility_path(self): 60 return os.path.join(self._autotest_install_dir, 61 'scheduler', 'drone_utility.py') 62 63 64 def used_capacity(self): 65 """Gets the capacity used by this drone 66 67 Returns a tuple of (percentage_full, -max_capacity). This is to aid 68 direct comparisons, so that a 0/10 drone is considered less heavily 69 loaded than a 0/2 drone. 70 71 This value should never be used directly. It should only be used in 72 direct comparisons using the basic comparison operators, or using the 73 cmp() function. 74 """ 75 if self.max_processes == 0: 76 return (1.0, 0) 77 return (float(self.active_processes) / self.max_processes, 78 -self.max_processes) 79 80 81 def usable_by(self, user): 82 if self.allowed_users is None: 83 return True 84 return user in self.allowed_users 85 86 87 def _execute_calls_impl(self, calls): 88 if not self._host: 89 raise ValueError('Drone cannot execute calls without a host.') 90 drone_utility_cmd = self._drone_utility_path 91 if self.timestamp_remote_calls: 92 drone_utility_cmd = '%s --call_time %s' % ( 93 drone_utility_cmd, time.time()) 94 logging.info("Running drone_utility on %s", self.hostname) 95 result = self._host.run('python %s' % drone_utility_cmd, 96 stdin=cPickle.dumps(calls), stdout_tee=None, 97 connect_timeout=300) 98 try: 99 return cPickle.loads(result.stdout) 100 except Exception: # cPickle.loads can throw all kinds of exceptions 101 logging.critical('Invalid response:\n---\n%s\n---', result.stdout) 102 raise 103 104 105 def _execute_calls(self, 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 """Execute queued calls. 137 138 If there are any processes queued to kill, kill them then process the 139 remaining queued up calls. 140 """ 141 if self._processes_to_kill: 142 self.queue_call('kill_processes', self._processes_to_kill) 143 self.clear_processes_to_kill() 144 145 if not self._calls: 146 return 147 results = self._execute_calls(self._calls) 148 self.clear_call_queue() 149 return results 150 151 152 def set_autotest_install_dir(self, path): 153 pass 154 155 156 @property 157 def support_ssp(self): 158 """Check if the drone supports server-side packaging with container. 159 160 @return: True if the drone supports server-side packaging with container 161 """ 162 if not self._host: 163 raise ValueError('Can not determine if drone supports server-side ' 164 'packaging before host is set.') 165 if self._support_ssp is None: 166 try: 167 # TODO(crbug.com/471316): We need a better way to check if drone 168 # supports container, and install/upgrade base container. The 169 # check of base container folder is not reliable and shall be 170 # obsoleted once that bug is fixed. 171 self._host.run('which lxc-start') 172 # Test if base container is setup. 173 base_container_name = CONFIG.get_config_value( 174 'AUTOSERV', 'container_base_name') 175 base_container = os.path.join(DEFAULT_CONTAINER_PATH, 176 base_container_name) 177 # SSP uses privileged containers, sudo access is required. If 178 # the process can't run sudo command without password, SSP can't 179 # work properly. sudo command option -n will avoid user input. 180 # If password is required, the command will fail and raise 181 # AutoservRunError exception. 182 self._host.run('sudo -n ls "%s"' % base_container) 183 self._support_ssp = True 184 except (error.AutoservRunError, error.AutotestHostRunError): 185 # Local drone raises AutotestHostRunError, while remote drone 186 # raises AutoservRunError. 187 logging.exception('Drone %s does not support server-side ' 188 'packaging.', self.hostname) 189 self._support_ssp = False 190 if SSP_REQUIRED: 191 raise 192 return self._support_ssp 193 194 195 def queue_kill_process(self, process): 196 """Queue a process to kill/abort. 197 198 @param process: Process to kill/abort. 199 """ 200 self._processes_to_kill.append(process) 201 202 203 def clear_processes_to_kill(self): 204 """Reset the list of processes to kill for this tick.""" 205 self._processes_to_kill = [] 206 207 208class _AbstractDrone(SiteDrone): 209 pass 210 211 212class _LocalDrone(_AbstractDrone): 213 def __init__(self, timestamp_remote_calls=True): 214 super(_LocalDrone, self).__init__( 215 timestamp_remote_calls=timestamp_remote_calls) 216 self.hostname = 'localhost' 217 self._host = local_host.LocalHost() 218 219 220 def send_file_to(self, drone, source_path, destination_path, 221 can_fail=False): 222 if drone.hostname == self.hostname: 223 self.queue_call('copy_file_or_directory', source_path, 224 destination_path) 225 else: 226 self.queue_call('send_file_to', drone.hostname, source_path, 227 destination_path, can_fail) 228 229 230class _RemoteDrone(_AbstractDrone): 231 def __init__(self, hostname, timestamp_remote_calls=True): 232 super(_RemoteDrone, self).__init__( 233 timestamp_remote_calls=timestamp_remote_calls) 234 self.hostname = hostname 235 self._host = drone_utility.create_host(hostname) 236 if not self._host.is_up(): 237 logging.error('Drone %s is unpingable, kicking out', hostname) 238 raise DroneUnreachable 239 240 241 def set_autotest_install_dir(self, path): 242 self._autotest_install_dir = path 243 244 245 def shutdown(self): 246 super(_RemoteDrone, self).shutdown() 247 self._host.close() 248 249 250 def send_file_to(self, drone, source_path, destination_path, 251 can_fail=False): 252 if drone.hostname == self.hostname: 253 self.queue_call('copy_file_or_directory', source_path, 254 destination_path) 255 elif isinstance(drone, _LocalDrone): 256 drone.queue_call('get_file_from', self.hostname, source_path, 257 destination_path) 258 else: 259 self.queue_call('send_file_to', drone.hostname, source_path, 260 destination_path, can_fail) 261 262 263def get_drone(hostname): 264 """ 265 Use this factory method to get drone objects. 266 """ 267 if hostname == 'localhost': 268 return _LocalDrone() 269 try: 270 return _RemoteDrone(hostname) 271 except DroneUnreachable: 272 return None 273