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

Make command reference publication race-safe

Browse source
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
Download patch file Download diff file Expand all files Collapse all files

170
tests/test_command_reference_tool.py
View file

@ -6,10 +6,13 @@ import hashlib
import io
import os
import tempfile
import threading
import unittest
from pathlib import Path
from unittest import mock
from docforge.mcp_server import ALL_TOOLS, APPLICATION_TOOLS
from tools import generate_command_reference as command_reference_tool
from tools.generate_command_reference import (
EXPECTED_CLI_ROWS,
EXPECTED_MCP_ROWS,
@ -183,6 +186,173 @@ class CommandReferenceToolTests(unittest.TestCase):
check=False,
)
def test_existing_target_race_is_restored_without_overwrite(self) -> None:
with tempfile.TemporaryDirectory() as directory:
repository = Path(directory).resolve()
generated = repository / "generated"
generated.mkdir()
output = Path("generated/reference.md")
target = repository / output
target.write_bytes(b"initial\n")
original_exchange = command_reference_tool._rename_exchange
injected = False
def exchange_after_race(directory_fd: int, first: str, second: str) -> None:
nonlocal injected
if not injected:
injected = True
target.write_bytes(b"concurrent\n")
original_exchange(directory_fd, first, second)
with (
mock.patch.object(
command_reference_tool,
"_rename_exchange",
side_effect=exchange_after_race,
),
self.assertRaisesRegex(
CommandReferenceToolError,
"changed during atomic publication",
),
):
publish_or_check_command_reference(
b"generated\n",
repository_root=repository,
project_root=ALPHA,
output=output,
check=False,
)
self.assertEqual(b"concurrent\n", target.read_bytes())
self.assertEqual(
[],
list(generated.glob(".reference.md.docforge-command-reference-*")),
)
def test_missing_target_race_is_never_clobbered(self) -> None:
with tempfile.TemporaryDirectory() as directory:
repository = Path(directory).resolve()
generated = repository / "generated"
generated.mkdir()
output = Path("generated/reference.md")
target = repository / output
original_link = command_reference_tool._link_no_replace
injected = False
def link_after_race(directory_fd: int, source: str, destination: str) -> None:
nonlocal injected
if not injected:
injected = True
target.write_bytes(b"concurrent\n")
original_link(directory_fd, source, destination)
with (
mock.patch.object(
command_reference_tool,
"_link_no_replace",
side_effect=link_after_race,
),
self.assertRaisesRegex(CommandReferenceToolError, "appeared"),
):
publish_or_check_command_reference(
b"generated\n",
repository_root=repository,
project_root=ALPHA,
output=output,
check=False,
)
self.assertEqual(b"concurrent\n", target.read_bytes())
self.assertEqual(
[],
list(generated.glob(".reference.md.docforge-command-reference-*")),
)
def test_concurrent_generators_are_serialized_around_inspection_and_write(self) -> None:
with tempfile.TemporaryDirectory() as directory:
repository = Path(directory).resolve()
(repository / "generated").mkdir()
output = Path("generated/reference.md")
target = repository / output
target.write_bytes(b"initial\n")
first_entered = threading.Event()
release_first = threading.Event()
second_started = threading.Event()
second_entered = threading.Event()
results: list[str] = []
errors: list[BaseException] = []
original_replace = command_reference_tool._atomic_replace
invocation_count = 0
invocation_lock = threading.Lock()
def blocking_replace(
directory_fd: int,
name: str,
content: bytes,
*,
expected: object,
expected_content: bytes | None,
) -> None:
nonlocal invocation_count
with invocation_lock:
invocation_count += 1
invocation = invocation_count
if invocation == 1:
first_entered.set()
if not release_first.wait(timeout=5):
raise AssertionError("Timed out waiting to release first generator")
else:
second_entered.set()
original_replace(
directory_fd,
name,
content,
expected=expected, # type: ignore[arg-type]
expected_content=expected_content,
)
def publish(content: bytes, *, started: threading.Event | None = None) -> None:
if started is not None:
started.set()
try:
results.append(
publish_or_check_command_reference(
content,
repository_root=repository,
project_root=ALPHA,
output=output,
check=False,
)
)
except BaseException as error:
errors.append(error)
with mock.patch.object(
command_reference_tool,
"_atomic_replace",
side_effect=blocking_replace,
):
first = threading.Thread(target=publish, args=(b"first\n",))
second = threading.Thread(
target=publish,
args=(b"second\n",),
kwargs={"started": second_started},
)
first.start()
self.assertTrue(first_entered.wait(timeout=5))
second.start()
self.assertTrue(second_started.wait(timeout=5))
self.assertFalse(second_entered.wait(timeout=0.1))
release_first.set()
first.join(timeout=5)
second.join(timeout=5)
self.assertFalse(first.is_alive())
self.assertFalse(second.is_alive())
self.assertEqual([], errors)
self.assertEqual(["written", "written"], results)
self.assertEqual(b"second\n", target.read_bytes())
def test_command_runs_with_explicit_repository_and_project_roots(self) -> None:
with tempfile.TemporaryDirectory() as directory:
repository = Path(directory).resolve()

215
tools/generate_command_reference.py
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)