diff --git a/dol/__init__.py b/dol/__init__.py index 548d790f..feeba4ae 100644 --- a/dol/__init__.py +++ b/dol/__init__.py @@ -116,6 +116,7 @@ def ihead(store, n=1): ) from dol.atomic import atomic_write # crash-safe file write (temp file + os.replace) +from dol.atomic import file_lock, FileLock # cross-process advisory lock (sidecar file) from dol.util import ( non_colliding_key, # make a key that does not collide with existing keys diff --git a/dol/atomic.py b/dol/atomic.py index 11ada846..fd8594ac 100644 --- a/dol/atomic.py +++ b/dol/atomic.py @@ -3,6 +3,9 @@ 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. +``file_lock`` is its companion: an atomic replace is only safe against concurrent +appenders and readers when they all hold the same advisory lock. + >>> import os, tempfile >>> path = os.path.join(tempfile.mkdtemp(), "f.txt") >>> atomic_write(path, b"hello") @@ -13,10 +16,15 @@ 'ab' """ +import errno import os import stat +import threading +import time import uuid +import weakref from collections.abc import Iterable +from contextlib import contextmanager _BYTES_LIKE = (bytes, bytearray, memoryview) @@ -116,3 +124,218 @@ def atomic_write( raise if fsync: _fsync_dir(dirname) + + +# --------------------------------------------------------------------------- +# file_lock + +DFLT_LOCK_SIDECAR = ".lock" +DFLT_POLL_INTERVAL = 0.001 # seconds between attempts on a contended lock (Windows) + +_LOCK_UNSUPPORTED = { # e.g. Lustre without its flock option, some FUSE/SMB mounts + getattr(errno, name) + for name in ("ENOLCK", "ENOTSUP", "EOPNOTSUPP", "ENOSYS") + if hasattr(errno, name) +} + +if os.name == "nt": + import msvcrt + + def _lock(fh, poll_interval: float) -> None: + # not LK_LOCK: it sleeps 1s on any contention and raises after ~10s. + # No timeout: Windows releases a dead process's locks. + fh.seek(0) + while True: + try: + msvcrt.locking(fh.fileno(), msvcrt.LK_NBLCK, 1) + return + except OSError as e: + if e.errno not in (errno.EACCES, errno.EDEADLOCK): + raise + time.sleep(poll_interval) + + def _unlock(fh) -> None: + fh.seek(0) + msvcrt.locking(fh.fileno(), msvcrt.LK_UNLCK, 1) + +else: + import fcntl + + def _lock(fh, poll_interval: float) -> None: + fcntl.flock(fh.fileno(), fcntl.LOCK_EX) # excludes other descriptors too + + def _unlock(fh) -> None: + fcntl.flock(fh.fileno(), fcntl.LOCK_UN) + + +def _open_lock_file(path: str): + """``path`` opened for locking: read-only if it exists but is not writable (both + ``flock`` and ``msvcrt.locking`` work on a read-only handle); ``None`` if the + folder is read-only, so locking is unavailable. Any other failure (missing folder, + too many open files...) is a real error and propagates.""" + try: + return open(path, "a+b") + except PermissionError: + try: + return open(path, "rb") + except OSError: + return None + except OSError as e: + if e.errno == errno.EROFS: + return None + raise + + +_live_locks = weakref.WeakSet() + + +def _renew_in_child() -> None: + # after a fork the child must not use, nor release, what the parent holds + for lock in list(_live_locks): + lock._thread_lock = threading.Lock() + lock._owner = None # the parent's section; __exit__ in the child is a no-op + lock.held = False + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_renew_in_child) + + +class FileLock: + """A reusable, exclusive, advisory lock on a sidecar file ````. + + Use it as ``with lock:``; the body runs either way, and ``lock.held`` says whether + the lock is actually held (it is not when the lock folder is read-only or the + filesystem does not support locking: advisory locks are best effort, and the caller + decides whether that matters). + + The handle is kept open between uses (cheaper for many short sections); ``close()`` + it when done. It is fork-safe: ``flock`` locks belong to the open file description, + so a forked child that inherited the handle would not exclude its parent; the pid + is remembered and the handle reopened on a mismatch (and a child never releases + its parent's lock). Threads of one process exclude each other too. Not reentrant: + entering again from the holding thread raises ``RuntimeError`` (a second lock on + the same path from the same thread would deadlock). For a one-off section, use + :func:`file_lock`. + + The lock file is a separate sidecar, not ``path`` itself, so that + :func:`atomic_write` replacing ``path`` does not swap the file being locked. + + >>> import tempfile + >>> lock = FileLock(os.path.join(tempfile.mkdtemp(), "data.json")) + >>> with lock: + ... lock.held + True + >>> lock.close() + """ + + def __init__( + self, + path, + *, + sidecar: str = DFLT_LOCK_SIDECAR, + poll_interval: float = DFLT_POLL_INTERVAL, + ) -> None: + """ + :param path: The file being protected; the lock lives next to it. + :param sidecar: Non-empty suffix of the lock file. + :param poll_interval: Seconds between attempts on a contended lock (Windows + only; POSIX blocks in the kernel). + """ + if not sidecar: + raise ValueError("sidecar must be non-empty: the lock file cannot be `path`") + self.path = os.fspath(path) + sidecar + self.poll_interval = poll_interval + self.held = False + self._fh = None + self._closer = None # weakref.finalize closing the handle on collection + self._pid = None # the process the handle belongs to + self._owner = None # thread currently inside the lock + self._thread_lock = threading.Lock() + _live_locks.add(self) + + def _ensure_handle(self) -> None: + if self._pid != os.getpid(): + self._fh = self._closer = None # inherited over a fork: leave it be + self._pid = os.getpid() + if self._fh is None: + self._fh = _open_lock_file(self.path) + if self._fh is not None: + self._closer = weakref.finalize(self, self._fh.close) + + def _try_lock(self) -> bool: + if self._fh is None: + return False + try: + _lock(self._fh, self.poll_interval) + return True + except OSError as e: + readonly = "+" not in self._fh.mode + # NFS emulates flock with fcntl locks, refused on a read-only descriptor + if e.errno in _LOCK_UNSUPPORTED or (readonly and e.errno == errno.EBADF): + return False + raise + + def __enter__(self) -> "FileLock": + if self._owner == threading.get_ident(): + raise RuntimeError(f"FileLock({self.path!r}) is not reentrant") + self._thread_lock.acquire() + try: + self._ensure_handle() + self.held = self._try_lock() + self._owner = threading.get_ident() + except BaseException: + self._thread_lock.release() + raise + return self + + def __exit__(self, *exc) -> None: + if self._owner != threading.get_ident(): + return # a forked child leaving the section its parent entered + try: + if self.held: + self.held = False + _unlock(self._fh) + finally: + self._owner = None + self._thread_lock.release() + + def close(self) -> None: + """Release the kept handle (not from inside the ``with`` block).""" + if self._owner == threading.get_ident(): + raise RuntimeError("close() called while holding the lock") + with self._thread_lock: + if self._closer is not None: + self._closer() # closes the handle, once + self._fh = self._closer = None + + +@contextmanager +def file_lock( + path, + *, + sidecar: str = DFLT_LOCK_SIDECAR, + poll_interval: float = DFLT_POLL_INTERVAL, +): + """Hold an exclusive advisory lock on ````; yields whether it is held. + + Cross-process and cross-platform (``flock`` on POSIX, ``msvcrt.locking`` on + Windows). Take it around every append, read and :func:`atomic_write` of ``path`` + so they cannot interleave. The body still runs when locking is unavailable + (read-only folder, filesystem without lock support): the yielded value is then + ``False``. A lock file that exists but is not writable is opened read-only. + Nesting it on the same path in one thread deadlocks; use one section. + + >>> import tempfile + >>> path = os.path.join(tempfile.mkdtemp(), "data.json") + >>> with file_lock(path) as held: + ... atomic_write(path, b"{}") + >>> held, os.path.exists(path + ".lock") + (True, True) + """ + lock = FileLock(path, sidecar=sidecar, poll_interval=poll_interval) + try: + with lock: + yield lock.held + finally: + lock.close() diff --git a/dol/tests/test_file_lock.py b/dol/tests/test_file_lock.py new file mode 100644 index 00000000..5694bbb1 --- /dev/null +++ b/dol/tests/test_file_lock.py @@ -0,0 +1,153 @@ +"""Tests for ``dol.file_lock`` / ``dol.FileLock``.""" + +import multiprocessing as mp +import os +import stat +import sys +import threading +import time + +import pytest + +from dol import FileLock, atomic_write, file_lock + +posix_only = pytest.mark.skipif(sys.platform == "win32", reason="POSIX only") + + +def _increment_n(path, n): + for _ in range(n): + with file_lock(path): # read-modify-write: only correct if truly exclusive + try: + count = int(open(path, "rb").read() or 0) + except FileNotFoundError: + count = 0 + atomic_write(path, str(count + 1).encode()) + + +def test_creates_sidecar_and_reports_held(tmp_path): + p = tmp_path / "f" + with file_lock(p) as held: + assert held is True + assert (tmp_path / "f.lock").exists() + with pytest.raises(ValueError): + FileLock(p, sidecar="") + + +def test_processes_exclude_each_other(tmp_path): + p = str(tmp_path / "f") + procs = [mp.Process(target=_increment_n, args=(p, 50)) for _ in range(4)] + [pr.start() for pr in procs] + [pr.join(60) for pr in procs] + assert all(pr.exitcode == 0 for pr in procs) + assert int(open(p, "rb").read()) == 200 + + +def test_threads_share_one_filelock(tmp_path): + lock = FileLock(tmp_path / "f") + inside, overlaps = [0], [0] + + def work(): + for _ in range(20): + with lock: + inside[0] += 1 + overlaps[0] += inside[0] > 1 + time.sleep(0.0005) + inside[0] -= 1 + + ts = [threading.Thread(target=work) for _ in range(4)] + [t.start() for t in ts] + [t.join() for t in ts] + lock.close() + assert overlaps[0] == 0 + + +def test_lock_released_on_exception(tmp_path): + lock = FileLock(tmp_path / "f") + with pytest.raises(ValueError): + with lock: + raise ValueError + with lock: + assert lock.held + + +def test_works_with_atomic_write(tmp_path): + p = tmp_path / "f" + with file_lock(p): + atomic_write(p, b"ok") + assert p.read_bytes() == b"ok" + + +@posix_only +def test_readonly_lock_file_still_locks(tmp_path): + p = tmp_path / "f" + lock_path = tmp_path / "f.lock" + lock_path.write_bytes(b"") + lock_path.chmod(stat.S_IRUSR) + if os.access(lock_path, os.W_OK): # running as root + pytest.skip("cannot make a file unwritable for this user") + with file_lock(p) as held: + assert held is True + + +def test_missing_folder_is_an_error_not_a_silent_no_lock(tmp_path): + with pytest.raises(FileNotFoundError): + with file_lock(tmp_path / "nodir" / "f"): + pass + + +def test_reentry_raises_instead_of_deadlocking(tmp_path): + lock = FileLock(tmp_path / "f") + with lock: + with pytest.raises(RuntimeError, match="reentrant"): + with lock: + pass + lock.close() + + +@posix_only +def test_fork_child_waits_for_parents_lock(tmp_path): + lock = FileLock(tmp_path / "f") + with lock: # the child inherits this very object, handle included + r, w = os.pipe() + pid = os.fork() + if pid == 0: + try: + t0 = time.monotonic() + with lock: + os.write(w, b"%.3f" % (time.monotonic() - t0)) + finally: + os._exit(0) + time.sleep(0.4) + waited = float(os.read(r, 16)) + os.waitpid(pid, 0) + lock.close() + assert waited >= 0.3 # a shared (inherited) lock would have let it straight in + + +@posix_only +def test_forked_child_leaving_the_block_keeps_parents_lock(tmp_path): + import fcntl + + lock = FileLock(tmp_path / "f") + with lock: + pid = os.fork() + if pid == 0: + code = 1 + try: + lock.__exit__(None, None, None) # the child leaves the inherited block + code = 0 + finally: + os._exit(code) + assert os.waitpid(pid, 0)[1] == 0 # the child exited cleanly + with open(lock.path, "rb") as fh: + with pytest.raises(OSError): + fcntl.flock(fh.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + lock.close() + + +def test_close_inside_the_block_raises(tmp_path): + lock = FileLock(tmp_path / "f") + with lock: + with pytest.raises(RuntimeError, match="close"): + lock.close() + lock.close()