Minor refactoring in task_runner.py in preparation for adding more code.

appengine/swarming/swarming_bot/bot_code/task_runner.py

 # Copyright 2013 The LUCI Authors. All rights reserved.
 # Use of this source code is governed under the Apache License, Version 2.0
 # that can be found in the LICENSE file.
 
 """Runs a Swarming task.
 
 Downloads all the necessary files to run the task, executes the command and
 streams results back to the Swarming server.
 
 The process exit code is 0 when the task was executed, even if the task itself
@@ skipping 13 matching lines @@
 import time
 import traceback
 
 from utils import file_path
 from utils import net
 from utils import on_error
 from utils import subprocess42
 from utils import zip_package
 
 import bot_auth
-import file_reader
 
 
 # Path to this file or the zip containing this file.
 THIS_FILE = os.path.abspath(zip_package.get_main_script_path())
 
 
 # Sends a maximum of 100kb of stdout per task_update packet.
 MAX_CHUNK_SIZE = 102400
 
 
@@ skipping 100 matching lines @@
   cmd.append('--')
   if task_details.command:
     cmd.extend(task_details.command)
   elif task_details.extra_args:
     cmd.extend(task_details.extra_args)
   return cmd
 
 
 class TaskDetails(object):
   def __init__(self, data):
-    """Loads the raw data from a manifest file specified by --in-file."""
     logging.info('TaskDetails(%s)', data)
     if not isinstance(data, dict):
-      raise ValueError('Expected dict, got %r' % data)
+      raise InternalError('Expected dict in task_runner_in.json, got %r' % data)
 
     # Get all the data first so it fails early if the task details is invalid.
     self.bot_id = data['bot_id']
 
     # Raw command. Only self.command or self.isolated.input can be set.
     self.command = data['command'] or []
 
     # Isolated command. Is a serialized version of task_request.FilesRef.
     self.isolated = data['isolated']
     self.extra_args = data['extra_args']
 
     self.cipd_input = data.get('cipd_input')
 
     self.env = {
       k.encode('utf-8'): v.encode('utf-8') for k, v in data['env'].iteritems()
     }
     self.grace_period = data['grace_period']
     self.hard_timeout = data['hard_timeout']
     self.io_timeout = data['io_timeout']
     self.task_id = data['task_id']
 
+  @staticmethod
+  def load(path):
+    """Loads the TaskDetails from a file on disk (specified via --in-file).
 
-class MustExit(Exception):
-  """Raised on signal that the process must exit immediately."""
+    Raises InternalError if the file can't be read or parsed.
+    """
+    try:
+      with open(path, 'rb') as f:
+        return TaskDetails(json.load(f))
+    except (IOError, ValueError) as e:
+      raise InternalError('Cannot load task_runner_in.json: %s' % e)
+
+
+class ExitSignal(Exception):
+  """Raised on a signal that the process must exit immediately."""
   def __init__(self, sig):
-    super(MustExit, self).__init__()
+    super(ExitSignal, self).__init__(u'task_runner received signal %s' % sig)
     self.signal = sig
 
 
+class InternalError(Exception):
+  """Raised on unrecoverable errors that abort task with 'internal error'."""
+
+
 def load_and_run(
     in_file, swarming_server, cost_usd_hour, start, out_file, min_free_space,
     bot_file, auth_params_file):
-  """Loads the task's metadata and execute it.
+  """Loads the task's metadata, prepares auth environment and executes the task.
 
   This may throw all sorts of exceptions in case of failure. It's up to the
   caller to trap them. These shall be considered 'internal_failure' instead of
   'failure' from a TaskRunResult standpoint.
   """
-  # The work directory is guaranteed to exist since it was created by
-  # bot_main.py and contains the manifest. Temporary files will be downloaded
-  # there. It's bot_main.py that will delete the directory afterward. Tests are
-  # not run from there.
+  auth_system = None
   task_result = None
+  work_dir = os.path.dirname(out_file)
+
   def handler(sig, _):
     logging.info('Got signal %s', sig)
-    raise MustExit(sig)
-  work_dir = os.path.dirname(out_file)
+    raise ExitSignal(sig)
+
   try:
     with subprocess42.set_signal_handler([SIG_BREAK_OR_TERM], handler):
+      # The work directory is guaranteed to exist since it was created by
+      # bot_main.py and contains the manifest. Temporary files will be
+      # downloaded there. It's bot_main.py that will delete the directory
+      # afterward. Tests are not run from there.
       if not os.path.isdir(work_dir):
-        raise ValueError('%s expected to exist' % work_dir)
+        raise InternalError('%s expected to exist' % work_dir)
 
-      with open(in_file, 'rb') as f:
-        task_details = TaskDetails(json.load(f))
+      # Raises InternalError on errors.
+      task_details = TaskDetails.load(in_file)
 
+      # This will start a thread that occasionally reads bot authentication
+      # headers from 'auth_params_file'.
+      if auth_params_file:
+        try:
+          auth_system = bot_auth.AuthSystem()
+          auth_system.start(auth_params_file)
+        except bot_auth.AuthSystemError as e:
+          raise InternalError('Failed to init auth: %s' % e)
+
+      # Returns bot authentication headers dict or raises InternalError.
+      def headers_cb():
+        try:
+          return auth_system.bot_headers if auth_system else {}
+        except bot_auth.AuthSystemError as e:
+          raise InternalError('Failed to grab bot auth headers: %s' % e)
+
+      # Auth environment is up, start the command. task_result is dumped to
+      # disk in 'finally' block.
       task_result = run_command(
           swarming_server, task_details, work_dir,
-          cost_usd_hour, start, min_free_space, bot_file, auth_params_file)
-  except MustExit as e:
+          cost_usd_hour, start, min_free_space, bot_file, headers_cb)
+
+  except (ExitSignal, InternalError) as e:
     # This normally means run_command() didn't get the chance to run, as it
-    # itself trap MustExit and will report accordingly. In this case, we want
+    # itself traps exceptions and will report accordingly. In this case, we want
     # the parent process to send the message instead.
     if not task_result:
       task_result = {
-        u'exit_code': None,
+        u'exit_code': -1,
         u'hard_timeout': False,
         u'io_timeout': False,
-        u'must_signal_internal_failure':
-            u'task_runner received signal %s' % e.signal,
+        u'must_signal_internal_failure': str(e.message or 'unknown error'),
         u'version': OUT_VERSION,
       }
+
   finally:
     # We've found tests to delete the working directory work_dir when quitting,
     # causing an exception here. Try to recreate the directory if necessary.
     if not os.path.isdir(work_dir):
       os.mkdir(work_dir)
+    if auth_system:
+      auth_system.stop()
     with open(out_file, 'wb') as f:
       json.dump(task_result, f)
 
 
 def post_update(
     swarming_server, auth_headers, params, exit_code,
     stdout, output_chunk_start):
   """Posts task update to task_update.
 
   Arguments:
     swarming_server: Base URL to Swarming server.
     auth_headers: dict with HTTP authentication headers.
     params: Default JSON parameters for the POST.
     exit_code: Process exit code, only when a command completed.
     stdout: Incremental output since last call, if any.
     output_chunk_start: Total number of stdout previously sent, for coherency
         with the server.
 
   Returns:
     False if the task should stop.
+
+  Raises:
+    InternalError if can't contact the server after many attempts or the server
+    replies with an error.
   """
   params = params.copy()
   if exit_code is not None:
     params['exit_code'] = exit_code
   if stdout:
     # The output_chunk_start is used by the server to make sure that the stdout
     # chunks are processed and saved in the DB in order.
     params['output'] = base64.b64encode(stdout)
     params['output_chunk_start'] = output_chunk_start
   # TODO(maruel): Support early cancellation.
   # https://code.google.com/p/swarming/issues/detail?id=62
   resp = net.url_read_json(
       swarming_server+'/swarming/api/v1/bot/task_update/%s' % params['task_id'],
       data=params,
       headers=auth_headers,
       follow_redirects=False)
   logging.debug('post_update() = %s', resp)
   if not resp or resp.get('error'):
     # Abandon it. This will force a process exit.
-    raise ValueError(resp.get('error') if resp else 'Failed to contact server')
+    raise InternalError(
+        resp.get('error') if resp else 'Failed to contact server')
   return not resp.get('must_stop', False)
 
 
 def should_post_update(stdout, now, last_packet):
   """Returns True if it's time to send a task_update packet via post_update().
 
   Sends a packet when one of this condition is met:
   - more than MAX_CHUNK_SIZE of stdout is buffered.
   - last packet was sent more than MIN_PACKET_INTERNAL seconds ago and there was
     stdout.
@@ skipping 24 matching lines @@
 
 def kill_and_wait(proc, grace_period, reason):
   logging.warning('SIGTERM finally due to %s', reason)
   proc.terminate()
   try:
     proc.wait(grace_period)
   except subprocess42.TimeoutError:
     logging.warning('SIGKILL finally due to %s', reason)
     proc.kill()
   exit_code = proc.wait()
-  logging.info('Waiting for proces exit in finally - done')
+  logging.info('Waiting for process exit in finally - done')
   return exit_code
 
 
-def start_reading_headers(auth_params_file):
-  """Spawns a thread that rereads headers from --auth-params-file path.
-
-  Returns:
-    Tuple (callback that returns the last known headers, stop callback).
-
-  Raises:
-    file_reader.FatalReadError if headers file can't be read.
-    ValueError if it can be read, but its body is invalid.
-  """
-  # Read headers more often than bot_main writes them, to reduce maximum
-  # possible latency between headers are updated and read.
-  reader = file_reader.FileReaderThread(auth_params_file, interval_sec=30)
-
-  def read_and_validate_headers():
-    val = bot_auth.process_auth_params_json(reader.last_value or {})
-    return val.swarming_http_headers
-
-  reader.start()
-  read_and_validate_headers() # initial validation, may raise ValueError
-  return read_and_validate_headers, reader.stop
-
-
 def run_command(
     swarming_server, task_details, work_dir, cost_usd_hour,
-    task_start, min_free_space, bot_file, auth_params_file):
+    task_start, min_free_space, bot_file, headers_cb):
   """Runs a command and sends packets to the server to stream results back.
 
   Implements both I/O and hard timeouts. Sends the packets numbered, so the
   server can ensure they are processed in order.
 
   Returns:
-    Metadata about the command.
+    Metadata dict with the execution result.
+
+  Raises:
+    ExitSignal if caught some signal when starting or stopping.
+    InternalError on unexpected internal errors.
   """
   # TODO(maruel): This function is incomprehensible, split and refactor.
 
-  # Grab initial auth headers and start rereading them in parallel thread. They
-  # MUST be there already. It's fatal internal error if they are not.
-  headers_cb = lambda: {}
-  stop_headers_reader = lambda: None
-  if auth_params_file:
-    try:
-      headers_cb, stop_headers_reader = start_reading_headers(auth_params_file)
-    except (ValueError, file_reader.FatalReadError) as e:
-      return {
-        u'exit_code': 1,
-        u'hard_timeout': False,
-        u'io_timeout': False,
-        u'must_signal_internal_failure': str(e),
-        u'version': OUT_VERSION,
-      }
-
   # Signal the command is about to be started.
   last_packet = start = now = monotonic_time()
   params = {
     'cost_usd': cost_usd_hour * (now - task_start) / 60. / 60.,
     'id': task_details.bot_id,
     'task_id': task_details.task_id,
   }
   if not post_update(swarming_server, headers_cb(), params, None, '', 0):
     # Don't even bother, the task was already canceled.
     return {
@@ skipping 74 matching lines @@
           last_io = now
 
         # Post update if necessary.
         if should_post_update(stdout, now, last_packet):
           last_packet = monotonic_time()
           params['cost_usd'] = (
               cost_usd_hour * (last_packet - task_start) / 60. / 60.)
           if not post_update(
               swarming_server, headers_cb(), params, None, stdout,
               output_chunk_start):
-            # Server is telling us to stop. Normally task cancelation.
+            # Server is telling us to stop. Normally task cancellation.
             if not kill_sent:
               logging.warning('Server induced stop; sending SIGKILL')
               proc.kill()
               kill_sent = True
 
           output_chunk_start += len(stdout)
           stdout = ''
 
         # Send signal on timeout if necessary. Both are failures, not
         # internal_failures.
@@ skipping 14 matching lines @@
             # states:
             # - signal but process exited within grace period,
             #   (hard_|io_)_timed_out will be set but the process exit code will
             #   be script provided.
             # - processed exited late, exit code will be -9 on posix.
             logging.warning(
                 'Grace of %.3fs exhausted at %.3fs; sending SIGKILL',
                 task_details.grace_period, now - timed_out)
             proc.kill()
             kill_sent = True
-      logging.info('Waiting for proces exit')
+      logging.info('Waiting for process exit')
       exit_code = proc.wait()
-    except MustExit as e:
-      # TODO(maruel): Do the send SIGTERM to child process and give it
-      # task_details.grace_period to terminate.
-      must_signal_internal_failure = (
-          u'task_runner received signal %s' % e.signal)
-      exit_code = kill_and_wait(
-          proc, task_details.grace_period, 'signal %d' % e.signal)
-    except (IOError, OSError):
+    except (ExitSignal, InternalError, IOError, OSError) as e:
       # Something wrong happened, try to kill the child process.
-      exit_code = kill_and_wait(
-          proc, task_details.grace_period, 'exception %s' % e)
+      must_signal_internal_failure = str(e.message or 'unknown error')
+      exit_code = kill_and_wait(proc, task_details.grace_period, e.message)
 
     # This is the very last packet for this command. It if was an isolated task,
     # include the output reference to the archived .isolated file.
     now = monotonic_time()
     params['cost_usd'] = cost_usd_hour * (now - task_start) / 60. / 60.
     params['duration'] = now - start
     params['io_timeout'] = had_io_timeout
     had_hard_timeout = False
     try:
       if not os.path.isfile(isolated_result):
@@ skipping 52 matching lines @@
         if cipd_stats:
           params['cipd_stats'] = cipd_stats
         cipd_pins = run_isolated_result.get('cipd_pins')
         if cipd_pins:
           params['cipd_pins'] = cipd_pins
     except (IOError, OSError, ValueError) as e:
       logging.error('Swallowing error: %s', e)
       if not must_signal_internal_failure:
         must_signal_internal_failure = '%s\n%s' % (
             e, traceback.format_exc()[-2048:])
+
     # TODO(maruel): Send the internal failure here instead of sending it through
     # bot_main, this causes a race condition.
     if exit_code is None:
       exit_code = -1
     params['hard_timeout'] = had_hard_timeout
-    # Ignore server reply to stop.
-    post_update(
-        swarming_server, headers_cb(), params, exit_code,
-        stdout, output_chunk_start)
+
+    # Ignore server reply to stop. Also ignore internal errors here if we are
+    # already handling some.
+    try:
+      post_update(
+          swarming_server, headers_cb(), params, exit_code,
+          stdout, output_chunk_start)
+    except InternalError as e:
+      logging.error('Internal error while finishing the task: %s', e)
+      if not must_signal_internal_failure:
+        must_signal_internal_failure = str(e.message or 'unknown error')
+
     return {
       u'exit_code': exit_code,
       u'hard_timeout': had_hard_timeout,
       u'io_timeout': had_io_timeout,
       u'must_signal_internal_failure': must_signal_internal_failure,
       u'version': OUT_VERSION,
     }
   finally:
     file_path.try_remove(unicode(isolated_result))
-    stop_headers_reader()
 
 
 def main(args):
   subprocess42.inhibit_os_error_reporting()
   parser = optparse.OptionParser(description=sys.modules[__name__].__doc__)
   parser.add_option('--in-file', help='Name of the request file')
   parser.add_option(
       '--out-file', help='Name of the JSON file to write a task summary to')
   parser.add_option(
       '--swarming-server', help='Swarming server to send data back')
@@ skipping 21 matching lines @@
     options.start = now
 
   try:
     load_and_run(
         options.in_file, options.swarming_server, options.cost_usd_hour,
         options.start, options.out_file, options.min_free_space,
         options.bot_file, options.auth_params_file)
     return 0
   finally:
     logging.info('quitting')