diff --git a/changes/4168.bugfix.md b/changes/4168.bugfix.md new file mode 100644 index 0000000000..4bb538cfbe --- /dev/null +++ b/changes/4168.bugfix.md @@ -0,0 +1,5 @@ +`ZipStore` now rejects pickling unless it was opened in read-only mode. Pickling +a writable store previously allowed multiple processes to reopen the same ZIP +archive for writing, so each writer could replace the archive's central +directory when it closed and corrupt the result. Read-only `ZipStore` instances +remain pickleable. diff --git a/docs/user-guide/performance.md b/docs/user-guide/performance.md index 52c1cf0d71..d1875fcf2e 100644 --- a/docs/user-guide/performance.md +++ b/docs/user-guide/performance.md @@ -307,8 +307,12 @@ When writing to the same chunks from multiple processes, you should use external ## Pickle support -Zarr arrays and groups can be pickled, as long as the underlying store object can be -pickled. All of the storage classes provided in the `zarr.storage` module can be pickled. +Zarr arrays and groups can be pickled as long as the underlying store object can be +pickled. All storage classes provided in `zarr.storage` support pickling for read-only +workloads. A `ZipStore` opened in an archive-writing mode is the exception: independent +ZIP writers can corrupt the archive, so writable `ZipStore` objects reject pickling. +Use a store that supports parallel writes, such as `LocalStore`, and create the ZIP +archive only after writing is complete. If an array or group is backed by a persistent store such as a `zarr.storage.LocalStore`, `zarr.storage.ZipStore` or `zarr.storage.FsspecStore` then the store data diff --git a/src/zarr/storage/_zip.py b/src/zarr/storage/_zip.py index 430b0c3e2a..840c24f206 100644 --- a/src/zarr/storage/_zip.py +++ b/src/zarr/storage/_zip.py @@ -43,6 +43,14 @@ class ZipStore(Store): will raise an exception when the ZIP file would require ZIP64 extensions. + Notes + ----- + A ``ZipStore`` opened in an archive-writing mode cannot be pickled. ZIP + archives do not support independent concurrent writers, and reopening a + serialized writer can corrupt the archive. For parallel writes, use a + store that supports them and create the ZIP archive after writing is + complete. + Attributes ---------- allowed_exceptions @@ -107,6 +115,13 @@ async def _open(self) -> None: self._sync_open() def __getstate__(self) -> dict[str, Any]: + if self._zmode != "r": + raise TypeError( + "ZipStore instances opened in an archive-writing mode cannot be pickled, " + "because independent ZIP writers can corrupt the archive. Use a LocalStore " + "for parallel writes and create the ZIP archive after writing is complete, " + "or reopen the ZipStore with mode='r' before pickling." + ) # We need a copy to not modify the state of the original store state = self.__dict__.copy() for attr in ["_zf", "_lock"]: diff --git a/src/zarr/testing/store.py b/src/zarr/testing/store.py index 46287ccffb..7a2e51a603 100644 --- a/src/zarr/testing/store.py +++ b/src/zarr/testing/store.py @@ -127,8 +127,8 @@ async def test_serializable_store(self, store: S) -> None: # quickly roundtrip data to a key to test that new store works data_buf = self.buffer_cls.from_bytes(b"\x01\x02\x03\x04") key = "foo" - await store.set(key, data_buf) - observed = await store.get(key, prototype=default_buffer_prototype()) + await new_store.set(key, data_buf) + observed = await new_store.get(key, prototype=default_buffer_prototype()) assert_bytes_equal(observed, data_buf) def test_store_read_only(self, store: S) -> None: diff --git a/tests/test_store/test_zip.py b/tests/test_store/test_zip.py index ed69114b51..15c30cc47c 100644 --- a/tests/test_store/test_zip.py +++ b/tests/test_store/test_zip.py @@ -1,6 +1,7 @@ from __future__ import annotations import os +import pickle import shutil import tempfile import zipfile @@ -84,6 +85,34 @@ def test_store_supports_writes(self, store: ZipStore) -> None: def test_store_supports_listing(self, store: ZipStore) -> None: assert store.supports_listing + async def test_serializable_store(self, store: ZipStore) -> None: + data = cpu.Buffer.from_bytes(b"preserve me") + await store.set("sentinel", data) + + with pytest.raises(TypeError, match="archive-writing mode cannot be pickled"): + pickle.dumps(store) + + store.close() + with zipfile.ZipFile(store.path, mode="r") as archive: + assert archive.read("sentinel") == data.to_bytes() + + async def test_read_only_store_is_serializable(self, tmp_path: Path) -> None: + path = tmp_path / "data.zip" + data = cpu.Buffer.from_bytes(b"preserve me") + + writable = await ZipStore.open(path, mode="w") + await writable.set("sentinel", data) + writable.close() + + read_only = await ZipStore.open(path, mode="r") + restored = pickle.loads(pickle.dumps(read_only)) + observed = await restored.get("sentinel", prototype=default_buffer_prototype()) + + assert observed is not None + assert observed.to_bytes() == data.to_bytes() + read_only.close() + restored.close() + # TODO: fix this warning @pytest.mark.filterwarnings("ignore:Unclosed client session:ResourceWarning") def test_api_integration(self, store: ZipStore) -> None: