diff --git a/changes/4157.bugfix.md b/changes/4157.bugfix.md new file mode 100644 index 0000000000..6b0d0fcc67 --- /dev/null +++ b/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. diff --git a/src/zarr/core/buffer/core.py b/src/zarr/core/buffer/core.py index 497543f88f..b8f7f11cd4 100644 --- a/src/zarr/core/buffer/core.py +++ b/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): diff --git a/src/zarr/storage/_memory.py b/src/zarr/storage/_memory.py index 5f5b632d4a..97dd355515 100644 --- a/src/zarr/storage/_memory.py +++ b/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 diff --git a/tests/test_store/test_memory.py b/tests/test_store/test_memory.py index 35504718a7..36265423e6 100644 --- a/tests/test_store/test_memory.py +++ b/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