Return specific files, not just those in $(ISOLATED_OUTDIR)

client/run_isolated.py

 #!/usr/bin/env python
 # Copyright 2012 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 command with optional isolated input/output.
 
 Despite name "run_isolated", can run a generic non-isolated command specified as
 args.
 
@@ skipping 273 matching lines @@
   return bundle, {
     'duration': time.time() - start,
     'initial_number_items': cache.initial_number_items,
     'initial_size': cache.initial_size,
     'items_cold': base64.b64encode(large.pack(sorted(cache.added))),
     'items_hot': base64.b64encode(
         large.pack(sorted(set(cache.used) - set(cache.added)))),
   }
 
 
+def link_outputs_to_outdir(run_dir, out_dir, outputs):
+  """Links any named outputs to out_dir so they can be uploaded.
+
+  Raises an error if the file already exists in that directory.
+  """
+  if not outputs:
+    return
+  isolateserver.create_directories(out_dir, outputs)
+  for o in outputs:
+    try:
+      file_path.link_file(
+          os.path.join(out_dir, o),
+          os.path.join(run_dir, o),
+          file_path.HARDLINK_WITH_FALLBACK)
+    except OSError as e:
+      # TODO(aludwin): surface this error
+      sys.stderr.write('<Could not return file %s: %s>' % (o, e))
+
+
 def delete_and_upload(storage, out_dir, leak_temp_dir):
   """Deletes the temporary run directory and uploads results back.
 
   Returns:
     tuple(outputs_ref, success, stats)
     - outputs_ref: a dict referring to the results archived back to the isolated
           server, if applicable.
     - success: False if something occurred that means that the task must
           forcibly be considered a failure, e.g. zombie processes were left
           behind.
     - stats: uploading stats.
   """
-
   # Upload out_dir and generate a .isolated file out of this directory. It is
   # only done if files were written in the directory.
   outputs_ref = None
   cold = []
   hot = []
   start = time.time()
 
   if fs.isdir(out_dir) and fs.listdir(out_dir):
     with tools.Profiler('ArchiveOutput'):
       try:
@@ skipping 29 matching lines @@
     logging.exception('Had difficulties removing out_dir %s: %s', out_dir, e)
   stats = {
     'duration': time.time() - start,
     'items_cold': base64.b64encode(large.pack(cold)),
     'items_hot': base64.b64encode(large.pack(hot)),
   }
   return outputs_ref, success, stats
 
 
 def map_and_run(
-    command, isolated_hash, storage, isolate_cache, init_name_caches,
+    command, isolated_hash, storage, isolate_cache, outputs, init_name_caches,
     leak_temp_dir, root_dir, hard_timeout, grace_period, bot_file, extra_args,
     install_packages_fn, use_symlinks):
   """Runs a command with optional isolated input/output.
 
   See run_tha_test for argument documentation.
 
   Returns metadata about the result.
   """
   assert root_dir or root_dir is None
   assert bool(command) ^ bool(isolated_hash)
@@ skipping 67 matching lines @@
         if os.environ.get('SWARMING_TASK_ID'):
           # Give an additional hint when running as a swarming task.
           sys.stderr.write('<This occurs at the \'isolate\' step>\n')
         result['exit_code'] = 1
         return result
 
       change_tree_read_only(run_dir, bundle.read_only)
       cwd = os.path.normpath(os.path.join(cwd, bundle.relative_cwd))
       command = bundle.command + extra_args
 
+    # If we have an explicit list of files to return, make sure their
+    # directories exist now.
+    if storage and outputs:
+      isolateserver.create_directories(run_dir, outputs)
+
     command = tools.fix_python_path(command)
     command = process_command(command, out_dir, bot_file)
     file_path.ensure_command_has_abs_path(command, cwd)
 
     init_name_caches(run_dir)
 
     sys.stdout.flush()
     start = time.time()
     try:
       result['exit_code'], result['had_hard_timeout'] = run_command(
           command, cwd, tmp_dir, hard_timeout, grace_period)
     finally:
       result['duration'] = max(time.time() - start, 0)
   except Exception as e:
     # An internal error occurred. Report accordingly so the swarming task will
     # be retried automatically.
     logging.exception('internal failure: %s', e)
     result['internal_failure'] = str(e)
     on_error.report(None)
+
+  # Clean up
   finally:
     try:
+      # Try to link files to the output directory, if specified.
+      if out_dir:
+        link_outputs_to_outdir(run_dir, out_dir, outputs)
+
       success = False
       if leak_temp_dir:
         success = True
         logging.warning(
             'Deliberately leaking %s for later examination', run_dir)
       else:
         # On Windows rmtree(run_dir) call above has a synchronization effect: it
         # finishes only when all task child processes terminate (since a running
         # process locks *.exe file). Examine out_dir only after that call
         # completes (since child processes may write to out_dir too and we need
@@ skipping 36 matching lines @@
         result['exit_code'] = 1
     except Exception as e:
       # Swallow any exception in the main finally clause.
       if out_dir:
         logging.exception('Leaking out_dir %s: %s', out_dir, e)
       result['internal_failure'] = str(e)
   return result
 
 
 def run_tha_test(
-    command, isolated_hash, storage, isolate_cache, init_name_caches,
+    command, isolated_hash, storage, isolate_cache, outputs, init_name_caches,
     leak_temp_dir, result_json, root_dir, hard_timeout, grace_period, bot_file,
     extra_args, install_packages_fn, use_symlinks):
   """Runs an executable and records execution metadata.
 
   Either command or isolated_hash must be specified.
 
   If isolated_hash is specified, downloads the dependencies in the cache,
   hardlinks them into a temporary directory and runs the command specified in
   the .isolated.
 
@@ skipping 46 matching lines @@
       'exit_code': None,
       'had_hard_timeout': False,
       'internal_failure': 'Was terminated before completion',
       'outputs_ref': None,
       'version': 5,
     }
     tools.write_json(result_json, result, dense=True)
 
   # run_isolated exit code. Depends on if result_json is used or not.
   result = map_and_run(
-      command, isolated_hash, storage, isolate_cache, init_name_caches,
+      command, isolated_hash, storage, isolate_cache, outputs, init_name_caches,
       leak_temp_dir, root_dir, hard_timeout, grace_period, bot_file, extra_args,
       install_packages_fn, use_symlinks)
   logging.info('Result:\n%s', tools.format_json(result, dense=True))
 
   if result_json:
     # We've found tests to delete 'work' when quitting, causing an exception
     # here. Try to recreate the directory if necessary.
     file_path.ensure_tree(os.path.dirname(result_json))
     tools.write_json(result_json, result, dense=True)
     # Only return 1 if there was an internal error.
@@ skipping 158 matching lines @@
   parser.add_option(
       '--hard-timeout', type='float', help='Enforce hard timeout in execution')
   parser.add_option(
       '--grace-period', type='float',
       help='Grace period between SIGTERM and SIGKILL')
   parser.add_option(
       '--bot-file',
       help='Path to a file describing the state of the host. The content is '
            'defined by on_before_task() in bot_config.')
   parser.add_option(
+      '--output', action='append',
+      help='Specifies an output to return. If no outputs are specified, all '
+           'files located in $(ISOLATED_OUTDIR) will be returned; '
+           'otherwise, outputs in both $(ISOLATED_OUTDIR) and those '
+           'specified by --output option (there can be multiple) will be '
+           'returned. Note that if a file in OUT_DIR has the same path '
+           'as an --output option, the --output version will be returned.')
+  parser.add_option(
       '-a', '--argsfile',
       # This is actually handled in parse_args; it's included here purely so it
       # can make it into the help text.
       help='Specify a file containing a JSON array of arguments to this '
            'script. If --argsfile is provided, no other argument may be '
            'provided on the command line.')
   data_group = optparse.OptionGroup(parser, 'Data source')
   data_group.add_option(
       '-s', '--isolated',
       help='Hash of the .isolated to grab from the isolate server.')
@@ skipping 108 matching lines @@
       storage = isolateserver.get_storage(
           options.isolate_server, options.namespace)
       with storage:
         # Hashing schemes used by |storage| and |isolate_cache| MUST match.
         assert storage.hash_algo == isolate_cache.hash_algo
         return run_tha_test(
             command,
             options.isolated,
             storage,
             isolate_cache,
+            options.output,
             init_named_caches,
             options.leak_temp_dir,
             options.json, options.root_dir,
             options.hard_timeout,
             options.grace_period,
             options.bot_file, args,
             install_packages_fn,
             options.use_symlinks)
     return run_tha_test(
         command,
         options.isolated,
         None,
         isolate_cache,
+        options.output,
         init_named_caches,
         options.leak_temp_dir,
         options.json,
         options.root_dir,
         options.hard_timeout,
         options.grace_period,
         options.bot_file, args,
         install_packages_fn,
         options.use_symlinks)
   except (cipd.Error, named_cache.Error) as ex:
     print >> sys.stderr, ex.message
     return 1
 
 
 if __name__ == '__main__':
   subprocess42.inhibit_os_error_reporting()
   # Ensure that we are always running with the correct encoding.
   fix_encoding.fix_encoding()
   file_path.enable_symlink()
 
   sys.exit(main(sys.argv[1:]))