Cache/retrieve extracted CIPD packages in local isolate cache

client/cipd.py

 # Copyright 2016 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.
 
 """Fetches CIPD client and installs packages."""
 
 import contextlib
 import hashlib
 import json
 import logging
 import optparse
 import os
 import platform
+import re
+import shutil
 import sys
 import tempfile
 import time
 import urllib
 
 from utils import file_path
 from utils import fs
 from utils import net
 from utils import subprocess42
 from utils import tools
+import isolate

[kjlubick, 2017/05/03 18:28:33]

This introduces a dependency cycle that I have no idea how to break.

 import isolated_format
 import isolateserver
 
 
 # .exe on Windows.
 EXECUTABLE_SUFFIX = '.exe' if sys.platform == 'win32' else ''
 
 
 if sys.platform == 'win32':
   def _ensure_batfile(client_path):
@@ skipping 99 matching lines @@
       package_name (str): the CIPD package name for the client itself.
       instance_id (str): the CIPD instance_id for the client itself.
       service_url (str): if not None, URL of the CIPD backend that overrides
         the default one.
     """
     self.binary_path = binary_path
     self.package_name = package_name
     self.instance_id = instance_id
     self.service_url = service_url
 
+  def _ensure_from_isolate(self, root, subdir, isolated_cipd, isolate_cache):
+    if not isolate_cache:
+      logging.info('Not ensuring cipd from isolate cache isolate_cache is not'
+                   'defined: %s', isolate_cache)
+      return False
+    try:
+      with open(isolated_cipd , 'r') as f:
+        digest = str(f.read())
+      try:
+        content = isolate_cache.getfileobj(digest).read()
+      except Exception as e:
+        logging.warning('Could not find isolated file in cache with digest '
+                        '%s: %s', digest, e)
+        return False
+
+      #FIXME(kjlubick): Don't assume sha1
+      sha_1 = isolated_format.SUPPORTED_ALGOS['sha-1']
+      ifile = isolated_format.IsolatedFile(digest, sha_1)
+      ifile.load(content)
+
+      subdir = os.path.join(root, subdir)
+      file_path.ensure_tree(subdir)
+      files = ifile.data.get(u'files', {})
+      for f in files.keys():
+        props = files.get(f, None)
+        if not props:
+          logging.warning('Problem getting info for %s', f)
+          return False
+        file_mode = props.get('m', None)
+        if file_mode:
+          # Ignore all bits apart from the user
+          file_mode &= 0700
+
+        dstpath = os.path.join(subdir, f)
+        file_path.ensure_tree(os.path.dirname(dstpath))
+        digest = props.get('h', None)
+        if not digest:
+          logging.warning('Hash can\'t be empty %s', f)
+          return False
+        srcpath = isolate_cache.getfileobj(digest).name
+
+        file_path.link_file(unicode(dstpath), unicode(srcpath),
+                            file_path.HARDLINK_WITH_FALLBACK)
+
+        if file_mode is not None:
+          fs.chmod(dstpath, file_mode)
+    except Exception as e:
+      logging.warning('Could not ensure cipd package from isolate %s', e)
+
+    return True
+
+
+  def _isolate_cipd(self, root, packages, isolate_cache, cipd_cache):
+    logging.debug('_isolate_cipd(%s, %s, %s, %s)', root, packages,
+                  isolate_cache, cipd_cache)
+    if not isolate_cache or not os.path.isdir(cipd_cache):
+      logging.info('Not putting cipd into isolate cache because one of the'
+                   'caches is empty: %s, %s', isolate_cache, cipd_cache)
+      return
+    for subdir, shafile in packages.iteritems():
+      isolated_file = os.path.join(cipd_cache, shafile[:-5])
+      complete_state = isolate.CompleteState.load_files(isolated_file)
+
+      subdir = os.path.join(root, subdir)
+
+      infiles = isolated_format.expand_directories_and_symlinks(
+        subdir, ['./'], [], False, True)
+
+      complete_state.saved_state.update_isolated(None, infiles, True, None)
+      complete_state.saved_state.root_dir = subdir
+      # Collapse symlinks to remove CIPD symlinking AND to simplify other code.
+      complete_state.files_to_metadata('', True)
+      complete_state.save_files()
+
+      for infile in infiles:
+        digest = complete_state.saved_state.files[infile].get('h', '')
+        if not digest:
+          logging.warning('No digest found in saved state %s:%s', infile,
+                          complete_state.saved_state.files[infile])
+          continue
+        with open(os.path.join(subdir, infile) , 'r') as f:
+          isolate_cache.write(digest, f)
+
+      with open(isolated_file , 'r') as f:
+        content = f.read()
+        digest = complete_state.saved_state.algo(content).hexdigest()
+        isolate_cache.write(digest, content)
+
+        with open(os.path.join(cipd_cache, shafile), 'w') as sf:
+          sf.write(digest)
+
+
+
   def ensure(
-      self, site_root, packages, cache_dir=None, tmp_dir=None, timeout=None):
+      self, site_root, packages, cache_dir=None, tmp_dir=None, timeout=None,
+      isolate_cache=None):
     """Ensures that packages installed in |site_root| equals |packages| set.
 
     Blocking call.
 
+    Attempts to use the isolate cache to store the unzipped cipd files, keeping
+    a .isolated file in the cipd cache_dir
+
     Args:
       site_root (str): where to install packages.
       packages: dict of subdir -> list of (package_template, version) tuples.
       cache_dir (str): if set, cache dir for cipd binary own cache.
         Typically contains packages and tags.
       tmp_dir (str): if not None, dir for temp files.
       timeout (int): if not None, timeout in seconds for this function to run.
 
     Returns:
       Pinned packages in the form of {subdir: [(package_name, package_id)]},
       which correspond 1:1 with the input packages argument.
 
     Raises:
       Error if could not install packages or timed out.
     """
     timeoutfn = tools.sliding_timeout(timeout)
     logging.info('Installing packages %r into %s', packages, site_root)
+    logging.info('Cache dir %s', cache_dir)
 
     ensure_file_handle, ensure_file_path = tempfile.mkstemp(
         dir=tmp_dir, prefix=u'cipd-ensure-file-', suffix='.txt')
     json_out_file_handle, json_file_path = tempfile.mkstemp(
       dir=tmp_dir, prefix=u'cipd-ensure-result-', suffix='.json')
     os.close(json_out_file_handle)
-
+    if cache_dir:
+      file_path.ensure_tree(unicode(cache_dir))
+    to_isolate = {}
+    from_isolate = {}
     try:
       try:
         for subdir, pkgs in sorted(packages.iteritems()):
           if '\n' in subdir:
             raise Error(
               'Could not install packages; subdir %r contains newline' % subdir)
+
+          versions = [p[1] for p in pkgs]
+          isolated_cipd = '%s.%s.isolated.sha1' % (subdir,
+                                                   '_'.join(versions))
+          abs_isolated_cipd = os.path.join(cache_dir, isolated_cipd)
+          if (os.path.isfile(abs_isolated_cipd) and
+              self._ensure_from_isolate(site_root, subdir, abs_isolated_cipd,
+                                        isolate_cache)):
+            from_isolate[unicode(subdir)] = pkgs
+            continue
+          to_isolate[subdir] = isolated_cipd
           os.write(ensure_file_handle, '@Subdir %s\n' % (subdir,))
           for pkg, version in pkgs:
             pkg = render_package_name_template(pkg)
             os.write(ensure_file_handle, '%s %s\n' % (pkg, version))
+
       finally:
         os.close(ensure_file_handle)
 
+      # to_isolate is the packages that we need to ensure from CIPD and then
+      # isolate.  Thus, if this is empty, we don't need to get anything from
+      # CIPD because they were successfully pulled from isolate. Thus return
+      # from_isolate, the pinned packages that we pulled from_isolate
+      if not to_isolate:
+        return from_isolate
+
       cmd = [
         self.binary_path, 'ensure',
         '-root', site_root,
         '-ensure-file', ensure_file_path,
         '-verbose',  # this is safe because cipd-ensure does not print a lot
         '-json-output', json_file_path,
       ]
       if cache_dir:
         cmd += ['-cache-dir', cache_dir]
       if self.service_url:
@@ skipping 16 matching lines @@
         if pipe_name == 'stderr':
           logging.debug('cipd client: %s', line)
         else:
           logging.info('cipd client: %s', line)
 
       exit_code = process.wait(timeout=timeoutfn())
       if exit_code != 0:
         raise Error(
             'Could not install packages; exit code %d\noutput:%s' % (
             exit_code, '\n'.join(output)))
+
+      self._isolate_cipd(site_root, to_isolate, isolate_cache, cache_dir)
+
       with open(json_file_path) as jfile:
         result_json = json.load(jfile)
-      return {
+      from_isolate.update({
         subdir: [(x['package'], x['instance_id']) for x in pins]
         for subdir, pins in result_json['result'].iteritems()
-      }
+      })
+      return from_isolate
     finally:
       fs.remove(ensure_file_path)
       fs.remove(json_file_path)
 
 
 def get_platform():
   """Returns ${platform} parameter value.
 
   Borrowed from
   https://chromium.googlesource.com/infra/infra/+/aaf9586/build/build.py#204
@@ skipping 234 matching lines @@
   """
   result = []
   for pkg in packages:
     path, name, version = pkg.split(':', 2)
     if not name:
       raise Error('Invalid package "%s": package name is not specified' % pkg)
     if not version:
       raise Error('Invalid package "%s": version is not specified' % pkg)
     result.append((path, name, version))
   return result