Fix/memorystore buffer aliasing

changes/4157.bugfix.md

@@ -0,0 +1,11 @@
+`MemoryStore` now copies buffers as they are written, so it never retains the
+caller's memory. Previously an uncompressed write handed the store a zero-copy
+view of the user's array, and mutating that array afterwards would silently
+rewrite chunks already committed to the store.
+
+Only `MemoryStore` is affected: stores that serialize on write, such as
+`LocalStore` and `ZipStore`, never aliased the caller's memory. Uncompressed
+writes to a `MemoryStore` are correspondingly slower, since the copy that makes
+the stored data independent is now actually performed; compressed writes are
+unchanged. Buffers supplied through the `store_dict` argument remain the
+caller's responsibility and are stored as-is.

src/zarr/core/buffer/core.py

@@ -45,6 +45,8 @@ def __getitem__(self, key: slice) -> Self: ...
 
     def __setitem__(self, key: slice, value: Any) -> None: ...
 
+    def copy(self) -> Self: ...
+
 
 @runtime_checkable
 class NDArrayLike(Protocol):

src/zarr/storage/_memory.py

@@ -26,6 +26,18 @@
 logger = getLogger(__name__)
 
 
+def _copy_buffer(value: Buffer) -> Buffer:
+    """Copy `value` so the store does not retain the caller's memory.
+
+    Encoding a chunk can hand the store a zero-copy view of the user's array
+    (an uncompressed write is the common case), and unlike stores that
+    serialize on write, this one keeps whatever it is given alive in a dict.
+    Without this copy a later mutation of the user's array would rewrite
+    chunks already committed to the store.
+    """
+    return type(value).from_array_like(value.as_array_like().copy())
+
+
 class MemoryStore(Store):
     """
     Store for local memory.
@@ -42,6 +54,12 @@ class MemoryStore(Store):
     supports_writes
     supports_deletes
     supports_listing
+
+    Notes
+    -----
+    Writes copy the buffer they are given, so the store never aliases the
+    caller's memory. Buffers passed via `store_dict` are the caller's
+    responsibility and are stored as-is.
     """
 
     supports_writes: bool = True
@@ -117,7 +135,7 @@ def set_sync(self, key: str, value: Buffer) -> None:
             raise TypeError(
                 f"MemoryStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead."
             )
-        self._store_dict[key] = value
+        self._store_dict[key] = _copy_buffer(value)
 
     def delete_sync(self, key: str) -> None:
         self._check_writable()
@@ -178,13 +196,13 @@ async def set(self, key: str, value: Buffer, byte_range: tuple[int, int] | None
             buf[byte_range[0] : byte_range[1]] = value
             self._store_dict[key] = buf
         else:
-            self._store_dict[key] = value
+            self._store_dict[key] = _copy_buffer(value)
 
     async def set_if_not_exists(self, key: str, value: Buffer) -> None:
         # docstring inherited
         self._check_writable()
         await self._ensure_open()
-        self._store_dict.setdefault(key, value)
+        self._store_dict.setdefault(key, _copy_buffer(value))
 
     async def delete(self, key: str) -> None:
         # docstring inherited

tests/test_store/test_memory.py

@@ -8,7 +8,7 @@
 import pytest
 
 import zarr
-from zarr.core.buffer import Buffer, cpu, gpu
+from zarr.core.buffer import Buffer, cpu, default_buffer_prototype, gpu
 from zarr.errors import ZarrUserWarning
 from zarr.storage import GpuMemoryStore, ManagedMemoryStore, MemoryStore
 from zarr.testing.store import StoreTests
@@ -76,6 +76,55 @@ async def test_deterministic_size(
         np.testing.assert_array_equal(a[:3], 1)
         np.testing.assert_array_equal(a[3:], 0)
 
+    @pytest.mark.parametrize("method", ["set", "set_sync", "set_if_not_exists"])
+    async def test_set_does_not_retain_caller_buffer(self, store: MemoryStore, method: str) -> None:
+        """Writing a buffer must not alias the caller's memory.
+
+        MemoryStore keeps whatever it is handed alive in a dict, so retaining
+        the caller's buffer lets a later mutation of that buffer rewrite data
+        already committed to the store.
+        """
+        source = np.frombuffer(bytearray(b"\x01\x02\x03\x04"), dtype="B")
+        value = cpu.Buffer.from_array_like(source)
+
+        if method == "set_sync":
+            store.set_sync("k", value)
+        else:
+            await getattr(store, method)("k", value)
+
+        source[:] = 0xF  # mutate the caller's memory after the write
+        stored = await store.get("k", prototype=default_buffer_prototype())
+        assert stored is not None
+        assert stored.to_bytes() == b"\x01\x02\x03\x04"
+
+    @pytest.mark.parametrize(
+        "pipeline",
+        [
+            "zarr.core.codec_pipeline.BatchedCodecPipeline",
+            "zarr.core.codec_pipeline.FusedCodecPipeline",
+        ],
+    )
+    @pytest.mark.parametrize(("shape", "chunks"), [((30,), (10,)), ((8,), (4,)), ((4,), (4,))])
+    def test_write_does_not_alias_source_array(
+        self, pipeline: str, shape: tuple[int], chunks: tuple[int]
+    ) -> None:
+        """Mutating the source array after a write must not corrupt stored chunks.
+
+        Without compression the encoded buffer is a zero-copy view of the
+        caller's array all the way down to the store, so this covers both the
+        single-chunk and multi-chunk write paths.
+        """
+        with zarr.config.set({"codec_pipeline.path": pipeline}):
+            array = zarr.create_array(
+                store=MemoryStore(), shape=shape, chunks=chunks, dtype="i4", compressors=None
+            )
+            source = np.arange(shape[0], dtype="i4")
+            expected = source.copy()
+            array[:] = source
+            source[:] = -1
+
+            np.testing.assert_array_equal(array[:], expected)
+
     # --- byte-range-write tests: disabled ---
     # Byte-range-write support (set_range / set_range_sync / SupportsSetRange)
     # was removed from this PR pending a decision on the store interface. These