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

[Vadim Sh., 2016/08/31 03:53:33]

moved comment to closer where it is relevant

+      # 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:

[Vadim Sh., 2016/08/31 03:53:33]

InternalError here is new. I believe it's more clean to catch it here and
transform it into task_result rather than let task_runner.py crash dumping the
stack trace. At least we'll see exception message in Swarming event log.

     # 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': e.message or 'unknown error',

[M-A Ruel, 2016/08/31 15:03:47]

str(e.message) if e.message else 'unknown error'

so that if message is not a string, it won't crash json encoder.

[Vadim Sh., 2016/09/01 00:15:56]

Done.

         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:

[M-A Ruel, 2016/08/31 15:03:47]

One thing that could help is to use contextlib. contextmanager more. No need to
change this CL but as a potential example:

@contextlib.contextmanager
def auth_system(auth_params_file):
  <lines 231-248>
  yield headers_cb
  <lines 271-272>

then use:
with auth_system(auth_params_files) as headers_cb:

wdyt?

[Vadim Sh., 2016/09/01 00:15:56]

I've considered this. AuthSystem isn't supposed to be used in many places. And
the place where it is used already has 'finally' block. So I decided just to
reuse this 'finally' block instead of adding additional layer of indentation.
Deep levels of try: and with: are hard to read.

+      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.

[M-A Ruel, 2016/08/31 15:03:47]

"Rarely used outside the U.S. According to graphs that show the popularity of
cancelation vs. cancellation in the U.S., cancelation is almost three times as
common in the U.S. after the 80s."  :D

But as a non-American, that's fine, you can keep it.

[Vadim Sh., 2016/09/01 00:15:56]

Oh. My spellchecker is British then or something :) I read the code, see a word
highlighted by spellchecker, and change it to suggested replacement. Good way to
find typos in comments.

             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:

[Vadim Sh., 2016/08/31 03:53:33]

same here... 

Was there a reason IOError and OSError are just killing the task, but not record
internal_error?

IIUC they never happen if task dies on its own (even if with error).

Also tentatively added InternalError here. It can happen if connectivity with
the server is lost. I think before we just crashed with ValueError without doing
a proper cleanup.

[M-A Ruel, 2016/08/31 15:03:47]

- If the file is missing, subprocess.Popen(['missing_file'], ...) will raise
OSError. This is not an internal error, it's a user error, the person forgot to
map the file.
- For IOError, I don't recall exactly. It _could_ be file mode on POSIX or
something else but I forgot to note why it could happen.
- Both are now moot because now task_runner always starts [sys.executable,
'swarming_bot.N.zip', 'run_isolated', ...], which wasn't the case before when
this code was first written. So now OSError *is* an internal error, as this now
means that swarming_bot.N.zip somehow disappeared.

[Vadim Sh., 2016/09/01 00:15:56]

I believe this is caught sooner (on line 434) when subprocess is starting (and
it is real error, not internal one).

       # Something wrong happened, try to kill the child process.
-      exit_code = kill_and_wait(
-          proc, task_details.grace_period, 'exception %s' % e)

[Vadim Sh., 2016/08/31 03:53:32]

'e' here was undefined :)

[M-A Ruel, 2016/08/31 15:03:47]

I tried to be defensive but sucked at it.

+      must_signal_internal_failure = e.message or 'unknown error'

[M-A Ruel, 2016/08/31 15:03:47]

str(e.message) if e.message else 'unknown error'

I prefer to be defensive here

[Vadim Sh., 2016/09/01 00:15:56]

Done.

+      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 = e.message or 'unknown error'

[M-A Ruel, 2016/08/31 15:03:47]

same

[Vadim Sh., 2016/09/01 00:15:56]

Done.

+
     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')