diff --git a/dol/__init__.py b/dol/__init__.py index ac855089..22d310f3 100644 --- a/dol/__init__.py +++ b/dol/__init__.py @@ -114,6 +114,8 @@ def ihead(store, n=1): subfolder_stores, # a store of stores, each store corresponding to a subfolder ) +from dol.atomic import atomic_write # crash-safe file write (temp file + os.replace) + from dol.util import ( non_colliding_key, # make a key that does not collide with existing keys get_app_folder, # get the path to a folder for application data diff --git a/dol/atomic.py b/dol/atomic.py new file mode 100644 index 00000000..11ada846 --- /dev/null +++ b/dol/atomic.py @@ -0,0 +1,118 @@ +"""Atomic file writes: write to a sibling temp file, then ``os.replace`` it into place. + +A crash (or an exception while producing the data) leaves the old content or the new, +never a torn or empty file. Dependency-free; ``atomic_write`` is the one entry point. + + >>> import os, tempfile + >>> path = os.path.join(tempfile.mkdtemp(), "f.txt") + >>> atomic_write(path, b"hello") + >>> open(path).read(), os.listdir(os.path.dirname(path)) + ('hello', ['f.txt']) + >>> atomic_write(path, (c.encode() for c in "ab")) # chunks are streamed + >>> open(path).read() + 'ab' +""" + +import os +import stat +import uuid +from collections.abc import Iterable + +_BYTES_LIKE = (bytes, bytearray, memoryview) + + +def _fsync_dir(dirpath: str) -> None: + """Best-effort fsync of a directory, so a rename survives power loss (POSIX only).""" + try: + fd = os.open(dirpath or ".", os.O_RDONLY) + except OSError: # e.g. Windows cannot open directories + return + try: + os.fsync(fd) + except OSError: + pass + finally: + os.close(fd) + + +def _mode_to_apply(path: str, perms: "int | None") -> "int | None": + """``perms`` if given, else the mode of the file being replaced (``None`` if new).""" + if perms is not None: + return perms + try: + return stat.S_IMODE(os.stat(path).st_mode) + except OSError: + return None + + +def atomic_write( + path, + data: "bytes | Iterable[bytes]", + *, + fsync: bool = False, + perms: "int | None" = None, +) -> None: + """Write ``data`` to ``path`` atomically (temp file in the same folder, then ``os.replace``). + + :param path: Target file; its parent folder must exist. A symlink is followed: + its target is what gets replaced. + :param data: Bytes-like, or an iterable of bytes-like chunks (written one at a + time, so a large payload need not be held in memory). If the iterable raises, + ``path`` is untouched. Text must be encoded first. + :param fsync: Flush the file (and, on POSIX, the folder) to disk before returning: + durable across power loss, not just across a process crash. Slower. (On macOS + ``os.fsync`` is not ``F_FULLFSYNC``: best effort.) + :param perms: Mode (e.g. ``0o600``) for the file, set at creation so the data is + never readable more widely, even briefly, and not narrowed by the umask. + Default: the mode of the file being replaced, else the usual umask mode. + + The temp file is a hidden, uniquely named sibling (``...tmp``), created + with ``O_EXCL``, so concurrent writers never share one and key listings skip it. + On any failure it is removed and the old content is untouched. Atomic on POSIX, + and on Windows within one volume. Hard links to the old file are detached, and its owner is not preserved. + + >>> import tempfile + >>> d = tempfile.mkdtemp() + >>> atomic_write(os.path.join(d, "a"), [b"x", b"y"], perms=0o600, fsync=True) + >>> open(os.path.join(d, "a"), "rb").read() + b'xy' + >>> def boom(): + ... yield b"partial" + ... raise RuntimeError("producer died") + >>> atomic_write(os.path.join(d, "a"), boom()) + Traceback (most recent call last): + ... + RuntimeError: producer died + >>> open(os.path.join(d, "a"), "rb").read(), os.listdir(d) # old content, no litter + (b'xy', ['a']) + """ + if isinstance(data, str): + raise TypeError( + "atomic_write takes bytes (or chunks of bytes); encode text first" + ) + path = os.path.realpath(os.fspath(path)) + chunks = [data] if isinstance(data, _BYTES_LIKE) else data + dirname, name = os.path.split(path) + tmp = os.path.join(dirname, f".{name[:100]}.{uuid.uuid4().hex[:8]}.tmp") + mode = _mode_to_apply(path, perms) + try: + fd = os.open( + tmp, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o666 if mode is None else mode + ) + with os.fdopen(fd, "wb") as fh: + if mode is not None and hasattr(os, "fchmod"): + os.fchmod(fh.fileno(), mode) # exact, not narrowed by the umask + for chunk in chunks: + fh.write(chunk) + if fsync: + fh.flush() + os.fsync(fh.fileno()) + os.replace(tmp, path) + except BaseException: + try: + os.unlink(tmp) + except OSError: + pass + raise + if fsync: + _fsync_dir(dirname) diff --git a/dol/tests/test_atomic.py b/dol/tests/test_atomic.py new file mode 100644 index 00000000..3785063a --- /dev/null +++ b/dol/tests/test_atomic.py @@ -0,0 +1,129 @@ +"""Tests for ``dol.atomic.atomic_write``.""" + +import os +import stat +import sys + +import pytest + +from dol import atomic_write + +posix_only = pytest.mark.skipif( + sys.platform == "win32", reason="POSIX mode bits/symlinks" +) + + +def _mode(p): + return stat.S_IMODE(os.stat(p).st_mode) + + +def test_bytes_and_chunks_and_overwrite(tmp_path): + p = tmp_path / "f" + atomic_write(p, b"one") + atomic_write(p, iter([b"t", bytearray(b"w"), memoryview(b"o")])) + assert p.read_bytes() == b"two" + assert os.listdir(tmp_path) == ["f"] + + +def test_empty_data(tmp_path): + atomic_write(tmp_path / "f", iter([])) + assert (tmp_path / "f").read_bytes() == b"" + + +def test_str_is_refused_clearly(tmp_path): + with pytest.raises(TypeError, match="encode"): + atomic_write(tmp_path / "f", "text") + assert os.listdir(tmp_path) == [] + + +def test_failure_keeps_old_content_and_cleans_tmp(tmp_path): + p = tmp_path / "f" + atomic_write(p, b"old") + + def gen(): + yield b"new" + raise ValueError + + with pytest.raises(ValueError): + atomic_write(p, gen()) + assert p.read_bytes() == b"old" and os.listdir(tmp_path) == ["f"] + + +def test_missing_parent_raises_and_leaves_nothing(tmp_path): + with pytest.raises(OSError): + atomic_write(tmp_path / "nope" / "f", b"x") + assert os.listdir(tmp_path) == [] + + +def test_tmp_is_hidden_unique_and_in_same_folder(tmp_path, monkeypatch): + seen = [] + real = os.replace + monkeypatch.setattr(os, "replace", lambda a, b: (seen.append(a), real(a, b))) + atomic_write(tmp_path / "f", b"x") + atomic_write(tmp_path / "f", b"x") + assert len(set(seen)) == 2 + for s in seen: + assert os.path.dirname(s) == os.path.realpath(tmp_path) + assert os.path.basename(s).startswith(".f.") + + +def test_long_names_do_not_overflow(tmp_path): + name = "n" * 240 + atomic_write(tmp_path / name, b"x") + assert (tmp_path / name).read_bytes() == b"x" + + +def test_fsync_only_when_asked(tmp_path, monkeypatch): + calls = [] + real = os.fsync + monkeypatch.setattr(os, "fsync", lambda fd: (calls.append(fd), real(fd))) + atomic_write(tmp_path / "f", b"x") + assert calls == [] + atomic_write(tmp_path / "f", b"x", fsync=True) + assert len(calls) >= 1 and (tmp_path / "f").read_bytes() == b"x" + + +@posix_only +def test_perms_exact_despite_umask(tmp_path): + old = os.umask(0o077) + try: + atomic_write(tmp_path / "f", b"x", perms=0o644) + finally: + os.umask(old) + assert _mode(tmp_path / "f") == 0o644 + + +@posix_only +def test_perms_set_before_data_is_written(tmp_path, monkeypatch): + monkeypatch.setattr(os, "fchmod", lambda *a: None) # creation alone must do it + seen = [] + + def gen(): + (tmp,) = [p for p in os.listdir(tmp_path) if p.startswith(".f.")] + seen.append(_mode(tmp_path / tmp)) + yield b"secret" + + old = os.umask(0) + try: + atomic_write(tmp_path / "f", gen(), perms=0o600) + finally: + os.umask(old) + assert seen == [0o600] + + +@posix_only +def test_existing_mode_preserved(tmp_path): + p = tmp_path / "f" + p.write_bytes(b"old") + os.chmod(p, 0o600) + atomic_write(p, b"new") + assert _mode(p) == 0o600 and p.read_bytes() == b"new" + + +@posix_only +def test_symlink_is_followed(tmp_path): + target, link = tmp_path / "target", tmp_path / "link" + target.write_bytes(b"old") + link.symlink_to(target) + atomic_write(link, b"new") + assert link.is_symlink() and target.read_bytes() == b"new"