136 lines
3.6 KiB
Python
136 lines
3.6 KiB
Python
"""Durable local-file primitives shared by host and sandbox file mutations.
|
|
|
|
Writes are staged in the destination directory, fsynced, then published with
|
|
``os.replace`` so readers see either the old complete file or the new complete
|
|
file. ``interprocess_file_lock`` uses the operating system's advisory lock;
|
|
the lock file may remain on disk, but the lock itself is released automatically
|
|
when a process exits.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import tempfile
|
|
import time
|
|
from contextlib import contextmanager
|
|
from pathlib import Path
|
|
from typing import Iterator, Optional
|
|
|
|
|
|
class FileLockBusy(RuntimeError):
|
|
"""A non-blocking or timed inter-process lock could not be acquired."""
|
|
|
|
|
|
def _atomic_replace(path: Path, data: bytes) -> None:
|
|
path = Path(path)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
old_mode: Optional[int] = None
|
|
try:
|
|
old_mode = path.stat().st_mode
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
fd, raw_tmp = tempfile.mkstemp(
|
|
prefix=f".{path.name}.", suffix=".tmp", dir=str(path.parent)
|
|
)
|
|
tmp = Path(raw_tmp)
|
|
try:
|
|
with os.fdopen(fd, "wb") as f:
|
|
f.write(data)
|
|
f.flush()
|
|
os.fsync(f.fileno())
|
|
if old_mode is not None:
|
|
os.chmod(tmp, old_mode)
|
|
os.replace(tmp, path)
|
|
# Persist the directory entry on POSIX. Windows cannot open directories
|
|
# this way; os.replace still gives atomic visibility there.
|
|
if os.name != "nt":
|
|
try:
|
|
dir_fd = os.open(path.parent, os.O_RDONLY)
|
|
try:
|
|
os.fsync(dir_fd)
|
|
finally:
|
|
os.close(dir_fd)
|
|
except OSError:
|
|
# Some network/virtual filesystems reject directory fsync.
|
|
# The file itself is already fsynced and atomically visible.
|
|
pass
|
|
finally:
|
|
try:
|
|
tmp.unlink()
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
|
|
def atomic_write_bytes(path: Path, data: bytes) -> None:
|
|
_atomic_replace(Path(path), data)
|
|
|
|
|
|
def atomic_write_text(path: Path, text: str, encoding: str = "utf-8") -> None:
|
|
_atomic_replace(Path(path), text.encode(encoding))
|
|
|
|
|
|
def _try_lock(f) -> bool:
|
|
f.seek(0)
|
|
if os.name == "nt":
|
|
import msvcrt
|
|
|
|
try:
|
|
msvcrt.locking(f.fileno(), msvcrt.LK_NBLCK, 1)
|
|
return True
|
|
except OSError:
|
|
return False
|
|
|
|
import fcntl
|
|
|
|
try:
|
|
fcntl.flock(f.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
return True
|
|
except BlockingIOError:
|
|
return False
|
|
|
|
|
|
def _unlock(f) -> None:
|
|
f.seek(0)
|
|
if os.name == "nt":
|
|
import msvcrt
|
|
|
|
msvcrt.locking(f.fileno(), msvcrt.LK_UNLCK, 1)
|
|
return
|
|
|
|
import fcntl
|
|
|
|
fcntl.flock(f.fileno(), fcntl.LOCK_UN)
|
|
|
|
|
|
@contextmanager
|
|
def interprocess_file_lock(
|
|
path: Path,
|
|
*,
|
|
timeout_seconds: Optional[float] = 0,
|
|
poll_seconds: float = 0.05,
|
|
) -> Iterator[None]:
|
|
"""Acquire an advisory exclusive lock.
|
|
|
|
``timeout_seconds=0`` is non-blocking; ``None`` waits indefinitely.
|
|
"""
|
|
path = Path(path)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with open(path, "a+b") as f:
|
|
f.seek(0, os.SEEK_END)
|
|
if f.tell() == 0:
|
|
f.write(b"\0")
|
|
f.flush()
|
|
|
|
deadline = (
|
|
None if timeout_seconds is None
|
|
else time.monotonic() + max(0.0, timeout_seconds)
|
|
)
|
|
while not _try_lock(f):
|
|
if deadline is not None and time.monotonic() >= deadline:
|
|
raise FileLockBusy(str(path))
|
|
time.sleep(poll_seconds)
|
|
try:
|
|
yield
|
|
finally:
|
|
_unlock(f)
|