Skip to content
7 changes: 7 additions & 0 deletions changes/4353.bugfix.md
Original file line number Diff line number Diff line change
@@ -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.
137 changes: 82 additions & 55 deletions src/zarr/storage/_local.py
Original file line number Diff line number Diff line change
Expand Up @@ -131,6 +131,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.
Expand Down Expand Up @@ -211,18 +263,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()}"
Expand All @@ -239,10 +298,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(
Expand Down Expand Up @@ -275,11 +331,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,
Expand All @@ -290,8 +342,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()
path = self.root / key

try:
Expand Down Expand Up @@ -323,8 +374,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()
if not isinstance(value, Buffer):
raise TypeError(
Expand All @@ -348,24 +398,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
Expand All @@ -374,40 +412,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
140 changes: 140 additions & 0 deletions tests/test_store/test_local.py
Original file line number Diff line number Diff line change
@@ -1,25 +1,165 @@
from __future__ import annotations

import asyncio
import inspect
import io
import os
import pathlib
import re
import time
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 _RETRY_DELAYS, _atomic_write, _move_with_retry
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())

Expand Down
Loading