fix(store): run every LocalStore filesystem call off the event loop

changes/4353.bugfix.md

@@ -0,0 +1,7 @@
+`LocalStore` no longer blocks the event loop on filesystem calls. `open`,
+`clear`, `delete`, `delete_dir`, `list`, `list_prefix`, `list_dir`,
+`move` and `getsize` now run their disk I/O in a worker thread via
+`asyncio.to_thread`, as `get` and `set` already did, so other tasks
+sharing the loop are not stalled while the store walks or modifies a directory
+tree. The `LocalStore` test suite now fails if any async store method touches
+the filesystem from the event loop thread.

src/zarr/storage/_local.py

@@ -85,6 +85,58 @@ def _put(path: Path, value: Buffer, exclusive: bool = False) -> int:
         return f.write(view)
 
 
+# The helpers below do the blocking filesystem work behind LocalStore's async methods.
+# Each async method runs exactly one of them via `asyncio.to_thread` so that the event
+# loop is never stalled on disk I/O; the synchronous methods call them directly.
+
+
+def _ensure_root(root: Path, *, create: bool) -> None:
+    if create:
+        root.mkdir(parents=True, exist_ok=True)
+    if not root.exists():
+        raise FileNotFoundError(f"{root} does not exist")
+
+
+def _clear(root: Path) -> None:
+    shutil.rmtree(root)
+    root.mkdir()
+
+
+def _delete(path: Path) -> None:
+    if path.is_dir():
+        shutil.rmtree(path)
+    else:
+        path.unlink(missing_ok=True)
+
+
+def _delete_dir(path: Path, prefix: str) -> None:
+    if path.is_dir():
+        shutil.rmtree(path)
+    elif path.is_file():
+        raise ValueError(f"delete_dir was passed a {prefix=!r} that is a file")
+    # A non-existent directory is a no-op; test_group:test_create_creates_parents relies on it.
+
+
+def _list_files(root: Path, prefix: str) -> list[str]:
+    """Keys (paths relative to `root`, POSIX style) of every file under `root / prefix`."""
+    to_strip = root.as_posix() + "/"
+    return [p.as_posix().removeprefix(to_strip) for p in (root / prefix).rglob("*") if p.is_file()]
+
+
+def _list_dir(base: Path) -> list[str]:
+    try:
+        return [p.name for p in base.iterdir()]
+    except (FileNotFoundError, NotADirectoryError):
+        return []
+
+
+def _move(src: Path, dest_root: Path) -> None:
+    dest_root.parent.mkdir(parents=True, exist_ok=True)
+    if dest_root.exists():
+        raise FileExistsError(f"Destination root {dest_root} already exists.")
+    shutil.move(src, dest_root)
+
+
 class LocalStore(Store):
     """
     Store for the local file system.
@@ -165,18 +217,25 @@ async def open(
         return store
 
     async def _open(self, *, mode: AccessModeLiteral | None = None) -> None:
-        if not self.read_only:
-            self.root.mkdir(parents=True, exist_ok=True)
-
-        if not self.root.exists():
-            raise FileNotFoundError(f"{self.root} does not exist")
+        await asyncio.to_thread(_ensure_root, self.root, create=not self.read_only)
         return await super()._open()
 
+    async def _ensure_open(self) -> None:
+        # docstring inherited
+        if not self._is_open:
+            await asyncio.to_thread(_ensure_root, self.root, create=not self.read_only)
+            # Concurrent lazy opens (every `set` of a `set_many`, say) all pass the check
+            # above and each verifies the root, which is idempotent; only the first may
+            # flip the flag, since `Store._open` refuses to open an open store.
+            # Note this calls `Store._open` directly, so a subclass's `_open` override is
+            # bypassed on the lazy-open path.
+            if not self._is_open:
+                await super()._open()
+
     async def clear(self) -> None:
         # docstring inherited
         self._check_writable()
-        shutil.rmtree(self.root)
-        self.root.mkdir()
+        await asyncio.to_thread(_clear, self.root)
 
     def __str__(self) -> str:
         return f"file://{self.root.as_posix()}"
@@ -193,10 +252,7 @@ def __eq__(self, other: object) -> bool:
 
     def _ensure_open_sync(self) -> None:
         if not self._is_open:
-            if not self.read_only:
-                self.root.mkdir(parents=True, exist_ok=True)
-            if not self.root.exists():
-                raise FileNotFoundError(f"{self.root} does not exist")
+            _ensure_root(self.root, create=not self.read_only)
             self._is_open = True
 
     def get_sync(
@@ -231,11 +287,7 @@ def set_sync(self, key: str, value: Buffer) -> None:
     def delete_sync(self, key: str) -> None:
         self._ensure_open_sync()
         self._check_writable()
-        path = self.root / key
-        if path.is_dir():
-            shutil.rmtree(path)
-        else:
-            path.unlink(missing_ok=True)
+        _delete(self.root / key)
 
     async def get(
         self,
@@ -246,8 +298,7 @@ async def get(
         # docstring inherited
         if prototype is None:
             prototype = default_buffer_prototype()
-        if not self._is_open:
-            await self._open()
+        await self._ensure_open()
         assert isinstance(key, str)
         path = self.root / key
 
@@ -281,8 +332,7 @@ async def set_if_not_exists(self, key: str, value: Buffer) -> None:
             pass
 
     async def _set(self, key: str, value: Buffer, exclusive: bool = False) -> None:
-        if not self._is_open:
-            await self._open()
+        await self._ensure_open()
         self._check_writable()
         assert isinstance(key, str)
         if not isinstance(value, Buffer):
@@ -307,24 +357,12 @@ async def delete(self, key: str) -> None:
         """
         # docstring inherited
         self._check_writable()
-        path = self.root / key
-        if path.is_dir():  # TODO: support deleting directories? shutil.rmtree?
-            shutil.rmtree(path)
-        else:
-            await asyncio.to_thread(path.unlink, True)  # Q: we may want to raise if path is missing
+        await asyncio.to_thread(_delete, self.root / key)
 
     async def delete_dir(self, prefix: str) -> None:
         # docstring inherited
         self._check_writable()
-        path = self.root / prefix
-        if path.is_dir():
-            shutil.rmtree(path)
-        elif path.is_file():
-            raise ValueError(f"delete_dir was passed a {prefix=!r} that is a file")
-        else:
-            # Non-existent directory
-            # This path is tested by test_group:test_create_creates_parents for one
-            pass
+        await asyncio.to_thread(_delete_dir, self.root / prefix, prefix)
 
     async def exists(self, key: str) -> bool:
         # docstring inherited
@@ -333,40 +371,29 @@ async def exists(self, key: str) -> bool:
 
     async def list(self) -> AsyncIterator[str]:
         # docstring inherited
-        to_strip = self.root.as_posix() + "/"
-        for p in list(self.root.rglob("*")):
-            if p.is_file():
-                yield p.as_posix().replace(to_strip, "")
+        for key in await asyncio.to_thread(_list_files, self.root, ""):
+            yield key
 
     async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
-        to_strip = self.root.as_posix() + "/"
-        prefix = prefix.rstrip("/")
-        for p in (self.root / prefix).rglob("*"):
-            if p.is_file():
-                yield p.as_posix().replace(to_strip, "")
+        for key in await asyncio.to_thread(_list_files, self.root, prefix.rstrip("/")):
+            yield key
 
     async def list_dir(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
-        base = self.root / prefix
-        try:
-            key_iter = base.iterdir()
-            for key in key_iter:
-                yield key.relative_to(base).as_posix()
-        except (FileNotFoundError, NotADirectoryError):
-            pass
+        for name in await asyncio.to_thread(_list_dir, self.root / prefix):
+            yield name
 
     async def move(self, dest_root: Path | str) -> None:
         """
         Move the store to another path. The old root directory is deleted.
         """
         if isinstance(dest_root, str):
             dest_root = Path(dest_root)
-        os.makedirs(dest_root.parent, exist_ok=True)
-        if dest_root.exists():
-            raise FileExistsError(f"Destination root {dest_root} already exists.")
-        shutil.move(self.root, dest_root)
+        await asyncio.to_thread(_move, self.root, dest_root)
         self.root = dest_root
 
     async def getsize(self, key: str) -> int:
-        return (self.root / key).stat().st_size
+        # docstring inherited
+        stat = await asyncio.to_thread((self.root / key).stat)
+        return stat.st_size

tests/test_store/test_local.py

@@ -1,24 +1,164 @@
 from __future__ import annotations
 
+import asyncio
+import inspect
+import io
+import os
 import pathlib
 import re
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 import pytest
 
 import zarr
+import zarr.storage._local
 from zarr import create_array
 from zarr.core.buffer import Buffer, cpu
 from zarr.storage import LocalStore
 from zarr.storage._local import _atomic_write
 from zarr.testing.store import StoreTests
 from zarr.testing.utils import assert_bytes_equal
 
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+_LOCAL_STORE_FILE = zarr.storage._local.__file__
+_ASYNC_CODE_FLAGS = inspect.CO_COROUTINE | inspect.CO_ASYNC_GENERATOR
+
+try:
+    import coverage as _coverage
+except ImportError:  # pragma: no cover
+    _COVERAGE_DIR = None
+else:
+    # coverage.py canonicalizes a source file's path (os.path.realpath, hence os.lstat) the
+    # first time its tracer sees code from that file, and it does so on whatever thread is
+    # running that code. Such calls carry coverage's own frames and are not LocalStore's.
+    _COVERAGE_DIR = os.path.dirname(_coverage.__file__) + os.sep
+
+# The syscall-level entry points that pathlib / os.path / shutil helpers bottom out in.
+# Patching these, rather than each Path method, catches a blocking call no matter which
+# helper made it. The list is deliberately wider than what LocalStore uses today.
+_FILESYSTEM_CALLS: tuple[tuple[Any, str], ...] = (
+    (os, "stat"),
+    (os, "lstat"),
+    (os, "access"),
+    (os, "scandir"),
+    (os, "listdir"),
+    (os, "mkdir"),
+    (os, "rmdir"),
+    (os, "unlink"),
+    (os, "remove"),
+    (os, "link"),
+    (os, "rename"),
+    (os, "replace"),
+    (os, "utime"),
+    (os, "open"),
+    (os, "fsync"),
+    (io, "open"),
+)
+
+
+@dataclass
+class _FilesystemCalls:
+    """What the patched filesystem entry points saw from LocalStore code during one test."""
+
+    off_loop: set[tuple[str, str]] = field(default_factory=set)
+    """``(outermost LocalStore function, op)`` pairs made from a thread with no running loop."""
+    on_loop: list[str] = field(default_factory=list)
+    """Calls made on an event loop's thread from inside a LocalStore coroutine: violations."""
+
+    def record(self, op: str) -> None:
+        frame = inspect.currentframe()
+        innermost = outermost = None
+        while frame is not None:
+            if _COVERAGE_DIR is not None and frame.f_code.co_filename.startswith(_COVERAGE_DIR):
+                return  # the coverage tracer resolving a filename, not LocalStore doing I/O
+            if frame.f_code.co_filename == _LOCAL_STORE_FILE:
+                if innermost is None:
+                    innermost = frame
+                outermost = frame
+            frame = frame.f_back
+        if outermost is None or innermost is None:
+            return  # not LocalStore's doing (pytest, tmp_path, the test body, ...)
+        try:
+            asyncio.get_running_loop()
+        except RuntimeError:
+            # A worker thread, such as the one asyncio.to_thread uses: blocking is fine here.
+            self.off_loop.add((outermost.f_code.co_name, op))
+            return
+        if not outermost.f_code.co_flags & _ASYNC_CODE_FLAGS:
+            return  # a synchronous LocalStore method: its caller chose to block the loop
+        site = (
+            f"LocalStore.{outermost.f_code.co_name} called {op} at _local.py:{innermost.f_lineno}"
+        )
+        if site not in self.on_loop:
+            self.on_loop.append(site)
+
 
 class TestLocalStore(StoreTests[LocalStore, cpu.Buffer]):
     store_cls = LocalStore
     buffer_cls = cpu.Buffer
 
+    async def test_delete_dir_on_a_file_raises(self, tmp_path: pathlib.Path) -> None:
+        """`delete_dir` refuses a prefix that names a file rather than a directory."""
+        store = await LocalStore.open(tmp_path)
+        await store.set("file", self.buffer_cls.from_bytes(b"x"))
+        with pytest.raises(ValueError, match="that is a file"):
+            await store.delete_dir("file")
+        assert await store.exists("file")
+
+    @pytest.fixture(autouse=True)
+    def filesystem_calls(self, monkeypatch: pytest.MonkeyPatch) -> Iterator[_FilesystemCalls]:
+        """Fail if any LocalStore coroutine does filesystem I/O on the event loop thread.
+
+        Every async method must hand its filesystem work to ``asyncio.to_thread``;
+        doing it inline stalls every other task sharing the loop.
+        """
+        calls = _FilesystemCalls()
+
+        def patch(module: Any, name: str) -> None:
+            original = getattr(module, name)
+
+            def wrapper(*args: Any, **kwargs: Any) -> Any:
+                calls.record(f"{module.__name__}.{name}")
+                return original(*args, **kwargs)
+
+            monkeypatch.setattr(module, name, wrapper)
+
+        for module, name in _FILESYSTEM_CALLS:
+            patch(module, name)
+        yield calls
+        assert not calls.on_loop, "filesystem calls on the event loop thread:\n" + "\n".join(
+            calls.on_loop
+        )
+
+    async def test_filesystem_calls_are_observed(
+        self, store: LocalStore, filesystem_calls: _FilesystemCalls
+    ) -> None:
+        """The detector must actually see LocalStore's I/O, or its silence means nothing."""
+        await store.set("foo", self.buffer_cls.from_bytes(b"x"))
+        await store.get("foo")
+        assert ("_put", "io.open") in filesystem_calls.off_loop
+        assert ("_get", "io.open") in filesystem_calls.off_loop
+
+    async def test_concurrent_lazy_open(self, store_not_open: LocalStore) -> None:
+        """Concurrent first calls on an unopened store all succeed.
+
+        Opening now suspends (the root check runs in a thread), so every caller that
+        finds the store closed races to open it; none of them may hit
+        ``Store._open``'s "already open" error.
+        """
+        data = self.buffer_cls.from_bytes(b"x")
+        keys = [f"k{i}" for i in range(8)]
+        await asyncio.gather(
+            store_not_open.get("missing"), *(store_not_open.set(k, data) for k in keys)
+        )
+        assert store_not_open._is_open
+        for key in keys:
+            assert_bytes_equal(await store_not_open.get(key), data)
+
     async def get(self, store: LocalStore, key: str) -> Buffer:
         return self.buffer_cls.from_bytes((store.root / key).read_bytes())