diff --git a/tests/test_command_reference_tool.py b/tests/test_command_reference_tool.py index a6c4feb..57d88f7 100644 --- a/tests/test_command_reference_tool.py +++ b/tests/test_command_reference_tool.py @@ -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() diff --git a/tools/generate_command_reference.py b/tools/generate_command_reference.py index e448467..0a605f7 100644 --- a/tools/generate_command_reference.py +++ b/tools/generate_command_reference.py @@ -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)