src/third_party/chromite/lib/cache.py

   1 # Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4
   5 """Contains on-disk caching functionality."""
   6
   7 from __future__ import print_function
   8
   9 import logging
  10 import os
  11 import shutil
  12 import urlparse
  13
  14 from chromite.lib import cros_build_lib
  15 from chromite.lib import locking
  16 from chromite.lib import osutils
  17 from chromite.lib import retry_util
  18
  19 # pylint: disable=W0212
  20
  21 def EntryLock(f):
  22   """Decorator that provides monitor access control."""
  23   def new_f(self, *args, **kwargs):
  24     # Ensure we don't have a read lock before potentially blocking while trying
  25     # to access the monitor.
  26     if self.read_locked:
  27       raise AssertionError(
  28           'Cannot call %s while holding a read lock.' % f.__name__)
  29
  30     with self._entry_lock:
  31       self._entry_lock.write_lock()
  32       return f(self, *args, **kwargs)
  33   return new_f
  34
  35
  36 def WriteLock(f):
  37   """Decorator that takes a write lock."""
  38   def new_f(self, *args, **kwargs):
  39     with self._lock.write_lock():
  40       return f(self, *args, **kwargs)
  41   return new_f
  42
  43
  44 class CacheReference(object):
  45   """Encapsulates operations on a cache key reference.
  46
  47   CacheReferences are returned by the DiskCache.Lookup() function.  They are
  48   used to read from and insert into the cache.
  49
  50   A typical example of using a CacheReference:
  51
  52   @contextlib.contextmanager
  53   def FetchFromCache()
  54     with cache.Lookup(key) as ref:
  55        # If entry doesn't exist in cache already, generate it ourselves, and
  56        # insert it into the cache, acquiring a read lock on it in the process.
  57        # If the entry does exist, we grab a read lock on it.
  58       if not ref.Exists(lock=True):
  59         path = PrepareItem()
  60         ref.SetDefault(path, lock=True)
  61
  62       # yield the path to the cached entry to consuming code.
  63       yield ref.path
  64   """
  65
  66   def __init__(self, cache, key):
  67     self._cache = cache
  68     self.key = key
  69     self.acquired = False
  70     self.read_locked = False
  71     self._lock = cache._LockForKey(key)
  72     self._entry_lock = cache._LockForKey(key, suffix='.entry_lock')
  73
  74   @property
  75   def path(self):
  76     """Returns on-disk path to the cached item."""
  77     return self._cache._GetKeyPath(self.key)
  78
  79   def Acquire(self):
  80     """Prepare the cache reference for operation.
  81
  82     This must be called (either explicitly or through entering a 'with'
  83     context) before calling any methods that acquire locks, or mutates
  84     reference.
  85     """
  86     if self.acquired:
  87       raise AssertionError(
  88           'Attempting to acquire an already acquired reference.')
  89
  90     self.acquired = True
  91     self._lock.__enter__()
  92
  93   def Release(self):
  94     """Release the cache reference.  Causes any held locks to be released."""
  95     if not self.acquired:
  96       raise AssertionError(
  97           'Attempting to release an unacquired reference.')
  98
  99     self.acquired = False
 100     self._lock.__exit__(None, None, None)
 101
 102   def __enter__(self):
 103     self.Acquire()
 104     return self
 105
 106   def __exit__(self, *args):
 107     self.Release()
 108
 109   def _ReadLock(self):
 110     self._lock.read_lock()
 111     self.read_locked = True
 112
 113   @WriteLock
 114   def _Assign(self, path):
 115     self._cache._Insert(self.key, path)
 116
 117   @WriteLock
 118   def _AssignText(self, text):
 119     self._cache._InsertText(self.key, text)
 120
 121   @WriteLock
 122   def _Remove(self, key):
 123     self._cache._Remove(key)
 124
 125   def _Exists(self):
 126     return self._cache._KeyExists(self.key)
 127
 128   @EntryLock
 129   def Assign(self, path):
 130     """Insert a file or a directory into the cache at the referenced key."""
 131     self._Assign(path)
 132
 133   @EntryLock
 134   def AssignText(self, text):
 135     """Create a file containing |text| and assign it to the key.
 136
 137     Args:
 138       text: Can be a string or an iterable.
 139     """
 140     self._AssignText(text)
 141
 142   @EntryLock
 143   def Remove(self, key):
 144     """Removes the key entry from the cache."""
 145     self._Remove(key)
 146
 147   @EntryLock
 148   def Exists(self, lock=False):
 149     """Tests for existence of entry.
 150
 151     Args:
 152       lock: If the entry exists, acquire and maintain a read lock on it.
 153     """
 154     if self._Exists():
 155       if lock:
 156         self._ReadLock()
 157       return True
 158     return False
 159
 160   @EntryLock
 161   def SetDefault(self, default_path, lock=False):
 162     """Assigns default_path if the entry doesn't exist.
 163
 164     Args:
 165       default_path: The path to assign if the entry doesn't exist.
 166       lock: Acquire and maintain a read lock on the entry.
 167     """
 168     if not self._Exists():
 169       self._Assign(default_path)
 170     if lock:
 171       self._ReadLock()
 172
 173   def Unlock(self):
 174     """Release read lock on the reference."""
 175     self._lock.unlock()
 176
 177
 178 class DiskCache(object):
 179   """Locked file system cache keyed by tuples.
 180
 181   Key entries can be files or directories.  Access to the cache is provided
 182   through CacheReferences, which are retrieved by using the cache Lookup()
 183   method.
 184   """
 185   # TODO(rcui): Add LRU cleanup functionality.
 186
 187   _STAGING_DIR = 'staging'
 188
 189   def __init__(self, cache_dir):
 190     self._cache_dir = cache_dir
 191     self.staging_dir = os.path.join(cache_dir, self._STAGING_DIR)
 192
 193     osutils.SafeMakedirsNonRoot(self._cache_dir)
 194     osutils.SafeMakedirsNonRoot(self.staging_dir)
 195
 196   def _KeyExists(self, key):
 197     return os.path.exists(self._GetKeyPath(key))
 198
 199   def _GetKeyPath(self, key):
 200     """Get the on-disk path of a key."""
 201     return os.path.join(self._cache_dir, '+'.join(key))
 202
 203   def _LockForKey(self, key, suffix='.lock'):
 204     """Returns an unacquired lock associated with a key."""
 205     key_path = self._GetKeyPath(key)
 206     osutils.SafeMakedirsNonRoot(os.path.dirname(key_path))
 207     lock_path = os.path.join(self._cache_dir, os.path.dirname(key_path),
 208                              os.path.basename(key_path) + suffix)
 209     return locking.FileLock(lock_path)
 210
 211   def _TempDirContext(self):
 212     return osutils.TempDir(base_dir=self.staging_dir)
 213
 214   def _Insert(self, key, path):
 215     """Insert a file or a directory into the cache at a given key."""
 216     self._Remove(key)
 217     key_path = self._GetKeyPath(key)
 218     osutils.SafeMakedirsNonRoot(os.path.dirname(key_path))
 219     shutil.move(path, key_path)
 220
 221   def _InsertText(self, key, text):
 222     """Inserts a file containing |text| into the cache."""
 223     with self._TempDirContext() as tempdir:
 224       file_path = os.path.join(tempdir, 'tempfile')
 225       osutils.WriteFile(file_path, text)
 226       self._Insert(key, file_path)
 227
 228   def _Remove(self, key):
 229     """Remove a key from the cache."""
 230     if self._KeyExists(key):
 231       with self._TempDirContext() as tempdir:
 232         shutil.move(self._GetKeyPath(key), tempdir)
 233
 234   def Lookup(self, key):
 235     """Get a reference to a given key."""
 236     return CacheReference(self, key)
 237
 238
 239 def Untar(path, cwd, sudo=False):
 240   """Untar a tarball."""
 241   functor = cros_build_lib.SudoRunCommand if sudo else cros_build_lib.RunCommand
 242   functor(['tar', '-xpf', path], cwd=cwd, debug_level=logging.DEBUG)
 243
 244
 245 class TarballCache(DiskCache):
 246   """Supports caching of extracted tarball contents."""
 247
 248   def __init__(self, cache_dir):
 249     DiskCache.__init__(self, cache_dir)
 250
 251   def _Fetch(self, url, local_path):
 252     """Fetch a remote file."""
 253     # We have to nest the import because gs.GSContext uses us to cache its own
 254     # gsutil tarball.  We know we won't get into a recursive loop though as it
 255     # only fetches files via non-gs URIs.
 256     from chromite.lib import gs
 257
 258     if url.startswith(gs.BASE_GS_URL):
 259       ctx = gs.GSContext()
 260       ctx.Copy(url, local_path)
 261     else:
 262       retry_util.RunCurl([url, '-o', local_path], debug_level=logging.DEBUG)
 263
 264   def _Insert(self, key, tarball_path):
 265     """Insert a tarball and its extracted contents into the cache.
 266
 267     Download the tarball first if a URL is provided as tarball_path.
 268     """
 269     with osutils.TempDir(prefix='tarball-cache',
 270                          base_dir=self.staging_dir) as tempdir:
 271
 272       o = urlparse.urlsplit(tarball_path)
 273       if o.scheme:
 274         url = tarball_path
 275         tarball_path = os.path.join(tempdir, os.path.basename(o.path))
 276         self._Fetch(url, tarball_path)
 277
 278       extract_path = os.path.join(tempdir, 'extract')
 279       os.mkdir(extract_path)
 280       Untar(tarball_path, extract_path)
 281       DiskCache._Insert(self, key, extract_path)