Make command reference publication race-safe

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()

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,
     *,
@@ -132,6 +153,8 @@ def publish_or_check_command_reference(
     parent_parts = relative.parts[:-1]
     target_name = relative.name
     parent_fd = _open_relative_directory(root, parent_parts)
+    try:
+        fcntl.flock(parent_fd, fcntl.LOCK_EX)
         try:
             initial = _file_identity(parent_fd, target_name)
             existing = (
@@ -149,8 +172,16 @@ def publish_or_check_command_reference(
                 return "current"
             if existing == content:
                 return "unchanged"
-        _atomic_replace(parent_fd, target_name, content, expected=initial)
+            _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,
+            limit=MAX_REFERENCE_BYTES,
-            src_dir_fd=directory_fd,
-            dst_dir_fd=directory_fd,
         )
+        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:
+        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)