284 lines
11 KiB
Python
284 lines
11 KiB
Python
"""Small, strict, process-owned advisory locks for harness coordination."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import errno
|
|
import fcntl
|
|
import json
|
|
import os
|
|
import stat
|
|
import time
|
|
from contextlib import contextmanager
|
|
from pathlib import Path
|
|
from typing import Callable
|
|
|
|
|
|
class LockError(RuntimeError):
|
|
"""Base class for managed lock failures."""
|
|
|
|
|
|
class LockTimeoutError(LockError, TimeoutError):
|
|
"""The lock could not be acquired before the configured deadline."""
|
|
|
|
|
|
class LockValidationError(LockError):
|
|
"""The persistent lock path or held lock failed a safety check."""
|
|
|
|
|
|
class LockOrderError(LockError):
|
|
"""A lock operation would violate the canonical acquisition order."""
|
|
|
|
|
|
def _canonical(path: os.PathLike[str] | str) -> Path:
|
|
raw_value = os.fspath(path)
|
|
raw = Path(raw_value)
|
|
if any(part in {".", ".."} for part in str(raw_value).split("/")):
|
|
raise LockValidationError(f"lock path must be absolute and traversal-free: {raw_value}")
|
|
supplied = Path(os.path.abspath(os.fspath(path)))
|
|
if not supplied.is_absolute():
|
|
raise LockValidationError(f"lock path must be absolute and traversal-free: {supplied}")
|
|
# Resolve ancestors through directory FDs. Path.is_symlink() followed by
|
|
# a later open is a check-then-use race.
|
|
try:
|
|
parent_fd, name = _open_parent(supplied, create_parent=False)
|
|
except FileNotFoundError:
|
|
return supplied
|
|
except OSError as exc:
|
|
raise LockValidationError(f"cannot inspect lock path {supplied}") from exc
|
|
try:
|
|
try:
|
|
st = os.stat(name, dir_fd=parent_fd, follow_symlinks=False)
|
|
if stat.S_ISLNK(st.st_mode):
|
|
raise LockValidationError(f"lock path must not traverse a symlink: {supplied}")
|
|
except FileNotFoundError:
|
|
pass
|
|
finally:
|
|
os.close(parent_fd)
|
|
return supplied
|
|
|
|
|
|
def _open_parent(path: Path, *, create_parent: bool) -> tuple[int, str]:
|
|
"""Open *path*'s parent one component at a time, without symlinks."""
|
|
parts = path.parts
|
|
fd = os.open(parts[0], os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW | os.O_CLOEXEC)
|
|
try:
|
|
for part in parts[1:-1]:
|
|
try:
|
|
nxt = os.open(part, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW | os.O_CLOEXEC,
|
|
dir_fd=fd)
|
|
except FileNotFoundError:
|
|
if not create_parent:
|
|
raise
|
|
os.mkdir(part, 0o700, dir_fd=fd)
|
|
nxt = os.open(part, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW | os.O_CLOEXEC,
|
|
dir_fd=fd)
|
|
os.close(fd)
|
|
fd = nxt
|
|
return fd, parts[-1]
|
|
except BaseException:
|
|
os.close(fd)
|
|
raise
|
|
|
|
|
|
def _validate_file(fd: int, path: Path) -> None:
|
|
try:
|
|
fd_stat = os.fstat(fd)
|
|
except OSError as exc:
|
|
raise LockValidationError(f"cannot inspect lock file {path}: {exc}") from exc
|
|
if not stat.S_ISREG(fd_stat.st_mode):
|
|
raise LockValidationError(f"lock is not a regular file: {path}")
|
|
if fd_stat.st_uid != os.getuid():
|
|
raise LockValidationError(f"lock is not owned by the current uid: {path}")
|
|
if stat.S_IMODE(fd_stat.st_mode) != 0o600:
|
|
raise LockValidationError(f"lock mode must be exactly 0600: {path}")
|
|
if fd_stat.st_nlink != 1:
|
|
raise LockValidationError(f"lock must have exactly one link: {path}")
|
|
|
|
|
|
def _open_lock(path: Path) -> int:
|
|
flags = os.O_RDWR | os.O_NOFOLLOW | os.O_CLOEXEC
|
|
parent_fd, name = _open_parent(path, create_parent=True)
|
|
try:
|
|
fd = os.open(name, flags | os.O_CREAT, 0o600, dir_fd=parent_fd)
|
|
except OSError as exc:
|
|
if exc.errno in (errno.ELOOP, errno.ENOTDIR):
|
|
raise LockValidationError(f"lock path is not a safe regular file: {path}") from exc
|
|
raise
|
|
try:
|
|
_validate_file(fd, path)
|
|
return fd
|
|
except BaseException:
|
|
os.close(fd)
|
|
raise
|
|
finally:
|
|
os.close(parent_fd)
|
|
|
|
|
|
def validate_held_flock(fd: int, path: os.PathLike[str] | str) -> bool:
|
|
"""Verify *fd* has this process's exclusive whole-file kernel FLOCK."""
|
|
canonical = _canonical(path)
|
|
_validate_file(fd, canonical)
|
|
try:
|
|
with open(f"/proc/self/fdinfo/{fd}", encoding="ascii") as info:
|
|
records = info.read().splitlines()
|
|
except (OSError, ValueError) as exc:
|
|
raise LockValidationError(f"cannot inspect fdinfo for fd {fd}") from exc
|
|
st = os.fstat(fd)
|
|
device = f"{os.major(st.st_dev):02x}:{os.minor(st.st_dev):02x}:{st.st_ino}"
|
|
for line in records:
|
|
fields = line.split()
|
|
if len(fields) == 9 and fields[0].startswith("lock:"):
|
|
if (fields[2:7] == ["FLOCK", "ADVISORY", "WRITE", str(os.getpid()), device]
|
|
and fields[7:] == ["0", "EOF"]):
|
|
return True
|
|
return False
|
|
|
|
|
|
class ManagedLock:
|
|
"""A persistent lock file whose exclusive FLOCK lasts until release."""
|
|
|
|
def __init__(self, path: os.PathLike[str] | str, *, timeout: float = 0.0,
|
|
poll_interval: float = 0.01) -> None:
|
|
if timeout < 0 or poll_interval <= 0:
|
|
raise ValueError("timeout must be non-negative and poll_interval positive")
|
|
self.path = _canonical(path)
|
|
self.timeout = timeout
|
|
self.poll_interval = poll_interval
|
|
self.fd: int | None = None
|
|
|
|
@property
|
|
def held(self) -> bool:
|
|
return self.fd is not None
|
|
|
|
def acquire(self) -> "ManagedLock":
|
|
if self.fd is not None:
|
|
raise LockError("lock is already acquired")
|
|
fd = _open_lock(self.path)
|
|
deadline = time.monotonic() + self.timeout
|
|
try:
|
|
while True:
|
|
try:
|
|
fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
if not validate_held_flock(fd, self.path):
|
|
raise LockValidationError("kernel FLOCK was not recorded as expected")
|
|
self.fd = fd
|
|
return self
|
|
except BlockingIOError:
|
|
if time.monotonic() >= deadline:
|
|
raise LockTimeoutError(f"timed out acquiring {self.path}")
|
|
time.sleep(min(self.poll_interval, max(0.0, deadline - time.monotonic())))
|
|
except BaseException:
|
|
os.close(fd)
|
|
raise
|
|
|
|
def release(self) -> None:
|
|
if self.fd is None:
|
|
return
|
|
fd, self.fd = self.fd, None
|
|
try:
|
|
fcntl.flock(fd, fcntl.LOCK_UN)
|
|
finally:
|
|
os.close(fd)
|
|
|
|
def write_metadata(self, metadata: object) -> None:
|
|
if self.fd is None or not validate_held_flock(self.fd, self.path):
|
|
raise RuntimeError("metadata can only be written while this lock is held")
|
|
payload = (json.dumps(metadata, sort_keys=True, separators=(",", ":")) + "\n").encode()
|
|
os.ftruncate(self.fd, 0)
|
|
os.lseek(self.fd, 0, os.SEEK_SET)
|
|
written = 0
|
|
while written < len(payload):
|
|
written += os.write(self.fd, payload[written:])
|
|
os.fsync(self.fd)
|
|
|
|
def __enter__(self) -> "ManagedLock":
|
|
return self.acquire()
|
|
|
|
def __exit__(self, exc_type, exc_value, traceback) -> None:
|
|
self.release()
|
|
|
|
|
|
def probe_free(path: os.PathLike[str] | str, *, poll_interval: float = 0.01) -> bool:
|
|
"""Probe a lock without removing or repairing its persistent file."""
|
|
try:
|
|
lock = ManagedLock(path, timeout=0.0, poll_interval=poll_interval).acquire()
|
|
except LockTimeoutError:
|
|
return False
|
|
lock.release()
|
|
return True
|
|
|
|
|
|
class ManagedLockSet:
|
|
"""Manage an Issue lock with a short-lived nested repository scope."""
|
|
|
|
@staticmethod
|
|
def _component(value: object, name: str) -> str:
|
|
if not isinstance(value, str) or not value or value in {".", ".."} or "\x00" in value:
|
|
raise ValueError(f"{name} must be one safe path component")
|
|
if "/" in value or "\\" in value:
|
|
raise ValueError(f"{name} must be one safe path component")
|
|
return value
|
|
|
|
@classmethod
|
|
def paths_for_issue(cls, harness_root: os.PathLike[str] | str,
|
|
owner: str, repository: str, issue: int) -> tuple[Path, Path]:
|
|
root = Path(os.fspath(harness_root))
|
|
if not root.is_absolute() or "." in root.parts or ".." in root.parts:
|
|
raise ValueError("harness_root must be an absolute, traversal-free path")
|
|
owner = cls._component(owner, "owner")
|
|
repository = cls._component(repository, "repository")
|
|
if isinstance(issue, bool) or not isinstance(issue, int) or issue <= 0:
|
|
raise ValueError("issue must be a positive integer")
|
|
locks = root / "repositories" / owner / repository / "locks"
|
|
return locks / f"issue-{issue}.lock", locks / "repository.lock"
|
|
|
|
@classmethod
|
|
def for_issue(cls, harness_root: os.PathLike[str] | str,
|
|
owner: str, repository: str, issue: int, **kwargs) -> "ManagedLockSet":
|
|
issue_path, repository_path = cls.paths_for_issue(harness_root, owner, repository, issue)
|
|
return cls(issue_path, repository_path, **kwargs)
|
|
|
|
def __init__(self, issue: os.PathLike[str] | str, repository: os.PathLike[str] | str,
|
|
*, timeout: float = 0.0, poll_interval: float = 0.01,
|
|
event_sink: Callable[[str], None] | None = None) -> None:
|
|
self.issue = ManagedLock(issue, timeout=timeout, poll_interval=poll_interval)
|
|
self.repository = ManagedLock(repository, timeout=timeout, poll_interval=poll_interval)
|
|
self._event_sink = event_sink or (lambda event: None)
|
|
|
|
def __enter__(self) -> "ManagedLockSet":
|
|
try:
|
|
self.issue.acquire()
|
|
self._event_sink("acquire:issue")
|
|
return self
|
|
except BaseException:
|
|
if self.issue.held:
|
|
self.issue.release()
|
|
self._event_sink("release:issue")
|
|
raise
|
|
|
|
def __exit__(self, exc_type, exc_value, traceback) -> None:
|
|
if self.repository.held:
|
|
self.repository.release()
|
|
self._event_sink("release:repository")
|
|
self.release_issue()
|
|
|
|
@contextmanager
|
|
def repository_scope(self):
|
|
if not self.issue.held:
|
|
raise LockOrderError("repository lock requires the Issue lock")
|
|
self.repository.acquire()
|
|
try:
|
|
self._event_sink("acquire:repository")
|
|
yield self
|
|
finally:
|
|
if self.repository.held:
|
|
self.repository.release()
|
|
self._event_sink("release:repository")
|
|
|
|
def release_issue(self) -> None:
|
|
if self.repository.held:
|
|
raise LockOrderError("release the repository lock before the Issue lock")
|
|
if self.issue.held:
|
|
self.issue.release()
|
|
self._event_sink("release:issue")
|