LRUStoreCache: cache "contains" by contains checks (zarr-developers#1499)

madsbk · web-flow · commit 7017ce16954f · 2023-08-16T15:45:13.000-04:00
* cache "contains" by contains checks

* updated the test_cache_keys counts

* release notes
diff --git a/docs/release.rst b/docs/release.rst
@@ -27,6 +27,9 @@ Maintenance
 * Add ``docs`` requirements to ``pyproject.toml``
   By :user:`John A. Kirkham <jakirkham>` :issue:`1494`.
 
+* Fixed caching issue in ``LRUStoreCache``.
+  By :user:`Mads R. B. Kristensen <madsbk>` :issue:`1499`.
+
 .. _release_2.16.0:
 
 2.16.0
diff --git a/zarr/_storage/v3.py b/zarr/_storage/v3.py
@@ -509,7 +509,7 @@ def __init__(self, store, max_size: int):
         self._max_size = max_size
         self._current_size = 0
         self._keys_cache = None
-        self._contains_cache = None
+        self._contains_cache = {}
         self._listdir_cache: Dict[Path, Any] = dict()
         self._values_cache: Dict[Path, Any] = OrderedDict()
         self._mutex = Lock()
diff --git a/zarr/storage.py b/zarr/storage.py
@@ -2393,7 +2393,7 @@ def __init__(self, store: StoreLike, max_size: int):
         self._max_size = max_size
         self._current_size = 0
         self._keys_cache = None
-        self._contains_cache = None
+        self._contains_cache: Dict[Any, Any] = {}
         self._listdir_cache: Dict[Path, Any] = dict()
         self._values_cache: Dict[Path, Any] = OrderedDict()
         self._mutex = Lock()
@@ -2434,9 +2434,9 @@ def __iter__(self):
 
     def __contains__(self, key):
         with self._mutex:
-            if self._contains_cache is None:
-                self._contains_cache = set(self._keys())
-            return key in self._contains_cache
+            if key not in self._contains_cache:
+                self._contains_cache[key] = key in self._store
+            return self._contains_cache[key]
 
     def clear(self):
         self._store.clear()
@@ -2506,7 +2506,7 @@ def invalidate_keys(self):
 
     def _invalidate_keys(self):
         self._keys_cache = None
-        self._contains_cache = None
+        self._contains_cache.clear()
         self._listdir_cache.clear()
 
     def _invalidate_value(self, key):
diff --git a/zarr/tests/test_storage.py b/zarr/tests/test_storage.py
@@ -2196,7 +2196,10 @@ def test_cache_keys(self):
         assert keys == sorted(cache.keys())
         assert 1 == store.counter["keys"]
         assert foo_key in cache
-        assert 0 == store.counter["__contains__", foo_key]
+        assert 1 == store.counter["__contains__", foo_key]
+        # the next check for `foo_key` is cached
+        assert foo_key in cache
+        assert 1 == store.counter["__contains__", foo_key]
         assert keys == sorted(cache)
         assert 0 == store.counter["__iter__"]
         assert 1 == store.counter["keys"]
@@ -2215,23 +2218,23 @@ def test_cache_keys(self):
         keys = sorted(cache.keys())
         assert keys == [bar_key, baz_key, foo_key]
         assert 3 == store.counter["keys"]
-        assert 0 == store.counter["__contains__", foo_key]
+        assert 1 == store.counter["__contains__", foo_key]
         assert 0 == store.counter["__iter__"]
         cache.invalidate_keys()
         keys = sorted(cache)
         assert keys == [bar_key, baz_key, foo_key]
         assert 4 == store.counter["keys"]
-        assert 0 == store.counter["__contains__", foo_key]
+        assert 1 == store.counter["__contains__", foo_key]
         assert 0 == store.counter["__iter__"]
         cache.invalidate_keys()
         assert foo_key in cache
-        assert 5 == store.counter["keys"]
-        assert 0 == store.counter["__contains__", foo_key]
+        assert 4 == store.counter["keys"]
+        assert 2 == store.counter["__contains__", foo_key]
         assert 0 == store.counter["__iter__"]
 
         # check these would get counted if called directly
         assert foo_key in store
-        assert 1 == store.counter["__contains__", foo_key]
+        assert 3 == store.counter["__contains__", foo_key]
         assert keys == sorted(store)
         assert 1 == store.counter["__iter__"]
 
