~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/fifo_cache.py

  • Committer: Vincent Ladeuil
  • Date: 2011-12-21 14:25:26 UTC
  • mto: This revision was merged to the branch mainline in revision 6397.
  • Revision ID: v.ladeuil+lp@free.fr-20111221142526-pnwau0xnalimujts
Provides MemoryStack to simplify configuration setup in tests

Show diffs side-by-side

added added

removed removed

Lines of Context:
 
1
# Copyright (C) 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 first-in-first-out (FIFO) cache."""
 
18
 
 
19
from __future__ import absolute_import
 
20
 
 
21
from collections import deque
 
22
 
 
23
 
 
24
class FIFOCache(dict):
 
25
    """A class which manages a cache of entries, removing old ones."""
 
26
 
 
27
    def __init__(self, max_cache=100, after_cleanup_count=None):
 
28
        dict.__init__(self)
 
29
        self._max_cache = max_cache
 
30
        if after_cleanup_count is None:
 
31
            self._after_cleanup_count = self._max_cache * 8 / 10
 
32
        else:
 
33
            self._after_cleanup_count = min(after_cleanup_count,
 
34
                                            self._max_cache)
 
35
        self._cleanup = {} # map to cleanup functions when items are removed
 
36
        self._queue = deque() # Track when things are accessed
 
37
 
 
38
    def __setitem__(self, key, value):
 
39
        """Add a value to the cache, there will be no cleanup function."""
 
40
        self.add(key, value, cleanup=None)
 
41
 
 
42
    def __delitem__(self, key):
 
43
        # Remove the key from an arbitrary location in the queue
 
44
        remove = getattr(self._queue, 'remove', None)
 
45
        # Python2.5's has deque.remove, but Python2.4 does not
 
46
        if remove is not None:
 
47
            remove(key)
 
48
        else:
 
49
            # TODO: It would probably be faster to pop()/popleft() until we get to the
 
50
            #       key, and then insert those back into the queue. We know
 
51
            #       the key should only be present in one position, and we
 
52
            #       wouldn't need to rebuild the whole queue.
 
53
            self._queue = deque([k for k in self._queue if k != key])
 
54
        self._remove(key)
 
55
 
 
56
    def add(self, key, value, cleanup=None):
 
57
        """Add a new value to the cache.
 
58
 
 
59
        Also, if the entry is ever removed from the queue, call cleanup.
 
60
        Passing it the key and value being removed.
 
61
 
 
62
        :param key: The key to store it under
 
63
        :param value: The object to store
 
64
        :param cleanup: None or a function taking (key, value) to indicate
 
65
                        'value' should be cleaned up
 
66
        """
 
67
        if key in self:
 
68
            # Remove the earlier reference to this key, adding it again bumps
 
69
            # it to the end of the queue
 
70
            del self[key]
 
71
        self._queue.append(key)
 
72
        dict.__setitem__(self, key, value)
 
73
        if cleanup is not None:
 
74
            self._cleanup[key] = cleanup
 
75
        if len(self) > self._max_cache:
 
76
            self.cleanup()
 
77
 
 
78
    def cache_size(self):
 
79
        """Get the number of entries we will cache."""
 
80
        return self._max_cache
 
81
 
 
82
    def cleanup(self):
 
83
        """Clear the cache until it shrinks to the requested size.
 
84
 
 
85
        This does not completely wipe the cache, just makes sure it is under
 
86
        the after_cleanup_count.
 
87
        """
 
88
        # Make sure the cache is shrunk to the correct size
 
89
        while len(self) > self._after_cleanup_count:
 
90
            self._remove_oldest()
 
91
        if len(self._queue) != len(self):
 
92
            raise AssertionError('The length of the queue should always equal'
 
93
                ' the length of the dict. %s != %s'
 
94
                % (len(self._queue), len(self)))
 
95
 
 
96
    def clear(self):
 
97
        """Clear out all of the cache."""
 
98
        # Clean up in FIFO order
 
99
        while self:
 
100
            self._remove_oldest()
 
101
 
 
102
    def _remove(self, key):
 
103
        """Remove an entry, making sure to call any cleanup function."""
 
104
        cleanup = self._cleanup.pop(key, None)
 
105
        # We override self.pop() because it doesn't play well with cleanup
 
106
        # functions.
 
107
        val = dict.pop(self, key)
 
108
        if cleanup is not None:
 
109
            cleanup(key, val)
 
110
        return val
 
111
 
 
112
    def _remove_oldest(self):
 
113
        """Remove the oldest entry."""
 
114
        key = self._queue.popleft()
 
115
        self._remove(key)
 
116
 
 
117
    def resize(self, max_cache, after_cleanup_count=None):
 
118
        """Increase/decrease the number of cached entries.
 
119
 
 
120
        :param max_cache: The maximum number of entries to cache.
 
121
        :param after_cleanup_count: After cleanup, we should have at most this
 
122
            many entries. This defaults to 80% of max_cache.
 
123
        """
 
124
        self._max_cache = max_cache
 
125
        if after_cleanup_count is None:
 
126
            self._after_cleanup_count = max_cache * 8 / 10
 
127
        else:
 
128
            self._after_cleanup_count = min(max_cache, after_cleanup_count)
 
129
        if len(self) > self._max_cache:
 
130
            self.cleanup()
 
131
 
 
132
    # raise NotImplementedError on dict functions that would mutate the cache
 
133
    # which have not been properly implemented yet.
 
134
    def copy(self):
 
135
        raise NotImplementedError(self.copy)
 
136
 
 
137
    def pop(self, key, default=None):
 
138
        # If there is a cleanup() function, than it is unclear what pop()
 
139
        # should do. Specifically, we would have to call the cleanup on the
 
140
        # value before we return it, which should cause whatever resources were
 
141
        # allocated to be removed, which makes the return value fairly useless.
 
142
        # So instead, we just don't implement it.
 
143
        raise NotImplementedError(self.pop)
 
144
 
 
145
    def popitem(self):
 
146
        # See pop()
 
147
        raise NotImplementedError(self.popitem)
 
148
 
 
149
    def setdefault(self, key, defaultval=None):
 
150
        """similar to dict.setdefault"""
 
151
        if key in self:
 
152
            return self[key]
 
153
        self[key] = defaultval
 
154
        return defaultval
 
155
 
 
156
    def update(self, *args, **kwargs):
 
157
        """Similar to dict.update()"""
 
158
        if len(args) == 1:
 
159
            arg = args[0]
 
160
            if isinstance(arg, dict):
 
161
                for key, val in arg.iteritems():
 
162
                    self.add(key, val)
 
163
            else:
 
164
                for key, val in args[0]:
 
165
                    self.add(key, val)
 
166
        elif len(args) > 1:
 
167
            raise TypeError('update expected at most 1 argument, got %d'
 
168
                            % len(args))
 
169
        if kwargs:
 
170
            for key, val in kwargs.iteritems():
 
171
                self.add(key, val)
 
172
 
 
173
 
 
174
class FIFOSizeCache(FIFOCache):
 
175
    """An FIFOCache that removes things based on the size of the values.
 
176
 
 
177
    This differs in that it doesn't care how many actual items there are,
 
178
    it restricts the cache to be cleaned based on the size of the data.
 
179
    """
 
180
 
 
181
    def __init__(self, max_size=1024*1024, after_cleanup_size=None,
 
182
                 compute_size=None):
 
183
        """Create a new FIFOSizeCache.
 
184
 
 
185
        :param max_size: The max number of bytes to store before we start
 
186
            clearing out entries.
 
187
        :param after_cleanup_size: After cleaning up, shrink everything to this
 
188
            size (defaults to 80% of max_size).
 
189
        :param compute_size: A function to compute the size of a value. If
 
190
            not supplied we default to 'len'.
 
191
        """
 
192
        # Arbitrary, we won't really be using the value anyway.
 
193
        FIFOCache.__init__(self, max_cache=max_size)
 
194
        self._max_size = max_size
 
195
        if after_cleanup_size is None:
 
196
            self._after_cleanup_size = self._max_size * 8 / 10
 
197
        else:
 
198
            self._after_cleanup_size = min(after_cleanup_size, self._max_size)
 
199
 
 
200
        self._value_size = 0
 
201
        self._compute_size = compute_size
 
202
        if compute_size is None:
 
203
            self._compute_size = len
 
204
 
 
205
    def add(self, key, value, cleanup=None):
 
206
        """Add a new value to the cache.
 
207
 
 
208
        Also, if the entry is ever removed from the queue, call cleanup.
 
209
        Passing it the key and value being removed.
 
210
 
 
211
        :param key: The key to store it under
 
212
        :param value: The object to store, this value by itself is >=
 
213
            after_cleanup_size, then we will not store it at all.
 
214
        :param cleanup: None or a function taking (key, value) to indicate
 
215
                        'value' sohuld be cleaned up.
 
216
        """
 
217
        # Even if the new value won't be stored, we need to remove the old
 
218
        # value
 
219
        if key in self:
 
220
            # Remove the earlier reference to this key, adding it again bumps
 
221
            # it to the end of the queue
 
222
            del self[key]
 
223
        value_len = self._compute_size(value)
 
224
        if value_len >= self._after_cleanup_size:
 
225
            return
 
226
        self._queue.append(key)
 
227
        dict.__setitem__(self, key, value)
 
228
        if cleanup is not None:
 
229
            self._cleanup[key] = cleanup
 
230
        self._value_size += value_len
 
231
        if self._value_size > self._max_size:
 
232
            # Time to cleanup
 
233
            self.cleanup()
 
234
 
 
235
    def cache_size(self):
 
236
        """Get the number of bytes we will cache."""
 
237
        return self._max_size
 
238
 
 
239
    def cleanup(self):
 
240
        """Clear the cache until it shrinks to the requested size.
 
241
 
 
242
        This does not completely wipe the cache, just makes sure it is under
 
243
        the after_cleanup_size.
 
244
        """
 
245
        # Make sure the cache is shrunk to the correct size
 
246
        while self._value_size > self._after_cleanup_size:
 
247
            self._remove_oldest()
 
248
 
 
249
    def _remove(self, key):
 
250
        """Remove an entry, making sure to maintain the invariants."""
 
251
        val = FIFOCache._remove(self, key)
 
252
        self._value_size -= self._compute_size(val)
 
253
        return val
 
254
 
 
255
    def resize(self, max_size, after_cleanup_size=None):
 
256
        """Increase/decrease the amount of cached data.
 
257
 
 
258
        :param max_size: The maximum number of bytes to cache.
 
259
        :param after_cleanup_size: After cleanup, we should have at most this
 
260
            many bytes cached. This defaults to 80% of max_size.
 
261
        """
 
262
        FIFOCache.resize(self, max_size)
 
263
        self._max_size = max_size
 
264
        if after_cleanup_size is None:
 
265
            self._after_cleanup_size = max_size * 8 / 10
 
266
        else:
 
267
            self._after_cleanup_size = min(max_size, after_cleanup_size)
 
268
        if self._value_size > self._max_size:
 
269
            self.cleanup()
 
270