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,

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

please start comment on the first line
generally prefer public methods first then private methods

[kjlubick, 2017/05/10 19:41:00]

Done.

+    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.

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

inconsistent dots, remove

[kjlubick, 2017/05/10 19:41:00]

Done.

+      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 , 'rb') as f:
+        digest = str(f.read())

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

it's already an str.

[kjlubick, 2017/05/10 19:41:01]

Done.

+      try:
+        content = isolate_cache.getfileobj(digest).read()
+      except Exception as e:

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

that's too broad, please scope down

[kjlubick, 2017/05/10 19:41:00]

Done.

+        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)
+      ifile.load(content)
+
+      if not isolateserver.is_cached(ifile, isolate_cache):
+        logging.info('Cached CIPD asset(s) %s are incomplete', cipd_isolated)
+        return False
+
+      file_path.ensure_tree(target_dir)
+      return isolateserver.extract(ifile, target_dir, isolate_cache)
+    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):
+    """
+    Puts the content of the CIPD subdirectories into the isolate cache,

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

single sentence
empty line
rest

[kjlubick, 2017/05/10 19:41:01]

Done.

+    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_version where name_version is
+        [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) , 'rb') as f:
+          isolate_cache.write(infile.digest, f)
+
+      with open(isolated_file , 'rb') as f:
+        content = f.read()
+        digest = self._cache_hash(content).hexdigest()
+        isolate_cache.write(digest, content)
+
+      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)
-
+    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]
+          cipd_isolated = u'%s.%s.isolated.%s' % (subdir, '_'.join(versions),
+                                                  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