isolateserver.py

#!/usr/bin/env python
# Copyright 2013 The Swarming 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.5'

import base64
import functools
import logging
import optparse
import os
import re
import signal
import sys
import tempfile
import threading
import time
import types
import urllib
import urlparse
import zlib

from third_party import colorama
from third_party.depot_tools import fix_encoding
from third_party.depot_tools import subcommand

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 threading_utils
from utils import tools

import auth
import isolated_format


# Version of isolate protocol passed to the server in /handshake request.
ISOLATE_PROTOCOL_VERSION = '1.0'


# The file size to be used when we don't know the correct file size,
# generally used for .isolated files.
UNKNOWN_FILE_SIZE = None


# Maximum expected delay (in seconds) between successive file fetches or uploads
# in Storage. If it takes longer than that, a deadlock might be happening
# and all stack frames for all threads are dumped to log.
DEADLOCK_TIMEOUT = 5 * 60


# The number of files to check the isolate server per /pre-upload query.
# All files are sorted by likelihood of a change in the file content
# (currently file size is used to estimate this: larger the file -> larger the
# possibility it has changed). Then first ITEMS_PER_CONTAINS_QUERIES[0] files
# are taken and send to '/pre-upload', then next ITEMS_PER_CONTAINS_QUERIES[1],
# and so on. Numbers here is a trade-off; the more per request, the lower the
# effect of HTTP round trip latency and TCP-level chattiness. On the other hand,
# larger values cause longer lookups, increasing the initial latency to start
# uploading, which is especially an issue for large files. This value is
# optimized for the "few thousands files to look up with minimal number of large
# files missing" case.
ITEMS_PER_CONTAINS_QUERIES = (20, 20, 50, 50, 50, 100)


# A list of already compressed extension types that should not receive any
# compression before being uploaded.
ALREADY_COMPRESSED_TYPES = [
    '7z', 'avi', 'cur', 'gif', 'h264', 'jar', 'jpeg', 'jpg', 'mp4', 'pdf',
    'png', 'wav', 'zip',
]


# Chunk size to use when reading from network stream.
NET_IO_FILE_CHUNK = 16 * 1024


# Read timeout in seconds for downloads from isolate storage. If there's no
# response from the server within this timeout whole download will be aborted.
DOWNLOAD_READ_TIMEOUT = 60


# The delay (in seconds) to wait between logging statements when retrieving
# the required files. This is intended to let the user (or buildbot) know that
# the program is still running.
DELAY_BETWEEN_UPDATES_IN_SECS = 30


DEFAULT_BLACKLIST = (
  # Temporary vim or python files.
  r'^.+\.(?:pyc|swp)$',
  # .git or .svn directory.
  r'^(?:.+' + re.escape(os.path.sep) + r'|)\.(?:git|svn)$',
)


# A class to use to communicate with the server by default. Can be changed by
# 'set_storage_api_class'. Default is IsolateServer.
_storage_api_cls = None


class Error(Exception):
  """Generic runtime error."""
  pass


class Aborted(Error):
  """Operation aborted."""
  pass


def stream_read(stream, chunk_size):
  """Reads chunks from |stream| and yields them."""
  while True:
    data = stream.read(chunk_size)
    if not data:
      break
    yield data


def file_read(path, chunk_size=isolated_format.DISK_FILE_CHUNK, offset=0):
  """Yields file content in chunks of |chunk_size| starting from |offset|."""
  with fs.open(path, 'rb') as f:
    if offset:
      f.seek(offset)
    while True:
      data = f.read(chunk_size)
      if not data:
        break
      yield data


def file_write(path, content_generator):
  """Writes file content as generated by content_generator.

  Creates the intermediary directory as needed.

  Returns the number of bytes written.

  Meant to be mocked out in unit tests.
  """
  filedir = os.path.dirname(path)
  if not fs.isdir(filedir):
    fs.makedirs(filedir)
  total = 0
  with fs.open(path, 'wb') as f:
    for d in content_generator:
      total += len(d)
      f.write(d)
  return total


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


def zip_decompress(
    content_generator, chunk_size=isolated_format.DISK_FILE_CHUNK):
  """Reads zipped data from |content_generator| and yields decompressed data.

  Decompresses data in small chunks (no larger than |chunk_size|) so that
  zip bomb file doesn't cause zlib to preallocate huge amount of memory.

  Raises IOError if data is corrupted or incomplete.
  """
  decompressor = zlib.decompressobj()
  compressed_size = 0
  try:
    for chunk in content_generator:
      compressed_size += len(chunk)
      data = decompressor.decompress(chunk, chunk_size)
      if data:
        yield data
      while decompressor.unconsumed_tail:
        data = decompressor.decompress(decompressor.unconsumed_tail, chunk_size)
        if data:
          yield data
    tail = decompressor.flush()
    if tail:
      yield tail
  except zlib.error as e:
    raise IOError(
        'Corrupted zip stream (read %d bytes) - %s' % (compressed_size, e))
  # Ensure all data was read and decompressed.
  if decompressor.unused_data or decompressor.unconsumed_tail:
    raise IOError('Not all data was decompressed')


def get_zip_compression_level(filename):
  """Given a filename calculates the ideal zip compression level to use."""
  file_ext = os.path.splitext(filename)[1].lower()
  # TODO(csharp): Profile to find what compression level works best.
  return 0 if file_ext in ALREADY_COMPRESSED_TYPES else 7


def create_directories(base_directory, files):
  """Creates the directory structure needed by the given list of files."""
  logging.debug('create_directories(%s, %d)', base_directory, len(files))
  # Creates the tree of directories to create.
  directories = set(os.path.dirname(f) for f in files)
  for item in list(directories):
    while item:
      directories.add(item)
      item = os.path.dirname(item)
  for d in sorted(directories):
    if d:
      fs.mkdir(os.path.join(base_directory, d))


def create_symlinks(base_directory, files):
  """Creates any symlinks needed by the given set of files."""
  for filepath, properties in files:
    if 'l' not in properties:
      continue
    if sys.platform == 'win32':
      # TODO(maruel): Create symlink via the win32 api.
      logging.warning('Ignoring symlink %s', filepath)
      continue
    outfile = os.path.join(base_directory, filepath)
    # os.symlink() doesn't exist on Windows.
    os.symlink(properties['l'], outfile)  # pylint: disable=E1101


def is_valid_file(path, size):
  """Determines if the given files appears valid.

  Currently it just checks the file's size.
  """
  if size == UNKNOWN_FILE_SIZE:
    return fs.isfile(path)
  actual_size = fs.stat(path).st_size
  if size != actual_size:
    logging.warning(
        'Found invalid item %s; %d != %d',
        os.path.basename(path), actual_size, size)
    return False
  return True


class Item(object):
  """An item to push to Storage.

  Its digest and size may be provided in advance, if known. Otherwise they will
  be derived from content(). If digest is provided, it MUST correspond to
  hash algorithm used by Storage.

  When used with Storage, Item starts its life in a main thread, travels
  to 'contains' thread, then to 'push' thread and then finally back to
  the main thread. It is never used concurrently from multiple threads.
  """

  def __init__(self, digest=None, size=None, high_priority=False):
    self.digest = digest
    self.size = size
    self.high_priority = high_priority
    self.compression_level = 6

  def content(self):
    """Iterable with content of this item as byte string (str) chunks."""
    raise NotImplementedError()

  def prepare(self, hash_algo):
    """Ensures self.digest and self.size are set.

    Uses content() as a source of data to calculate them. Does nothing if digest
    and size is already known.

    Arguments:
      hash_algo: hash algorithm to use to calculate digest.
    """
    if self.digest is None or self.size is None:
      digest = hash_algo()
      total = 0
      for chunk in self.content():
        digest.update(chunk)
        total += len(chunk)
      self.digest = digest.hexdigest()
      self.size = total


class FileItem(Item):
  """A file to push to Storage.

  Its digest and size may be provided in advance, if known. Otherwise they will
  be derived from the file content.
  """

  def __init__(self, path, digest=None, size=None, high_priority=False):
    super(FileItem, self).__init__(
        digest,
        size if size is not None else fs.stat(path).st_size,
        high_priority)
    self.path = path
    self.compression_level = get_zip_compression_level(path)

  def content(self):
    return file_read(self.path)


class BufferItem(Item):
  """A byte buffer to push to Storage."""

  def __init__(self, buf, high_priority=False):
    super(BufferItem, self).__init__(None, len(buf), high_priority)
    self.buffer = buf

  def content(self):
    return [self.buffer]


class Storage(object):
  """Efficiently downloads or uploads large set of files via StorageApi.

  Implements compression support, parallel 'contains' checks, parallel uploads
  and more.

  Works only within single namespace (and thus hashing algorithm and compression
  scheme are fixed).

  Spawns multiple internal threads. Thread safe, but not fork safe. Modifies
  signal handlers table to handle Ctrl+C.
  """

  def __init__(self, storage_api):
    self._storage_api = storage_api
    self._use_zip = isolated_format.is_namespace_with_compression(
        storage_api.namespace)
    self._hash_algo = isolated_format.get_hash_algo(storage_api.namespace)
    self._cpu_thread_pool = None
    self._net_thread_pool = None
    self._aborted = False
    self._prev_sig_handlers = {}

  @property
  def hash_algo(self):
    """Hashing algorithm used to name files in storage based on their content.

    Defined by |namespace|. See also isolated_format.get_hash_algo().
    """
    return self._hash_algo

  @property
  def location(self):
    """URL of the backing store that this class is using."""
    return self._storage_api.location

  @property
  def namespace(self):
    """Isolate namespace used by this storage.

    Indirectly defines hashing scheme and compression method used.
    """
    return self._storage_api.namespace

  @property
  def cpu_thread_pool(self):
    """ThreadPool for CPU-bound tasks like zipping."""
    if self._cpu_thread_pool is None:
      threads = max(threading_utils.num_processors(), 2)
      if sys.maxsize <= 2L**32:
        # On 32 bits userland, do not try to use more than 16 threads.
        threads = min(threads, 16)
      self._cpu_thread_pool = threading_utils.ThreadPool(2, threads, 0, 'zip')
    return self._cpu_thread_pool

  @property
  def net_thread_pool(self):
    """AutoRetryThreadPool for IO-bound tasks, retries IOError."""
    if self._net_thread_pool is None:
      self._net_thread_pool = threading_utils.IOAutoRetryThreadPool()
    return self._net_thread_pool

  def close(self):
    """Waits for all pending tasks to finish."""
    logging.info('Waiting for all threads to die...')
    if self._cpu_thread_pool:
      self._cpu_thread_pool.join()
      self._cpu_thread_pool.close()
      self._cpu_thread_pool = None
    if self._net_thread_pool:
      self._net_thread_pool.join()
      self._net_thread_pool.close()
      self._net_thread_pool = None
    logging.info('Done.')

  def abort(self):
    """Cancels any pending or future operations."""
    # This is not strictly theadsafe, but in the worst case the logging message
    # will be printed twice. Not a big deal. In other places it is assumed that
    # unprotected reads and writes to _aborted are serializable (it is true
    # for python) and thus no locking is used.
    if not self._aborted:
      logging.warning('Aborting... It can take a while.')
      self._aborted = True

  def __enter__(self):
    """Context manager interface."""
    assert not self._prev_sig_handlers, self._prev_sig_handlers
    for s in (signal.SIGINT, signal.SIGTERM):
      self._prev_sig_handlers[s] = signal.signal(s, lambda *_args: self.abort())
    return self

  def __exit__(self, _exc_type, _exc_value, _traceback):
    """Context manager interface."""
    self.close()
    while self._prev_sig_handlers:
      s, h = self._prev_sig_handlers.popitem()
      signal.signal(s, h)
    return False

  def upload_items(self, items):
    """Uploads a bunch of items to the isolate server.

    It figures out what items are missing from the server and uploads only them.

    Arguments:
      items: list of Item instances that represents data to upload.

    Returns:
      List of items that were uploaded. All other items are already there.
    """
    logging.info('upload_items(items=%d)', len(items))

    # Ensure all digests are calculated.
    for item in items:
      item.prepare(self._hash_algo)

    # For each digest keep only first Item that matches it. All other items
    # are just indistinguishable copies from the point of view of isolate
    # server (it doesn't care about paths at all, only content and digests).
    seen = {}
    duplicates = 0
    for item in items:
      if seen.setdefault(item.digest, item) is not item:
        duplicates += 1
    items = seen.values()
    if duplicates:
      logging.info('Skipped %d files with duplicated content', duplicates)

    # Enqueue all upload tasks.
    missing = set()
    uploaded = []
    channel = threading_utils.TaskChannel()
    for missing_item, push_state in self.get_missing_items(items):
      missing.add(missing_item)
      self.async_push(channel, missing_item, push_state)

    # No need to spawn deadlock detector thread if there's nothing to upload.
    if missing:
      with threading_utils.DeadlockDetector(DEADLOCK_TIMEOUT) as detector:
        # Wait for all started uploads to finish.
        while len(uploaded) != len(missing):
          detector.ping()
          item = channel.pull()
          uploaded.append(item)
          logging.debug(
              'Uploaded %d / %d: %s', len(uploaded), len(missing), item.digest)
    logging.info('All files are uploaded')

    # Print stats.
    total = len(items)
    total_size = sum(f.size for f in items)
    logging.info(
        'Total:      %6d, %9.1fkb',
        total,
        total_size / 1024.)
    cache_hit = set(items) - missing
    cache_hit_size = sum(f.size for f in cache_hit)
    logging.info(
        'cache hit:  %6d, %9.1fkb, %6.2f%% files, %6.2f%% size',
        len(cache_hit),
        cache_hit_size / 1024.,
        len(cache_hit) * 100. / total,
        cache_hit_size * 100. / total_size if total_size else 0)
    cache_miss = missing
    cache_miss_size = sum(f.size for f in cache_miss)
    logging.info(
        'cache miss: %6d, %9.1fkb, %6.2f%% files, %6.2f%% size',
        len(cache_miss),
        cache_miss_size / 1024.,
        len(cache_miss) * 100. / total,
        cache_miss_size * 100. / total_size if total_size else 0)

    return uploaded

  def async_push(self, channel, item, push_state):
    """Starts asynchronous push to the server in a parallel thread.

    Can be used only after |item| was checked for presence on a server with
    'get_missing_items' call. 'get_missing_items' returns |push_state| object
    that contains storage specific information describing how to upload
    the item (for example in case of cloud storage, it is signed upload URLs).

    Arguments:
      channel: TaskChannel that receives back |item| when upload ends.
      item: item to upload as instance of Item class.
      push_state: push state returned by 'get_missing_items' call for |item|.

    Returns:
      None, but |channel| later receives back |item| when upload ends.
    """
    # Thread pool task priority.
    priority = (
        threading_utils.PRIORITY_HIGH if item.high_priority
        else threading_utils.PRIORITY_MED)

    def push(content):
      """Pushes an Item and returns it to |channel|."""
      if self._aborted:
        raise Aborted()
      item.prepare(self._hash_algo)
      self._storage_api.push(item, push_state, content)
      return item

    # If zipping is not required, just start a push task.
    if not self._use_zip:
      self.net_thread_pool.add_task_with_channel(
          channel, priority, push, item.content())
      return

    # If zipping is enabled, zip in a separate thread.
    def zip_and_push():
      # TODO(vadimsh): Implement streaming uploads. Before it's done, assemble
      # content right here. It will block until all file is zipped.
      try:
        if self._aborted:
          raise Aborted()
        stream = zip_compress(item.content(), item.compression_level)
        data = ''.join(stream)
      except Exception as exc:
        logging.error('Failed to zip \'%s\': %s', item, exc)
        channel.send_exception()
        return
      self.net_thread_pool.add_task_with_channel(
          channel, priority, push, [data])
    self.cpu_thread_pool.add_task(priority, zip_and_push)

  def push(self, item, push_state):
    """Synchronously pushes a single item to the server.

    If you need to push many items at once, consider using 'upload_items' or
    'async_push' with instance of TaskChannel.

    Arguments:
      item: item to upload as instance of Item class.
      push_state: push state returned by 'get_missing_items' call for |item|.

    Returns:
      Pushed item (same object as |item|).
    """
    channel = threading_utils.TaskChannel()
    with threading_utils.DeadlockDetector(DEADLOCK_TIMEOUT):
      self.async_push(channel, item, push_state)
      pushed = channel.pull()
      assert pushed is item
    return item

  def async_fetch(self, channel, priority, digest, size, sink):
    """Starts asynchronous fetch from the server in a parallel thread.

    Arguments:
      channel: TaskChannel that receives back |digest| when download ends.
      priority: thread pool task priority for the fetch.
      digest: hex digest of an item to download.
      size: expected size of the item (after decompression).
      sink: function that will be called as sink(generator).
    """
    def fetch():
      try:
        # Prepare reading pipeline.
        stream = self._storage_api.fetch(digest)
        if self._use_zip:
          stream = zip_decompress(stream, isolated_format.DISK_FILE_CHUNK)
        # Run |stream| through verifier that will assert its size.
        verifier = FetchStreamVerifier(stream, size)
        # Verified stream goes to |sink|.
        sink(verifier.run())
      except Exception as err:
        logging.error('Failed to fetch %s: %s', digest, err)
        raise
      return digest

    # Don't bother with zip_thread_pool for decompression. Decompression is
    # really fast and most probably IO bound anyway.
    self.net_thread_pool.add_task_with_channel(channel, priority, fetch)

  def get_missing_items(self, items):
    """Yields items that are missing from the server.

    Issues multiple parallel queries via StorageApi's 'contains' method.

    Arguments:
      items: a list of Item objects to check.

    Yields:
      For each missing item it yields a pair (item, push_state), where:
        * item - Item object that is missing (one of |items|).
        * push_state - opaque object that contains storage specific information
            describing how to upload the item (for example in case of cloud
            storage, it is signed upload URLs). It can later be passed to
            'async_push'.
    """
    channel = threading_utils.TaskChannel()
    pending = 0

    # Ensure all digests are calculated.
    for item in items:
      item.prepare(self._hash_algo)

    def contains(batch):
      if self._aborted:
        raise Aborted()
      return self._storage_api.contains(batch)

    # Enqueue all requests.
    for batch in batch_items_for_check(items):
      self.net_thread_pool.add_task_with_channel(
          channel, threading_utils.PRIORITY_HIGH, contains, batch)
      pending += 1

    # Yield results as they come in.
    for _ in xrange(pending):
      for missing_item, push_state in channel.pull().iteritems():
        yield missing_item, push_state


def batch_items_for_check(items):
  """Splits list of items to check for existence on the server into batches.

  Each batch corresponds to a single 'exists?' query to the server via a call
  to StorageApi's 'contains' method.

  Arguments:
    items: a list of Item objects.

  Yields:
    Batches of items to query for existence in a single operation,
    each batch is a list of Item objects.
  """
  batch_count = 0
  batch_size_limit = ITEMS_PER_CONTAINS_QUERIES[0]
  next_queries = []
  for item in sorted(items, key=lambda x: x.size, reverse=True):
    next_queries.append(item)
    if len(next_queries) == batch_size_limit:
      yield next_queries
      next_queries = []
      batch_count += 1
      batch_size_limit = ITEMS_PER_CONTAINS_QUERIES[
          min(batch_count, len(ITEMS_PER_CONTAINS_QUERIES) - 1)]
  if next_queries:
    yield next_queries


class FetchQueue(object):
  """Fetches items from Storage and places them into LocalCache.

  It manages multiple concurrent fetch operations. Acts as a bridge between
  Storage and LocalCache so that Storage and LocalCache don't depend on each
  other at all.
  """

  def __init__(self, storage, cache):
    self.storage = storage
    self.cache = cache
    self._channel = threading_utils.TaskChannel()
    self._pending = set()
    self._accessed = set()
    self._fetched = cache.cached_set()

  def add(
      self,
      digest,
      size=UNKNOWN_FILE_SIZE,
      priority=threading_utils.PRIORITY_MED):
    """Starts asynchronous fetch of item |digest|."""
    # Fetching it now?
    if digest in self._pending:
      return

    # Mark this file as in use, verify_all_cached will later ensure it is still
    # in cache.
    self._accessed.add(digest)

    # Already fetched? Notify cache to update item's LRU position.
    if digest in self._fetched:
      # 'touch' returns True if item is in cache and not corrupted.
      if self.cache.touch(digest, size):
        return
      # Item is corrupted, remove it from cache and fetch it again.
      self._fetched.remove(digest)
      self.cache.evict(digest)

    # TODO(maruel): It should look at the free disk space, the current cache
    # size and the size of the new item on every new item:
    # - Trim the cache as more entries are listed when free disk space is low,
    #   otherwise if the amount of data downloaded during the run > free disk
    #   space, it'll crash.
    # - Make sure there's enough free disk space to fit all dependencies of
    #   this run! If not, abort early.

    # Start fetching.
    self._pending.add(digest)
    self.storage.async_fetch(
        self._channel, priority, digest, size,
        functools.partial(self.cache.write, digest))

  def wait(self, digests):
    """Starts a loop that waits for at least one of |digests| to be retrieved.

    Returns the first digest retrieved.
    """
    # Flush any already fetched items.
    for digest in digests:
      if digest in self._fetched:
        return digest

    # Ensure all requested items are being fetched now.
    assert all(digest in self._pending for digest in digests), (
        digests, self._pending)

    # Wait for some requested item to finish fetching.
    while self._pending:
      digest = self._channel.pull()
      self._pending.remove(digest)
      self._fetched.add(digest)
      if digest in digests:
        return digest

    # Should never reach this point due to assert above.
    raise RuntimeError('Impossible state')

  def inject_local_file(self, path, algo):
    """Adds local file to the cache as if it was fetched from storage."""
    with fs.open(path, 'rb') as f:
      data = f.read()
    digest = algo(data).hexdigest()
    self.cache.write(digest, [data])
    self._fetched.add(digest)
    return digest

  @property
  def pending_count(self):
    """Returns number of items to be fetched."""
    return len(self._pending)

  def verify_all_cached(self):
    """True if all accessed items are in cache."""
    return self._accessed.issubset(self.cache.cached_set())


class FetchStreamVerifier(object):
  """Verifies that fetched file is valid before passing it to the LocalCache."""

  def __init__(self, stream, expected_size):
    assert stream is not None
    self.stream = stream
    self.expected_size = expected_size
    self.current_size = 0

  def run(self):
    """Generator that yields same items as |stream|.

    Verifies |stream| is complete before yielding a last chunk to consumer.

    Also wraps IOError produced by consumer into MappingError exceptions since
    otherwise Storage will retry fetch on unrelated local cache errors.
    """
    # Read one chunk ahead, keep it in |stored|.
    # That way a complete stream can be verified before pushing last chunk
    # to consumer.
    stored = None
    for chunk in self.stream:
      assert chunk is not None
      if stored is not None:
        self._inspect_chunk(stored, is_last=False)
        try:
          yield stored
        except IOError as exc:
          raise isolated_format.MappingError(
              'Failed to store an item in cache: %s' % exc)
      stored = chunk
    if stored is not None:
      self._inspect_chunk(stored, is_last=True)
      try:
        yield stored
      except IOError as exc:
        raise isolated_format.MappingError(
            'Failed to store an item in cache: %s' % exc)

  def _inspect_chunk(self, chunk, is_last):
    """Called for each fetched chunk before passing it to consumer."""
    self.current_size += len(chunk)
    if (is_last and
        (self.expected_size != UNKNOWN_FILE_SIZE) and
        (self.expected_size != self.current_size)):
      raise IOError('Incorrect file size: expected %d, got %d' % (
          self.expected_size, self.current_size))


class StorageApi(object):
  """Interface for classes that implement low-level storage operations.

  StorageApi is oblivious of compression and hashing scheme used. This details
  are handled in higher level Storage class.

  Clients should generally not use StorageApi directly. Storage class is
  preferred since it implements compression and upload optimizations.
  """

  @property
  def location(self):
    """URL of the backing store that this class is using."""
    raise NotImplementedError()

  @property
  def namespace(self):
    """Isolate namespace used by this storage.

    Indirectly defines hashing scheme and compression method used.
    """
    raise NotImplementedError()

  def fetch(self, digest, offset=0):
    """Fetches an object and yields its content.

    Arguments:
      digest: hash digest of item to download.
      offset: offset (in bytes) from the start of the file to resume fetch from.

    Yields:
      Chunks of downloaded item (as str objects).
    """
    raise NotImplementedError()

  def push(self, item, push_state, content=None):
    """Uploads an |item| with content generated by |content| generator.

    |item| MUST go through 'contains' call to get |push_state| before it can
    be pushed to the storage.

    To be clear, here is one possible usage:
      all_items = [... all items to push as Item subclasses ...]
      for missing_item, push_state in storage_api.contains(all_items).items():
        storage_api.push(missing_item, push_state)

    When pushing to a namespace with compression, data that should be pushed
    and data provided by the item is not the same. In that case |content| is
    not None and it yields chunks of compressed data (using item.content() as
    a source of original uncompressed data). This is implemented by Storage
    class.

    Arguments:
      item: Item object that holds information about an item being pushed.
      push_state: push state object as returned by 'contains' call.
      content: a generator that yields chunks to push, item.content() if None.

    Returns:
      None.
    """
    raise NotImplementedError()

  def contains(self, items):
    """Checks for |items| on the server, prepares missing ones for upload.

    Arguments:
      items: list of Item objects to check for presence.

    Returns:
      A dict missing Item -> opaque push state object to be passed to 'push'.
      See doc string for 'push'.
    """
    raise NotImplementedError()


class _IsolateServerPushState(object):
  """Per-item state passed from IsolateServer.contains to IsolateServer.push.

  Note this needs to be a global class to support pickling.
  """

  def __init__(self, preupload_status, size):
    self.preupload_status = preupload_status
    gs_upload_url = preupload_status.get('gs_upload_url') or None
    if gs_upload_url:
      self.upload_url = gs_upload_url
      self.finalize_url = '_ah/api/isolateservice/v1/finalize_gs_upload'
    else:
      self.upload_url = '_ah/api/isolateservice/v1/store_inline'
      self.finalize_url = None
    self.uploaded = False
    self.finalized = False
    self.size = size


class IsolateServer(StorageApi):
  """StorageApi implementation that downloads and uploads to Isolate Server.

  It uploads and downloads directly from Google Storage whenever appropriate.
  Works only within single namespace.
  """

  def __init__(self, base_url, namespace):
    super(IsolateServer, self).__init__()
    assert file_path.is_url(base_url), base_url
    self._base_url = base_url.rstrip('/')
    self._namespace = namespace
    self._namespace_dict = {
        'compression': 'flate' if namespace.endswith(
            ('-gzip', '-flate')) else '',
        'digest_hash': 'sha-1',
        'namespace': namespace,
    }
    self._lock = threading.Lock()
    self._server_caps = None
    self._memory_use = 0

  @property
  def _server_capabilities(self):
    """Gets server details.

    Returns:
      Server capabilities dictionary as returned by /server_details endpoint.
    """
    # TODO(maruel): Make this request much earlier asynchronously while the
    # files are being enumerated.

    # TODO(vadimsh): Put |namespace| in the URL so that server can apply
    # namespace-level ACLs to this call.

    with self._lock:
      if self._server_caps is None:
        self._server_caps = net.url_read_json(
            url='%s/_ah/api/isolateservice/v1/server_details' % self._base_url,
            data={})
      return self._server_caps

  @property
  def location(self):
    return self._base_url

  @property
  def namespace(self):
    return self._namespace

  def fetch(self, digest, offset=0):
    assert offset >= 0
    source_url = '%s/_ah/api/isolateservice/v1/retrieve' % (
        self._base_url)
    logging.debug('download_file(%s, %d)', source_url, offset)
    response = self.do_fetch(source_url, digest, offset)

    if not response:
      raise IOError(
          'Attempted to fetch from %s; no data exist: %s / %s.' % (
            source_url, self._namespace, digest))

    # for DB uploads
    content = response.get('content')
    if content is not None:
      return base64.b64decode(content)

    # for GS entities
    connection = net.url_open(response['url'])
    if not connection:
      raise IOError('Failed to download %s / %s' % (self._namespace, digest))

    # If |offset|, verify server respects it by checking Content-Range.
    if offset:
      content_range = connection.get_header('Content-Range')
      if not content_range:
        raise IOError('Missing Content-Range header')

      # 'Content-Range' format is 'bytes <offset>-<last_byte_index>/<size>'.
      # According to a spec, <size> can be '*' meaning "Total size of the file
      # is not known in advance".
      try:
        match = re.match(r'bytes (\d+)-(\d+)/(\d+|\*)', content_range)
        if not match:
          raise ValueError()
        content_offset = int(match.group(1))
        last_byte_index = int(match.group(2))
        size = None if match.group(3) == '*' else int(match.group(3))
      except ValueError:
        raise IOError('Invalid Content-Range header: %s' % content_range)

      # Ensure returned offset equals requested one.
      if offset != content_offset:
        raise IOError('Expecting offset %d, got %d (Content-Range is %s)' % (
            offset, content_offset, content_range))

      # Ensure entire tail of the file is returned.
      if size is not None and last_byte_index + 1 != size:
        raise IOError('Incomplete response. Content-Range: %s' % content_range)

    return stream_read(connection, NET_IO_FILE_CHUNK)

  def push(self, item, push_state, content=None):
    assert isinstance(item, Item)
    assert item.digest is not None
    assert item.size is not None
    assert isinstance(push_state, _IsolateServerPushState)
    assert not push_state.finalized

    # Default to item.content().
    content = item.content() if content is None else content
    logging.info('Push state size: %d', push_state.size)
    if isinstance(content, (basestring, list)):
      # Memory is already used, too late.
      with self._lock:
        self._memory_use += push_state.size
    else:
      # TODO(vadimsh): Do not read from |content| generator when retrying push.
      # If |content| is indeed a generator, it can not be re-winded back to the
      # beginning of the stream. A retry will find it exhausted. A possible
      # solution is to wrap |content| generator with some sort of caching
      # restartable generator. It should be done alongside streaming support
      # implementation.
      #
      # In theory, we should keep the generator, so that it is not serialized in
      # memory. Sadly net.HttpService.request() requires the body to be
      # serialized.
      assert isinstance(content, types.GeneratorType), repr(content)
      slept = False
      # HACK HACK HACK. Please forgive me for my sins but OMG, it works!
      # One byte less than 512mb. This is to cope with incompressible content.
      max_size = int(sys.maxsize * 0.25)
      while True:
        with self._lock:
          # This is due to 32 bits python when uploading very large files. The
          # problem is that it's comparing uncompressed sizes, while we care
          # about compressed sizes since it's what is serialized in memory.
          # The first check assumes large files are compressible and that by
          # throttling one upload at once, we can survive. Otherwise, kaboom.
          memory_use = self._memory_use
          if ((push_state.size >= max_size and not memory_use) or
              (memory_use + push_state.size <= max_size)):
            self._memory_use += push_state.size
            memory_use = self._memory_use
            break
        time.sleep(0.1)
        slept = True
      if slept:
        logging.info('Unblocked: %d %d', memory_use, push_state.size)

    try:
      # This push operation may be a retry after failed finalization call below,
      # no need to reupload contents in that case.
      if not push_state.uploaded:
        # PUT file to |upload_url|.
        success = self.do_push(push_state, content)
        if not success:
          raise IOError('Failed to upload file with hash %s to URL %s' % (
              item.digest, push_state.upload_url))
        push_state.uploaded = True
      else:
        logging.info(
            'A file %s already uploaded, retrying finalization only',
            item.digest)

      # Optionally notify the server that it's done.
      if push_state.finalize_url:
        # TODO(vadimsh): Calculate MD5 or CRC32C sum while uploading a file and
        # send it to isolated server. That way isolate server can verify that
        # the data safely reached Google Storage (GS provides MD5 and CRC32C of
        # stored files).
        # TODO(maruel): Fix the server to accept properly data={} so
        # url_read_json() can be used.
        response = net.url_read_json(
            url='%s/%s' % (self._base_url, push_state.finalize_url),
            data={
                'upload_ticket': push_state.preupload_status['upload_ticket'],
            })
        if not response or not response['ok']:
          raise IOError('Failed to finalize file with hash %s.' % item.digest)
      push_state.finalized = True
    finally:
      with self._lock:
        self._memory_use -= push_state.size

  def contains(self, items):
    # Ensure all items were initialized with 'prepare' call. Storage does that.
    assert all(i.digest is not None and i.size is not None for i in items)

    # Request body is a json encoded list of dicts.
    body = {
        'items': [
          {
            'digest': item.digest,
            'is_isolated': bool(item.high_priority),
            'size': item.size,
          } for item in items
        ],
        'namespace': self._namespace_dict,
    }

    query_url = '%s/_ah/api/isolateservice/v1/preupload' % self._base_url

    # Response body is a list of push_urls (or null if file is already present).
    response = None
    try:
      response = net.url_read_json(url=query_url, data=body)
      if response is None:
        raise isolated_format.MappingError(
            'Failed to execute preupload query')
    except ValueError as err:
      raise isolated_format.MappingError(
          'Invalid response from server: %s, body is %s' % (err, response))

    # Pick Items that are missing, attach _PushState to them.
    missing_items = {}
    for preupload_status in response.get('items', []):
      assert 'upload_ticket' in preupload_status, (
          preupload_status, '/preupload did not generate an upload ticket')
      index = int(preupload_status['index'])
      missing_items[items[index]] = _IsolateServerPushState(
          preupload_status, items[index].size)
    logging.info('Queried %d files, %d cache hit',
        len(items), len(items) - len(missing_items))
    return missing_items

  def do_fetch(self, url, digest, offset):
    """Fetches isolated data from the URL.

    Used only for fetching files, not for API calls. Can be overridden in
    subclasses.

    Args:
      url: URL to fetch the data from, can possibly return http redirect.
      offset: byte offset inside the file to start fetching from.

    Returns:
      net.HttpResponse compatible object, with 'read' and 'get_header' calls.
    """
    assert isinstance(offset, int)
    data = {
        'digest': digest.encode('utf-8'),
        'namespace': self._namespace_dict,
        'offset': offset,
    }
    # TODO(maruel): url + '?' + urllib.urlencode(data) once a HTTP GET endpoint
    # is added.
    return net.url_read_json(
        url=url,
        data=data,
        read_timeout=DOWNLOAD_READ_TIMEOUT)

  def do_push(self, push_state, content):
    """Uploads isolated file to the URL.

    Used only for storing files, not for API calls. Can be overridden in
    subclasses.

    Args:
      url: URL to upload the data to.
      push_state: an _IsolateServicePushState instance
      item: the original Item to be uploaded
      content: an iterable that yields 'str' chunks.
    """
    # A cheezy way to avoid memcpy of (possibly huge) file, until streaming
    # upload support is implemented.
    if isinstance(content, list) and len(content) == 1:
      content = content[0]
    else:
      content = ''.join(content)

    # DB upload
    if not push_state.finalize_url:
      url = '%s/%s' % (self._base_url, push_state.upload_url)
      content = base64.b64encode(content)
      data = {
          'upload_ticket': push_state.preupload_status['upload_ticket'],
          'content': content,
      }
      response = net.url_read_json(url=url, data=data)
      return response is not None and response['ok']

    # upload to GS
    url = push_state.upload_url
    response = net.url_read(
        content_type='application/octet-stream',
        data=content,
        method='PUT',
        url=url)
    return response is not None


class LocalCache(object):
  """Local cache that stores objects fetched via Storage.

  It can be accessed concurrently from multiple threads, so it should protect
  its internal state with some lock.
  """
  cache_dir = None

  def __enter__(self):
    """Context manager interface."""
    return self

  def __exit__(self, _exc_type, _exec_value, _traceback):
    """Context manager interface."""
    return False

  def cached_set(self):
    """Returns a set of all cached digests (always a new object)."""
    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."""
    raise NotImplementedError()

  def write(self, digest, content):
    """Reads data from |content| generator and stores it in cache."""
    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__()
    self._file_mode_mask = file_mode_mask
    # Let's not assume dict is thread safe.
    self._lock = threading.Lock()
    self._contents = {}

  def cached_set(self):
    with self._lock:
      return set(self._contents)

  def touch(self, digest, size):
    with self._lock:
      return digest in self._contents

  def evict(self, digest):
    with self._lock:
      self._contents.pop(digest, None)

  def read(self, digest):
    with self._lock:
      return self._contents[digest]

  def write(self, digest, content):
    # Assemble whole stream before taking the lock.
    data = ''.join(content)
    with self._lock:
      self._contents[digest] = data

  def hardlink(self, digest, dest, file_mode):
    """Since data is kept in memory, there is no filenode to hardlink."""
    file_write(dest, [self.read(digest)])
    if file_mode is not None:
      fs.chmod(dest, file_mode & self._file_mode_mask)


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
                 enforce a limit.
    """
    self.max_cache_size = max_cache_size
    self.min_free_space = min_free_space
    self.max_items = max_items


class DiskCache(LocalCache):
  """Stateful LRU cache in a flat hash table in a directory.

  Saves its state as json file.
  """
  STATE_FILE = u'state.json'

  def __init__(self, cache_dir, policies, hash_algo):
    """
    Arguments:
      cache_dir: directory where to place the cache.
      policies: cache retention policies.
      algo: hashing algorithm used.
    """
    super(DiskCache, self).__init__()
    self.cache_dir = cache_dir
    self.policies = policies
    self.hash_algo = hash_algo
    self.state_file = os.path.join(cache_dir, self.STATE_FILE)

    # All protected methods (starting with '_') except _path should be called
    # with this lock locked.
    self._lock = threading_utils.LockWithAssert()
    self._lru = lru.LRUDict()

    # Profiling values.
    self._added = []
    self._removed = []
    self._free_disk = 0

    with tools.Profiler('Setup'):
      with self._lock:
        self._load()

  def __enter__(self):
    return self

  def __exit__(self, _exc_type, _exec_value, _traceback):
    with tools.Profiler('CleanupTrimming'):
      with self._lock:
        self._trim()

        logging.info(
            '%5d (%8dkb) added',
            len(self._added), sum(self._added) / 1024)
        logging.info(
            '%5d (%8dkb) current',
            len(self._lru),
            sum(self._lru.itervalues()) / 1024)
        logging.info(
            '%5d (%8dkb) removed',
            len(self._removed), sum(self._removed) / 1024)
        logging.info(
            '       %8dkb free',
            self._free_disk / 1024)
    return False

  def cached_set(self):
    with self._lock:
      return self._lru.keys_set()

  def touch(self, digest, size):
    """Verifies an actual file is valid.

    Note that is doesn't compute the hash so it could still be corrupted if the
    file size didn't change.

    TODO(maruel): More stringent verification while keeping the check fast.
    """
    # Do the check outside the lock.
    if not is_valid_file(self._path(digest), size):
      return False

    # Update it's LRU position.
    with self._lock:
      if digest not in self._lru:
        return False
      self._lru.touch(digest)
    return True

  def evict(self, digest):
    with self._lock:
      self._lru.pop(digest)
      self._delete_file(digest, UNKNOWN_FILE_SIZE)

  def read(self, digest):
    with fs.open(self._path(digest), 'rb') as f:
      return f.read()

  def write(self, digest, content):
    assert content is not None
    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
    # in write mode. Take no chance here.
    file_path.try_remove(path)
    try:
      size = file_write(path, content)
    except:
      # There are two possible places were an exception can occur:
      #   1) Inside |content| generator in case of network or unzipping errors.
      #   2) Inside file_write itself in case of disk IO errors.
      # In any case delete an incomplete file and propagate the exception to
      # caller, it will be logged there.
      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)

  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)

  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.
      file_path.make_tree_read_only(self.cache_dir)

    # Load state of the cache.
    if fs.isfile(self.state_file):
      try:
        self._lru = lru.LRUDict.load(self.state_file)
      except ValueError as err:
        logging.error('Failed to load cache state: %s' % (err,))
        # Don't want to keep broken state file.
        file_path.try_remove(self.state_file)

    # Ensure that all files listed in the state still exist and add new ones.
    previous = self._lru.keys_set()
    unknown = []
    for filename in fs.listdir(self.cache_dir):
      if filename == self.STATE_FILE:
        continue
      if filename in previous:
        previous.remove(filename)
        continue
      # An untracked file.
      if not isolated_format.is_valid_hash(filename, self.hash_algo):
        logging.warning('Removing unknown file %s from cache', filename)
        p = self._path(filename)
        if fs.isdir(p):
          try:
            file_path.rmtree(p)
          except OSError:
            pass
        else:
          file_path.try_remove(p)
        continue
      # File that's not referenced in 'state.json'.
      # TODO(vadimsh): Verify its SHA1 matches file name.
      logging.warning('Adding unknown file %s to cache', filename)
      unknown.append(filename)

    if unknown:
      # Add as oldest files. They will be deleted eventually if not accessed.
      self._add_oldest_list(unknown)
      logging.warning('Added back %d unknown files', len(unknown))

    if previous:
      # Filter out entries that were not found.
      logging.warning('Removed %d lost files', len(previous))
      for filename in previous:
        self._lru.pop(filename)
    self._trim()

  def _save(self):
    """Saves the LRU ordering."""
    self._lock.assert_locked()
    if sys.platform != 'win32':
      d = os.path.dirname(self.state_file)
      if fs.isdir(d):
        # Necessary otherwise the file can't be created.
        file_path.set_read_only(d, False)
    if fs.isfile(self.state_file):
      file_path.set_read_only(self.state_file, False)
    self._lru.save(self.state_file)

  def _trim(self):
    """Trims anything we don't know, make sure enough free space exists."""
    self._lock.assert_locked()

    # Ensure maximum cache size.
    if self.policies.max_cache_size:
      total_size = sum(self._lru.itervalues())
      while total_size > self.policies.max_cache_size:
        total_size -= self._remove_lru_file()

    # Ensure maximum number of items in the cache.
    if self.policies.max_items and len(self._lru) > self.policies.max_items:
      for _ in xrange(len(self._lru) - self.policies.max_items):
        self._remove_lru_file()

    # Ensure enough free space.
    self._free_disk = file_path.get_free_space(self.cache_dir)
    trimmed_due_to_space = False
    while (
        self.policies.min_free_space and
        self._lru and
        self._free_disk < self.policies.min_free_space):
      trimmed_due_to_space = True
      self._remove_lru_file()
      self._free_disk = file_path.get_free_space(self.cache_dir)
    if trimmed_due_to_space:
      total_usage = sum(self._lru.itervalues())
      usage_percent = 0.
      if total_usage:
        usage_percent = 100. * self.policies.max_cache_size / float(total_usage)
      logging.warning(
          'Trimmed due to not enough free disk space: %.1fkb free, %.1fkb '
          'cache (%.1f%% of its maximum capacity)',
          self._free_disk / 1024.,
          total_usage / 1024.,
          usage_percent)
    self._save()

  def _path(self, digest):
    """Returns the path to one item."""
    return os.path.join(self.cache_dir, digest)

  def _remove_lru_file(self):
    """Removes the last recently used file and returns its size."""
    self._lock.assert_locked()
    digest, size = self._lru.pop_oldest()
    self._delete_file(digest, size)
    return size

  def _add(self, digest, size=UNKNOWN_FILE_SIZE):
    """Adds an item into LRU cache marking it as a newest one."""
    self._lock.assert_locked()
    if size == UNKNOWN_FILE_SIZE:
      size = fs.stat(self._path(digest)).st_size
    self._added.append(size)
    self._lru.add(digest, size)

  def _add_oldest_list(self, digests):
    """Adds a bunch of items into LRU cache marking them as oldest ones."""
    self._lock.assert_locked()
    pairs = []
    for digest in digests:
      size = fs.stat(self._path(digest)).st_size
      self._added.append(size)
      pairs.append((digest, size))
    self._lru.batch_insert_oldest(pairs)

  def _delete_file(self, digest, size=UNKNOWN_FILE_SIZE):
    """Deletes cache file from the file system."""
    self._lock.assert_locked()
    try:
      if size == UNKNOWN_FILE_SIZE:
        size = fs.stat(self._path(digest)).st_size
      file_path.try_remove(self._path(digest))
      self._removed.append(size)
    except OSError as e:
      logging.error('Error attempting to delete a file %s:\n%s' % (digest, e))


class IsolatedBundle(object):
  """Fetched and parsed .isolated file with all dependencies."""

  def __init__(self):
    self.command = []
    self.files = {}
    self.read_only = None
    self.relative_cwd = None
    # The main .isolated file, a IsolatedFile instance.
    self.root = None

  def fetch(self, fetch_queue, root_isolated_hash, algo):
    """Fetches the .isolated and all the included .isolated.

    It enables support for "included" .isolated files. They are processed in
    strict order but fetched asynchronously from the cache. This is important so
    that a file in an included .isolated file that is overridden by an embedding
    .isolated file is not fetched needlessly. The includes are fetched in one
    pass and the files are fetched as soon as all the ones on the left-side
    of the tree were fetched.

    The prioritization is very important here for nested .isolated files.
    'includes' have the highest priority and the algorithm is optimized for both
    deep and wide trees. A deep one is a long link of .isolated files referenced
    one at a time by one item in 'includes'. A wide one has a large number of
    'includes' in a single .isolated file. 'left' is defined as an included
    .isolated file earlier in the 'includes' list. So the order of the elements
    in 'includes' is important.

    As a side effect this method starts asynchronous fetch of all data files
    by adding them to |fetch_queue|. It doesn't wait for data files to finish
    fetching though.
    """
    self.root = isolated_format.IsolatedFile(root_isolated_hash, algo)

    # Isolated files being retrieved now: hash -> IsolatedFile instance.
    pending = {}
    # Set of hashes of already retrieved items to refuse recursive includes.
    seen = set()
    # Set of IsolatedFile's whose data files have already being fetched.
    processed = set()

    def retrieve_async(isolated_file):
      h = isolated_file.obj_hash
      if h in seen:
        raise isolated_format.IsolatedError(
            'IsolatedFile %s is retrieved recursively' % h)
      assert h not in pending
      seen.add(h)
      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))

      # 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:
          # Not visited, and not yet loaded -> wait for it to load.
          if not node.is_loaded:
            break
          # Not visited and loaded -> process it and continue the traversal.
          self._start_fetching_files(node, fetch_queue)
          processed.add(node)

    # All *.isolated files should be processed by now and only them.
    all_isolateds = set(isolated_format.walk_includes(self.root))
    assert all_isolateds == processed, (all_isolateds, processed)

    # Extract 'command' and other bundle properties.
    for node in isolated_format.walk_includes(self.root):
      self._update_self(node)
    self.relative_cwd = self.relative_cwd or ''

  def _start_fetching_files(self, isolated, fetch_queue):
    """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
        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).
    """
    # Grabs properties.
    if not self.command and node.data.get('command'):
      # Ensure paths are correctly separated on windows.
      self.command = node.data['command']
      if self.command:
        self.command[0] = self.command[0].replace('/', os.path.sep)
        self.command = tools.fix_python_path(self.command)
    if self.read_only is None and node.data.get('read_only') is not None:
      self.read_only = node.data['read_only']
    if (self.relative_cwd is None and
        node.data.get('relative_cwd') is not None):
      self.relative_cwd = node.data['relative_cwd']


def set_storage_api_class(cls):
  """Replaces StorageApi implementation used by default."""
  global _storage_api_cls
  assert _storage_api_cls is None
  assert issubclass(cls, StorageApi)
  _storage_api_cls = cls


def get_storage_api(url, namespace):
  """Returns an object that implements low-level StorageApi interface.

  It is used by Storage to work with single isolate |namespace|. It should
  rarely be used directly by clients, see 'get_storage' for
  a better alternative.

  Arguments:
    url: URL of isolate service to use shared cloud based storage.
    namespace: isolate namespace to operate in, also defines hashing and
        compression scheme used, i.e. namespace names that end with '-gzip'
        store compressed data.

  Returns:
    Instance of StorageApi subclass.
  """
  cls = _storage_api_cls or IsolateServer
  return cls(url, namespace)


def get_storage(url, namespace):
  """Returns Storage class that can upload and download from |namespace|.

  Arguments:
    url: URL of isolate service to use shared cloud based storage.
    namespace: isolate namespace to operate in, also defines hashing and
        compression scheme used, i.e. namespace names that end with '-gzip'
        store compressed data.

  Returns:
    Instance of Storage.
  """
  return Storage(get_storage_api(url, namespace))


def upload_tree(base_url, infiles, namespace):
  """Uploads the given tree to the given url.

  Arguments:
    base_url:  The url of the isolate server to upload to.
    infiles:   iterable of pairs (absolute path, metadata dict) of files.
    namespace: The namespace to use on the server.
  """
  # Convert |infiles| into a list of FileItem objects, skip duplicates.
  # Filter out symlinks, since they are not represented by items on isolate
  # server side.
  items = []
  seen = set()
  skipped = 0
  for filepath, metadata in infiles:
    assert isinstance(filepath, unicode), filepath
    if 'l' not in metadata and filepath not in seen:
      seen.add(filepath)
      item = FileItem(
          path=filepath,
          digest=metadata['h'],
          size=metadata['s'],
          high_priority=metadata.get('priority') == '0')
      items.append(item)
    else:
      skipped += 1

  logging.info('Skipped %d duplicated entries', skipped)
  with get_storage(base_url, namespace) as storage:
    storage.upload_items(items)


def fetch_isolated(isolated_hash, storage, cache, outdir, require_command):
  """Aggressively downloads the .isolated file(s), then download all the files.

  Arguments:
    isolated_hash: hash of the root *.isolated file.
    storage: Storage class that communicates with isolate storage.
    cache: LocalCache class that knows how to store and map files locally.
    outdir: Output directory to map file tree to.
    require_command: Ensure *.isolated specifies a command to run.

  Returns:
    IsolatedBundle object that holds details about loaded *.isolated file.
  """
  logging.debug(
      'fetch_isolated(%s, %s, %s, %s, %s)',
      isolated_hash, storage, cache, outdir, require_command)
  # Hash algorithm to use, defined by namespace |storage| is using.
  algo = storage.hash_algo
  with cache:
    fetch_queue = FetchQueue(storage, cache)
    bundle = IsolatedBundle()

    with tools.Profiler('GetIsolateds'):
      # Optionally support local files by manually adding them to cache.
      if not isolated_format.is_valid_hash(isolated_hash, algo):
        logging.debug('%s is not a valid hash, assuming a file', isolated_hash)
        path = unicode(os.path.abspath(isolated_hash))
        try:
          isolated_hash = fetch_queue.inject_local_file(path, algo)
        except IOError:
          raise isolated_format.MappingError(
              '%s doesn\'t seem to be a valid file. Did you intent to pass a '
              'valid hash?' % isolated_hash)

      # Load all *.isolated and start loading rest of the files.
      bundle.fetch(fetch_queue, isolated_hash, algo)
      if require_command and not bundle.command:
        # TODO(vadimsh): All fetch operations are already enqueue and there's no
        # easy way to cancel them.
        raise isolated_format.IsolatedError('No command to run')

    with tools.Profiler('GetRest'):
      # Create file system hierarchy.
      if not fs.isdir(outdir):
        fs.makedirs(outdir)
      create_directories(outdir, bundle.files)
      create_symlinks(outdir, bundle.files.iteritems())

      # Ensure working directory exists.
      cwd = os.path.normpath(os.path.join(outdir, bundle.relative_cwd))
      if not fs.isdir(cwd):
        fs.makedirs(cwd)

      # Multimap: digest -> list of pairs (path, props).
      remaining = {}
      for filepath, props in bundle.files.iteritems():
        if 'h' in props:
          remaining.setdefault(props['h'], []).append((filepath, props))

      # Now block on the remaining files to be downloaded and mapped.
      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.
          for filepath, props in remaining.pop(digest):
            cache.hardlink(
                digest, os.path.join(outdir, filepath), props.get('m'))

          # 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.
  if not fetch_queue.verify_all_cached():
    raise isolated_format.MappingError(
        'Cache is too small to hold all requested files')
  return bundle


def directory_to_metadata(root, algo, blacklist):
  """Returns the FileItem list and .isolated metadata for a directory."""
  root = file_path.get_native_path_case(root)
  paths = isolated_format.expand_directory_and_symlink(
      root, '.' + os.path.sep, blacklist, sys.platform != 'win32')
  metadata = {
    relpath: isolated_format.file_to_metadata(
        os.path.join(root, relpath), {}, 0, algo)
    for relpath in paths
  }
  for v in metadata.itervalues():
    v.pop('t')
  items = [
      FileItem(
          path=os.path.join(root, relpath),
          digest=meta['h'],
          size=meta['s'],
          high_priority=relpath.endswith('.isolated'))
      for relpath, meta in metadata.iteritems() if 'h' in meta
  ]
  return items, metadata


def archive_files_to_storage(storage, files, blacklist):
  """Stores every entries and returns the relevant data.

  Arguments:
    storage: a Storage object that communicates with the remote object store.
    files: list of file paths to upload. If a directory is specified, a
           .isolated file is created and its hash is returned.
    blacklist: function that returns True if a file should be omitted.
  """
  assert all(isinstance(i, unicode) for i in files), files
  if len(files) != len(set(map(os.path.abspath, files))):
    raise Error('Duplicate entries found.')

  results = []
  # The temporary directory is only created as needed.
  tempdir = None
  try:
    # TODO(maruel): Yield the files to a worker thread.
    items_to_upload = []
    for f in files:
      try:
        filepath = os.path.abspath(f)
        if fs.isdir(filepath):
          # Uploading a whole directory.
          items, metadata = directory_to_metadata(
              filepath, storage.hash_algo, blacklist)

          # Create the .isolated file.
          if not tempdir:
            tempdir = tempfile.mkdtemp(prefix=u'isolateserver')
          handle, isolated = tempfile.mkstemp(dir=tempdir, suffix=u'.isolated')
          os.close(handle)
          data = {
              'algo':
                  isolated_format.SUPPORTED_ALGOS_REVERSE[storage.hash_algo],
              'files': metadata,
              'version': isolated_format.ISOLATED_FILE_VERSION,
          }
          isolated_format.save_isolated(isolated, data)
          h = isolated_format.hash_file(isolated, storage.hash_algo)
          items_to_upload.extend(items)
          items_to_upload.append(
              FileItem(
                  path=isolated,
                  digest=h,
                  size=fs.stat(isolated).st_size,
                  high_priority=True))
          results.append((h, f))

        elif fs.isfile(filepath):
          h = isolated_format.hash_file(filepath, storage.hash_algo)
          items_to_upload.append(
            FileItem(
                path=filepath,
                digest=h,
                size=fs.stat(filepath).st_size,
                high_priority=f.endswith('.isolated')))
          results.append((h, f))
        else:
          raise Error('%s is neither a file or directory.' % f)
      except OSError:
        raise Error('Failed to process %s.' % f)
    # Technically we would care about which files were uploaded but we don't
    # much in practice.
    _uploaded_files = storage.upload_items(items_to_upload)
    return results
  finally:
    if tempdir and fs.isdir(tempdir):
      file_path.rmtree(tempdir)


def archive(out, namespace, files, blacklist):
  if files == ['-']:
    files = sys.stdin.readlines()

  if not files:
    raise Error('Nothing to upload')

  files = [f.decode('utf-8') for f in files]
  blacklist = tools.gen_blacklist(blacklist)
  with get_storage(out, namespace) as storage:
    results = archive_files_to_storage(storage, files, blacklist)
  print('\n'.join('%s %s' % (r[0], r[1]) for r in results))


@subcommand.usage('<file1..fileN> or - to read from stdin')
def CMDarchive(parser, args):
  """Archives data to the server.

  If a directory is specified, a .isolated file is created the whole directory
  is uploaded. Then this .isolated file can be included in another one to run
  commands.

  The commands output each file that was processed with its content hash. For
  directories, the .isolated generated for the directory is listed as the
  directory entry itself.
  """
  add_isolate_server_options(parser)
  add_archive_options(parser)
  options, files = parser.parse_args(args)
  process_isolate_server_options(parser, options, True)
  try:
    archive(options.isolate_server, options.namespace, files, options.blacklist)
  except Error as e:
    parser.error(e.args[0])
  return 0


def CMDdownload(parser, args):
  """Download data from the server.

  It can either download individual files or a complete tree from a .isolated
  file.
  """
  add_isolate_server_options(parser)
  parser.add_option(
      '-s', '--isolated', metavar='HASH',
      help='hash of an isolated file, .isolated file content is discarded, use '
           '--file if you need it')
  parser.add_option(
      '-f', '--file', metavar='HASH DEST', default=[], action='append', nargs=2,
      help='hash and destination of a file, can be used multiple times')
  parser.add_option(
      '-t', '--target', metavar='DIR', default='download',
      help='destination directory')
  add_cache_options(parser)
  options, args = parser.parse_args(args)
  if args:
    parser.error('Unsupported arguments: %s' % args)

  process_isolate_server_options(parser, options, True)
  if bool(options.isolated) == bool(options.file):
    parser.error('Use one of --isolated or --file, and only one.')

  cache = process_cache_options(options)
  options.target = unicode(os.path.abspath(options.target))
  if options.isolated:
    if (fs.isfile(options.target) or
        (fs.isdir(options.target) and fs.listdir(options.target))):
      parser.error(
          '--target \'%s\' exists, please use another target' % options.target)
  with get_storage(options.isolate_server, options.namespace) as storage:
    # Fetching individual files.
    if options.file:
      # TODO(maruel): Enable cache in this case too.
      channel = threading_utils.TaskChannel()
      pending = {}
      for digest, dest in options.file:
        pending[digest] = dest
        storage.async_fetch(
            channel,
            threading_utils.PRIORITY_MED,
            digest,
            UNKNOWN_FILE_SIZE,
            functools.partial(file_write, os.path.join(options.target, dest)))
      while pending:
        fetched = channel.pull()
        dest = pending.pop(fetched)
        logging.info('%s: %s', fetched, dest)

    # Fetching whole isolated tree.
    if options.isolated:
      with cache:
        bundle = fetch_isolated(
            isolated_hash=options.isolated,
            storage=storage,
            cache=cache,
            outdir=options.target,
            require_command=False)
      if bundle.command:
        rel = os.path.join(options.target, bundle.relative_cwd)
        print('To run this test please run from the directory %s:' %
              os.path.join(options.target, rel))
        print('  ' + ' '.join(bundle.command))

  return 0


def add_archive_options(parser):
  parser.add_option(
      '--blacklist',
      action='append', default=list(DEFAULT_BLACKLIST),
      help='List of regexp to use as blacklist filter when uploading '
           'directories')


def add_isolate_server_options(parser):
  """Adds --isolate-server and --namespace options to parser."""
  parser.add_option(
      '-I', '--isolate-server',
      metavar='URL', default=os.environ.get('ISOLATE_SERVER', ''),
      help='URL of the Isolate Server to use. Defaults to the environment '
           'variable ISOLATE_SERVER if set. No need to specify https://, this '
           'is assumed.')
  parser.add_option(
      '--namespace', default='default-gzip',
      help='The namespace to use on the Isolate Server, default: %default')


def process_isolate_server_options(parser, options, set_exception_handler):
  """Processes the --isolate-server option and aborts if not specified.

  Returns the identity as determined by the server.
  """
  if not options.isolate_server:
    parser.error('--isolate-server is required.')
  try:
    options.isolate_server = net.fix_url(options.isolate_server)
  except ValueError as e:
    parser.error('--isolate-server %s' % e)
  if set_exception_handler:
    on_error.report_on_exception_exit(options.isolate_server)
  try:
    return auth.ensure_logged_in(options.isolate_server)
  except ValueError as e:
    parser.error(str(e))


def add_cache_options(parser):
  cache_group = optparse.OptionGroup(parser, 'Cache management')
  cache_group.add_option(
      '--cache', metavar='DIR',
      help='Directory to keep a local cache of the files. Accelerates download '
           'by reusing already downloaded files. Default=%default')
  cache_group.add_option(
      '--max-cache-size',
      type='int',
      metavar='NNN',
      default=50*1024*1024*1024,
      help='Trim if the cache gets larger than this value, default=%default')
  cache_group.add_option(
      '--min-free-space',
      type='int',
      metavar='NNN',
      default=2*1024*1024*1024,
      help='Trim if disk free space becomes lower than this value, '
           'default=%default')
  cache_group.add_option(
      '--max-items',
      type='int',
      metavar='NNN',
      default=100000,
      help='Trim if more than this number of items are in the cache '
           'default=%default')
  parser.add_option_group(cache_group)


def process_cache_options(options):
  if options.cache:
    policies = CachePolicies(
        options.max_cache_size, options.min_free_space, options.max_items)

    # |options.cache| path may not exist until DiskCache() instance is created.
    return DiskCache(
        unicode(os.path.abspath(options.cache)),
        policies,
        isolated_format.get_hash_algo(options.namespace))
  else:
    return MemoryCache()


class OptionParserIsolateServer(logging_utils.OptionParserWithLogging):
  def __init__(self, **kwargs):
    logging_utils.OptionParserWithLogging.__init__(
        self,
        version=__version__,
        prog=os.path.basename(sys.modules[__name__].__file__),
        **kwargs)
    auth.add_auth_options(self)

  def parse_args(self, *args, **kwargs):
    options, args = logging_utils.OptionParserWithLogging.parse_args(
        self, *args, **kwargs)
    auth.process_auth_options(self, options)
    return options, args


def main(args):
  dispatcher = subcommand.CommandDispatcher(__name__)
  return dispatcher.execute(OptionParserIsolateServer(), args)


if __name__ == '__main__':
  fix_encoding.fix_encoding()
  tools.disable_buffering()
  colorama.init()
  sys.exit(main(sys.argv[1:]))