Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# This file is part of Hypothesis, which may be found at
2# https://github.com/HypothesisWorks/hypothesis/
3#
4# Copyright the Hypothesis Authors.
5# Individual contributors are listed in AUTHORS.rst and the git log.
6#
7# This Source Code Form is subject to the terms of the Mozilla Public License,
8# v. 2.0. If a copy of the MPL was not distributed with this file, You can
9# obtain one at https://mozilla.org/MPL/2.0/.
11import threading
12from collections import OrderedDict
13from dataclasses import dataclass
14from typing import Any, Generic, TypeVar
16from hypothesis.errors import InvalidArgument
18K = TypeVar("K")
19V = TypeVar("V")
22@dataclass(slots=True, frozen=False)
23class Entry(Generic[K, V]):
24 key: K
25 value: V
26 score: int
27 pins: int = 0
29 @property
30 def sort_key(self) -> tuple[int, ...]:
31 if self.pins == 0:
32 # Unpinned entries are sorted by score.
33 return (0, self.score)
34 else:
35 # Pinned entries sort after unpinned ones. Beyond that, we don't
36 # worry about their relative order.
37 return (1,)
40class GenericCache(Generic[K, V]):
41 """Generic supertype for cache implementations.
43 Defines a dict-like mapping with a maximum size, where as well as mapping
44 to a value, each key also maps to a score. When a write would cause the
45 dict to exceed its maximum size, it first evicts the existing key with
46 the smallest score, then adds the new key to the map. If due to pinning
47 no key can be evicted, ValueError is raised.
49 A key has the following lifecycle:
51 1. key is written for the first time, the key is given the score
52 self.new_entry(key, value)
53 2. whenever an existing key is read or written, self.on_access(key, value,
54 score) is called. This returns a new score for the key.
55 3. After a key is evicted, self.on_evict(key, value, score) is called.
57 The cache will be in a valid state in all of these cases.
59 Implementations are expected to implement new_entry and optionally
60 on_access and on_evict to implement a specific scoring strategy.
61 """
63 __slots__ = ("_threadlocal", "max_size")
65 def __init__(self, max_size: int):
66 if max_size <= 0:
67 raise InvalidArgument("Cache size must be at least one.")
69 self.max_size = max_size
71 # Implementation: We store a binary heap of Entry objects in self.data,
72 # with the heap property requiring that a parent's score is <= that of
73 # its children. keys_to_index then maps keys to their index in the
74 # heap. We keep these two in sync automatically - the heap is never
75 # reordered without updating the index.
76 self._threadlocal = threading.local()
78 @property
79 def keys_to_indices(self) -> dict[K, int]:
80 try:
81 return self._threadlocal.keys_to_indices
82 except AttributeError:
83 self._threadlocal.keys_to_indices = {}
84 return self._threadlocal.keys_to_indices
86 @property
87 def data(self) -> list[Entry[K, V]]:
88 try:
89 return self._threadlocal.data
90 except AttributeError:
91 self._threadlocal.data = []
92 return self._threadlocal.data
94 def _repair_if_interrupted(self) -> None:
95 """An exception raised partway through a mutation - most plausibly a
96 RecursionError below a deeply recursive strategy - can leave the heap
97 property or the key index broken. We note the interruption and
98 rebuild both here, on the next operation, from the surviving entries
99 (a sorted list is heap-ordered), rather than attempting a repair at
100 the moment the stack overflowed.
101 """
102 if getattr(self._threadlocal, "interrupted", False):
103 self.data.sort(key=lambda e: e.sort_key)
104 self.keys_to_indices.clear()
105 for i, e in enumerate(self.data):
106 self.keys_to_indices[e.key] = i
107 self._threadlocal.interrupted = False
109 def __len__(self) -> int:
110 self._repair_if_interrupted()
111 assert len(self.keys_to_indices) == len(self.data)
112 return len(self.data)
114 def __contains__(self, key: K) -> bool:
115 self._repair_if_interrupted()
116 return key in self.keys_to_indices
118 def __getitem__(self, key: K) -> V:
119 self._repair_if_interrupted()
120 i = self.keys_to_indices[key]
121 result = self.data[i]
122 try:
123 self.__entry_was_accessed(i)
124 except BaseException:
125 self._threadlocal.interrupted = True
126 raise
127 return result.value
129 def __setitem__(self, key: K, value: V) -> None:
130 self._repair_if_interrupted()
131 evicted = None
132 try:
133 i = self.keys_to_indices[key]
134 except KeyError:
135 entry = Entry(key, value, self.new_entry(key, value))
136 if len(self.data) >= self.max_size:
137 evicted = self.data[0]
138 if evicted.pins > 0:
139 raise ValueError(
140 "Cannot increase size of cache where all keys have been pinned."
141 ) from None
142 del self.keys_to_indices[evicted.key]
143 i = 0
144 self.data[0] = entry
145 else:
146 i = len(self.data)
147 self.data.append(entry)
148 try:
149 self.keys_to_indices[key] = i
150 self.__balance(i)
151 except BaseException:
152 self._threadlocal.interrupted = True
153 raise
154 else:
155 entry = self.data[i]
156 assert entry.key == key
157 entry.value = value
158 try:
159 self.__entry_was_accessed(i)
160 except BaseException:
161 self._threadlocal.interrupted = True
162 raise
164 if evicted is not None:
165 if self.data[0] is not entry:
166 assert evicted.sort_key <= self.data[0].sort_key
167 self.on_evict(evicted.key, evicted.value, evicted.score)
169 def __iter__(self):
170 return iter(self.keys_to_indices)
172 def pin(self, key: K, value: V) -> None:
173 """Mark ``key`` as pinned (with the given value). That is, it may not
174 be evicted until ``unpin(key)`` has been called. The same key may be
175 pinned multiple times, possibly changing its value, and will not be
176 unpinned until the same number of calls to unpin have been made.
177 """
178 self[key] = value
180 i = self.keys_to_indices[key]
181 entry = self.data[i]
182 entry.pins += 1
183 if entry.pins == 1:
184 try:
185 self.__balance(i)
186 except BaseException:
187 self._threadlocal.interrupted = True
188 raise
190 def unpin(self, key: K) -> None:
191 """Undo one previous call to ``pin(key)``. The value stays the same.
192 Once all calls are undone this key may be evicted as normal."""
193 self._repair_if_interrupted()
194 i = self.keys_to_indices[key]
195 entry = self.data[i]
196 if entry.pins == 0:
197 raise ValueError(f"Key {key!r} has not been pinned")
198 entry.pins -= 1
199 if entry.pins == 0:
200 try:
201 self.__balance(i)
202 except BaseException:
203 self._threadlocal.interrupted = True
204 raise
206 def is_pinned(self, key: K) -> bool:
207 """Returns True if the key is currently pinned."""
208 i = self.keys_to_indices[key]
209 return self.data[i].pins > 0
211 def clear(self) -> None:
212 """Remove all keys, regardless of their pinned status."""
213 del self.data[:]
214 self.keys_to_indices.clear()
216 def __repr__(self) -> str:
217 return "{" + ", ".join(f"{e.key!r}: {e.value!r}" for e in self.data) + "}"
219 def new_entry(self, key: K, value: V) -> int:
220 """Called when a key is written that does not currently appear in the
221 map.
223 Returns the score to associate with the key.
224 """
225 raise NotImplementedError
227 def on_access(self, key: K, value: V, score: Any) -> Any:
228 """Called every time a key that is already in the map is read or
229 written.
231 Returns the new score for the key.
232 """
233 return score
235 def on_evict(self, key: K, value: V, score: Any) -> Any:
236 """Called after a key has been evicted, with the score it had had at
237 the point of eviction."""
239 def check_valid(self) -> None:
240 """Debugging method for use in tests.
242 Asserts that all of the cache's invariants hold. When everything
243 is working correctly this should be an expensive no-op.
244 """
245 assert len(self.keys_to_indices) == len(self.data)
246 for i, e in enumerate(self.data):
247 assert self.keys_to_indices[e.key] == i
248 for j in [i * 2 + 1, i * 2 + 2]:
249 if j < len(self.data):
250 assert e.sort_key <= self.data[j].sort_key, self.data
252 def __entry_was_accessed(self, i: int) -> None:
253 entry = self.data[i]
254 new_score = self.on_access(entry.key, entry.value, entry.score)
255 if new_score != entry.score:
256 entry.score = new_score
257 # changing the score of a pinned entry cannot unbalance the heap, as
258 # we place all pinned entries after unpinned ones, regardless of score.
259 if entry.pins == 0:
260 self.__balance(i)
262 def __swap(self, i: int, j: int) -> None:
263 assert i < j
264 assert self.data[j].sort_key < self.data[i].sort_key
265 self.data[i], self.data[j] = self.data[j], self.data[i]
266 self.keys_to_indices[self.data[i].key] = i
267 self.keys_to_indices[self.data[j].key] = j
269 def __balance(self, i: int) -> None:
270 """When we have made a modification to the heap such that
271 the heap property has been violated locally around i but previously
272 held for all other indexes (and no other values have been modified),
273 this fixes the heap so that the heap property holds everywhere."""
274 # bubble up (if score is too low for current position)
275 while (parent := (i - 1) // 2) >= 0:
276 if self.__out_of_order(parent, i):
277 self.__swap(parent, i)
278 i = parent
279 else:
280 break
281 # or bubble down (if score is too high for current position)
282 while children := [j for j in (2 * i + 1, 2 * i + 2) if j < len(self.data)]:
283 smallest_child = min(children, key=lambda j: self.data[j].sort_key)
284 if self.__out_of_order(i, smallest_child):
285 self.__swap(i, smallest_child)
286 i = smallest_child
287 else:
288 break
290 def __out_of_order(self, i: int, j: int) -> bool:
291 """Returns True if the indices i, j are in the wrong order.
293 i must be the parent of j.
294 """
295 assert i == (j - 1) // 2
296 return self.data[j].sort_key < self.data[i].sort_key
299class LRUReusedCache(GenericCache[K, V]):
300 """The only concrete implementation of GenericCache we use outside of tests
301 currently.
303 Adopts a modified least-recently used eviction policy: It evicts the key
304 that has been used least recently, but it will always preferentially evict
305 keys that have never been accessed after insertion. Among keys that have been
306 accessed, it ignores the number of accesses.
308 This retains most of the benefits of an LRU cache, but adds an element of
309 scan-resistance to the process: If we end up scanning through a large
310 number of keys without reusing them, this does not evict the existing
311 entries in preference for the new ones.
312 """
314 __slots__ = ("__tick",)
316 def __init__(self, max_size: int):
317 super().__init__(max_size)
318 self.__tick: int = 0
320 def tick(self) -> int:
321 self.__tick += 1
322 return self.__tick
324 def new_entry(self, key: K, value: V) -> Any:
325 return (1, self.tick())
327 def on_access(self, key: K, value: V, score: Any) -> Any:
328 return (2, self.tick())
331class LRUCache(Generic[K, V]):
332 """
333 This is a drop-in replacement for a GenericCache (despite the lack of inheritance)
334 in performance critical environments. It turns out that GenericCache's heap
335 balancing for arbitrary scores can be quite expensive compared to the doubly
336 linked list approach of lru_cache or OrderedDict.
338 This class is a pure LRU and does not provide any sort of affininty towards
339 the number of accesses beyond recency. If soft-pinning entries which have been
340 accessed at least once is important, use LRUReusedCache.
341 """
343 # Here are some nice performance references for lru_cache vs OrderedDict:
344 # https://github.com/python/cpython/issues/72426#issuecomment-1093727671
345 # https://discuss.python.org/t/simplify-lru-cache/18192/6
346 #
347 # We use OrderedDict here because it is unclear to me we can provide the same
348 # api as GenericCache using @lru_cache without messing with lru_cache internals.
349 #
350 # Anecdotally, OrderedDict seems quite competitive with lru_cache, but perhaps
351 # that is localized to our access patterns.
353 def __init__(self, max_size: int) -> None:
354 assert max_size > 0
355 self.max_size = max_size
356 self._threadlocal = threading.local()
358 @property
359 def cache(self) -> OrderedDict[K, V]:
360 try:
361 return self._threadlocal.cache
362 except AttributeError:
363 self._threadlocal.cache = OrderedDict()
364 return self._threadlocal.cache
366 def __setitem__(self, key: K, value: V) -> None:
367 self.cache[key] = value
368 self.cache.move_to_end(key)
370 while len(self.cache) > self.max_size:
371 self.cache.popitem(last=False)
373 def __getitem__(self, key: K) -> V:
374 val = self.cache[key]
375 self.cache.move_to_end(key)
376 return val
378 def __iter__(self):
379 return iter(self.cache)
381 def __len__(self) -> int:
382 return len(self.cache)
384 def __contains__(self, key: K) -> bool:
385 return key in self.cache
387 # implement GenericCache interface, for tests
388 def check_valid(self) -> None:
389 pass