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