Merge pull request #4064 from tybug/ir-fast-cache

Use a more performant `LRUCache` for ir kwargs caches

Author: Zac-HD

hypothesis-python/RELEASE.rst

@@ -0,0 +1,3 @@
+RELEASE_TYPE: patch
+
+This patch improves our example generation performance by adjusting our internal cache implementation.

hypothesis-python/src/hypothesis/internal/cache.py

@@ -9,6 +9,7 @@
 # obtain one at https://mozilla.org/MPL/2.0/.
 
 import threading
+from collections import OrderedDict
 
 import attr
 
@@ -282,3 +283,64 @@ def new_entry(self, key, value):
 
     def on_access(self, key, value, score):
         return (2, self.tick())
+
+
+class LRUCache:
+    """
+    This is a drop-in replacement for a GenericCache (despite the lack of inheritance)
+    in performance critical environments. It turns out that GenericCache's heap
+    balancing for arbitrary scores can be quite expensive compared to the doubly
+    linked list approach of lru_cache or OrderedDict.
+
+    This class is a pure LRU and does not provide any sort of affininty towards
+    the number of accesses beyond recency. If soft-pinning entries which have been
+    accessed at least once is important, use LRUReusedCache.
+    """
+
+    # Here are some nice performance references for lru_cache vs OrderedDict:
+    # https://github.com/python/cpython/issues/72426#issuecomment-1093727671
+    # https://discuss.python.org/t/simplify-lru-cache/18192/6
+    #
+    # We use OrderedDict here because it is unclear to me we can provide the same
+    # api as GenericCache without messing with @lru_cache internals.
+    #
+    # Anecdotally, OrderedDict seems quite competitive with lru_cache, but perhaps
+    # that is localized to our access patterns.
+
+    def __init__(self, max_size):
+        assert max_size > 0
+        self.max_size = max_size
+        self._threadlocal = threading.local()
+
+    @property
+    def cache(self):
+        try:
+            return self._threadlocal.cache
+        except AttributeError:
+            self._threadlocal.cache = OrderedDict()
+            return self._threadlocal.cache
+
+    def __setitem__(self, key, value):
+        self.cache[key] = value
+        self.cache.move_to_end(key)
+
+        while len(self.cache) > self.max_size:
+            self.cache.popitem(last=False)
+
+    def __getitem__(self, key):
+        val = self.cache[key]
+        self.cache.move_to_end(key)
+        return val
+
+    def __iter__(self):
+        return iter(self.cache)
+
+    def __len__(self):
+        return len(self.cache)
+
+    def __contains__(self, key):
+        return key in self.cache
+
+    # implement GenericCache interface, for tests
+    def check_valid(self):
+        pass

hypothesis-python/src/hypothesis/internal/conjecture/data.py

@@ -41,7 +41,7 @@
 import attr
 
 from hypothesis.errors import Frozen, InvalidArgument, StopTest
-from hypothesis.internal.cache import LRUReusedCache
+from hypothesis.internal.cache import LRUCache
 from hypothesis.internal.compat import add_note, floor, int_from_bytes, int_to_bytes
 from hypothesis.internal.conjecture.floats import float_to_lex, lex_to_float
 from hypothesis.internal.conjecture.junkdrawer import (
@@ -200,9 +200,11 @@ def structural_coverage(label: int) -> StructuralCoverageTag:
 NASTY_FLOATS = list(map(float, NASTY_FLOATS))
 NASTY_FLOATS.extend([-x for x in NASTY_FLOATS])
 
-FLOAT_INIT_LOGIC_CACHE = LRUReusedCache(4096)
-
-POOLED_KWARGS_CACHE = LRUReusedCache(4096)
+# These caches, especially the kwargs cache, can be quite hot and so we prefer
+# LRUCache over LRUReusedCache for performance. We lose scan resistance, but
+# that's probably fine here.
+FLOAT_INIT_LOGIC_CACHE = LRUCache(4096)
+POOLED_KWARGS_CACHE = LRUCache(4096)
 
 DRAW_STRING_DEFAULT_MAX_SIZE = 10**10  # "arbitrarily large"
 

hypothesis-python/tests/cover/test_cache_implementation.py

@@ -24,12 +24,12 @@
     strategies as st,
 )
 from hypothesis.errors import InvalidArgument
-from hypothesis.internal.cache import GenericCache, LRUReusedCache
+from hypothesis.internal.cache import GenericCache, LRUCache, LRUReusedCache
 
 from tests.common.utils import skipif_emscripten
 
 
-class LRUCache(GenericCache):
+class LRUCacheAlternative(GenericCache):
     __slots__ = ("__tick",)
 
     def __init__(self, max_size):
@@ -88,7 +88,8 @@ def on_access(self, key, value, score):
 
 
 @pytest.mark.parametrize(
-    "implementation", [LRUCache, LFUCache, LRUReusedCache, ValueScored, RandomCache]
+    "implementation",
+    [LRUCache, LFUCache, LRUReusedCache, ValueScored, RandomCache, LRUCacheAlternative],
 )
 @example(writes=[(0, 0), (3, 0), (1, 0), (2, 0), (2, 0), (1, 0)], size=4)
 @example(writes=[(0, 0)], size=1)
@@ -308,6 +309,15 @@ def test_iterates_over_remaining_keys():
     assert sorted(cache) == [1, 2]
 
 
+def test_lru_cache_is_actually_lru():
+    cache = LRUCache(2)
+    cache[1] = 1  # [1]
+    cache[2] = 2  # [1, 2]
+    cache[1]  # [2, 1]
+    cache[3] = 2  # [2, 1, 3] -> drop least recently used -> [1, 3]
+    assert list(cache) == [1, 3]
+
+
 @skipif_emscripten
 def test_cache_is_threadsafe_issue_2433_regression():
     errors = []