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