Skip to content

fix(cache_store): stop CacheStore serving data the source no longer holds #298

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
d-v-b wants to merge 2 commits into
base: main
Choose a base branch
from
Open

fix(cache_store): stop CacheStore serving data the source no longer holds #298

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
eff15e0
fix(cache_store): stop CacheStore serving data the source no longer h…
d-v-b Aug 14, 2026
1a12bd9
chore: name the changelog fragment after the PR number
d-v-b Aug 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions changes/298.bugfix.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
Fixed several ways `zarr.experimental.cache_store.CacheStore` could serve data the source store no longer holds. `set_if_not_exists`, `_set_many`, `delete_dir` and `clear` were not overridden, so they were forwarded straight to the source and the cache went on serving the superseded value — most visibly, creating an array with `overwrite=True` through a `CacheStore` left the *old* array readable through the cache. Each now invalidates the keys it affects. `delete` also failed to reclaim its entry's bytes, so every delete permanently inflated `current_size` and ate into the `max_size` budget, and a value too large to cache was left in the backing cache untracked — uncounted against `max_size`, never eviction-eligible, but still served. Finally, every backing-store mutation is now published in the same locked section as its tracking update, so a concurrent write can no longer be observed half-applied.

`CacheStore.open()` now works: it inherited `WrapperStore.open()`, which builds the wrapped store from a `store_cls` argument and so could not supply the required `cache_store`.

`CacheStore` now requires its `cache_store` to support listing (in addition to deletes), since prefix deletions need it. This is checked in the constructor rather than failing later mid-write.
234 changes: 159 additions & 75 deletions src/zarr/experimental/cache_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,8 @@
logger = logging.getLogger(__name__)

if TYPE_CHECKING:
from collections.abc import Iterable

from zarr.core.buffer.core import Buffer, BufferPrototype

# A cache entry identifier. Plain ``str`` for full-key entries that live in
Expand Down Expand Up @@ -53,7 +55,8 @@ class CacheStore(WrapperStore[Store]):
The underlying store to wrap with caching
cache_store : Store
The store to use for caching (can be any Store implementation that
supports deletes)
supports deletes and listing). Listing is required so that ``delete_dir``
and ``clear`` can drop the cached keys they invalidate.
max_age_seconds : int or "infinity", optional
Maximum age of cached entries in seconds. The string "infinity" means
entries never expire. Default is "infinity".
Expand Down Expand Up @@ -112,6 +115,14 @@ def __init__(
)
raise ValueError(msg)

if not cache_store.supports_listing:
msg = (
f"The provided cache store {cache_store} does not support listing. "
"The cache_store must support listing so that prefix deletions "
"(delete_dir, clear) can invalidate the keys they remove."
)
raise ValueError(msg)

self._cache = cache_store
# Validate and set max_age_seconds
if isinstance(max_age_seconds, str):
Expand All @@ -124,6 +135,20 @@ def __init__(
self.cache_set_data = cache_set_data
self._state = _CacheState()

@classmethod
async def open(cls, *args: Any, **kwargs: Any) -> Self:
"""Create and open a ``CacheStore``.

``WrapperStore.open`` builds the wrapped store from a ``store_cls``
argument and then wraps it, which cannot supply the required
``cache_store``; inheriting it makes ``CacheStore.open`` raise. Take
``Store.open``'s behaviour instead: construct from the arguments this
class actually accepts, then open.
"""
store = cls(*args, **kwargs)
await store._open()
return store

def _with_store(self, store: Store) -> Self:
# Cannot support this operation because it would share a cache, but have a new store
# So cache keys would conflict
Expand Down Expand Up @@ -175,22 +200,11 @@ async def _evict_key(self, entry_key: _CacheEntryKey) -> None:
For ``(str, ByteRequest)`` keys the entry is removed from the
in-memory range cache.
"""
key_size = self._state.key_sizes.get(entry_key, 0)

if isinstance(entry_key, str):
await self._cache.delete(entry_key)
self._remove_from_tracking(entry_key)
else:
base_key, byte_range = entry_key
per_key = self._state.range_cache.get(base_key)
if per_key is not None:
per_key.pop(byte_range, None)
if not per_key:
del self._state.range_cache[base_key]

self._state.cache_order.pop(entry_key, None)
self._state.key_insert_times.pop(entry_key, None)
self._state.key_sizes.pop(entry_key, None)
self._state.current_size = max(0, self._state.current_size - key_size)
self._drop_range(*entry_key)
self._state.evictions += 1

async def _track_entry(self, entry_key: _CacheEntryKey, value: Buffer) -> bool:
Expand All @@ -200,45 +214,60 @@ async def _track_entry(self, entry_key: _CacheEntryKey, value: Buffer) -> bool:
exceeds ``max_size`` and was skipped. Callers should roll back any
data they already stored when this returns ``False``.

This method holds the lock for the entire operation to ensure atomicity.
Must be called while holding ``self._state.lock``, which the callers hold
across the backing-store write as well, so the value and its tracking
entry are published together.
"""
value_size = len(value)

# Drop any prior tracking for this entry first, reclaiming its bytes.
# This also removes it from the eviction candidates, so
# ``_accommodate_value`` cannot select the very entry being (re)tracked
# -- which would double-subtract its size, stop the eviction loop early,
# and delete the value the caller just wrote to the backing store.
self._remove_from_tracking(entry_key)

# Check if value exceeds max size
if self.max_size is not None and value_size > self.max_size:
return False

async with self._state.lock:
# If key already exists, subtract old size first
if entry_key in self._state.key_sizes:
old_size = self._state.key_sizes[entry_key]
self._state.current_size -= old_size

# Make room for the new value
await self._accommodate_value(value_size)

# Update tracking atomically
self._state.cache_order[entry_key] = None
self._state.current_size += value_size
self._state.key_sizes[entry_key] = value_size
self._state.key_insert_times[entry_key] = time.monotonic()

# Make room for the new value, then track it (appended as most-recent).
await self._accommodate_value(value_size)
self._state.cache_order[entry_key] = None
self._state.current_size += value_size
self._state.key_sizes[entry_key] = value_size
self._state.key_insert_times[entry_key] = time.monotonic()
return True

async def _update_access_order(self, entry_key: _CacheEntryKey) -> None:
"""Update the access order for LRU tracking."""
if entry_key in self._state.cache_order:
async with self._state.lock:
async with self._state.lock:
# Re-check membership under the lock: the entry may have been evicted
# by a concurrent operation between the call and acquiring the lock.
if entry_key in self._state.cache_order:
self._state.cache_order.move_to_end(entry_key)

def _remove_from_tracking(self, entry_key: _CacheEntryKey) -> None:
"""Remove an entry from all tracking structures.
"""Remove an entry from all tracking structures, reclaiming its bytes.

Must be called while holding self._state.lock.
"""
self._state.cache_order.pop(entry_key, None)
self._state.key_insert_times.pop(entry_key, None)
self._state.key_sizes.pop(entry_key, None)
entry_size = self._state.key_sizes.pop(entry_key, 0)
self._state.current_size = max(0, self._state.current_size - entry_size)

def _drop_range(self, key: str, byte_range: ByteRequest) -> None:
"""Remove one byte-range entry from the range cache and tracking.

Must be called while holding self._state.lock.
"""
per_key = self._state.range_cache.get(key)
if per_key is not None:
per_key.pop(byte_range, None)
if not per_key:
del self._state.range_cache[key]
self._remove_from_tracking((key, byte_range))

def _invalidate_range_entries(self, key: str) -> None:
"""Remove all byte-range entries for *key* from the range cache and tracking.
Expand All @@ -248,11 +277,7 @@ def _invalidate_range_entries(self, key: str) -> None:
per_key = self._state.range_cache.pop(key, None)
if per_key is not None:
for byte_range in per_key:
entry_key: _CacheEntryKey = (key, byte_range)
entry_size = self._state.key_sizes.pop(entry_key, 0)
self._state.cache_order.pop(entry_key, None)
self._state.key_insert_times.pop(entry_key, None)
self._state.current_size = max(0, self._state.current_size - entry_size)
self._remove_from_tracking((key, byte_range))

# ------------------------------------------------------------------
# get helpers
Expand All @@ -261,36 +286,34 @@ def _invalidate_range_entries(self, key: str) -> None:
async def _cache_miss(
self, key: str, byte_range: ByteRequest | None, result: Buffer | None
) -> None:
"""Handle a cache miss by storing or cleaning up after a source-store fetch."""
"""Handle a cache miss by storing or cleaning up after a source-store fetch.

Each branch takes ``self._state.lock`` for the whole of its backing-store
mutation *and* the matching tracking mutation, so the two can never be
observed out of step by a concurrent operation.
"""
if result is None:
if byte_range is None:
await self._cache.delete(key)
async with self._state.lock:
await self._cache.delete(key)
self._remove_from_tracking(key)
else:
entry_key: _CacheEntryKey = (key, byte_range)
async with self._state.lock:
per_key = self._state.range_cache.get(key)
if per_key is not None:
per_key.pop(byte_range, None)
if not per_key:
del self._state.range_cache[key]
self._remove_from_tracking(entry_key)
self._drop_range(key, byte_range)
else:
if byte_range is None:
await self._cache.set(key, result)
await self._track_entry(key, result)
async with self._state.lock:
await self._cache.set(key, result)
if not await self._track_entry(key, result):
# Value too large for the cache — roll back so the backing
# cache holds no untracked (uncounted, unevictable) orphan.
await self._cache.delete(key)
else:
entry_key = (key, byte_range)
self._state.range_cache.setdefault(key, {})[byte_range] = result
tracked = await self._track_entry(entry_key, result)
if not tracked:
# Value too large for the cache — roll back the insertion
per_key = self._state.range_cache.get(key)
if per_key is not None:
per_key.pop(byte_range, None)
if not per_key:
del self._state.range_cache[key]
async with self._state.lock:
self._state.range_cache.setdefault(key, {})[byte_range] = result
if not await self._track_entry((key, byte_range), result):
# Value too large for the cache — roll back the insertion
self._drop_range(key, byte_range)

async def _get_try_cache(
self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None
Expand Down Expand Up @@ -380,17 +403,49 @@ async def set(self, key: str, value: Buffer) -> None:
The data to store
"""
await super().set(key, value)
# Invalidate all cached byte-range entries (source data changed)
async with self._state.lock:
# The value just written supersedes any cached byte ranges for the key.
self._invalidate_range_entries(key)
if self.cache_set_data:
await self._cache.set(key, value)
await self._track_entry(key, value)
else:
await self._cache.delete(key)
async with self._state.lock:
if self.cache_set_data:
await self._cache.set(key, value)
if not await self._track_entry(key, value):
# Value too large for the cache — roll back so the backing cache
# holds no untracked (uncounted, unevictable) orphan.
await self._cache.delete(key)
else:
await self._cache.delete(key)
self._remove_from_tracking(key)

async def set_if_not_exists(self, key: str, value: Buffer) -> None:
"""
Store data only if the key does not already exist in the source store.

Parameters
----------
key : str
The key to store under
value : Buffer
The data to store
"""
await super().set_if_not_exists(key, value)
# Whether or not the write happened, any cached state for this key may now
# be stale (we may have just written a new value, or it already existed).
# Drop all of it so the next read reflects the source. Invalidating
# unconditionally is always safe. We do not populate the cache here: there
# is no guaranteed-fresh value to store (the write may have been a no-op).
async with self._state.lock:
self._invalidate_range_entries(key)
await self._cache.delete(key)
self._remove_from_tracking(key)

async def _set_many(self, values: Iterable[tuple[str, Buffer]]) -> None:
"""Bulk writes routed through ``self.set``, so they invalidate the cache.

``WrapperStore._set_many`` forwards straight to the source store, so
without this override a key written in bulk keeps its stale cached value.
"""
await Store._set_many(self, values)

async def delete(self, key: str) -> None:
"""
Delete data from both the underlying store and cache.
Expand All @@ -401,13 +456,42 @@ async def delete(self, key: str) -> None:
The key to delete
"""
await super().delete(key)
# Invalidate all cached byte-range entries
async with self._state.lock:
self._invalidate_range_entries(key)
await self._cache.delete(key)
async with self._state.lock:
await self._cache.delete(key)
self._remove_from_tracking(key)

async def delete_dir(self, prefix: str) -> None:
"""
Delete a prefix from the underlying store and drop its cached keys.

``WrapperStore.delete_dir`` delegates straight to the source store, so
without this override no ``delete`` runs for the keys under *prefix* and
the cache keeps serving them after e.g. an ``overwrite=True`` array
creation.
"""
await super().delete_dir(prefix)
if prefix != "" and not prefix.endswith("/"):
prefix += "/"
async with self._state.lock:
await self._cache.delete_dir(prefix)
for base_key in [k for k in self._state.range_cache if k.startswith(prefix)]:
self._invalidate_range_entries(base_key)
for entry_key in [
k for k in self._state.cache_order if isinstance(k, str) and k.startswith(prefix)
]:
self._remove_from_tracking(entry_key)

async def clear(self) -> None:
"""Clear the underlying store, and with it everything cached from it.

Same bypass as ``delete_dir``: ``WrapperStore.clear`` delegates to the
source store, which would leave the cache serving values for keys that no
longer exist anywhere.
"""
await super().clear()
await self.clear_cache()

def cache_info(self) -> dict[str, Any]:
"""Return information about the cache state."""
return {
Expand Down Expand Up @@ -436,12 +520,12 @@ def cache_stats(self) -> dict[str, Any]:

async def clear_cache(self) -> None:
"""Clear all cached data and tracking information."""
# Clear the cache store if it supports clear
if hasattr(self._cache, "clear"):
await self._cache.clear()

# Reset tracking
# The backing-store wipe and the tracking reset share one locked section:
# a ``set`` landing between them would otherwise leave a value in the
# backing cache with no tracking entry — uncounted against ``max_size``,
# never eviction-eligible, and served as a hit indefinitely.
async with self._state.lock:
await self._cache.clear()
self._state.key_insert_times.clear()
self._state.cache_order.clear()
self._state.key_sizes.clear()
Expand Down
Loading
Loading