isolateserver: DiskCache format v2

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.6.0'
+__version__ = '0.7.0'
 
 import base64
+import collections
 import errno
 import functools
+import json
 import io
 import logging
 import optparse
 import os
 import re
 import signal
 import stat
 import sys
 import tempfile
 import threading
 import time
 import types
 import zlib
 
 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
 
 import auth
 import isolated_format
 
 
@@ skipping 1428 matching lines @@
                       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.
+  """Stateful LRU cache in a semi-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):
+  # All protected methods (starting with '_') except _path should be called
+  # with self._lock held.
+
+  def __init__(self, cache_dir, policies, hash_algo, time_fn=None):
     """
     Arguments:
       cache_dir: directory where to place the cache.
       policies: cache retention policies.
       algo: hashing algorithm used.
+      time_fn: function to take current timestamp when adding new items.
+        Defaults to time.time.
     """
-    # All protected methods (starting with '_') except _path should be called
-    # with self._lock held.
     super(DiskCache, self).__init__()
+    self.time_fn = time_fn or time.time
     self.cache_dir = cache_dir
     self.policies = policies
     self.hash_algo = hash_algo
     self.state_file = os.path.join(cache_dir, self.STATE_FILE)
-    # Items in a LRU lookup dict(digest: size).
-    self._lru = lru.LRUDict()
+    # Items in a LRU lookup dict(digest: [size, timestamp]).
+    # We use lists instead of tuples because JSON arrays are parsed to lists.
+    # Mutations must be followed by `self._dirty = True`.
+    self._lru = collections.OrderedDict()
+    self._dirty = False  # if True, self._lru was modified since loading/saving.
     # Current cached free disk space. It is updated by self._trim().
     self._free_disk = 0
     # The first item in the LRU cache that must not be evicted during this run
     # since it was referenced. All items more recent that _protected in the LRU
     # cache are also inherently protected. It could be a set() of all items
     # referenced but this increases memory usage without a use case.
     self._protected = None
     # Cleanup operations done by self._load(), if any.
     self._operations = []
+
+    self._free_disk = file_path.get_free_space(self.cache_dir)
+
     with tools.Profiler('Setup'):
       with self._lock:
-        # self._load() calls self._trim() which initializes self._free_disk.
         self._load()
 
   def __contains__(self, digest):
     with self._lock:
       return digest in self._lru
 
   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)
+            self._cache_disk_size() / 1024)
         logging.info(
             '%5d (%8dkb) evicted',
             len(self._evicted), sum(self._evicted) / 1024)
         logging.info(
             '       %8dkb free',
             self._free_disk / 1024)
     return False
 
+  def _sizes(self):
+    """Returns an iterator of pairs (digest, size)."""
+    return (
+      (digest, size)
+      for digest, (size, _) in self._lru.iteritems())
+
   def cached_set(self):
     with self._lock:
-      return self._lru.keys_set()
+      return set(self._lru)
 
   def cleanup(self):
     """Cleans up the cache directory.
 
     Ensures there is no unknown files in cache_dir.
     Ensures the read-only bits are set correctly.
 
     At that point, the cache was already loaded, trimmed to respect cache
     policies.
     """
-    fs.chmod(self.cache_dir, 0700)
-    # Ensure that all files listed in the state still exist and add new ones.
-    previous = self._lru.keys_set()
-    # It'd be faster if there were a readdir() function.
-    for filename in fs.listdir(self.cache_dir):
-      if filename == self.STATE_FILE:
-        fs.chmod(os.path.join(self.cache_dir, filename), 0600)
-        continue
-      if filename in previous:
-        fs.chmod(os.path.join(self.cache_dir, filename), 0400)
-        previous.remove(filename)
-        continue
-
-      # An untracked file. Delete it.
+    def try_remove(filename):
+      assert os.path.isabs(filename)
       logging.warning('Removing unknown file %s from cache', filename)
-      p = self._path(filename)
-      if fs.isdir(p):
+      if fs.isdir(filename):
         try:
-          file_path.rmtree(p)
+          file_path.rmtree(filename)
         except OSError:
           pass
       else:
-        file_path.try_remove(p)
-      continue
+        file_path.try_remove(filename)
+
+    fs.chmod(self.cache_dir, 0700)
+    # Ensure that all files listed in the state still exist and add new ones.
+    previous = set(self._lru)

[M-A Ruel, 2016/10/13 20:51:12]

self.cached_set() ? just for the form. :)

+    # It'd be faster if there were a readdir() function.
+    for filename_l1 in fs.listdir(self.cache_dir):
+      full_name_l1 = os.path.join(self.cache_dir, filename_l1)
+      if filename_l1 == self.STATE_FILE:
+        fs.chmod(full_name_l1, 0600)
+        continue
+      if len(filename_l1) != 2:
+        try_remove(full_name_l1)
+        continue
+      for filename_l2 in fs.listdir(full_name_l1):
+        digest = filename_l1 + filename_l2
+        full_name_l2 = os.path.join(full_name_l1, filename_l2)
+        if digest in previous:
+          fs.chmod(full_name_l2, 0400)
+          previous.remove(digest)
+        else:
+          try_remove(full_name_l2)
 
     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)
+      for digest in previous:
+        self._lru.pop(digest)
+      self._dirty = True
 
     # What remains to be done is to hash every single item to
     # detect corruption, then save to ensure state.json is up to date.
     # Sadly, on a 50Gb cache with 100mib/s I/O, this is still over 8 minutes.
     # TODO(maruel): Let's revisit once directory metadata is stored in
     # state.json so only the files that had been mapped since the last cleanup()
     # call are manually verified.
     #
     #with self._lock:
     #  for digest in self._lru:
     #    if not isolated_format.is_valid_hash(
     #        self._path(digest), self.hash_algo):
     #      self.evict(digest)
     #      logging.info('Deleted corrupted item: %s', digest)
 
   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:
+      if self._lru.pop(digest, None) is None:
         return False
-      self._lru.touch(digest)
+      self._lru[digest] = [size, self.time_fn()]
+      self._dirty = True
       self._protected = self._protected or digest
     return True
 
   def evict(self, digest):
     with self._lock:
       # Do not check for 'digest == self._protected' since it could be because
       # the object is corrupted.
-      self._lru.pop(digest)
+      self._lru.pop(digest, None)
+      self._dirty = True
       self._delete_file(digest, UNKNOWN_FILE_SIZE)
 
   def getfileobj(self, digest):
     try:
       f = fs.open(self._path(digest), 'rb')
       with self._lock:
-        self._used.append(self._lru[digest])
+        self._used.append(self._lru[digest][0])
       return f
     except IOError:
       raise CacheMiss(digest)
 
   def write(self, digest, content):
     assert content is not None
     with self._lock:
       self._protected = self._protected or digest
     path = self._path(digest)
     # A stale broken file may remain. It is possible for the file to have write
@@ skipping 11 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 _cache_disk_size(self):

[M-A Ruel, 2016/10/13 20:51:12]

can you move it aside cache_set? I feel it'd make more sense to have both close.

+    """Returns number of bytes that cache files take."""
+    return sum(size for (size, _) in self._lru.itervalues())
+
+  @staticmethod
+  def _validate_state(state):
+    """Raises ValueError if the given serialized state is invalid."""
+    if not isinstance(state, dict):
+      raise ValueError('not a JSON object')
+
+    version = state.get('version')
+    if not isinstance(version, int):
+      raise ValueError('version %r is not an integer' % (version,))
+
+    entries = state.get('entries')
+    if not isinstance(entries, list):
+      raise ValueError('entries is not an array')
+    for entry in entries:
+      try:
+        if not isinstance(entry, list) or len(entry) != 2:
+          raise ValueError('not a pair')
+        if not isinstance(entry[0], basestring):
+          raise ValueError('first item is not a string')
+
+        pair_of_numbers = (
+          isinstance(entry[1], list) and
+          len(entry[1]) == 2 and
+          all(isinstance(x, (int, float)) for x in entry[1])
+        )
+        if not pair_of_numbers:
+          raise ValueError('second item is not a pair of numbers')
+      except ValueError as ex:
+        raise ValueError('invalid entry %r: %s' % (entry, ex))
+
   def _load(self):
-    """Loads state of the cache from json file.
+    """Loads state of the cache from json file if it exists.
 
     If cache_dir does not exist on disk, it is created.
+
+    Raises ValueError if the state file is broken.
     """
     self._lock.assert_locked()
 
     if not fs.isfile(self.state_file):
       if not os.path.isdir(self.cache_dir):
         fs.makedirs(self.cache_dir)
     else:
       # Load state of the cache.
       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)
+        with open(self.state_file, 'r') as f:

[M-A Ruel, 2016/10/13 20:51:12]

'rb'

+          state = json.load(f)
+        if not isinstance(state, dict):
+          # Possibly old format.

[M-A Ruel, 2016/10/13 20:51:12]

# Possibly old format. Hard reset with a clear cache.

+          file_path.rmtree(self.cache_dir)
+          fs.makedirs(self.cache_dir)
+          self._lru = collections.OrderedDict()
+        else:
+          self._validate_state(state)
+          if state['version'] != 2:
+            raise ValueError(
+                'unsupported version %s, supported is 2' % (state['version']))
+          self._lru = collections.OrderedDict(state['entries'])
+        self._dirty = False
+      except (IOError, ValueError) as e:
+        raise ValueError('Cannot load state file %s: %s' % (self.state_file, e))
+
     self._trim()
     # We want the initial cache size after trimming, i.e. what is readily
     # avaiable.
     self._initial_number_items = len(self._lru)
-    self._initial_size = sum(self._lru.itervalues())
+    self._initial_size = self._cache_disk_size()
     if self._evicted:
       logging.info(
           'Trimming evicted items with the following sizes: %s',
           sorted(self._evicted))
 
   def _save(self):
-    """Saves the LRU ordering."""
+    """Saves the LRU ordering if changed."""
     self._lock.assert_locked()
+    if not self._dirty:
+      return
+
     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)
+    with open(self.state_file, 'w') as f:

[M-A Ruel, 2016/10/13 20:51:12]

'wb'

+      contents = {
+        'version': 2,
+        'entries': self._lru.items(),
+      }
+      json.dump(contents, f)

[M-A Ruel, 2016/10/13 20:51:12]

separators=(',',':')

will save a few valuable bytes.

+    self._dirty = False
 
   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())
+      total_size = self._cache_disk_size()
       while total_size > self.policies.max_cache_size:
         total_size -= self._remove_lru_file(True)
 
     # 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(True)
 
     # Ensure enough free space.
     self._free_disk = file_path.get_free_space(self.cache_dir)
     trimmed_due_to_space = 0
     while (
         self.policies.min_free_space and
         self._lru and
         self._free_disk < self.policies.min_free_space):
       trimmed_due_to_space += 1
       self._remove_lru_file(True)
 
     if trimmed_due_to_space:
-      total_usage = sum(self._lru.itervalues())
+      total_usage = self._cache_disk_size()
       usage_percent = 0.
       if total_usage:
         usage_percent = 100. * float(total_usage) / self.policies.max_cache_size
 
       logging.warning(
           'Trimmed %s file(s) due to not enough free disk space: %.1fkb free,'
           ' %.1fkb cache (%.1f%% of its maximum capacity of %.1fkb)',
           trimmed_due_to_space,
           self._free_disk / 1024.,
           total_usage / 1024.,
           usage_percent,
           self.policies.max_cache_size / 1024.)
     self._save()
 
   def _path(self, digest):
     """Returns the path to one item."""
-    return os.path.join(self.cache_dir, digest)
+    assert len(digest) > 2
+    return os.path.join(self.cache_dir, digest[:2], digest[2:])
 
   def _remove_lru_file(self, allow_protected):
     """Removes the lastest recently used file and returns its size."""

[M-A Ruel, 2016/10/13 20:51:12]

Probably a typo of mine; latest ?

     self._lock.assert_locked()
-    try:
-      digest, size = self._lru.get_oldest()
-      if not allow_protected and digest == self._protected:
-        raise Error('Not enough space to map the whole isolated tree')
-    except KeyError:
+
+    digest = None
+    for key in self._lru:
+      digest = key
+      break
+    if digest is None:
       raise Error('Nothing to remove')
-    digest, size = self._lru.pop_oldest()
+
+    if not allow_protected and digest == self._protected:
+      raise Error('Not enough space to map the whole isolated tree')
+
+    digest, (size, _) = self._lru.popitem(last=False)
+    self._dirty = True
     logging.debug("Removing LRU file %s", digest)
     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."""
+    """Adds an item into LRU cache marking it as a newest one.
+
+    Asumes the file exists.

[M-A Ruel, 2016/10/13 20:51:12]

Assumes

+    """
     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)
+    self._lru.pop(digest, None)
+    self._lru[digest] = [size, self.time_fn()]

[M-A Ruel, 2016/10/13 20:51:12]

self._lru[digest] = (size, self.time_fn())

it's slightly more memory efficient. It doesn't help once serialized but oh
well.

+    self._dirty = True
     self._free_disk -= size
     # Do a quicker version of self._trim(). It only enforces free disk space,
     # not cache size limits. It doesn't actually look at real free disk space,
     # only uses its cache values. self._trim() will be called later to enforce
     # real trimming but doing this quick version here makes it possible to map
     # an isolated that is larger than the current amount of free disk space when
     # the cache size is already large.
     while (
         self.policies.min_free_space and
         self._lru and
         self._free_disk < self.policies.min_free_space):
       self._remove_lru_file(False)
 
   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))
+      path = self._path(digest)
+      file_path.try_remove(path)
       self._evicted.append(size)
       self._free_disk += size
+
+      parent = os.path.dirname(path)
+      if len(os.listdir(parent)) == 0:

[M-A Ruel, 2016/10/13 20:51:12]

I don't think it should be done. It's better to leave the parents there.

+        fs.rmtree(parent)

[M-A Ruel, 2016/10/13 20:51:12]

I don't see a corresponding call to mkdir() to create the two letters parent
directories.

     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 = {}
@@ skipping 625 matching lines @@
   return dispatcher.execute(OptionParserIsolateServer(), args)
 
 
 if __name__ == '__main__':
   subprocess42.inhibit_os_error_reporting()
   fix_encoding.fix_encoding()
   tools.disable_buffering()
   colorama.init()
   file_path.enable_symlink()
   sys.exit(main(sys.argv[1:]))