showard | 170873e | 2009-01-07 00:22:26 +0000 | [diff] [blame] | 1 | #!/usr/bin/python2.4 |
| 2 | |
| 3 | import pickle, subprocess, os, shutil, socket, sys, time, signal, getpass |
showard | 78d4d97 | 2009-01-16 03:04:16 +0000 | [diff] [blame^] | 4 | import datetime, traceback, tempfile |
showard | 170873e | 2009-01-07 00:22:26 +0000 | [diff] [blame] | 5 | import common |
| 6 | from autotest_lib.client.common_lib import utils, global_config, error |
| 7 | from autotest_lib.server import hosts, subcommand |
showard | d1ee1dd | 2009-01-07 21:33:08 +0000 | [diff] [blame] | 8 | from autotest_lib.scheduler import email_manager, scheduler_config |
showard | 170873e | 2009-01-07 00:22:26 +0000 | [diff] [blame] | 9 | |
| 10 | _TEMPORARY_DIRECTORY = 'drone_tmp' |
| 11 | _TRANSFER_FAILED_FILE = '.transfer_failed' |
| 12 | |
| 13 | class _MethodCall(object): |
| 14 | def __init__(self, method, args, kwargs): |
| 15 | self._method = method |
| 16 | self._args = args |
| 17 | self._kwargs = kwargs |
| 18 | |
| 19 | |
| 20 | def execute_on(self, drone_utility): |
| 21 | method = getattr(drone_utility, self._method) |
| 22 | return method(*self._args, **self._kwargs) |
| 23 | |
| 24 | |
| 25 | def __str__(self): |
| 26 | args = ', '.join(repr(arg) for arg in self._args) |
| 27 | kwargs = ', '.join('%s=%r' % (key, value) for key, value in |
| 28 | self._kwargs.iteritems()) |
| 29 | full_args = ', '.join(item for item in (args, kwargs) if item) |
| 30 | return '%s(%s)' % (self._method, full_args) |
| 31 | |
| 32 | |
| 33 | def call(method, *args, **kwargs): |
| 34 | return _MethodCall(method, args, kwargs) |
| 35 | |
| 36 | |
| 37 | class DroneUtility(object): |
| 38 | """ |
| 39 | This class executes actual OS calls on the drone machine. |
| 40 | |
| 41 | All paths going into and out of this class are absolute. |
| 42 | """ |
| 43 | _PS_ARGS = ['pid', 'pgid', 'ppid', 'comm', 'args'] |
showard | 170873e | 2009-01-07 00:22:26 +0000 | [diff] [blame] | 44 | _WARNING_DURATION = 60 |
| 45 | |
| 46 | def __init__(self): |
| 47 | self.warnings = [] |
| 48 | self._subcommands = [] |
| 49 | |
| 50 | |
| 51 | def initialize(self, results_dir): |
| 52 | temporary_directory = os.path.join(results_dir, _TEMPORARY_DIRECTORY) |
| 53 | if os.path.exists(temporary_directory): |
| 54 | shutil.rmtree(temporary_directory) |
| 55 | self._ensure_directory_exists(temporary_directory) |
| 56 | |
| 57 | # make sure there are no old parsers running |
| 58 | os.system('killall parse') |
| 59 | |
| 60 | |
| 61 | def _warn(self, warning): |
| 62 | self.warnings.append(warning) |
| 63 | |
| 64 | |
showard | 0205a3e | 2009-01-16 03:03:50 +0000 | [diff] [blame] | 65 | def _refresh_processes(self, command_name): |
showard | 170873e | 2009-01-07 00:22:26 +0000 | [diff] [blame] | 66 | ps_proc = subprocess.Popen( |
| 67 | ['/bin/ps', 'x', '-o', ','.join(self._PS_ARGS)], |
| 68 | stdout=subprocess.PIPE) |
| 69 | ps_output = ps_proc.communicate()[0] |
| 70 | |
| 71 | # split each line into the columns output by ps |
| 72 | split_lines = [line.split(None, 4) for line in ps_output.splitlines()] |
| 73 | process_infos = [dict(zip(self._PS_ARGS, line_components)) |
| 74 | for line_components in split_lines] |
| 75 | |
| 76 | processes = [] |
| 77 | for info in process_infos: |
showard | 0205a3e | 2009-01-16 03:03:50 +0000 | [diff] [blame] | 78 | if info['comm'] == command_name: |
showard | 170873e | 2009-01-07 00:22:26 +0000 | [diff] [blame] | 79 | processes.append(info) |
| 80 | |
| 81 | return processes |
| 82 | |
| 83 | |
| 84 | def _read_pidfiles(self, pidfile_paths): |
| 85 | pidfiles = {} |
| 86 | for pidfile_path in pidfile_paths: |
| 87 | if not os.path.exists(pidfile_path): |
| 88 | continue |
| 89 | try: |
| 90 | file_object = open(pidfile_path, 'r') |
| 91 | pidfiles[pidfile_path] = file_object.read() |
| 92 | file_object.close() |
| 93 | except IOError: |
| 94 | continue |
| 95 | return pidfiles |
| 96 | |
| 97 | |
| 98 | def refresh(self, pidfile_paths): |
| 99 | results = { |
| 100 | 'pidfiles' : self._read_pidfiles(pidfile_paths), |
showard | 0205a3e | 2009-01-16 03:03:50 +0000 | [diff] [blame] | 101 | 'autoserv_processes' : self._refresh_processes('autoserv'), |
| 102 | 'parse_processes' : self._refresh_processes('parse'), |
showard | 170873e | 2009-01-07 00:22:26 +0000 | [diff] [blame] | 103 | 'pidfiles_second_read' : self._read_pidfiles(pidfile_paths), |
| 104 | } |
| 105 | return results |
| 106 | |
| 107 | |
| 108 | def _is_process_running(self, process): |
| 109 | # TODO: enhance this to check the process args |
| 110 | proc_path = os.path.join('/proc', str(process.pid)) |
| 111 | return os.path.exists(proc_path) |
| 112 | |
| 113 | |
| 114 | def kill_process(self, process): |
| 115 | if self._is_process_running(process): |
| 116 | os.kill(process.pid, signal.SIGCONT) |
| 117 | os.kill(process.pid, signal.SIGTERM) |
| 118 | |
| 119 | |
showard | 78d4d97 | 2009-01-16 03:04:16 +0000 | [diff] [blame^] | 120 | def _convert_old_host_log(self, log_path): |
| 121 | """ |
| 122 | For backwards compatibility only. This can safely be removed in the |
| 123 | future. |
| 124 | |
| 125 | The scheduler used to create files at results/hosts/<hostname>, and |
| 126 | append all host logs to that file. Now, it creates directories at |
| 127 | results/hosts/<hostname>, and places individual timestamped log files |
| 128 | into that directory. |
| 129 | |
| 130 | This can be a problem the first time the scheduler runs after upgrading. |
| 131 | To work around that, we'll look for a file at the path where the |
| 132 | directory should be, and if we find one, we'll automatically convert it |
| 133 | to a directory containing the old logfile. |
| 134 | """ |
| 135 | # move the file out of the way |
| 136 | temp_dir = tempfile.mkdtemp(suffix='.convert_host_log') |
| 137 | base_name = os.path.basename(log_path) |
| 138 | temp_path = os.path.join(temp_dir, base_name) |
| 139 | os.rename(log_path, temp_path) |
| 140 | |
| 141 | os.mkdir(log_path) |
| 142 | |
| 143 | # and move it into the new directory |
| 144 | os.rename(temp_path, os.path.join(log_path, 'old_log')) |
| 145 | os.rmdir(temp_dir) |
| 146 | |
| 147 | |
showard | 170873e | 2009-01-07 00:22:26 +0000 | [diff] [blame] | 148 | def _ensure_directory_exists(self, path): |
showard | 78d4d97 | 2009-01-16 03:04:16 +0000 | [diff] [blame^] | 149 | if os.path.isdir(path): |
| 150 | return |
| 151 | |
| 152 | if os.path.exists(path): |
| 153 | # path exists already, but as a file, not a directory |
| 154 | if '/hosts/' in path: |
| 155 | self._convert_old_host_log(path) |
| 156 | return |
| 157 | else: |
| 158 | raise IOError('Path %s exists as a file, not a directory') |
| 159 | |
| 160 | os.makedirs(path) |
showard | 170873e | 2009-01-07 00:22:26 +0000 | [diff] [blame] | 161 | |
| 162 | |
| 163 | def execute_command(self, command, working_directory, log_file, |
| 164 | pidfile_name): |
| 165 | out_file = None |
| 166 | if log_file: |
| 167 | self._ensure_directory_exists(os.path.dirname(log_file)) |
| 168 | try: |
| 169 | out_file = open(log_file, 'a') |
| 170 | separator = ('*' * 80) + '\n' |
| 171 | out_file.write('\n' + separator) |
| 172 | out_file.write("%s> %s\n" % (time.strftime("%X %x"), command)) |
| 173 | out_file.write(separator) |
| 174 | except (OSError, IOError): |
| 175 | email_manager.manager.log_stacktrace( |
| 176 | 'Error opening log file %s' % log_file) |
| 177 | |
| 178 | if not out_file: |
| 179 | out_file = open('/dev/null', 'w') |
| 180 | |
| 181 | in_devnull = open('/dev/null', 'r') |
| 182 | |
| 183 | self._ensure_directory_exists(working_directory) |
| 184 | pidfile_path = os.path.join(working_directory, pidfile_name) |
| 185 | if os.path.exists(pidfile_path): |
| 186 | self._warn('Pidfile %s already exists' % pidfile_path) |
| 187 | os.remove(pidfile_path) |
| 188 | |
| 189 | subprocess.Popen(command, stdout=out_file, stderr=subprocess.STDOUT, |
| 190 | stdin=in_devnull) |
| 191 | out_file.close() |
| 192 | in_devnull.close() |
| 193 | |
| 194 | |
| 195 | def write_to_file(self, file_path, contents): |
| 196 | self._ensure_directory_exists(os.path.dirname(file_path)) |
| 197 | try: |
| 198 | file_object = open(file_path, 'a') |
| 199 | file_object.write(contents) |
| 200 | file_object.close() |
| 201 | except IOError, exc: |
| 202 | self._warn('Error write to file %s: %s' % (file_path, exc)) |
| 203 | |
| 204 | |
| 205 | def copy_file(self, source_path, destination_path): |
| 206 | if source_path == destination_path: |
| 207 | return |
| 208 | self._ensure_directory_exists(os.path.dirname(destination_path)) |
| 209 | shutil.copy(source_path, destination_path) |
| 210 | |
| 211 | |
showard | c408c5e | 2009-01-08 23:30:53 +0000 | [diff] [blame] | 212 | def wait_for_all_async_commands(self): |
| 213 | for subproc in self._subcommands: |
| 214 | subproc.fork_waitfor() |
| 215 | self._subcommands = [] |
| 216 | |
| 217 | |
| 218 | def _poll_async_commands(self): |
| 219 | still_running = [] |
| 220 | for subproc in self._subcommands: |
| 221 | if subproc.poll() is None: |
| 222 | still_running.append(subproc) |
| 223 | self._subcommands = still_running |
| 224 | |
| 225 | |
| 226 | def _wait_for_some_async_commands(self): |
| 227 | self._poll_async_commands() |
| 228 | max_processes = scheduler_config.config.max_transfer_processes |
| 229 | while len(self._subcommands) >= max_processes: |
| 230 | time.sleep(1) |
| 231 | self._poll_async_commands() |
| 232 | |
| 233 | |
showard | 170873e | 2009-01-07 00:22:26 +0000 | [diff] [blame] | 234 | def run_async_command(self, function, args): |
| 235 | subproc = subcommand.subcommand(function, args) |
| 236 | self._subcommands.append(subproc) |
| 237 | subproc.fork_start() |
| 238 | |
| 239 | |
showard | 170873e | 2009-01-07 00:22:26 +0000 | [diff] [blame] | 240 | def _sync_get_file_from(self, hostname, source_path, destination_path): |
| 241 | self._ensure_directory_exists(os.path.dirname(destination_path)) |
| 242 | host = create_host(hostname) |
| 243 | host.get_file(source_path, destination_path, delete_dest=True) |
| 244 | |
| 245 | |
| 246 | def get_file_from(self, hostname, source_path, destination_path): |
| 247 | self.run_async_command(self._sync_get_file_from, |
| 248 | (hostname, source_path, destination_path)) |
| 249 | |
| 250 | |
| 251 | def _sync_send_file_to(self, hostname, source_path, destination_path, |
| 252 | can_fail): |
| 253 | host = create_host(hostname) |
| 254 | try: |
| 255 | host.run('mkdir -p ' + os.path.dirname(destination_path)) |
| 256 | host.send_file(source_path, destination_path, delete_dest=True) |
| 257 | except error.AutoservError: |
| 258 | if not can_fail: |
| 259 | raise |
| 260 | |
| 261 | if os.path.isdir(source_path): |
| 262 | failed_file = os.path.join(source_path, _TRANSFER_FAILED_FILE) |
| 263 | file_object = open(failed_file, 'w') |
| 264 | try: |
| 265 | file_object.write('%s:%s\n%s\n%s' % |
| 266 | (hostname, destination_path, |
| 267 | datetime.datetime.now(), |
| 268 | traceback.format_exc())) |
| 269 | finally: |
| 270 | file_object.close() |
| 271 | else: |
| 272 | copy_to = destination_path + _TRANSFER_FAILED_FILE |
| 273 | self._ensure_directory_exists(os.path.dirname(copy_to)) |
| 274 | self.copy_file(source_path, copy_to) |
| 275 | |
| 276 | |
| 277 | def send_file_to(self, hostname, source_path, destination_path, |
| 278 | can_fail=False): |
| 279 | self.run_async_command(self._sync_send_file_to, |
| 280 | (hostname, source_path, destination_path, |
| 281 | can_fail)) |
| 282 | |
| 283 | |
| 284 | def _report_long_execution(self, calls, duration): |
| 285 | call_count = {} |
| 286 | for call in calls: |
| 287 | call_count.setdefault(call._method, 0) |
| 288 | call_count[call._method] += 1 |
| 289 | call_summary = '\n'.join('%d %s' % (count, method) |
| 290 | for method, count in call_count.iteritems()) |
| 291 | self._warn('Execution took %f sec\n%s' % (duration, call_summary)) |
| 292 | |
| 293 | |
| 294 | def execute_calls(self, calls): |
| 295 | results = [] |
| 296 | start_time = time.time() |
showard | c408c5e | 2009-01-08 23:30:53 +0000 | [diff] [blame] | 297 | max_processes = scheduler_config.config.max_transfer_processes |
showard | 170873e | 2009-01-07 00:22:26 +0000 | [diff] [blame] | 298 | for method_call in calls: |
| 299 | results.append(method_call.execute_on(self)) |
showard | d1ee1dd | 2009-01-07 21:33:08 +0000 | [diff] [blame] | 300 | if len(self._subcommands) >= max_processes: |
showard | c408c5e | 2009-01-08 23:30:53 +0000 | [diff] [blame] | 301 | self._wait_for_some_async_commands() |
| 302 | self.wait_for_all_async_commands() |
showard | 170873e | 2009-01-07 00:22:26 +0000 | [diff] [blame] | 303 | |
| 304 | duration = time.time() - start_time |
| 305 | if duration > self._WARNING_DURATION: |
| 306 | self._report_long_execution(calls, duration) |
| 307 | |
| 308 | warnings = self.warnings |
| 309 | self.warnings = [] |
| 310 | return dict(results=results, warnings=warnings) |
| 311 | |
| 312 | |
| 313 | def create_host(hostname): |
| 314 | username = global_config.global_config.get_config_value( |
| 315 | 'SCHEDULER', hostname + '_username', default=getpass.getuser()) |
| 316 | return hosts.SSHHost(hostname, user=username) |
| 317 | |
| 318 | |
| 319 | def parse_input(): |
| 320 | input_chunks = [] |
| 321 | chunk_of_input = sys.stdin.read() |
| 322 | while chunk_of_input: |
| 323 | input_chunks.append(chunk_of_input) |
| 324 | chunk_of_input = sys.stdin.read() |
| 325 | pickled_input = ''.join(input_chunks) |
| 326 | |
| 327 | try: |
| 328 | return pickle.loads(pickled_input) |
| 329 | except Exception, exc: |
| 330 | separator = '*' * 50 |
| 331 | raise ValueError('Unpickling input failed\n' |
| 332 | 'Input: %r\n' |
| 333 | 'Exception from pickle:\n' |
| 334 | '%s\n%s\n%s' % |
| 335 | (pickled_input, separator, traceback.format_exc(), |
| 336 | separator)) |
| 337 | |
| 338 | |
| 339 | def return_data(data): |
| 340 | print pickle.dumps(data) |
| 341 | |
| 342 | |
| 343 | def main(): |
| 344 | calls = parse_input() |
| 345 | drone_utility = DroneUtility() |
| 346 | return_value = drone_utility.execute_calls(calls) |
| 347 | return_data(return_value) |
| 348 | |
| 349 | |
| 350 | if __name__ == '__main__': |
| 351 | main() |