run_isolated.py: install CIPD packages

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.7.0'
+__version__ = '0.8.0'
 
 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}'
 
 # 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.
 if zip_package.get_main_script_path():
   MAIN_DIR = os.path.dirname(
@@ skipping 19 matching lines @@
   # Building a zip package when running from another zip package is
   # unsupported and probably unneeded.
   assert not zip_package.is_zipped_module(sys.modules[__name__])
   assert THIS_FILE_PATH
   assert BASE_DIR
   package = zip_package.ZipPackage(root=BASE_DIR)
   package.add_python_file(THIS_FILE_PATH, '__main__.py' if executable else None)
   package.add_python_file(os.path.join(BASE_DIR, 'isolated_format.py'))
   package.add_python_file(os.path.join(BASE_DIR, 'isolateserver.py'))
   package.add_python_file(os.path.join(BASE_DIR, 'auth.py'))
+  package.add_python_file(os.path.join(BASE_DIR, 'cipd.py'))
   package.add_directory(os.path.join(BASE_DIR, 'third_party'))
   package.add_directory(os.path.join(BASE_DIR, 'utils'))
   return package
 
 
 def make_temp_dir(prefix, root_dir=None):
   """Returns a temporary directory.
 
   If root_dir is given and /tmp is on same file system as root_dir, uses /tmp.
   Otherwise makes a new temp directory under root_dir.
@@ skipping 36 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):
-  """Replaces isolated specific variables in a command line."""
+def process_command(command, out_dir, cipd_path):
+  """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:
-      assert out_dir
+      if not out_dir:
+        raise ValueError('out_dir is requested in command, but not provided')
       arg = arg.replace(ISOLATED_OUTDIR_PARAMETER, out_dir)
-      # Replace slashes only if ISOLATED_OUTDIR_PARAMETER is present
+      replace_slash = True
+    if replace_slash:
+      # Replace slashes only if parameters are present
       # because of arguments like '${ISOLATED_OUTDIR}/foo/bar'
       arg = arg.replace('/', os.sep)
     return arg
 
   return [fix(arg) for arg in command]
 
 
 def run_command(command, cwd, tmp_dir, hard_timeout, grace_period):
   """Runs the command.
 
@@ skipping 146 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, extra_args):
+    hard_timeout, grace_period, extra_args, cipd_path, cipd_stats):
   """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,
     'had_hard_timeout': False,
     'internal_failure': None,
     'stats': {
     # 'isolated': {
+    #    'cipd': {
+    #      'duration': 0.,
+    #      'get_client_duration': 0.,
+    #    },
     #    'download': {
     #      'duration': 0.,
     #      'initial_number_items': 0,
     #      'initial_size': 0,
     #      'items_cold': '<large.pack()>',
     #      'items_hot': '<large.pack()>',
     #    },
     #    'upload': {
     #      'duration': 0.,
     #      'items_cold': '<large.pack()>',
     #      'items_hot': '<large.pack()>',
     #    },
     #  },
     },
     'outputs_ref': None,
-    'version': 4,
+    'version': 5,
   }
+  if cipd_stats:
+    result['stats']['cipd'] = cipd_stats
+
   if root_dir:
     file_path.ensure_tree(root_dir, 0700)
-    prefix = u''
   else:
     root_dir = os.path.dirname(cache.cache_dir) if cache.cache_dir else None
-    prefix = u'isolated_'
-  run_dir = make_temp_dir(prefix + u'run', root_dir)
-  out_dir = make_temp_dir(prefix + u'out', root_dir) if storage else None
-  tmp_dir = make_temp_dir(prefix + u'tmp', root_dir)
+  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:
     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)
     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(
-          process_command(command, out_dir), cwd, tmp_dir, hard_timeout,
-          grace_period)
+          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.
     logging.exception('internal failure: %s', e)
     result['internal_failure'] = str(e)
     on_error.report(None)
   finally:
     try:
@@ skipping 45 matching lines @@
     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, cache, leak_temp_dir, result_json,
-    root_dir, hard_timeout, grace_period, extra_args):
+    root_dir, hard_timeout, grace_period, extra_args, cipd_path, cipd_stats):
   """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 occured.
+                 exit code is always 0 unless an internal error occurred.
     root_dir: directory to the path 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.
 
   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': 2,
+      '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, extra_args)
+      hard_timeout, grace_period, extra_args, cipd_path, cipd_stats)
   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']))
 
 
-def main(args):
+@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.
+
+  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.
+    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.
+    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
+  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)
+  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:
+      client.ensure(
+          site_root, packages,
+          cache_dir=os.path.join(cache_dir, 'cipd_internal'),
+          timeout=timeoutfn())
+
+      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)
+
+
+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 '
            'and returns without executing anything; use with -v to know what '
            'was done')
   parser.add_option(
       '--json',
       help='dump output metadata to json file. When used, run_isolated returns '
            'non-zero only on internal failure')
   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')
   data_group = optparse.OptionGroup(parser, 'Data source')
   data_group.add_option(
       '-s', '--isolated',
       help='Hash of the .isolated to grab from the isolate server.')
   isolateserver.add_isolate_server_options(data_group)
   parser.add_option_group(data_group)
 
   isolateserver.add_cache_options(parser)
-  parser.set_defaults(cache='cache')
+
+  cipd.add_cipd_options(parser)
 
   debug_group = optparse.OptionGroup(parser, 'Debugging')
   debug_group.add_option(
       '--leak-temp-dir',
       action='store_true',
-      help='Deliberately leak isolate\'s temp dir for later examination '
-           '[default: %default]')
+      help='Deliberately leak isolate\'s temp dir for later examination. '
+           'Default: %default')
   debug_group.add_option(
       '--root-dir', help='Use a directory instead of a random one')
   parser.add_option_group(debug_group)
 
   auth.add_auth_options(parser)
+
+  parser.set_defaults(cache='cache', cipd_cache='cipd_cache')
+  return parser
+
+
+def main(args):
+  parser = create_option_parser()
   options, args = parser.parse_args(args)
 
   cache = isolateserver.process_cache_options(options)
   if options.clean:
     if options.isolated:
       parser.error('Can\'t use --isolated with --clean.')
     if options.isolate_server:
       parser.error('Can\'t use --isolate-server with --clean.')
     if options.json:
       parser.error('Can\'t use --json with --clean.')
     cache.cleanup()
     return 0
 
   if not options.isolated and not args:
     parser.error('--isolated or command to run is required.')
 
   auth.process_auth_options(parser, options)
 
   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))
 
-  command = [] if options.isolated else args
-  if options.isolate_server:
-    storage = isolateserver.get_storage(
-        options.isolate_server, options.namespace)
-    # Hashing schemes used by |storage| and |cache| MUST match.
-    with storage:
-      assert storage.hash_algo == cache.hash_algo
-      return run_tha_test(
-          command, options.isolated, storage, cache, options.leak_temp_dir,
-          options.json, options.root_dir, options.hard_timeout,
-          options.grace_period, args)
+  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:
-    return run_tha_test(
-        command, options.isolated, None, cache, options.leak_temp_dir,
-        options.json, options.root_dir, options.hard_timeout,
-        options.grace_period, args)
+    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)
+
+  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, args, cipd_path, cipd_stats)
+      else:
+        return run_tha_test(
+            command, options.isolated, None, cache, options.leak_temp_dir,
+            options.json, root_dir, options.hard_timeout,
+            options.grace_period, args, cipd_path, cipd_stats)
+  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:]))