1
0
Fork 0
Code Issues Pull requests Projects Releases 2 Packages Wiki Activity Actions Pages

Make command reference publication race-safe

This commit is contained in:
Andraxion 2026-07-29 15:00:28 -04:00
parent efe4a443b7
commit 3bf3836ac3
2 changed files with 361 additions and 24 deletions

View file

@ -5,6 +5,9 @@ from __future__ import annotations
import argparse
import asyncio
import contextlib
import ctypes
import errno
import fcntl
import json
import os
import secrets
@ -15,7 +18,7 @@ import tempfile
from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Literal
from typing import Literal, Protocol, cast
from mcp.shared.memory import create_connected_server_and_client_session
@ -35,6 +38,7 @@ DEFAULT_PROJECT_ROOT = REPOSITORY_ROOT / "tests" / "fixtures" / "alpha"
EXPECTED_CLI_ROWS = 28
EXPECTED_MCP_ROWS = 36
MAX_REFERENCE_BYTES = 5_000_000
RENAME_EXCHANGE = 2
class CommandReferenceToolError(RuntimeError):
@ -50,11 +54,28 @@ class _FileIdentity:
device: int
inode: int
mode: int
owner: int
group: int
size: int
modified_ns: int
changed_ns: int
class _RenameAt2(Protocol):
argtypes: list[object]
restype: object
def __call__(
self,
old_directory_fd: int,
old_name: bytes,
new_directory_fd: int,
new_name: bytes,
flags: int,
/,
) -> int: ...
async def collect_command_reference_rows(
project_root: Path,
*,
@ -133,24 +154,34 @@ def publish_or_check_command_reference(
target_name = relative.name
parent_fd = _open_relative_directory(root, parent_parts)
try:
initial = _file_identity(parent_fd, target_name)
existing = (
_read_regular_file(parent_fd, target_name, limit=MAX_REFERENCE_BYTES)
if initial is not None
else None
)
if _file_identity(parent_fd, target_name) != initial:
raise CommandReferenceToolError("Output target changed during inspection")
if check:
if existing != content:
raise CommandReferenceDrift(
f"Command reference is missing or stale: {relative.as_posix()}"
)
return "current"
if existing == content:
return "unchanged"
_atomic_replace(parent_fd, target_name, content, expected=initial)
return "written"
fcntl.flock(parent_fd, fcntl.LOCK_EX)
try:
initial = _file_identity(parent_fd, target_name)
existing = (
_read_regular_file(parent_fd, target_name, limit=MAX_REFERENCE_BYTES)
if initial is not None
else None
)
if _file_identity(parent_fd, target_name) != initial:
raise CommandReferenceToolError("Output target changed during inspection")
if check:
if existing != content:
raise CommandReferenceDrift(
f"Command reference is missing or stale: {relative.as_posix()}"
)
return "current"
if existing == content:
return "unchanged"
_atomic_replace(
parent_fd,
target_name,
content,
expected=initial,
expected_content=existing,
)
return "written"
finally:
fcntl.flock(parent_fd, fcntl.LOCK_UN)
finally:
os.close(parent_fd)
@ -226,6 +257,8 @@ def _file_identity(directory_fd: int, name: str) -> _FileIdentity | None:
device=status.st_dev,
inode=status.st_ino,
mode=status.st_mode,
owner=status.st_uid,
group=status.st_gid,
size=status.st_size,
modified_ns=status.st_mtime_ns,
changed_ns=status.st_ctime_ns,
@ -263,6 +296,8 @@ def _identity_from_stat(status: os.stat_result) -> _FileIdentity:
device=status.st_dev,
inode=status.st_ino,
mode=status.st_mode,
owner=status.st_uid,
group=status.st_gid,
size=status.st_size,
modified_ns=status.st_mtime_ns,
changed_ns=status.st_ctime_ns,
@ -275,6 +310,7 @@ def _atomic_replace(
content: bytes,
*,
expected: _FileIdentity | None,
expected_content: bytes | None,
) -> None:
temporary_name = f".{name}.docforge-command-reference-{secrets.token_hex(12)}"
flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW
@ -296,21 +332,152 @@ def _atomic_replace(
raise
else:
os.close(descriptor)
temporary_contains_only_staged_content = True
try:
staged = _file_identity(directory_fd, temporary_name)
if staged is None:
raise CommandReferenceToolError("Atomic output temporary disappeared")
if _file_identity(directory_fd, name) != expected:
raise CommandReferenceToolError("Output target changed before atomic publication")
os.replace(
if expected is None:
_link_no_replace(directory_fd, temporary_name, name)
_unlink_at(directory_fd, temporary_name)
os.fsync(directory_fd)
return
if expected_content is None:
raise CommandReferenceToolError("Expected output content was not captured")
_rename_exchange(directory_fd, temporary_name, name)
temporary_contains_only_staged_content = False
displaced = _file_identity(directory_fd, temporary_name)
published = _file_identity(directory_fd, name)
displaced_content = _read_regular_file(
directory_fd,
temporary_name,
name,
src_dir_fd=directory_fd,
dst_dir_fd=directory_fd,
limit=MAX_REFERENCE_BYTES,
)
if (
displaced is None
or published is None
or not _same_identity_after_rename(displaced, expected)
or displaced_content != expected_content
or not _same_identity_after_rename(published, staged)
):
try:
_rename_exchange(directory_fd, temporary_name, name)
except Exception as error:
raise CommandReferenceToolError(
"Output target raced publication; displaced data was retained "
f"in {temporary_name}"
) from error
temporary_contains_only_staged_content = True
if (
not _same_identity_after_rename(
_file_identity(directory_fd, name),
displaced,
)
or not _same_identity_after_rename(
_file_identity(directory_fd, temporary_name),
staged,
)
):
raise CommandReferenceToolError(
"Output target raced publication and could not be safely restored"
)
raise CommandReferenceToolError("Output target changed during atomic publication")
_unlink_at(directory_fd, temporary_name)
os.fsync(directory_fd)
except Exception:
_unlink_at(directory_fd, temporary_name)
if temporary_contains_only_staged_content:
_unlink_at(directory_fd, temporary_name)
raise
def _same_identity_after_rename(
actual: _FileIdentity | None,
expected: _FileIdentity | None,
) -> bool:
if actual is None or expected is None:
return False
return (
actual.device,
actual.inode,
actual.mode,
actual.owner,
actual.group,
actual.size,
actual.modified_ns,
) == (
expected.device,
expected.inode,
expected.mode,
expected.owner,
expected.group,
expected.size,
expected.modified_ns,
)
def _link_no_replace(directory_fd: int, source: str, target: str) -> None:
try:
os.link(
source,
target,
src_dir_fd=directory_fd,
dst_dir_fd=directory_fd,
follow_symlinks=False,
)
except FileExistsError as error:
raise CommandReferenceToolError(
"Output target appeared during atomic publication"
) from error
except OSError as error:
raise CommandReferenceToolError("Could not publish atomic output") from error
def _rename_exchange(directory_fd: int, first: str, second: str) -> None:
rename_at2 = _load_rename_at2()
ctypes.set_errno(0)
result = rename_at2(
directory_fd,
os.fsencode(first),
directory_fd,
os.fsencode(second),
RENAME_EXCHANGE,
)
if result == 0:
return
error_number = ctypes.get_errno()
if error_number in {errno.ENOSYS, errno.EINVAL, errno.EOPNOTSUPP}:
raise CommandReferenceToolError(
"Atomic exchange publication is unavailable on this filesystem"
)
raise CommandReferenceToolError("Could not exchange atomic output") from OSError(
error_number,
os.strerror(error_number),
)
def _load_rename_at2() -> _RenameAt2:
library = ctypes.CDLL(None, use_errno=True)
try:
rename_at2 = cast(_RenameAt2, library.renameat2)
except AttributeError as error:
raise CommandReferenceToolError(
"Atomic exchange publication is unavailable on this platform"
) from error
rename_at2.argtypes = [
ctypes.c_int,
ctypes.c_char_p,
ctypes.c_int,
ctypes.c_char_p,
ctypes.c_uint,
]
rename_at2.restype = ctypes.c_int
return rename_at2
def _unlink_at(directory_fd: int, name: str) -> None:
with contextlib.suppress(FileNotFoundError):
os.unlink(name, dir_fd=directory_fd)