run_isolated.py: install CIPD packages

client/run_isolated.py

Index: client/run_isolated.py
diff --git a/client/run_isolated.py b/client/run_isolated.py
index 43276449eade74f6108754febbb740d555e68c8a..a94808947cb35a85f0c10f6d316f155f02bafa8d 100755
--- a/client/run_isolated.py
+++ b/client/run_isolated.py
@@ -13,15 +13,24 @@ 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
@@ -41,10 +50,13 @@ 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
@@ -84,6 +96,7 @@ def get_as_zip_package(executable=True):
   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
@@ -140,13 +153,28 @@ def change_tree_read_only(rootdir, read_only):
         (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
@@ -313,7 +341,7 @@ def delete_and_upload(storage, out_dir, leak_temp_dir):
 
 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.
@@ -328,6 +356,10 @@ def map_and_run(
     'internal_failure': None,
     'stats': {
     # 'isolated': {
+    #    'cipd': {
+    #      'duration': 0.,
+    #      'get_client_duration': 0.,
+    #    },
     #    'download': {
     #      'duration': 0.,
     #      'initial_number_items': 0,
@@ -343,17 +375,18 @@ def map_and_run(
     #  },
     },
     '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:
@@ -378,14 +411,16 @@ def map_and_run(
       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:
@@ -451,7 +486,7 @@ def map_and_run(
 
 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.
@@ -480,7 +515,7 @@ def run_tha_test(
     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
@@ -488,12 +523,16 @@ def run_tha_test(
     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'
 
@@ -504,14 +543,16 @@ def run_tha_test(
       '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
@@ -536,7 +577,71 @@ def run_tha_test(
   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__,
@@ -563,19 +668,27 @@ def main(args):
   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)
@@ -603,27 +716,48 @@ def main(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:
+    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:
-    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)
+    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__':