Cache/retrieve extracted CIPD packages in local isolate cache

client/cipd.py

Index: client/cipd.py
diff --git a/client/cipd.py b/client/cipd.py
index 927ec594354790a93fbc22c418c1af17fda29daf..a6ebd3043cd2b9a565f790e9d73467951b469f76 100644
--- a/client/cipd.py
+++ b/client/cipd.py
@@ -11,6 +11,8 @@ import logging
 import optparse
 import os
 import platform
+import re
+import shutil
 import sys
 import tempfile
 import time
@@ -21,6 +23,7 @@ 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
 
@@ -140,12 +143,109 @@ class CipdClient(object):
     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.
@@ -163,26 +263,49 @@ class CipdClient(object):
     """
     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,
@@ -219,12 +342,16 @@ class CipdClient(object):
         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)