luci-py/isolateserver.py: Add archive support when downloading.

client/isolateserver.py

 #!/usr/bin/env python
 # Copyright 2013 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.
 
 """Archives a set of files or directories to an Isolate Server."""
 
 __version__ = '0.4.9'
 
 import base64
 import functools
 import errno
 import logging
 import optparse
 import os
 import re
+import stat
 import signal
 import sys
 import tempfile
 import threading
 import time
 import types
 import zlib
 
+import cStringIO as StringIO
+
 from third_party import colorama
 from third_party.depot_tools import fix_encoding
 from third_party.depot_tools import subcommand
 
+from libs import arfile
 from utils import file_path
 from utils import fs
 from utils import logging_utils
 from utils import lru
 from utils import net
 from utils import on_error
 from utils import subprocess42
 from utils import threading_utils
 from utils import tools
 
@@ skipping 103 matching lines @@
   """
   file_path.ensure_tree(os.path.dirname(path))
   total = 0
   with fs.open(path, 'wb') as f:
     for d in content_generator:
       total += len(d)
       f.write(d)
   return total
 
 
+def fileobj_path(fileobj):
+  name = getattr(fileobj, "name", None)

[nodir, 2016/06/20 15:54:35]

s/"/'/

[mithro, 2016/06/21 06:42:33]

Done.

+  if name is None:
+    return
+
+  if not isinstance(name, unicode):
+    name = name.decode(sys.getfilesystemencoding())
+
+  if fs.exists(name):
+    return name
+
+
+def fileobj_copy(dstfileobj, srcfileobj, amount=-1):
+  written = 0
+  while written != amount:
+    readsize = NET_IO_FILE_CHUNK
+    if amount > 0:
+      readsize = min(readsize, amount-written)
+    data = srcfileobj.read(readsize)
+    if not data:
+      if amount == -1:
+        break
+      else:

[nodir, 2016/06/20 15:54:35]

unnecessary else

[mithro, 2016/06/21 06:42:33]

Done.

+        raise IOError('partial file, got %s, wanted %s' % (written, amount))
+    dstfileobj.write(data)
+    written += readsize

[nodir, 2016/06/20 15:54:35]

you might have to `written = destfileobj.tell()` because readsize is not
necessarily equal to # of bytes actually read from srcfileobject, which is not
necessarily equal to # of bytes actually written to descobjfile

[mithro, 2016/06/21 06:42:33]

len(data) should be enough.

If write doesn't write the full amount, it raises an IOError.

[nodir, 2016/06/21 22:20:05]

Acknowledged.

+
+
+def putfile(srcfileobj, dstpath, file_mode=None, limit=-1):

[nodir, 2016/06/20 15:54:35]

should self._used be updated in this method?

[mithro, 2016/06/21 06:42:33]

This function doesn't have access to self._used and I don't think should be
independent of that class too.

[nodir, 2016/06/21 22:20:04]

for some reason I thought it is a method, not a global func

+  srcpath = fileobj_path(srcfileobj)
+  if srcpath:
+    readonly = file_mode is None or (
+        file_mode & (stat.S_IWUSR | stat.S_IWGRP | stat.S_IWOTH))
+    if readonly:
+      # We can link the file
+      file_path.link_file(dstpath, srcpath, file_path.HARDLINK_WITH_FALLBACK)
+    else:
+      # We must copy the file
+      file_path.link_file(dstpath, srcpath, file_path.COPY)
+  else:
+    # Need to write out the file
+    with fs.open(dstpath, "wb") as dstfileobj:
+      fileobj_copy(dstfileobj, srcfileobj, limit)
+
+  # file_mode of 0 is actually valid, so need explicit check.
+  if file_mode is not None:
+    fs.chmod(dstpath, file_mode)
+
+
 def zip_compress(content_generator, level=7):
   """Reads chunks from |content_generator| and yields zip compressed chunks."""
   compressor = zlib.compressobj(level)
   for chunk in content_generator:
     compressed = compressor.compress(chunk)
     if compressed:
       yield compressed
   tail = compressor.flush(zlib.Z_FINISH)
   if tail:
     yield tail
@@ skipping 1054 matching lines @@
   """
   cache_dir = None
 
   def __init__(self):
     self._lock = threading_utils.LockWithAssert()
     # Profiling values.
     self._added = []
     self._initial_number_items = 0
     self._initial_size = 0
     self._evicted = []
-    self._linked = []
+    self._used = []
 
   def __contains__(self, digest):
     raise NotImplementedError()
 
   def __enter__(self):
     """Context manager interface."""
     return self
 
   def __exit__(self, _exc_type, _exec_value, _traceback):
     """Context manager interface."""
     return False
 
   @property
   def added(self):
     return self._added[:]
 
   @property
   def evicted(self):
     return self._evicted[:]
 
   @property
+  def used(self):
+    return self._used[:]
+
+  @property
   def initial_number_items(self):
     return self._initial_number_items
 
   @property
   def initial_size(self):
     return self._initial_size
 
-  @property
-  def linked(self):
-    return self._linked[:]
-
   def cached_set(self):
     """Returns a set of all cached digests (always a new object)."""
     raise NotImplementedError()
 
   def cleanup(self):
     """Deletes any corrupted item from the cache and trims it if necessary."""
     raise NotImplementedError()
 
   def touch(self, digest, size):
     """Ensures item is not corrupted and updates its LRU position.
 
     Arguments:
       digest: hash digest of item to check.
       size: expected size of this item.
 
     Returns:
       True if item is in cache and not corrupted.
     """
     raise NotImplementedError()
 
   def evict(self, digest):
     """Removes item from cache if it's there."""
     raise NotImplementedError()
 
-  def read(self, digest):
-    """Returns contents of the cached item as a single str."""
+  def getfileobj(self, digest):
+    """Returns a file like object.
+
+    If file exists on the file system it will have a .name attribute with an
+    absolute path to the file.
+    """

[nodir, 2016/06/20 15:54:34]

document that it is the responsibility of the caller to close the returned file

[mithro, 2016/06/21 06:42:33]

I updated the docs to include that it is a *readable* file object.

In Python the file will be closed when the object is garbage collected. 

There are only really two reasons you want to explicitly close file objects in
python;
 * You are opening / closing a large number of files and are running out of file
descriptors before garbage collection is occurring.
 * You are writing to a file and have finished with it.

[M-A Ruel, 2016/06/21 13:31:08]

We are potentially hitting this case.

[nodir, 2016/06/21 22:20:04]

IIRC windows requires a file handle to be closed for another process to access
the file, so we should close all file handles before we spawn the task command.

the general recommendation, on windows at least, i think, is to close a file
handle as soon as possible. it is orthogonal to the language

[mithro, 2016/06/22 12:03:10]

I believe this is correct if the file is opened in write or exclusive read mode
but not shared read mode (which is the default).
I'm doing this now.

     raise NotImplementedError()
 
   def write(self, digest, content):
     """Reads data from |content| generator and stores it in cache.
 
     Returns digest to simplify chaining.
     """
     raise NotImplementedError()
 
-  def hardlink(self, digest, dest, file_mode):
-    """Ensures file at |dest| has same content as cached |digest|.
-
-    If file_mode is provided, it is used to set the executable bit if
-    applicable.
-    """
-    raise NotImplementedError()
-
 
 class MemoryCache(LocalCache):
   """LocalCache implementation that stores everything in memory."""
 
   def __init__(self, file_mode_mask=0500):
     """Args:
       file_mode_mask: bit mask to AND file mode with. Default value will make
           all mapped files to be read only.
     """
     super(MemoryCache, self).__init__()
@@ skipping 14 matching lines @@
   def touch(self, digest, size):
     with self._lock:
       return digest in self._contents
 
   def evict(self, digest):
     with self._lock:
       v = self._contents.pop(digest, None)
       if v is not None:
         self._evicted.add(v)
 
-  def read(self, digest):
+  def getfileobj(self, digest):
     with self._lock:
       try:
-        return self._contents[digest]
+        d = self._contents[digest]
+        self._used.append(len(d))
+        return StringIO.StringIO(d)
       except KeyError:
         raise CacheMiss(digest)
 
   def write(self, digest, content):
     # Assemble whole stream before taking the lock.
     data = ''.join(content)
     with self._lock:
       self._contents[digest] = data
       self._added.append(len(data))
     return digest
 
-  def hardlink(self, digest, dest, file_mode):
-    """Since data is kept in memory, there is no filenode to hardlink."""
-    data = self.read(digest)
-    file_write(dest, [data])
-    if file_mode is not None:
-      fs.chmod(dest, file_mode & self._file_mode_mask)
-    with self._lock:
-      self._linked.append(len(data))
-
 
 class CachePolicies(object):
   def __init__(self, max_cache_size, min_free_space, max_items):
     """
     Arguments:
     - 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.
     - max_items: Maximum number of items to keep in the cache. If 0, do not
@@ skipping 111 matching lines @@
       self._protected.add(digest)
     return True
 
   def evict(self, digest):
     with self._lock:
       # Do not check for 'digest in self._protected' since it could be because
       # the object is corrupted.
       self._lru.pop(digest)
       self._delete_file(digest, UNKNOWN_FILE_SIZE)
 
-  def read(self, digest):
+  def getfileobj(self, digest):
     try:
-      with fs.open(self._path(digest), 'rb') as f:
-        return f.read()
+      f = fs.open(self._path(digest), 'rb')
+      with self._lock:
+        self._used.append(self._lru[digest])
+      return f
     except IOError:
       raise CacheMiss(digest)
 
   def write(self, digest, content):
     assert content is not None
     with self._lock:
       self._protected.add(digest)
     path = self._path(digest)
     # A stale broken file may remain. It is possible for the file to have write
     # access bit removed which would cause the file_write() call to fail to open
@@ skipping 10 matching lines @@
       file_path.try_remove(path)
       raise
     # Make the file read-only in the cache.  This has a few side-effects since
     # the file node is modified, so every directory entries to this file becomes
     # read-only. It's fine here because it is a new file.
     file_path.set_read_only(path, True)
     with self._lock:
       self._add(digest, size)
     return digest
 
-  def hardlink(self, digest, dest, file_mode):
-    """Hardlinks the file to |dest|.
-
-    Note that the file permission bits are on the file node, not the directory
-    entry, so changing the access bit on any of the directory entries for the
-    file node will affect them all.
-    """
-    path = self._path(digest)
-    if not file_path.link_file(dest, path, file_path.HARDLINK_WITH_FALLBACK):
-      # Report to the server that it failed with more details. We'll want to
-      # squash them all.
-      on_error.report('Failed to hardlink\n%s -> %s' % (path, dest))
-
-    if file_mode is not None:
-      # Ignores all other bits.
-      fs.chmod(dest, file_mode & 0500)
-    with self._lock:
-      self._linked.append(self._lru[digest])
-
   def _load(self):
     """Loads state of the cache from json file."""
     self._lock.assert_locked()
 
     if not os.path.isdir(self.cache_dir):
       fs.makedirs(self.cache_dir)
     else:
       # Make sure the cache is read-only.
       # TODO(maruel): Calculate the cost and optimize the performance
       # accordingly.
@@ skipping 210 matching lines @@
       pending[h] = isolated_file
       fetch_queue.add(h, priority=threading_utils.PRIORITY_HIGH)
 
     # Start fetching root *.isolated file (single file, not the whole bundle).
     retrieve_async(self.root)
 
     while pending:
       # Wait until some *.isolated file is fetched, parse it.
       item_hash = fetch_queue.wait(pending)
       item = pending.pop(item_hash)
-      item.load(fetch_queue.cache.read(item_hash))
+      item.load(fetch_queue.cache.getfileobj(item_hash).read())
 
       # Start fetching included *.isolated files.
       for new_child in item.children:
         retrieve_async(new_child)
 
       # Always fetch *.isolated files in traversal order, waiting if necessary
       # until next to-be-processed node loads. "Waiting" is done by yielding
       # back to the outer loop, that waits until some *.isolated is loaded.
       for node in isolated_format.walk_includes(self.root):
         if node not in processed:
@@ skipping 17 matching lines @@
     """Starts fetching files from |isolated| that are not yet being fetched.
 
     Modifies self.files.
     """
     logging.debug('fetch_files(%s)', isolated.obj_hash)
     for filepath, properties in isolated.data.get('files', {}).iteritems():
       # Root isolated has priority on the files being mapped. In particular,
       # overridden files must not be fetched.
       if filepath not in self.files:
         self.files[filepath] = properties
+
+        # Make sure if the isolated is read only, the mode doesn't have write
+        # bits.
+        if 'm' in properties and self.read_only:
+          properties['m'] &= ~(stat.S_IWUSR | stat.S_IWGRP | stat.S_IWOTH)
+
+        # Preemptively request hashed files.
         if 'h' in properties:
-          # Preemptively request files.
           logging.debug('fetching %s', filepath)
           fetch_queue.add(
               properties['h'], properties['s'], threading_utils.PRIORITY_MED)
 
   def _update_self(self, node):
     """Extracts bundle global parameters from loaded *.isolated file.
 
     Will be called with each loaded *.isolated file in order of traversal of
     isolated include graph (see isolated_format.walk_includes).
     """
@@ skipping 141 matching lines @@
       logging.info('Retrieving remaining files (%d of them)...',
           fetch_queue.pending_count)
       last_update = time.time()
       with threading_utils.DeadlockDetector(DEADLOCK_TIMEOUT) as detector:
         while remaining:
           detector.ping()
 
           # Wait for any item to finish fetching to cache.
           digest = fetch_queue.wait(remaining)
 
-          # Link corresponding files to a fetched item in cache.
+          # Create the files in the destination using item in cache as the
+          # source.
           for filepath, props in remaining.pop(digest):
-            dest = os.path.join(outdir, filepath)
-            if os.path.exists(dest):
-              raise AlreadyExists('File %s already exists' % dest)
-            cache.hardlink(digest, dest, props.get('m'))
+            fullpath = os.path.join(outdir, filepath)
+            srcfileobj = cache.getfileobj(digest)
+
+            isarchive = props.get('a')
+            if isarchive:
+              basedir = os.path.dirname(fullpath)
+
+              extractor = None
+              if isarchive == 'ar':

[nodir, 2016/06/20 15:54:35]

perhaps inverse the condition and remove `extractor = None`

[mithro, 2016/06/21 06:42:33]

extractor = None is actually not needed. Removed.

[nodir, 2016/06/21 22:20:04]

the goal was to reduce nesting. With condition inversed the nesting is less

[mithro, 2016/06/22 12:03:10]

I refactored this code a little.

+                extractor = arfile.ArFileReader(srcfileobj, fullparse=False)
+              else:
+                raise isolated_format.IsolatedError(
+                    'Unknown archive format %r', isarchive)
+
+              for ai, ifd in extractor:
+                fp = os.path.normpath(os.path.join(basedir, ai.name))
+                file_path.ensure_tree(os.path.dirname(fp))
+                putfile(ifd, fp, ai.mode, ai.size)
+
+            else:
+              file_mode = props.get('m')
+              if file_mode:
+                # Ignore all bits apart from the user
+                file_mode &= 0700
+              putfile(srcfileobj, fullpath, file_mode)
 
           # Report progress.
           duration = time.time() - last_update
           if duration > DELAY_BETWEEN_UPDATES_IN_SECS:
             msg = '%d files remaining...' % len(remaining)
             print msg
             logging.info(msg)
             last_update = time.time()
 
   # Cache could evict some items we just tried to fetch, it's a fatal error.
@@ skipping 316 matching lines @@
   dispatcher = subcommand.CommandDispatcher(__name__)
   return dispatcher.execute(OptionParserIsolateServer(), args)
 
 
 if __name__ == '__main__':
   subprocess42.inhibit_os_error_reporting()
   fix_encoding.fix_encoding()
   tools.disable_buffering()
   colorama.init()
   sys.exit(main(sys.argv[1:]))