~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/fifo_cache.py

  • Committer: John Arbash Meinel
  • Date: 2008-12-10 17:22:13 UTC
  • mfrom: (3890 +trunk)
  • mto: This revision was merged to the branch mainline in revision 3912.
  • Revision ID: john@arbash-meinel.com-20081210172213-h2b0auuil3qaz28u
Merge bzr.dev 3890, bring in the FIFOCache.

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