Support with read only in wrappers (#3700)

* Support with_read_only in LoggingStore and LatencyStore.

Fixes #3699

* Support CacheStore.with_read_only

* Add entry to changes/

* use a dataclass for mutable cache state

---------

Co-authored-by: Davis Bennett <davis.v.bennett@gmail.com>

Author: BorisTheBrave

changes/3700.bugfix.md

@@ -0,0 +1 @@
+CacheStore, LoggingStore and LatencyStore now support with_read_only.

src/zarr/experimental/cache_store.py

@@ -4,7 +4,8 @@
 import logging
 import time
 from collections import OrderedDict
-from typing import TYPE_CHECKING, Any, Literal
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Literal, Self
 
 from zarr.abc.store import ByteRequest, Store
 from zarr.storage._wrapper import WrapperStore
@@ -15,6 +16,18 @@
     from zarr.core.buffer.core import Buffer, BufferPrototype
 
 
+@dataclass(slots=True)
+class _CacheState:
+    cache_order: OrderedDict[str, None] = field(default_factory=OrderedDict)
+    current_size: int = 0
+    key_sizes: dict[str, int] = field(default_factory=dict)
+    lock: asyncio.Lock = field(default_factory=asyncio.Lock)
+    hits: int = 0
+    misses: int = 0
+    evictions: int = 0
+    key_insert_times: dict[str, float] = field(default_factory=dict)
+
+
 class CacheStore(WrapperStore[Store]):
     """
     A dual-store caching implementation for Zarr stores.
@@ -36,9 +49,6 @@ class CacheStore(WrapperStore[Store]):
         Maximum size of the cache in bytes. When exceeded, least recently used
         items are evicted. None means unlimited size. Default is None.
         Note: Individual values larger than max_size will not be cached.
-    key_insert_times : dict[str, float] | None, optional
-        Dictionary to track insertion times (using monotonic time).
-        Primarily for internal use. Default is None (creates new dict).
     cache_set_data : bool, optional
         Whether to cache data when it's written to the store. Default is True.
 
@@ -69,15 +79,8 @@ class CacheStore(WrapperStore[Store]):
     _cache: Store
     max_age_seconds: int | Literal["infinity"]
     max_size: int | None
-    key_insert_times: dict[str, float]
     cache_set_data: bool
-    _cache_order: OrderedDict[str, None]  # Track access order for LRU
-    _current_size: int  # Track current cache size
-    _key_sizes: dict[str, int]  # Track size of each cached key
-    _lock: asyncio.Lock
-    _hits: int  # Cache hit counter
-    _misses: int  # Cache miss counter
-    _evictions: int  # Cache eviction counter
+    _state: _CacheState
 
     def __init__(
         self,
@@ -86,7 +89,6 @@ def __init__(
         cache_store: Store,
         max_age_seconds: int | str = "infinity",
         max_size: int | None = None,
-        key_insert_times: dict[str, float] | None = None,
         cache_set_data: bool = True,
     ) -> None:
         super().__init__(store)
@@ -107,18 +109,25 @@ def __init__(
         else:
             self.max_age_seconds = max_age_seconds
         self.max_size = max_size
-        if key_insert_times is None:
-            self.key_insert_times = {}
-        else:
-            self.key_insert_times = key_insert_times
         self.cache_set_data = cache_set_data
-        self._cache_order = OrderedDict()
-        self._current_size = 0
-        self._key_sizes = {}
-        self._lock = asyncio.Lock()
-        self._hits = 0
-        self._misses = 0
-        self._evictions = 0
+        self._state = _CacheState()
+
+    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
+        raise NotImplementedError("CacheStore does not support this operation.")
+
+    def with_read_only(self, read_only: bool = False) -> Self:
+        # Create a new cache store that shares the same cache and mutable state
+        store = type(self)(
+            store=self._store.with_read_only(read_only),
+            cache_store=self._cache,
+            max_age_seconds=self.max_age_seconds,
+            max_size=self.max_size,
+            cache_set_data=self.cache_set_data,
+        )
+        store._state = self._state
+        return store
 
     def _is_key_fresh(self, key: str) -> bool:
         """Check if a cached key is still fresh based on max_age_seconds.
@@ -128,7 +137,7 @@ def _is_key_fresh(self, key: str) -> bool:
         if self.max_age_seconds == "infinity":
             return True
         now = time.monotonic()
-        elapsed = now - self.key_insert_times.get(key, 0)
+        elapsed = now - self._state.key_insert_times.get(key, 0)
         return elapsed < self.max_age_seconds
 
     async def _accommodate_value(self, value_size: int) -> None:
@@ -140,9 +149,9 @@ async def _accommodate_value(self, value_size: int) -> None:
             return
 
         # Remove least recently used items until we have enough space
-        while self._current_size + value_size > self.max_size and self._cache_order:
+        while self._state.current_size + value_size > self.max_size and self._state.cache_order:
             # Get the least recently used key (first in OrderedDict)
-            lru_key = next(iter(self._cache_order))
+            lru_key = next(iter(self._state.cache_order))
             await self._evict_key(lru_key)
 
     async def _evict_key(self, key: str) -> None:
@@ -152,15 +161,15 @@ async def _evict_key(self, key: str) -> None:
         Updates size tracking atomically with deletion.
         """
         try:
-            key_size = self._key_sizes.get(key, 0)
+            key_size = self._state.key_sizes.get(key, 0)
 
             # Delete from cache store
             await self._cache.delete(key)
 
             # Update tracking after successful deletion
             self._remove_from_tracking(key)
-            self._current_size = max(0, self._current_size - key_size)
-            self._evictions += 1
+            self._state.current_size = max(0, self._state.current_size - key_size)
+            self._state.evictions += 1
 
             logger.debug("_evict_key: evicted key %s, freed %d bytes", key, key_size)
         except Exception:
@@ -183,39 +192,39 @@ async def _cache_value(self, key: str, value: Buffer) -> None:
             )
             return
 
-        async with self._lock:
+        async with self._state.lock:
             # If key already exists, subtract old size first
-            if key in self._key_sizes:
-                old_size = self._key_sizes[key]
-                self._current_size -= old_size
+            if key in self._state.key_sizes:
+                old_size = self._state.key_sizes[key]
+                self._state.current_size -= old_size
                 logger.debug("_cache_value: updating existing key %s, old size %d", key, old_size)
 
             # Make room for the new value (this calls _evict_key_locked internally)
             await self._accommodate_value(value_size)
 
             # Update tracking atomically
-            self._cache_order[key] = None  # OrderedDict to track access order
-            self._current_size += value_size
-            self._key_sizes[key] = value_size
-            self.key_insert_times[key] = time.monotonic()
+            self._state.cache_order[key] = None  # OrderedDict to track access order
+            self._state.current_size += value_size
+            self._state.key_sizes[key] = value_size
+            self._state.key_insert_times[key] = time.monotonic()
 
             logger.debug("_cache_value: cached key %s with size %d bytes", key, value_size)
 
     async def _update_access_order(self, key: str) -> None:
         """Update the access order for LRU tracking."""
-        if key in self._cache_order:
-            async with self._lock:
+        if key in self._state.cache_order:
+            async with self._state.lock:
                 # Move to end (most recently used)
-                self._cache_order.move_to_end(key)
+                self._state.cache_order.move_to_end(key)
 
     def _remove_from_tracking(self, key: str) -> None:
         """Remove a key from all tracking structures.
 
-        Must be called while holding self._lock.
+        Must be called while holding self._state.lock.
         """
-        self._cache_order.pop(key, None)
-        self.key_insert_times.pop(key, None)
-        self._key_sizes.pop(key, None)
+        self._state.cache_order.pop(key, None)
+        self._state.key_insert_times.pop(key, None)
+        self._state.key_sizes.pop(key, None)
 
     async def _get_try_cache(
         self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None
@@ -224,20 +233,20 @@ async def _get_try_cache(
         maybe_cached_result = await self._cache.get(key, prototype, byte_range)
         if maybe_cached_result is not None:
             logger.debug("_get_try_cache: key %s found in cache (HIT)", key)
-            self._hits += 1
+            self._state.hits += 1
             # Update access order for LRU
             await self._update_access_order(key)
             return maybe_cached_result
         else:
             logger.debug(
                 "_get_try_cache: key %s not found in cache (MISS), fetching from store", key
             )
-            self._misses += 1
+            self._state.misses += 1
             maybe_fresh_result = await super().get(key, prototype, byte_range)
             if maybe_fresh_result is None:
                 # Key doesn't exist in source store
                 await self._cache.delete(key)
-                async with self._lock:
+                async with self._state.lock:
                     self._remove_from_tracking(key)
             else:
                 # Cache the newly fetched value
@@ -249,12 +258,12 @@ async def _get_no_cache(
         self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None
     ) -> Buffer | None:
         """Get data directly from source store and update cache."""
-        self._misses += 1
+        self._state.misses += 1
         maybe_fresh_result = await super().get(key, prototype, byte_range)
         if maybe_fresh_result is None:
             # Key doesn't exist in source, remove from cache and tracking
             await self._cache.delete(key)
-            async with self._lock:
+            async with self._state.lock:
                 self._remove_from_tracking(key)
         else:
             logger.debug("_get_no_cache: key %s found in store, setting in cache", key)
@@ -312,7 +321,7 @@ async def set(self, key: str, value: Buffer) -> None:
         else:
             logger.debug("set: deleting key %s from cache", key)
             await self._cache.delete(key)
-            async with self._lock:
+            async with self._state.lock:
                 self._remove_from_tracking(key)
 
     async def delete(self, key: str) -> None:
@@ -328,7 +337,7 @@ async def delete(self, key: str) -> None:
         await super().delete(key)
         logger.debug("delete: deleting key %s from cache", key)
         await self._cache.delete(key)
-        async with self._lock:
+        async with self._state.lock:
             self._remove_from_tracking(key)
 
     def cache_info(self) -> dict[str, Any]:
@@ -339,20 +348,20 @@ def cache_info(self) -> dict[str, Any]:
             if self.max_age_seconds == "infinity"
             else self.max_age_seconds,
             "max_size": self.max_size,
-            "current_size": self._current_size,
+            "current_size": self._state.current_size,
             "cache_set_data": self.cache_set_data,
-            "tracked_keys": len(self.key_insert_times),
-            "cached_keys": len(self._cache_order),
+            "tracked_keys": len(self._state.key_insert_times),
+            "cached_keys": len(self._state.cache_order),
         }
 
     def cache_stats(self) -> dict[str, Any]:
         """Return cache performance statistics."""
-        total_requests = self._hits + self._misses
-        hit_rate = self._hits / total_requests if total_requests > 0 else 0.0
+        total_requests = self._state.hits + self._state.misses
+        hit_rate = self._state.hits / total_requests if total_requests > 0 else 0.0
         return {
-            "hits": self._hits,
-            "misses": self._misses,
-            "evictions": self._evictions,
+            "hits": self._state.hits,
+            "misses": self._state.misses,
+            "evictions": self._state.evictions,
             "total_requests": total_requests,
             "hit_rate": hit_rate,
         }
@@ -364,11 +373,11 @@ async def clear_cache(self) -> None:
             await self._cache.clear()
 
         # Reset tracking
-        async with self._lock:
-            self.key_insert_times.clear()
-            self._cache_order.clear()
-            self._key_sizes.clear()
-            self._current_size = 0
+        async with self._state.lock:
+            self._state.key_insert_times.clear()
+            self._state.cache_order.clear()
+            self._state.key_sizes.clear()
+            self._state.current_size = 0
         logger.debug("clear_cache: cleared all cache data")
 
     def __repr__(self) -> str:
@@ -379,6 +388,6 @@ def __repr__(self) -> str:
             f"cache_store={self._cache!r}, "
             f"max_age_seconds={self.max_age_seconds}, "
             f"max_size={self.max_size}, "
-            f"current_size={self._current_size}, "
-            f"cached_keys={len(self._cache_order)})"
+            f"current_size={self._state.current_size}, "
+            f"cached_keys={len(self._state.cache_order)})"
         )

src/zarr/storage/_logging.py

@@ -77,6 +77,9 @@ def _default_handler(self) -> logging.Handler:
         )
         return handler
 
+    def _with_store(self, store: T_Store) -> Self:
+        return type(self)(store=store, log_level=self.log_level, log_handler=self.log_handler)
+
     @contextmanager
     def log(self, hint: Any = "") -> Generator[None, None, None]:
         """Context manager to log method calls

src/zarr/storage/_wrapper.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Generic, TypeVar
+from typing import TYPE_CHECKING, Generic, TypeVar, cast
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, AsyncIterator, Iterable
@@ -31,14 +31,23 @@ class WrapperStore(Store, Generic[T_Store]):
     def __init__(self, store: T_Store) -> None:
         self._store = store
 
+    def _with_store(self, store: T_Store) -> Self:
+        """
+        Constructs a new instance of the wrapper store with the same details but a new store.
+        """
+        return type(self)(store=store)
+
     @classmethod
     async def open(cls: type[Self], store_cls: type[T_Store], *args: Any, **kwargs: Any) -> Self:
         store = store_cls(*args, **kwargs)
         await store._open()
         return cls(store=store)
 
+    def with_read_only(self, read_only: bool = False) -> Self:
+        return self._with_store(cast(T_Store, self._store.with_read_only(read_only)))
+
     def __enter__(self) -> Self:
-        return type(self)(self._store.__enter__())
+        return self._with_store(self._store.__enter__())
 
     def __exit__(
         self,

src/zarr/testing/store.py

@@ -4,7 +4,7 @@
 import json
 import pickle
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Generic, TypeVar
+from typing import TYPE_CHECKING, Generic, Self, TypeVar
 
 from zarr.storage import WrapperStore
 
@@ -578,10 +578,13 @@ class LatencyStore(WrapperStore[Store]):
     get_latency: float
     set_latency: float
 
-    def __init__(self, cls: Store, *, get_latency: float = 0, set_latency: float = 0) -> None:
+    def __init__(self, store: Store, *, get_latency: float = 0, set_latency: float = 0) -> None:
         self.get_latency = float(get_latency)
         self.set_latency = float(set_latency)
-        self._store = cls
+        self._store = store
+
+    def _with_store(self, store: Store) -> Self:
+        return type(self)(store, get_latency=self.get_latency, set_latency=self.set_latency)
 
     async def set(self, key: str, value: Buffer) -> None:
         """