diff --git a/changes/4353.bugfix.md b/changes/4353.bugfix.md new file mode 100644 index 0000000000..e1ee57219c --- /dev/null +++ b/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. diff --git a/src/zarr/storage/_local.py b/src/zarr/storage/_local.py index 1627c1a6b5..6ade7e283d 100644 --- a/src/zarr/storage/_local.py +++ b/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,28 +371,18 @@ 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: """ @@ -362,11 +390,10 @@ async def move(self, dest_root: Path | str) -> None: """ 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 diff --git a/tests/test_store/test_local.py b/tests/test_store/test_local.py index 90d214ee2c..131312bfa2 100644 --- a/tests/test_store/test_local.py +++ b/tests/test_store/test_local.py @@ -1,12 +1,19 @@ 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 @@ -14,11 +21,144 @@ 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())