fix(storage): open ZipStore archive on first use from every method and reopen it without truncating

changes/4450.bugfix.md

@@ -0,0 +1 @@
+`ZipStore` now opens its archive on first use from every method, so calling `get`, `get_partial_values`, `set_if_not_exists`, or `clear` on a store that was never opened no longer raises `AttributeError`. A store created with `mode="w"` or `mode="x"` now keeps its entries when it is used again after `close()`, moved with `move()`, or unpickled, instead of truncating the archive or refusing to open it. Previously, even a read on a closed `mode="w"` store erased the archive, and such a store backed by a file object dropped its earlier entries when written to again after `close()`. A store backed by a readable, seekable file object now keeps them, and one backed by any other file object raises `io.UnsupportedOperation` instead, because it cannot read the archive back. Threads that use a store for the first time at the same moment no longer open the archive twice and lose writes, and neither does a thread that uses a store while another thread closes or moves it. If closing the archive fails, for example because the disk is full, every later use of the store that would reopen the archive now raises `RuntimeError`, instead of reopening an archive that may lack its central directory and silently starting a new one; for a writable store backed by a path, `clear()` replaces the archive and makes the store usable again. A `clear()` that fails partway now leaves the store closed instead of open on a closed archive. An unpickled store now opens its archive on first use, so a copy that is never used, such as those dask makes while building a graph, no longer touches the file; before, such a copy reopened the archive and corrupted it, so writing a dask array to a `ZipStore` produced an unreadable file. A store that reopens its archive and finds it is not a zip, as happens while another copy is still in its first writing session, raises `zipfile.BadZipFile` instead of starting a new archive; other concurrent writers are not detected. On a store that was never opened, `move()` now applies the first-open mode at the original path, so mode `"x"` refuses an existing file instead of moving it. The type of the `mode` parameter now includes `"x"`, which the store already accepted at runtime.

src/zarr/storage/_zip.py

@@ -21,7 +21,7 @@
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator, Iterable
 
-ZipStoreAccessModeLiteral = Literal["r", "w", "a"]
+ZipStoreAccessModeLiteral = Literal["r", "w", "a", "x"]
 
 
 class _RawReaderAdapter(io.RawIOBase):
@@ -72,6 +72,17 @@ def readinto(self, b: Any) -> int:
         return n
 
 
+def _is_unfinished_archive(path: Path) -> bool:
+    """Whether the file at `path` exists but is not a zip archive."""
+    try:
+        with path.open("rb") as f:
+            # is_zipfile swallows OSError, so open the file here to let a
+            # directory or a permission error surface as itself
+            return not zipfile.is_zipfile(f)
+    except FileNotFoundError:
+        return False
+
+
 class ZipStore(Store):
     """
     Store using a ZIP file.
@@ -85,10 +96,20 @@ class ZipStore(Store):
         can only be used for reading (`mode="r"`). The file object must stay
         open for the lifetime of the store, and operations that require a
         filesystem location (`clear`, `move`, pickling) are not supported.
+        Using the store again after `close()` reopens the archive, which
+        requires a file object that is readable and seekable; otherwise it
+        raises `io.UnsupportedOperation`.
     mode : str, optional
         One of 'r' to read an existing file, 'w' to truncate and write a new
         file, 'a' to append to an existing file, or 'x' to exclusively create
-        and write a new file.
+        and write a new file. 'w' and 'x' apply to the first open only; the
+        store reopens its archive with 'a' after `close()`, `move()`, or
+        unpickling, so the entries it already wrote are kept. An unpickled
+        store opens its archive on first use. If `close()` raises, the
+        archive may be incomplete, and every later use that would reopen it
+        raises `RuntimeError` instead; for a writable store backed
+        by a path, `clear()` replaces the archive and makes the store usable
+        again.
     compression : int, optional
         Compression method to use when writing to the archive.
     allowZip64 : bool, optional
@@ -162,24 +183,78 @@ def __init__(
         self._zmode = mode
         self.compression = compression
         self.allowZip64 = allowZip64
+        self._lock = threading.RLock()
+        self._was_opened = False
+        self._close_failed = False
 
     def _sync_open(self) -> None:
         if self._is_open:
             raise ValueError("store is already open")
-
-        self._lock = threading.RLock()
+        if self._close_failed:
+            # the central directory may be missing; appending to such a file
+            # makes zipfile start a new archive and drop the earlier entries
+            raise RuntimeError(
+                f"closing the archive of {self!r} failed, so it may be incomplete; "
+                "the store will not reopen it"
+            )
+        if (
+            self.path is not None
+            and self._was_opened
+            and self._zmode == "a"
+            and _is_unfinished_archive(self.path)
+        ):
+            # the file lacks a central directory, e.g. because another copy of
+            # this store still has it open for writing; appending would start
+            # a new archive and drop the entries already written
+            raise zipfile.BadZipFile(
+                f"{self!r} cannot reopen {self.path}: it is not a zip archive. It "
+                "may be incomplete, for example because another store still has it "
+                "open for writing."
+            )
+        if (
+            self.path is None
+            and self._fileobj is not None
+            and self._was_opened
+            and not (self._fileobj.readable() and self._fileobj.seekable())
+        ):
+            # reopening appends, which needs to read the archive back; on a
+            # write-only file zipfile would start a new archive and drop the
+            # earlier entries
+            raise io.UnsupportedOperation(
+                "a ZipStore backed by a file object that is not readable and "
+                "seekable cannot be used again after close(), because the "
+                "archive it wrote cannot be read back"
+            )
 
         self._zf = zipfile.ZipFile(
             self.path if self.path is not None else self._fileobj,  # type: ignore[arg-type]
             mode=self._zmode,
             compression=self.compression,
             allowZip64=self.allowZip64,
         )
+        # "w" truncates and "x" refuses an existing file. Both apply only to the
+        # first open: reopening after close(), move(), or unpickling must keep
+        # the entries already written.
+        if self._zmode in ("w", "x"):
+            self._zmode = "a"
 
+        self._was_opened = True
         self._is_open = True
 
+    def _zipfile(self) -> zipfile.ZipFile:
+        """Return the archive, opening it on first use."""
+        with self._lock:
+            if not self._is_open:
+                self._sync_open()
+            return self._zf
+
     async def _open(self) -> None:
-        self._sync_open()
+        with self._lock:
+            self._sync_open()
+
+    async def _ensure_open(self) -> None:
+        # the base class checks _is_open outside the lock
+        self._zipfile()
 
     def __getstate__(self) -> dict[str, Any]:
         if self.path is None:
@@ -189,24 +264,39 @@ def __getstate__(self) -> dict[str, Any]:
                 "cannot pickle a ZipStore backed by a file-like object; "
                 "construct the store from a path instead"
             )
-        # We need a copy to not modify the state of the original store
-        state = self.__dict__.copy()
+        # We need a copy to not modify the state of the original store. The
+        # lock keeps it from catching _sync_open between opening with "w" and
+        # switching to "a", which would make unpickling truncate the archive.
+        with self._lock:
+            state = self.__dict__.copy()
         for attr in ["_zf", "_lock"]:
             state.pop(attr, None)
         return state
 
     def __setstate__(self, state: dict[str, Any]) -> None:
         self.__dict__ = state
+        self.__dict__.setdefault("_close_failed", False)
+        self.__dict__.setdefault("_was_opened", False)
+        self._lock = threading.RLock()
+        # open on first use: a copy that is never used, like those dask makes
+        # while building a graph, must not open the file, because zipfile
+        # writes a central directory when it closes an archive it could not read
         self._is_open = False
-        self._sync_open()
 
     def close(self) -> None:
         # docstring inherited
-        if not self._is_open:
-            return
-        super().close()
+        # hold the lock until the archive is closed: a thread that reopened it
+        # before its central directory was written would lose the entries
         with self._lock:
-            self._zf.close()
+            if not self._is_open:
+                return
+            try:
+                self._zf.close()
+            except BaseException:
+                self._close_failed = True
+                raise
+            finally:
+                super().close()
 
     async def clear(self) -> None:
         # docstring inherited
@@ -216,11 +306,20 @@ async def clear(self) -> None:
                 raise NotImplementedError(
                     "clear() is not supported for a ZipStore backed by a file-like object"
                 )
-            self._zf.close()
-            os.remove(self.path)
-            self._zf = zipfile.ZipFile(
-                self.path, mode="w", compression=self.compression, allowZip64=self.allowZip64
-            )
+            if self._zmode == "x":
+                # "x" lasts only until the first open; opening now keeps it from
+                # deleting a file this store never claimed
+                self._zipfile()
+            self.close()
+            # the file may be damaged or gone; clear() replaces it either way
+            self.path.unlink(missing_ok=True)
+            # replacing the file cannot drop entries, so clear() is the one way
+            # to recover a store whose close() failed
+            self._close_failed = False
+            # if this open fails the store stays closed, and the next use
+            # creates the archive again
+            self._zmode = "w"
+            self._sync_open()
 
     def __str__(self) -> str:
         if self.path is None:
@@ -243,11 +342,9 @@ def _get(
         prototype: BufferPrototype,
         byte_range: ByteRequest | None = None,
     ) -> Buffer | None:
-        if not self._is_open:
-            self._sync_open()
         # docstring inherited
         try:
-            with self._zf.open(key) as f:  # will raise KeyError
+            with self._zipfile().open(key) as f:  # will raise KeyError
                 if byte_range is None:
                     return prototype.buffer.from_bytes(f.read())
                 elif isinstance(byte_range, RangeByteRequest):
@@ -288,8 +385,6 @@ async def get_partial_values(
         return out
 
     def _set(self, key: str, value: Buffer) -> None:
-        if not self._is_open:
-            self._sync_open()
         # generally, this should be called inside a lock
         keyinfo = zipfile.ZipInfo(filename=key, date_time=time.localtime(time.time())[:6])
         keyinfo.compress_type = self.compression
@@ -298,13 +393,11 @@ def _set(self, key: str, value: Buffer) -> None:
             keyinfo.external_attr |= 0x10  # MS-DOS directory flag
         else:
             keyinfo.external_attr = 0o644 << 16  # ?rw-r--r--
-        self._zf.writestr(keyinfo, value.to_bytes())
+        self._zipfile().writestr(keyinfo, value.to_bytes())
 
     async def set(self, key: str, value: Buffer) -> None:
         # docstring inherited
         self._check_writable()
-        if not self._is_open:
-            self._sync_open()
         if not isinstance(value, Buffer):
             raise TypeError(
                 f"ZipStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead."
@@ -315,7 +408,7 @@ async def set(self, key: str, value: Buffer) -> None:
     async def set_if_not_exists(self, key: str, value: Buffer) -> None:
         self._check_writable()
         with self._lock:
-            members = self._zf.namelist()
+            members = self._zipfile().namelist()
             if key not in members:
                 self._set(key, value)
 
@@ -337,23 +430,20 @@ async def delete(self, key: str) -> None:
 
     async def exists(self, key: str) -> bool:
         # docstring inherited
-        if not self._is_open:
-            self._sync_open()
         with self._lock:
             try:
-                self._zf.getinfo(key)
+                self._zipfile().getinfo(key)
             except KeyError:
                 return False
             else:
                 return True
 
     async def list(self) -> AsyncIterator[str]:
         # docstring inherited
-        if not self._is_open:
-            self._sync_open()
-        with self._lock:
-            for key in self._zf.namelist():
-                yield key
+        # namelist() is a copy; holding the lock across yield would block
+        # other threads for as long as the caller iterates
+        for key in self._zipfile().namelist():
+            yield key
 
     async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
@@ -363,11 +453,9 @@ async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
 
     async def list_dir(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
-        if not self._is_open:
-            self._sync_open()
         prefix = prefix.rstrip("/")
 
-        keys = self._zf.namelist()
+        keys = self._zipfile().namelist()
         seen = set()
         if prefix == "":
             keys_unique = {k.split("/")[0] for k in keys}
@@ -393,8 +481,12 @@ async def move(self, path: Path | str) -> None:
             )
         if isinstance(path, str):
             path = Path(path)
-        self.close()
-        os.makedirs(path.parent, exist_ok=True)
-        shutil.move(self.path, path)
-        self.path = path
-        await self._open()
+        # hold the lock so that no thread reopens the old path mid-move
+        with self._lock:
+            # opening first keeps mode "x" from moving a file it may not claim
+            self._zipfile()
+            self.close()
+            os.makedirs(path.parent, exist_ok=True)
+            shutil.move(self.path, path)
+            self.path = path
+            self._sync_open()