chemenu/tools/chemenu/filelock.py

"""An exclusive lock on an open file that holds on POSIX and on Windows alike.

Two writers share files across processes: the budget state (`commands/run_budget.py`),
where a lost race loses an increment and lets a session run past its gate, and the
telemetry trace (`telemetry/writer.py`), where it lets one line land inside another.
Both used `fcntl.flock` and skipped the lock where `fcntl` does not exist - which is
every native Windows install, silently. This module is the one place that knows the
platform difference, and the only module allowed to import `fcntl`
(`tests/test_portability_guards.py` holds that).

Windows has no `flock`. `msvcrt.locking` locks a byte range instead, and the lock is
mandatory: a locked byte cannot be read or written by any other process. So the lock
does not go on the file's data, which a reader such as `trace_ingest.py` would then fail
on, but on one byte far past it. Windows allows a lock beyond the end of a file, and
there it excludes the other lockers and nobody else.
"""
from __future__ import annotations

import os
from contextlib import contextmanager
from typing import IO, Iterator

# Where the Windows lock sits: past any size a trace or budget file reaches, and
# still inside a signed 32-bit offset.
WINDOWS_LOCK_OFFSET = 2**31 - 2


@contextmanager
def exclusive(handle: IO) -> Iterator[None]:
    """Hold an exclusive lock on `handle` for the duration of the block.

    Blocks until the lock is free, on both platforms. The handle's file position is
    left where it was, so an append handle keeps appending.
    """
    try:
        import fcntl
    except ImportError:
        fcntl = None

    if fcntl is not None:
        fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
        try:
            yield
        finally:
            fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
        return

    import msvcrt

    handle.flush()
    _windows_lock(handle.fileno(), msvcrt, msvcrt.LK_LOCK)
    try:
        yield
    finally:
        handle.flush()
        _windows_lock(handle.fileno(), msvcrt, msvcrt.LK_UNLCK)


def _windows_lock(fd: int, msvcrt, mode: int) -> None:
    """Lock or unlock the one byte at WINDOWS_LOCK_OFFSET.

    `msvcrt.locking` acts on the current position, so this seeks there and back.
    `LK_LOCK` gives up after ten one-second retries with an OSError; it is retried
    here until it succeeds, which is what `flock` does on POSIX.
    """
    position = os.lseek(fd, 0, os.SEEK_CUR)
    os.lseek(fd, WINDOWS_LOCK_OFFSET, os.SEEK_SET)
    try:
        while True:
            try:
                msvcrt.locking(fd, mode, 1)
                return
            except OSError:
                if mode != msvcrt.LK_LOCK:
                    raise
    finally:
        os.lseek(fd, position, os.SEEK_SET)