Add tool to use the manifest, fetch and cache dependencies and run the test.

tools/isolate/tree_creator.py

+#!/usr/bin/env python
 # Copyright (c) 2012 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
-"""File related utility functions.
+"""Reads a manifest, creates a tree of hardlinks and runs the test.
 
-Creates a tree of hardlinks, symlinks or copy the inputs files. Calculate files
-hash.
+Keeps a local cache.
 """
 
 import ctypes
-import hashlib
+import json
 import logging
+import optparse
 import os
+import re
 import shutil
-import stat
+import subprocess
 import sys
+import tempfile
 import time
+import urllib
 
 
 # Types of action accepted by recreate_tree().
 HARDLINK, SYMLINK, COPY = range(4)[1:]
 
 
 class MappingError(OSError):
   """Failed to recreate the tree."""
   pass
 
 
 def os_link(source, link_name):
   """Add support for os.link() on Windows."""
   if sys.platform == 'win32':
     if not ctypes.windll.kernel32.CreateHardLinkW(
         unicode(link_name), unicode(source), 0):
       raise OSError()
   else:
     os.link(source, link_name)
 
 
-def expand_directories(indir, infiles, blacklist):
-  """Expands the directories, applies the blacklist and verifies files exist."""
-  logging.debug('expand_directories(%s, %s, %s)' % (indir, infiles, blacklist))
-  outfiles = []
-  for relfile in infiles:
-    if os.path.isabs(relfile):
-      raise MappingError('Can\'t map absolute path %s' % relfile)
-    infile = os.path.normpath(os.path.join(indir, relfile))
-    if not infile.startswith(indir):
-      raise MappingError('Can\'t map file %s outside %s' % (infile, indir))
-
-    if relfile.endswith('/'):
-      if not os.path.isdir(infile):
-        raise MappingError(
-            'Input directory %s must have a trailing slash' % infile)
-      for dirpath, dirnames, filenames in os.walk(infile):
-        # Convert the absolute path to subdir + relative subdirectory.
-        relpath = dirpath[len(indir)+1:]
-        outfiles.extend(os.path.join(relpath, f) for f in filenames)
-        for index, dirname in enumerate(dirnames):
-          # Do not process blacklisted directories.
-          if blacklist(os.path.join(relpath, dirname)):
-            del dirnames[index]
-    else:
-      if not os.path.isfile(infile):
-        raise MappingError('Input file %s doesn\'t exist' % infile)
-      outfiles.append(relfile)
-  return outfiles
-
-
-def process_inputs(indir, infiles, need_hash, read_only):
-  """Returns a dictionary of input files, populated with the files' mode and
-  hash.
-
-  The file mode is manipulated if read_only is True. In practice, we only save
-  one of 4 modes: 0755 (rwx), 0644 (rw), 0555 (rx), 0444 (r).
-  """
-  outdict = {}
-  for infile in infiles:
-    filepath = os.path.join(indir, infile)
-    filemode = stat.S_IMODE(os.stat(filepath).st_mode)
-    # Remove write access for non-owner.
-    filemode &= ~(stat.S_IWGRP | stat.S_IWOTH)
-    if read_only:
-      filemode &= ~stat.S_IWUSR
-    if filemode & stat.S_IXUSR:
-      filemode |= (stat.S_IXGRP | stat.S_IXOTH)
-    else:
-      filemode &= ~(stat.S_IXGRP | stat.S_IXOTH)
-    outdict[infile] = {
-      'mode': filemode,
-    }
-    if need_hash:
-      h = hashlib.sha1()
-      with open(filepath, 'rb') as f:
-        h.update(f.read())
-      outdict[infile]['sha-1'] = h.hexdigest()
-  return outdict
-
-
 def link_file(outfile, infile, action):
   """Links a file. The type of link depends on |action|."""
   logging.debug('Mapping %s to %s' % (infile, outfile))
+  if action not in (HARDLINK, SYMLINK, COPY):
+    raise ValueError('Unknown mapping action %s' % action)
   if os.path.isfile(outfile):
     raise MappingError('%s already exist' % outfile)
 
   if action == COPY:
     shutil.copy(infile, outfile)
   elif action == SYMLINK and sys.platform != 'win32':
+    # On windows, symlink are converted to hardlink and fails over to copy.
     os.symlink(infile, outfile)
-  elif action == HARDLINK:
+  else:
     try:
       os_link(infile, outfile)
     except OSError:
       # Probably a different file system.
       logging.warn(
           'Failed to hardlink, failing back to copy %s to %s' % (
             infile, outfile))
       shutil.copy(infile, outfile)
-  else:
-    raise ValueError('Unknown mapping action %s' % action)
-
-
-def recreate_tree(outdir, indir, infiles, action):
-  """Creates a new tree with only the input files in it.
-
-  Arguments:
-    outdir:    Output directory to create the files in.
-    indir:     Root directory the infiles are based in.
-    infiles:   List of files to map from |indir| to |outdir|.
-    action:    See assert below.
-  """
-  logging.debug(
-      'recreate_tree(%s, %s, %s, %s)' % (outdir, indir, infiles, action))
-  logging.info('Mapping from %s to %s' % (indir, outdir))
-
-  assert action in (HARDLINK, SYMLINK, COPY)
-  outdir = os.path.normpath(outdir)
-  if not os.path.isdir(outdir):
-    logging.info ('Creating %s' % outdir)
-    os.makedirs(outdir)
-  # Do not call abspath until the directory exists.
-  outdir = os.path.abspath(outdir)
-
-  for relfile in infiles:
-    infile = os.path.join(indir, relfile)
-    outfile = os.path.join(outdir, relfile)
-    outsubdir = os.path.dirname(outfile)
-    if not os.path.isdir(outsubdir):
-      os.makedirs(outsubdir)
-    link_file(outfile, infile, action)
 
 
 def _set_write_bit(path, read_only):
   """Sets or resets the executable bit on a file or directory."""
   mode = os.stat(path).st_mode
   if read_only:
     mode = mode & 0500
   else:
     mode = mode | 0200
   if hasattr(os, 'lchmod'):
     os.lchmod(path, mode)  # pylint: disable=E1101
   else:
     # TODO(maruel): Implement proper DACL modification on Windows.
     os.chmod(path, mode)
 
 
 def make_writable(root, read_only):
   """Toggle the writable bit on a directory tree."""
   root = os.path.abspath(root)
   for dirpath, dirnames, filenames in os.walk(root, topdown=True):
     for filename in filenames:
       _set_write_bit(os.path.join(dirpath, filename), read_only)
 
     for dirname in dirnames:
       _set_write_bit(os.path.join(dirpath, dirname), read_only)
 
 
 def rmtree(root):
   """Wrapper around shutil.rmtree() to retry automatically on Windows."""
+  make_writable(root, False)
   if sys.platform == 'win32':
     for i in range(3):
       try:
         shutil.rmtree(root)
         break
       except WindowsError:  # pylint: disable=E0602
         delay = (i+1)*2
         print >> sys.stderr, (
             'The test has subprocess outliving it. Sleep %d seconds.' % delay)
         time.sleep(delay)
   else:
     shutil.rmtree(root)
+
+
+def open_remote(file_or_url):
+  """Reads a file or url."""
+  if re.match(r'^https?://.+$', file_or_url):
+    return urllib.urlopen(file_or_url)
+  return open(file_or_url, 'rb')
+
+
+def download_or_copy(file_or_url, dest):
+  """Copies a file or download an url."""
+  if re.match(r'^https?://.+$', file_or_url):
+    urllib.urlretrieve(file_or_url, dest)
+  else:
+    shutil.copy(file_or_url, dest)
+
+
+def get_free_space(path):
+  """Returns the number of free bytes."""
+  if sys.platform == 'win32':
+    free_bytes = ctypes.c_ulonglong(0)
+    ctypes.windll.kernel32.GetDiskFreeSpaceExW(
+        ctypes.c_wchar_p(path), None, None, ctypes.pointer(free_bytes))
+    return free_bytes.value
+  f = os.statvfs(path)
+  return f.f_bfree * f.f_frsize
+
+
+class Cache(object):
+  """Stateful LRU cache.
+
+  Saves its state as json file.
+  """
+  STATE_FILE = 'state.json'
+
+  def __init__(self, cache_dir, remote, max_cache_size, min_free_space):
+    """
+    Arguments:
+    - cache_dir: Directory where to place the cache.
+    - remote: Remote directory (NFS, SMB, etc) or HTTP url to fetch the objects
+              from
+    - max_cache_size: Trim if the cache gets larger than this value. If 0, the
+                      cache is effectively a leak.
+    - min_free_space: Trim if disk free space becomes lower than this value. If
+                      0, it unconditionally fill the disk.
+    """
+    self.cache_dir = cache_dir
+    self.remote = remote
+    self.max_cache_size = max_cache_size
+    self.min_free_space = min_free_space
+    self.state_file = os.path.join(cache_dir, self.STATE_FILE)
+    # The files are kept as an array in a LRU style. E.g. self.state[0] is the
+    # oldest item.
+    self.state = []
+
+    if not os.path.isdir(self.cache_dir):
+      os.makedirs(self.cache_dir)
+    if os.path.isfile(self.state_file):
+      try:
+        self.state = json.load(open(self.state_file, 'rb'))
+      except ValueError:
+        # Too bad. The file will be overwritten and the cache cleared.
+        pass
+    self.trim()
+
+  def trim(self):
+    """Trims anything we don't know, make sure enough free space exists."""
+    for f in os.listdir(self.cache_dir):
+      if f == self.STATE_FILE or f in self.state:
+        continue
+      logging.warn('Unknown file %s from cache' % f)
+      # Insert as the oldest file. It will be deleted eventually if not
+      # accessed.
+      self.state.insert(0, f)
+
+    # Ensure enough free space.
+    while (
+        self.min_free_space and
+        self.state and
+        get_free_space(self.cache_dir) < self.min_free_space):
+      os.remove(self.path(self.state.pop(0)))
+
+    # Ensure maximum cache size.
+    if self.max_cache_size and self.state:
+      sizes = [os.stat(self.path(f)).st_size for f in self.state]
+      while sizes and sum(sizes) > self.max_cache_size:
+        # Delete the oldest item.
+        os.remove(self.path(self.state.pop(0)))
+        sizes.pop(0)
+
+    self.save()
+
+  def retrieve(self, item):
+    """Retrieves a file from the remote and add it to the cache."""
+    assert not '/' in item
+    try:
+      index = self.state.index(item)
+      # Was already in cache. Update it's LRU value.
+      self.state.pop(index)
+      self.state.append(item)
+      return False
+    except ValueError:
+      out = self.path(item)
+      download_or_copy(os.path.join(self.remote, item), out)
+      self.state.append(item)
+      return True
+    finally:
+      self.save()
+
+  def path(self, item):
+    """Returns the path to one item."""
+    return os.path.join(self.cache_dir, item)
+
+  def save(self):
+    """Saves the LRU ordering."""
+    json.dump(self.state, open(self.state_file, 'wb'))
+
+
+def run_tha_test(manifest, cache_dir, remote, max_cache_size, min_free_space):
+  """Downloads the dependencies in the cache, hardlinks them into a temporary
+  directory and runs the executable.
+  """
+  cache = Cache(cache_dir, remote, max_cache_size, min_free_space)
+  outdir = tempfile.mkdtemp(prefix='run_tha_test')
+  try:
+    for filepath, properties in manifest['files'].iteritems():
+      infile = properties['sha-1']
+      outfile = os.path.join(outdir, filepath)
+      cache.retrieve(infile)
+      outfiledir = os.path.dirname(outfile)
+      if not os.path.isdir(outfiledir):
+        os.makedirs(outfiledir)
+      link_file(outfile, cache.path(infile), HARDLINK)
+      os.chmod(outfile, properties['mode'])
+
+    cwd = os.path.join(outdir, manifest['relative_cwd'])
+    if not os.path.isdir(cwd):
+      os.makedirs(cwd)
+    if manifest.get('read_only'):
+      make_writable(outdir, True)
+    cmd = manifest['command']
+    logging.info('Running %s, cwd=%s' % (cmd, cwd))
+    return subprocess.call(cmd, cwd=cwd)
+  finally:
+    # Save first, in case an exception occur in the following lines, then clean
+    # up.
+    cache.save()
+    rmtree(outdir)
+    cache.trim()
+
+
+def main():
+  parser = optparse.OptionParser(
+      usage='%prog <options>', description=sys.modules[__name__].__doc__)
+  parser.add_option(
+      '-v', '--verbose', action='count', default=0, help='Use multiple times')
+  parser.add_option(
+      '-m', '--manifest',
+      metavar='FILE',
+      help='File/url describing what to map or run')
+  parser.add_option('--no-run', action='store_true', help='Skip the run part')
+  parser.add_option(
+      '--cache',
+      default='cache',
+      metavar='DIR',
+      help='Cache directory, default=%default')
+  parser.add_option(
+      '-r', '--remote', metavar='URL', help='Remote where to get the items')
+  parser.add_option(
+      '--max-cache-size',
+      type='int',
+      metavar='NNN',
+      default=20*1024*1024*1024,
+      help='Trim if the cache gets larger than this value, default=%default')
+  parser.add_option(
+      '--min-free-space',
+      type='int',
+      metavar='NNN',
+      default=1*1024*1024*1024,
+      help='Trim if disk free space becomes lower than this value, '
+           'default=%default')
+
+  options, args = parser.parse_args()
+  level = [logging.ERROR, logging.INFO, logging.DEBUG][min(2, options.verbose)]
+  logging.basicConfig(
+      level=level,
+      format='%(levelname)5s %(module)15s(%(lineno)3d): %(message)s')
+
+  if not options.manifest:
+    parser.error('--manifest is required.')
+  if not options.remote:
+    parser.error('--remote is required.')
+  if args:
+    parser.error('Unsupported args %s' % ' '.join(args))
+
+  manifest = json.load(open_remote(options.manifest))
+  return run_tha_test(
+      manifest, os.path.abspath(options.cache), options.remote,
+      options.max_cache_size, options.min_free_space)
+
+
+if __name__ == '__main__':
+  sys.exit(main())