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..ca6d692e6e1682665571dc0fbded2d49fe96281b 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 ${EXE_SUFFIX} on the command line will be replaced with ".exe" string on

[M-A Ruel, 2016/06/06 23:34:51]

The isolate client uses EXECUTABLE_SUFFIX. It think it'd be less confusing to
use the same wording.

[nodir, 2016/06/07 18:46:35]

Done

+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,17 @@ 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}'
+EXE_SUFFIX_PARAMETER = '${EXE_SUFFIX}'
+
+# .exe on Windows.
+EXE_SUFFIX = '.exe' if sys.platform == 'win32' else ''
+
 
 # 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 +100,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 +157,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:
+    AssertionError if a parameter is requested in |command| but its value is not

[Vadim Sh., 2016/06/06 21:32:13]

ValueError

assertion errors must not happen in correctly constructed program (regardless of
what's given as its input).

[nodir, 2016/06/07 18:46:35]

the code is correct. The cmdline is checked before calling this method.

I've changed to ValueError

+      provided.
+  """
   def fix(arg):
+    arg = arg.replace(EXE_SUFFIX_PARAMETER, EXE_SUFFIX)
+    replace_slash = False
+    if CIPD_PATH_PARAMETER in arg:
+      assert cipd_path, (
+        'cipd_path is requested in command %r, '
+        'but its value is not provided is empty' % command)
+      arg = arg.replace(CIPD_PATH_PARAMETER, cipd_path)
+      replace_slash = True
     if ISOLATED_OUTDIR_PARAMETER in arg:
       assert out_dir
       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 +345,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.
@@ -343,17 +375,18 @@ def map_and_run(
     #  },

[M-A Ruel, 2016/06/06 23:34:51]

Add expectation about:
'cipd': {
  ...
},

[nodir, 2016/06/07 18:46:35]

Done.

     },
     '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,

[M-A Ruel, 2016/06/06 23:34:51]

oops!

+      '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,89 @@ 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,
+    client_cache=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.
+    client_cache (isolatedserver.DiskCache): if not None, overrides client
+      cache derived from |cache_dir|.
+    timeout: max duration in seconds that this function can take.
+
+  Yields:
+    CIPD stats as dict.
+  """
+  assert cache_dir
+  timeouter = tools.Timeouter(timeout)
+  assert not client_cache or isinstance(client_cache, isolateserver.DiskCache)
+  if not packages:
+    yield
+    return
+
+  start = time.time()
+
+  # Prepare caches.
+  cache_dir = os.path.abspath(cache_dir)
+  # version_cache is {version_digest -> instance id} mapping.
+  # It does not take a lot of disk space.
+  version_cache = isolateserver.DiskCache(
+      unicode(os.path.join(cache_dir, 'versions')),
+      isolateserver.CachePolicies(0, 0, 300),
+      hashlib.sha1)
+  # client_cache is {instance_id -> client binary} mapping.
+  # It is bounded by 5 client versions.
+  client_cache = client_cache or isolateserver.DiskCache(
+      unicode(os.path.join(cache_dir, 'clients')),
+      isolateserver.CachePolicies(0, 0, 5),
+      hashlib.sha1)
+
+  # Get CIPD client.
+  client_package_name, client_version = cipd.parse_package(client_package)
+  get_client_start = time.time()
+  client = cipd.get_client(
+      service_url, client_package_name, client_version,
+      version_cache=version_cache, client_cache=client_cache,
+      timeout=timeouter.left())
+  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, 0777)
+  file_path.make_tree_writeable(site_root)
+  try:
+    client.ensure(
+        site_root, packages,
+        cache_dir=os.path.join(cache_dir, 'cipd_internal'),
+        timeout=timeouter.left())
+
+    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 +686,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='isolate_cache', cipd_cache='cipd_cache')

[Vadim Sh., 2016/06/06 21:32:14]

I believe this will "leak" all existing isolate caches on all the bots when
deployed. 

Either come up with a deployment strategy that involves migrating the directory
cleanly, or just don't change it.

[nodir, 2016/06/07 18:46:35]

Done.

+  return parser
+
+
+def main(args):
+  parser = create_option_parser()
   options, args = parser.parse_args(args)
 
   cache = isolateserver.process_cache_options(options)
@@ -603,27 +734,47 @@ 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))
+  root_dir = options.root_dir
+  if root_dir:
+    root_dir = unicode(os.path.abspath(root_dir))
+    file_path.ensure_tree(root_dir, 0700)

[Vadim Sh., 2016/06/06 21:32:14]

nit: do it after all options are validated

[nodir, 2016/06/07 18:46:35]

Done.

+  else:
+    root_dir = os.path.dirname(cache.cache_dir) if cache.cache_dir else None
+
   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)
+  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 ex.message
+    return 1
 
 
 if __name__ == '__main__':