dulwich/lru_cache.py

   1 # Copyright (C) 2006, 2008 Canonical Ltd
   2 #
   3 # This program is free software; you can redistribute it and/or modify
   4 # it under the terms of the GNU General Public License as published by
   5 # the Free Software Foundation; either version 2 of the License, or
   6 # (at your option) any later version.
   7 #
   8 # This program is distributed in the hope that it will be useful,
   9 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  10 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  11 # GNU General Public License for more details.
  12 #
  13 # You should have received a copy of the GNU General Public License
  14 # along with this program; if not, write to the Free Software
  15 # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  16
  17 """A simple least-recently-used (LRU) cache."""
  18
  19 _null_key = object()
  20
  21 class _LRUNode(object):
  22     """This maintains the linked-list which is the lru internals."""
  23
  24     __slots__ = ('prev', 'next_key', 'key', 'value', 'cleanup', 'size')
  25
  26     def __init__(self, key, value, cleanup=None):
  27         self.prev = None
  28         self.next_key = _null_key
  29         self.key = key
  30         self.value = value
  31         self.cleanup = cleanup
  32         # TODO: We could compute this 'on-the-fly' like we used to, and remove
  33         #       one pointer from this object, we just need to decide if it
  34         #       actually costs us much of anything in normal usage
  35         self.size = None
  36
  37     def __repr__(self):
  38         if self.prev is None:
  39             prev_key = None
  40         else:
  41             prev_key = self.prev.key
  42         return '%s(%r n:%r p:%r)' % (self.__class__.__name__, self.key,
  43                                      self.next_key, prev_key)
  44
  45     def run_cleanup(self):
  46         if self.cleanup is not None:
  47             self.cleanup(self.key, self.value)
  48         self.cleanup = None
  49         # Just make sure to break any refcycles, etc
  50         self.value = None
  51
  52
  53 class LRUCache(object):
  54     """A class which manages a cache of entries, removing unused ones."""
  55
  56     def __init__(self, max_cache=100, after_cleanup_count=None):
  57         self._cache = {}
  58         # The "HEAD" of the lru linked list
  59         self._most_recently_used = None
  60         # The "TAIL" of the lru linked list
  61         self._least_recently_used = None
  62         self._update_max_cache(max_cache, after_cleanup_count)
  63
  64     def __contains__(self, key):
  65         return key in self._cache
  66
  67     def __getitem__(self, key):
  68         cache = self._cache
  69         node = cache[key]
  70         # Inlined from _record_access to decrease the overhead of __getitem__
  71         # We also have more knowledge about structure if __getitem__ is
  72         # succeeding, then we know that self._most_recently_used must not be
  73         # None, etc.
  74         mru = self._most_recently_used
  75         if node is mru:
  76             # Nothing to do, this node is already at the head of the queue
  77             return node.value
  78         # Remove this node from the old location
  79         node_prev = node.prev
  80         next_key = node.next_key
  81         # benchmarking shows that the lookup of _null_key in globals is faster
  82         # than the attribute lookup for (node is self._least_recently_used)
  83         if next_key is _null_key:
  84             # 'node' is the _least_recently_used, because it doesn't have a
  85             # 'next' item. So move the current lru to the previous node.
  86             self._least_recently_used = node_prev
  87         else:
  88             node_next = cache[next_key]
  89             node_next.prev = node_prev
  90         node_prev.next_key = next_key
  91         # Insert this node at the front of the list
  92         node.next_key = mru.key
  93         mru.prev = node
  94         self._most_recently_used = node
  95         node.prev = None
  96         return node.value
  97
  98     def __len__(self):
  99         return len(self._cache)
 100
 101     def _walk_lru(self):
 102         """Walk the LRU list, only meant to be used in tests."""
 103         node = self._most_recently_used
 104         if node is not None:
 105             if node.prev is not None:
 106                 raise AssertionError('the _most_recently_used entry is not'
 107                                      ' supposed to have a previous entry'
 108                                      ' %s' % (node,))
 109         while node is not None:
 110             if node.next_key is _null_key:
 111                 if node is not self._least_recently_used:
 112                     raise AssertionError('only the last node should have'
 113                                          ' no next value: %s' % (node,))
 114                 node_next = None
 115             else:
 116                 node_next = self._cache[node.next_key]
 117                 if node_next.prev is not node:
 118                     raise AssertionError('inconsistency found, node.next.prev'
 119                                          ' != node: %s' % (node,))
 120             if node.prev is None:
 121                 if node is not self._most_recently_used:
 122                     raise AssertionError('only the _most_recently_used should'
 123                                          ' not have a previous node: %s'
 124                                          % (node,))
 125             else:
 126                 if node.prev.next_key != node.key:
 127                     raise AssertionError('inconsistency found, node.prev.next'
 128                                          ' != node: %s' % (node,))
 129             yield node
 130             node = node_next
 131
 132     def add(self, key, value, cleanup=None):
 133         """Add a new value to the cache.
 134
 135         Also, if the entry is ever removed from the cache, call
 136         cleanup(key, value).
 137
 138         :param key: The key to store it under
 139         :param value: The object to store
 140         :param cleanup: None or a function taking (key, value) to indicate
 141                         'value' should be cleaned up.
 142         """
 143         if key is _null_key:
 144             raise ValueError('cannot use _null_key as a key')
 145         if key in self._cache:
 146             node = self._cache[key]
 147             node.run_cleanup()
 148             node.value = value
 149             node.cleanup = cleanup
 150         else:
 151             node = _LRUNode(key, value, cleanup=cleanup)
 152             self._cache[key] = node
 153         self._record_access(node)
 154
 155         if len(self._cache) > self._max_cache:
 156             # Trigger the cleanup
 157             self.cleanup()
 158
 159     def cache_size(self):
 160         """Get the number of entries we will cache."""
 161         return self._max_cache
 162
 163     def get(self, key, default=None):
 164         node = self._cache.get(key, None)
 165         if node is None:
 166             return default
 167         self._record_access(node)
 168         return node.value
 169
 170     def keys(self):
 171         """Get the list of keys currently cached.
 172
 173         Note that values returned here may not be available by the time you
 174         request them later. This is simply meant as a peak into the current
 175         state.
 176
 177         :return: An unordered list of keys that are currently cached.
 178         """
 179         return self._cache.keys()
 180
 181     def items(self):
 182         """Get the key:value pairs as a dict."""
 183         return dict((k, n.value) for k, n in self._cache.iteritems())
 184
 185     def cleanup(self):
 186         """Clear the cache until it shrinks to the requested size.
 187
 188         This does not completely wipe the cache, just makes sure it is under
 189         the after_cleanup_count.
 190         """
 191         # Make sure the cache is shrunk to the correct size
 192         while len(self._cache) > self._after_cleanup_count:
 193             self._remove_lru()
 194
 195     def __setitem__(self, key, value):
 196         """Add a value to the cache, there will be no cleanup function."""
 197         self.add(key, value, cleanup=None)
 198
 199     def _record_access(self, node):
 200         """Record that key was accessed."""
 201         # Move 'node' to the front of the queue
 202         if self._most_recently_used is None:
 203             self._most_recently_used = node
 204             self._least_recently_used = node
 205             return
 206         elif node is self._most_recently_used:
 207             # Nothing to do, this node is already at the head of the queue
 208             return
 209         # We've taken care of the tail pointer, remove the node, and insert it
 210         # at the front
 211         # REMOVE
 212         if node is self._least_recently_used:
 213             self._least_recently_used = node.prev
 214         if node.prev is not None:
 215             node.prev.next_key = node.next_key
 216         if node.next_key is not _null_key:
 217             node_next = self._cache[node.next_key]
 218             node_next.prev = node.prev
 219         # INSERT
 220         node.next_key = self._most_recently_used.key
 221         self._most_recently_used.prev = node
 222         self._most_recently_used = node
 223         node.prev = None
 224
 225     def _remove_node(self, node):
 226         if node is self._least_recently_used:
 227             self._least_recently_used = node.prev
 228         self._cache.pop(node.key)
 229         # If we have removed all entries, remove the head pointer as well
 230         if self._least_recently_used is None:
 231             self._most_recently_used = None
 232         node.run_cleanup()
 233         # Now remove this node from the linked list
 234         if node.prev is not None:
 235             node.prev.next_key = node.next_key
 236         if node.next_key is not _null_key:
 237             node_next = self._cache[node.next_key]
 238             node_next.prev = node.prev
 239         # And remove this node's pointers
 240         node.prev = None
 241         node.next_key = _null_key
 242
 243     def _remove_lru(self):
 244         """Remove one entry from the lru, and handle consequences.
 245
 246         If there are no more references to the lru, then this entry should be
 247         removed from the cache.
 248         """
 249         self._remove_node(self._least_recently_used)
 250
 251     def clear(self):
 252         """Clear out all of the cache."""
 253         # Clean up in LRU order
 254         while self._cache:
 255             self._remove_lru()
 256
 257     def resize(self, max_cache, after_cleanup_count=None):
 258         """Change the number of entries that will be cached."""
 259         self._update_max_cache(max_cache,
 260                                after_cleanup_count=after_cleanup_count)
 261
 262     def _update_max_cache(self, max_cache, after_cleanup_count=None):
 263         self._max_cache = max_cache
 264         if after_cleanup_count is None:
 265             self._after_cleanup_count = self._max_cache * 8 / 10
 266         else:
 267             self._after_cleanup_count = min(after_cleanup_count,
 268                                             self._max_cache)
 269         self.cleanup()
 270
 271
 272 class LRUSizeCache(LRUCache):
 273     """An LRUCache that removes things based on the size of the values.
 274
 275     This differs in that it doesn't care how many actual items there are,
 276     it just restricts the cache to be cleaned up after so much data is stored.
 277
 278     The size of items added will be computed using compute_size(value), which
 279     defaults to len() if not supplied.
 280     """
 281
 282     def __init__(self, max_size=1024*1024, after_cleanup_size=None,
 283                  compute_size=None):
 284         """Create a new LRUSizeCache.
 285
 286         :param max_size: The max number of bytes to store before we start
 287             clearing out entries.
 288         :param after_cleanup_size: After cleaning up, shrink everything to this
 289             size.
 290         :param compute_size: A function to compute the size of the values. We
 291             use a function here, so that you can pass 'len' if you are just
 292             using simple strings, or a more complex function if you are using
 293             something like a list of strings, or even a custom object.
 294             The function should take the form "compute_size(value) => integer".
 295             If not supplied, it defaults to 'len()'
 296         """
 297         self._value_size = 0
 298         self._compute_size = compute_size
 299         if compute_size is None:
 300             self._compute_size = len
 301         self._update_max_size(max_size, after_cleanup_size=after_cleanup_size)
 302         LRUCache.__init__(self, max_cache=max(int(max_size/512), 1))
 303
 304     def add(self, key, value, cleanup=None):
 305         """Add a new value to the cache.
 306
 307         Also, if the entry is ever removed from the cache, call
 308         cleanup(key, value).
 309
 310         :param key: The key to store it under
 311         :param value: The object to store
 312         :param cleanup: None or a function taking (key, value) to indicate
 313                         'value' should be cleaned up.
 314         """
 315         if key is _null_key:
 316             raise ValueError('cannot use _null_key as a key')
 317         node = self._cache.get(key, None)
 318         value_len = self._compute_size(value)
 319         if value_len >= self._after_cleanup_size:
 320             # The new value is 'too big to fit', as it would fill up/overflow
 321             # the cache all by itself
 322             if node is not None:
 323                 # We won't be replacing the old node, so just remove it
 324                 self._remove_node(node)
 325             if cleanup is not None:
 326                 cleanup(key, value)
 327             return
 328         if node is None:
 329             node = _LRUNode(key, value, cleanup=cleanup)
 330             self._cache[key] = node
 331         else:
 332             self._value_size -= node.size
 333         node.size = value_len
 334         self._value_size += value_len
 335         self._record_access(node)
 336
 337         if self._value_size > self._max_size:
 338             # Time to cleanup
 339             self.cleanup()
 340
 341     def cleanup(self):
 342         """Clear the cache until it shrinks to the requested size.
 343
 344         This does not completely wipe the cache, just makes sure it is under
 345         the after_cleanup_size.
 346         """
 347         # Make sure the cache is shrunk to the correct size
 348         while self._value_size > self._after_cleanup_size:
 349             self._remove_lru()
 350
 351     def _remove_node(self, node):
 352         self._value_size -= node.size
 353         LRUCache._remove_node(self, node)
 354
 355     def resize(self, max_size, after_cleanup_size=None):
 356         """Change the number of bytes that will be cached."""
 357         self._update_max_size(max_size, after_cleanup_size=after_cleanup_size)
 358         max_cache = max(int(max_size/512), 1)
 359         self._update_max_cache(max_cache)
 360
 361     def _update_max_size(self, max_size, after_cleanup_size=None):
 362         self._max_size = max_size
 363         if after_cleanup_size is None:
 364             self._after_cleanup_size = self._max_size * 8 / 10
 365         else:
 366             self._after_cleanup_size = min(after_cleanup_size, self._max_size)