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
@@ skipping 108 matching lines @@
       binary_path (str): path to the CIPD client binary.
       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
+    self._cache_hash = hashlib.sha1
+    self._cache_hash_algo = (
+        isolated_format.SUPPORTED_ALGOS_REVERSE[self._cache_hash])
+
+  def _ensure_from_isolate(self, target_dir, cipd_isolated, isolate_cache):
+    """
+    Retrieves the CIPD subdirectories from the isolate cache, if they exist,
+    hardlinking or copying the files into the provided directory.
+
+    Does the opposite of _isolate_cipd for a given cipd_isolated file.
+
+    Args:
+      target_dir: directory in which to put the files.
+      cipd_isolated: the isolated.[hash] file created previously in
+        _isolate_cipd
+      isolate_cache: the isolateserver.DiskCache in which the files may be
+        stored.
+
+    Returns:
+      True if the isolated file and all the isolated contents were successfully
+      found in the isolate cache and put into place using hardlinks with a copy
+      fallback. False otherwise.
+    """
+    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(cipd_isolated , 'r') as f:

[M-A Ruel, 2017/05/08 20:43:16]

'rb'
it makes it faster on windows.

[kjlubick, 2017/05/09 17:38:25]

Done.

+        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
+
+      ifile = isolated_format.IsolatedFile(digest, self._cache_hash)

[M-A Ruel, 2017/05/08 20:43:16]

This part should be a function not a method. I'm thinking in isolateserver.py,
as it's similar to
https://github.com/luci/luci-py/blob/master/client/isolateserver.py#L1769 or so.

[kjlubick, 2017/05/09 17:38:25]

Done.

+      ifile.load(content)
+
+      file_path.ensure_tree(target_dir)
+      files = ifile.data.get(u'files', {})
+      for f in files.keys():

[M-A Ruel, 2017/05/08 20:43:16]

for f, props in files.iteritems():

[kjlubick, 2017/05/09 17:38:25]

I can't - I get "Too many values to unpack"  Comment added

+        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(target_dir, 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 False
+
+    return True
+
+  def _isolate_cipd(self, root, subdirs, isolate_cache, cipd_cache):

[kjlubick, 2017/05/08 18:59:29]

This doesn't support includes like we mentioned.  I couldn't figure out how to
incorporate it or where the includes would come from, be accessed, etc.

Any advice?

+    """
+    Puts the content of the CIPD subdirectories into the isolate cache,
+    creating a .isolated file representing each subdirectory. This .isolated
+    file goes into the isolate_cache as well, and a .isolated.[hash] file
+    goes into cipd_cache for extraction in _ensure_from_isolate(). The suffix
+    will be related to self._cache_hash_algo (.sha-1 for SHA-1, etc)
+
+    This allows for disks with slow I/O (e.g. Raspberry Pis) to not have to
+    re-extract the CIPD zips every time, potentially saving a bunch of time.
+
+    Arg:
+      root: where packages are installed.
+      subdirs: dict of subdir -> name_verision where name_version is

[M-A Ruel, 2017/05/08 20:43:16]

name_version

[kjlubick, 2017/05/09 17:38:25]

Done.

+        [subdir].[pkg1version_pkg2version...].isolated.[hash] This way, if
+        any packages are updated, the cached files will also change.
+        isolate_cache: A isolateserver.DiskCache used to store files locally.
+        cipd_cache: A directory in which to put the *isolated.[hash] files.
+    """
+    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, hashfile in subdirs.iteritems():
+      if not subdir:
+        logging.info('not caching to %s because it extracts to ./', hashfile)
+        continue
+      # The subdirs given to us from cipd are always in foo/bar/baz format
+      # which freaks Windows out.
+      subdir = subdir.replace('/', os.path.sep)
+      subdir = os.path.join(root, subdir)
+
+      if not os.path.isdir(subdir):
+        logging.warning('%r is not a directory, so it can\'t be isolated',
+                        subdir)
+        continue
+
+      infiles, metadata = isolateserver.directory_to_metadata(
+          subdir, self._cache_hash, [], True)
+
+      # The .isolated file to be created on disk. hashfile represents the
+      # file that will also be created along with this.
+      # e.g. *.isolated.sha-1 if the sha1 algorithm is used
+      isolated_file = unicode(os.path.join(cipd_cache, hashfile[:-5]))
+      data = {
+        'algo': self._cache_hash_algo,
+        'files': metadata,
+        'version': isolated_format.ISOLATED_FILE_VERSION,
+      }
+      isolated_format.save_isolated(isolated_file, data)
+
+      for infile in infiles:
+        with open(os.path.join(subdir, infile.path) , 'r') as f:

[M-A Ruel, 2017/05/08 20:43:16]

'rb'

[kjlubick, 2017/05/09 17:38:25]

Done.

+          isolate_cache.write(infile.digest, f)
+
+      digest = isolated_format.hash_file(isolated_file, self._cache_hash)
+      with open(isolated_file , 'r') as f:

[M-A Ruel, 2017/05/08 20:43:16]

'rb'

[kjlubick, 2017/05/09 17:38:25]

Done

+        content = f.read()
+        isolate_cache.write(digest, content)

[M-A Ruel, 2017/05/08 20:43:16]

you can calculate hash here, so the file is not read twice.

[kjlubick, 2017/05/09 17:38:25]

Done.

+
+      with open(os.path.join(cipd_cache, hashfile), 'w') as f:
+        f.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.
+      isolate_cache (isolateserver.DiskCache): if not None, CIPD assets will
+        be unzipped and stored in this disk cache and extracted from there.
 
     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)
-
     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))

[M-A Ruel, 2017/05/08 20:43:16]

I'd prefer to update call sites instead.

[kjlubick, 2017/05/09 17:38:25]

Done.  Only run_isolated needed it.

+    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)
+
+          # Join all the versions together so as to cause a new cached isolated
+          # to be used if any of them change.
+          versions = [p[1] for p in pkgs]

[M-A Ruel, 2017/05/08 20:43:16]

why not each package be individually packed?

[kjlubick, 2017/05/09 17:38:25]

I had tried that originally, but that gets tricky when you have two packages
extracted into the same subdirectory. It ends up being cleaner to cache on a
directory by directory basis, at the expense of more cache churn if only one of
the packages updates.

[M-A Ruel, 2017/05/10 13:31:58]

Hummm I still think it's not a great idea. Since there can be collision, extract
in a temporary directory, archive this, then remap into the destination. It's 3x
more I/O initially but I think it's worth it because then you can detect if
there are files conflicts and it will improve cache hit rate when one package is
rolled frequently but not the others.

+          cipd_isolated = '%s.%s.isolated.%s' % (subdir, '_'.join(versions),

[M-A Ruel, 2017/05/08 20:43:16]

u

[kjlubick, 2017/05/09 17:38:25]

Done.

+                                                 self._cache_hash_algo)
+          cipd_isolated = os.path.join(cache_dir, cipd_isolated)
+          if (os.path.isfile(cipd_isolated) and
+              self._ensure_from_isolate(os.path.join(site_root, subdir),
+                                        cipd_isolated, isolate_cache)):
+            from_isolate[unicode(subdir)] = pkgs
+            continue
+          to_isolate[subdir] = cipd_isolated
           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 236 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