swarming: custom cipd package paths

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.
 
 If input isolated hash is provided, fetches it, creates a tree of hard links,
 appends args to the command in the fetched isolated and runs it.
 To improve performance, keeps a local cache.
 The local cache can safely be deleted.
 
 Any ${EXECUTABLE_SUFFIX} on the command line will be replaced with ".exe" string
 on Windows and "" on other platforms.
 
-If at least one CIPD package was specified, any ${CIPD_PATH} on the command line
-will be replaced by location of a temporary directory that contains installed
-packages.
-
 Any ${ISOLATED_OUTDIR} on the command line will be replaced by the location of a
 temporary directory upon execution of the command specified in the .isolated
 file. All content written to this directory will be uploaded upon termination
 and the .isolated file describing this directory will be printed to stdout.
 """
 
-__version__ = '0.8.0'
+__version__ = '0.8.1'
 
 import base64
-import contextlib
-import hashlib
 import logging
 import optparse
 import os
 import sys
 import tempfile
 import time
 
 from third_party.depot_tools import fix_encoding
 
 from utils import file_path
 from utils import fs
 from utils import large
 from utils import logging_utils
 from utils import on_error
 from utils import subprocess42
 from utils import tools
 from utils import zip_package
 
 import auth
 import cipd
 import isolateserver
 
 
 ISOLATED_OUTDIR_PARAMETER = '${ISOLATED_OUTDIR}'
-CIPD_PATH_PARAMETER = '${CIPD_PATH}'
 EXECUTABLE_SUFFIX_PARAMETER = '${EXECUTABLE_SUFFIX}'
 SWARMING_BOT_FILE_PARAMETER = '${SWARMING_BOT_FILE}'
 
 # Absolute path to this file (can be None if running from zip on Mac).
 THIS_FILE_PATH = os.path.abspath(__file__) if __file__ else None
 
 # Directory that contains this file (might be inside zip package).
 BASE_DIR = os.path.dirname(THIS_FILE_PATH) if __file__ else None
 
 # Directory that contains currently running script file.
@@ skipping 78 matching lines @@
     # is not yet changed to verify the hash of the content of the files it is
     # looking at, so that if a test modifies an input file, the file must be
     # deleted.
     file_path.make_tree_writeable(rootdir)
   else:
     raise ValueError(
         'change_tree_read_only(%s, %s): Unknown flag %s' %
         (rootdir, read_only, read_only))
 
 
-def process_command(command, out_dir, cipd_path, bot_file):
+def process_command(command, out_dir, bot_file):
   """Replaces variables in a command line.
 
   Raises:
     ValueError if a parameter is requested in |command| but its value is not
       provided.
   """
   def fix(arg):
     arg = arg.replace(EXECUTABLE_SUFFIX_PARAMETER, cipd.EXECUTABLE_SUFFIX)
     replace_slash = False
-    if CIPD_PATH_PARAMETER in arg:
-      if not cipd_path:
-        raise ValueError('cipd_path is requested in command, but not provided')
-      arg = arg.replace(CIPD_PATH_PARAMETER, cipd_path)
-      replace_slash = True
     if ISOLATED_OUTDIR_PARAMETER in arg:
       if not out_dir:
         raise ValueError('out_dir is requested in command, but not provided')
       arg = arg.replace(ISOLATED_OUTDIR_PARAMETER, out_dir)
       replace_slash = True
+    if SWARMING_BOT_FILE_PARAMETER in arg:
+      if bot_file:
+        arg = arg.replace(SWARMING_BOT_FILE_PARAMETER, bot_file)
+        replace_slash = True
+      else:
+        logging.warning('SWARMING_BOT_FILE_PARAMETER found in command, but no '
+                        'bot_file specified. Leaving parameter unchanged.')
     if replace_slash:
       # Replace slashes only if parameters are present
       # because of arguments like '${ISOLATED_OUTDIR}/foo/bar'
       arg = arg.replace('/', os.sep)
-    if SWARMING_BOT_FILE_PARAMETER in arg:
-      if bot_file:
-        arg = arg.replace(SWARMING_BOT_FILE_PARAMETER, bot_file)
-        arg = arg.replace('/', os.sep)
-      else:
-        logging.warning('SWARMING_BOT_FILE_PARAMETER found in command, but no '
-                        'bot_file specified. Leaving paramater unchanged.')
     return arg
 
   return [fix(arg) for arg in command]
 
 
 def run_command(command, cwd, tmp_dir, hard_timeout, grace_period):
   """Runs the command.
 
   Returns:
     tuple(process exit code, bool if had a hard timeout)
@@ skipping 144 matching lines @@
   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, cache, leak_temp_dir, root_dir,
-    hard_timeout, grace_period, bot_file, extra_args, cipd_path, cipd_stats):
+    hard_timeout, grace_period, bot_file, extra_args, install_packages_fn):
   """Runs a command with optional isolated input/output.
 
   See run_tha_test for argument documentation.
 
   Returns metadata about the result.
   """
   assert bool(command) ^ bool(isolated_hash)
   result = {
     'duration': None,
     'exit_code': None,
@@ skipping 15 matching lines @@
     #    'upload': {
     #      'duration': 0.,
     #      'items_cold': '<large.pack()>',
     #      'items_hot': '<large.pack()>',
     #    },
     #  },
     },
     'outputs_ref': None,
     'version': 5,
   }
-  if cipd_stats:
-    result['stats']['cipd'] = cipd_stats
 
   if root_dir:
     file_path.ensure_tree(root_dir, 0700)
   else:
     root_dir = os.path.dirname(cache.cache_dir) if cache.cache_dir else None
   run_dir = make_temp_dir(u'isolated_run', root_dir)
   out_dir = make_temp_dir(u'isolated_out', root_dir) if storage else None
   tmp_dir = make_temp_dir(u'isolated_tmp', root_dir)
   cwd = run_dir
 
   try:
+    cipd_stats = install_packages_fn(run_dir)
+    if cipd_stats:
+      result['stats']['cipd'] = cipd_stats
+
     if isolated_hash:
       isolated_stats = result['stats'].setdefault('isolated', {})
       bundle, isolated_stats['download'] = fetch_and_measure(
           isolated_hash=isolated_hash,
           storage=storage,
           cache=cache,
           outdir=run_dir)
       if not bundle.command:
         # Handle this as a task failure, not an internal failure.
         sys.stderr.write(
             '<The .isolated doesn\'t declare any command to run!>\n'
             '<Check your .isolate for missing \'command\' variable>\n')
         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
 
     command = tools.fix_python_path(command)
-    command = process_command(command, out_dir, cipd_path, bot_file)
+    command = process_command(command, out_dir, bot_file)
     file_path.ensure_command_has_abs_path(command, cwd)
 
     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 occured. Report accordingly so the swarming task will be
-    # retried automatically.
+    # 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)
   finally:
     try:
       if leak_temp_dir:
         logging.warning(
             'Deliberately leaking %s for later examination', run_dir)
       else:
         # On Windows rmtree(run_dir) call above has a synchronization effect: it
@@ skipping 41 matching lines @@
       # 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, cache, leak_temp_dir, result_json,
     root_dir, hard_timeout, grace_period, bot_file, extra_args,
-    cipd_path, cipd_stats):
+    install_packages_fn):
   """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.
 
   A temporary directory is created to hold the output files. The content inside
   this directory will be uploaded back to |storage| packaged as a .isolated
   file.
 
   Arguments:
     command: the command to run, a list of strings. Mutually exclusive with
              isolated_hash.
     isolated_hash: the SHA-1 of the .isolated file that must be retrieved to
                    recreate the tree of files to run the target executable.
                    The command specified in the .isolated is executed.
                    Mutually exclusive with command argument.
     storage: an isolateserver.Storage object to retrieve remote objects. This
              object has a reference to an isolateserver.StorageApi, which does
              the actual I/O.
     cache: an isolateserver.LocalCache to keep from retrieving the same objects
            constantly by caching the objects retrieved. Can be on-disk or
            in-memory.
     leak_temp_dir: if true, the temporary directory will be deliberately leaked
                    for later examination.
     result_json: file path to dump result metadata into. If set, the process
                  exit code is always 0 unless an internal error occurred.
-    root_dir: directory to the path to use to create the temporary directory. If
+    root_dir: path to the directory to use to create the temporary directory. If
               not specified, a random temporary directory is created.
     hard_timeout: kills the process if it lasts more than this amount of
                   seconds.
     grace_period: number of seconds to wait between SIGTERM and SIGKILL.
     extra_args: optional arguments to add to the command stated in the .isolate
                 file. Ignored if isolate_hash is empty.
-    cipd_path: value for CIPD_PATH_PARAMETER. If empty, command or extra_args
-               must not use CIPD_PATH_PARAMETER.
-    cipd_stats: CIPD stats to include in the metadata written to result_json.
+    install_packages_fn: function (dir) => cipd_stats. Installs packages.
 
   Returns:
     Process exit code that should be used.
   """
   assert bool(command) ^ bool(isolated_hash)
   extra_args = extra_args or []
 
   if any(ISOLATED_OUTDIR_PARAMETER in a for a in (command or extra_args)):
     assert storage is not None, 'storage is None although outdir is specified'
 
   if result_json:
     # Write a json output file right away in case we get killed.
     result = {
       'exit_code': None,
       'had_hard_timeout': False,
       'internal_failure': 'Was terminated before completion',
       'outputs_ref': None,
       'version': 5,
     }
-    if cipd_stats:
-      result['stats'] = {'cipd': cipd_stats}
     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, cache, leak_temp_dir, root_dir,
-      hard_timeout, grace_period, bot_file, extra_args, cipd_path, cipd_stats)
+      hard_timeout, grace_period, bot_file, extra_args, install_packages_fn)
   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.
     return int(bool(result['internal_failure']))
 
   # Marshall into old-style inline output.
   if result['outputs_ref']:
     data = {
       'hash': result['outputs_ref']['isolated'],
       'namespace': result['outputs_ref']['namespace'],
       'storage': result['outputs_ref']['isolatedserver'],
     }
     sys.stdout.flush()
     print(
         '[run_isolated_out_hack]%s[/run_isolated_out_hack]' %
         tools.format_json(data, dense=True))
     sys.stdout.flush()
   return result['exit_code'] or int(bool(result['internal_failure']))
 
 
-@contextlib.contextmanager
-def ensure_packages(
-    site_root, packages, service_url, client_package, cache_dir=None,
-    timeout=None):
-  """Returns a context manager that installs packages into site_root.
-
-  Creates/recreates site_root dir and install packages before yielding.
-  After yielding deletes site_root dir.
+def install_packages(
+    run_dir, package_list_file, service_url, client_package_name,
+    client_version, cache_dir=None, timeout=None):
+  """Installs packages. Returns stats.
 
   Args:
-    site_root (str): where to install the packages.
-    packages (list of str): list of package to ensure.
-      Package format is same as for "--cipd-package" option.
+    run_dir (str): root of installation.
+    package_list_file (str): path to a file with a list of packages to install.
     service_url (str): CIPD server url, e.g.
       "https://chrome-infra-packages.appspot.com."
-    client_package (str): CIPD package of CIPD client.
-      Format is same as for "--cipd-package" option.
+    client_package_name (str): CIPD package name of CIPD client.
+    client_version (str): Version of CIPD client.
     cache_dir (str): where to keep cache of cipd clients, packages and tags.
     timeout: max duration in seconds that this function can take.
-
-  Yields:
-    CIPD stats as dict.
   """
   assert cache_dir
+  if not package_list_file:
+    return None
+
   timeoutfn = tools.sliding_timeout(timeout)
-  if not packages:
-    yield
-    return
-
   start = time.time()
-
   cache_dir = os.path.abspath(cache_dir)
 
-  # Get CIPD client.
-  client_package_name, client_version = cipd.parse_package(client_package)
+  run_dir = os.path.abspath(run_dir)
+  package_list = cipd.parse_package_list_file(package_list_file)
+
   get_client_start = time.time()
   client_manager = cipd.get_client(
       service_url, client_package_name, client_version, cache_dir,
       timeout=timeoutfn())
   with client_manager as client:
     get_client_duration = time.time() - get_client_start
-    # Create site_root, install packages, yield, delete site_root.
-    if fs.isdir(site_root):
-      file_path.rmtree(site_root)
-    file_path.ensure_tree(site_root, 0770)
-    try:
+    for path, packages in package_list.iteritems():
+      site_root = os.path.abspath(os.path.join(run_dir, path))
+      if not site_root.startswith(run_dir):
+        raise cipd.Error('Invalid CIPD package path "%s"' % path)
+
+      # Do not clean site_root before installation because it may contain other
+      # site roots.
+      file_path.ensure_tree(site_root, 0770)
       client.ensure(
           site_root, packages,
           cache_dir=os.path.join(cache_dir, 'cipd_internal'),
           timeout=timeoutfn())
+      file_path.make_tree_files_read_only(site_root)
 
-      total_duration = time.time() - start
-      logging.info(
-          'Installing CIPD client and packages took %d seconds', total_duration)
+  total_duration = time.time() - start
+  logging.info(
+      'Installing CIPD client and packages took %d seconds', total_duration)
 
-      file_path.make_tree_files_read_only(site_root)
-      yield {
-        'duration': total_duration,
-        'get_client_duration': get_client_duration,
-      }
-    finally:
-      file_path.rmtree(site_root)
+  return {
+    'duration': total_duration,
+    'get_client_duration': get_client_duration,
+  }
 
 
 def create_option_parser():
   parser = logging_utils.OptionParserWithLogging(
       usage='%prog <options> [command to run or extra args]',
       version=__version__,
       log_file=RUN_ISOLATED_LOG_FILE)
   parser.add_option(
       '--clean', action='store_true',
       help='Cleans the cache, trimming it necessary and remove corrupted items '
@@ skipping 61 matching lines @@
 
   isolateserver.process_isolate_server_options(
     parser, options, True, False)
   if not options.isolate_server:
     if options.isolated:
       parser.error('--isolated requires --isolate-server')
     if ISOLATED_OUTDIR_PARAMETER in args:
       parser.error(
         '%s in args requires --isolate-server' % ISOLATED_OUTDIR_PARAMETER)
 
+  if options.root_dir:
+    options.root_dir = unicode(os.path.abspath(options.root_dir))
   if options.json:
     options.json = unicode(os.path.abspath(options.json))
 
   cipd.validate_cipd_options(parser, options)
 
-  root_dir = options.root_dir
-  if root_dir:
-    root_dir = unicode(os.path.abspath(root_dir))
-    file_path.ensure_tree(root_dir, 0700)
-  else:
-    root_dir = os.path.dirname(cache.cache_dir) if cache.cache_dir else None
-
-  cipd_path = None
-  if not options.cipd_package:
-    if CIPD_PATH_PARAMETER in args:
-      parser.error('%s in args requires --cipd-package' % CIPD_PATH_PARAMETER)
-  else:
-    cipd_path = make_temp_dir(u'cipd_site_root', root_dir)
+  install_packages_fn = lambda run_dir: install_packages(
+      run_dir, options.cipd_package_list, options.cipd_server,
+      options.cipd_client_package, options.cipd_client_version,
+      cache_dir=options.cipd_cache)
 
   try:
-    with ensure_packages(
-        cipd_path, options.cipd_package, options.cipd_server,
-        options.cipd_client_package, options.cipd_cache) as cipd_stats:
-      command = [] if options.isolated else args
-      if options.isolate_server:
-        storage = isolateserver.get_storage(
-            options.isolate_server, options.namespace)
-        with storage:
-          # Hashing schemes used by |storage| and |cache| MUST match.
-          assert storage.hash_algo == cache.hash_algo
-          return run_tha_test(
-              command, options.isolated, storage, cache, options.leak_temp_dir,
-              options.json, root_dir, options.hard_timeout,
-              options.grace_period, options.bot_file, args,
-              cipd_path, cipd_stats)
-      else:
+    command = [] if options.isolated else args
+    if options.isolate_server:
+      storage = isolateserver.get_storage(
+          options.isolate_server, options.namespace)
+      with storage:
+        # Hashing schemes used by |storage| and |cache| MUST match.
+        assert storage.hash_algo == cache.hash_algo
         return run_tha_test(
-            command, options.isolated, None, cache, options.leak_temp_dir,
-            options.json, root_dir, options.hard_timeout,
-            options.grace_period, options.bot_file, args,
-            cipd_path, cipd_stats)
+            command, options.isolated, storage, cache, options.leak_temp_dir,
+            options.json, options.root_dir, options.hard_timeout,
+            options.grace_period, options.bot_file, args, install_packages_fn)
+    else:
+      return run_tha_test(
+          command, options.isolated, None, cache, options.leak_temp_dir,
+          options.json, options.root_dir, options.hard_timeout,
+          options.grace_period, options.bot_file, args, install_packages_fn)
   except cipd.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()
   sys.exit(main(sys.argv[1:]))