diff --git a/harness/tests/test_dwh_preprocess_job.py b/harness/tests/test_dwh_preprocess_job.py index c2e810ad..1c5b2a13 100644 --- a/harness/tests/test_dwh_preprocess_job.py +++ b/harness/tests/test_dwh_preprocess_job.py @@ -120,7 +120,6 @@ def test_writer_rejects_dangling_legacy_symlinks_before_claim_or_callback(tmp_pa def test_owner_publication_remains_on_locked_root_when_path_is_swapped( monkeypatch, tmp_path, ): - import os import pytest import tht.jobs.dwh_pipeline as module @@ -135,6 +134,7 @@ def test_owner_publication_remains_on_locked_root_when_path_is_swapped( swapped = True replacement.rename(moved) replacement.mkdir(mode=0o700) + (moved / "generation.lock").rename(replacement / "generation.lock") return real_replace(source, destination, *args, **kwargs) monkeypatch.setattr(module.os, "replace", swapping_replace) @@ -149,7 +149,7 @@ def test_owner_publication_remains_on_locked_root_when_path_is_swapped( assert swapped assert (moved / "OWNER.json").is_file() assert not (replacement / "OWNER.json").exists() - assert os.stat(moved / "generation.lock").st_ino != os.stat(replacement).st_ino + assert (replacement / "generation.lock").is_file() def test_owner_requires_exact_read_only_owner_mode_and_active_requires_binding(tmp_path): diff --git a/harness/tht/jobs/dwh_pipeline.py b/harness/tht/jobs/dwh_pipeline.py index 8567737c..a4d8722e 100644 --- a/harness/tht/jobs/dwh_pipeline.py +++ b/harness/tht/jobs/dwh_pipeline.py @@ -49,9 +49,21 @@ def _binding_digest(binding: dict[str, str]) -> str: def _read_root_binding(workspace_root: Path) -> dict[str, str]: - marker = workspace_root / ".tht-dwh" / OWNER_MARKER try: - fd = os.open(marker, os.O_RDONLY | os.O_NOFOLLOW) + root_fd = os.open( + workspace_root / ".tht-dwh", os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW + ) + try: + return _read_root_binding_fd(root_fd) + finally: + os.close(root_fd) + except OSError as error: + raise CorruptCheckpointError("DWH workspace ownership marker is missing or invalid") from error + + +def _read_root_binding_fd(root_fd: int) -> dict[str, str]: + try: + fd = os.open(OWNER_MARKER, os.O_RDONLY | os.O_NOFOLLOW, dir_fd=root_fd) try: info = os.fstat(fd) if ( @@ -81,14 +93,26 @@ def _read_root_binding(workspace_root: Path) -> dict[str, str]: def _validate_root_binding(workspace_root: Path, expected: dict[str, str]) -> None: - if _read_root_binding(workspace_root) != expected: + try: + root_fd = os.open( + workspace_root / ".tht-dwh", os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW + ) + except OSError as error: + raise CorruptCheckpointError("DWH workspace ownership marker is missing or invalid") from error + try: + _validate_root_binding_fd(root_fd, expected) + finally: + os.close(root_fd) + + +def _validate_root_binding_fd(root_fd: int, expected: dict[str, str]) -> None: + if _read_root_binding_fd(root_fd) != expected: raise CorruptCheckpointError("DWH artifacts belong to a different workspace configuration") - root = workspace_root / ".tht-dwh" - generations = root / "generations" - active_exists = (root / "ACTIVE").exists() + entries = set(os.listdir(root_fd)) + active_exists = "ACTIVE" in entries try: generations_fd = os.open( - generations, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW + "generations", os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=root_fd ) except FileNotFoundError: generation_entries = [] @@ -107,22 +131,13 @@ def _validate_root_binding(workspace_root: Path, expected: dict[str, str]) -> No def _claim_or_validate_root_binding( - workspace_root: Path, binding: dict[str, str], lock_fd: int + root_fd: int, binding: dict[str, str] ) -> None: - root = workspace_root / ".tht-dwh" - root_fd = os.open(root, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW) temporary: str | None = None try: entries = set(os.listdir(root_fd)) - opened_lock = os.open("generation.lock", os.O_RDONLY | os.O_NOFOLLOW, dir_fd=root_fd) - try: - held_info, opened_info = os.fstat(lock_fd), os.fstat(opened_lock) - if (held_info.st_dev, held_info.st_ino) != (opened_info.st_dev, opened_info.st_ino): - raise OSError("DWH workspace root changed while locked") - finally: - os.close(opened_lock) if OWNER_MARKER in entries: - _validate_root_binding(workspace_root, binding) + _validate_root_binding_fd(root_fd, binding) return if not entries <= {"generation.lock", "generations"} or "generation.lock" not in entries: raise CorruptCheckpointError( @@ -173,7 +188,6 @@ def _claim_or_validate_root_binding( os.unlink(temporary, dir_fd=root_fd) except FileNotFoundError: pass - os.close(root_fd) @dataclass(frozen=True) @@ -183,39 +197,76 @@ class DwhArtifactSnapshot: lsh_dir: Path +@dataclass +class _GenerationLease: + root_fd: int + lock_fd: int + + def assert_root_identity(self, root: Path) -> None: + try: + path_info = os.stat(root, follow_symlinks=False) + opened_info = os.fstat(self.root_fd) + except OSError as error: + raise CorruptCheckpointError("DWH workspace root changed while locked") from error + if ( + not stat.S_ISDIR(path_info.st_mode) + or (path_info.st_dev, path_info.st_ino) + != (opened_info.st_dev, opened_info.st_ino) + ): + raise CorruptCheckpointError("DWH workspace root changed while locked") + + def close(self) -> None: + try: + fcntl.flock(self.lock_fd, fcntl.LOCK_UN) + finally: + try: + os.close(self.lock_fd) + finally: + os.close(self.root_fd) + + class DwhSnapshotLease: def __init__(self, cfg) -> None: self.cfg = cfg - self._fd: int | None = None + self._lease: _GenerationLease | None = None self.snapshot: DwhArtifactSnapshot | None = None def __enter__(self) -> DwhArtifactSnapshot: binding = config_dwh_binding(self.cfg) - self._fd = _acquire_existing_generation_lock( + self._lease = _acquire_existing_generation_lock( self.cfg.paths.artifacts.parent, exclusive=False ) try: - self.snapshot = _resolve_dwh_snapshot_locked(self.cfg, binding) + self.snapshot = _resolve_dwh_snapshot_locked( + self.cfg, binding, self._lease.root_fd + ) return self.snapshot except BaseException: self.__exit__(None, None, None) raise def __exit__(self, *_args) -> None: - if self._fd is not None: - fd, self._fd = self._fd, None - fcntl.flock(fd, fcntl.LOCK_UN) - os.close(fd) + if self._lease is not None: + lease, self._lease = self._lease, None + lease.close() def lease_dwh_snapshot(cfg) -> DwhSnapshotLease: return DwhSnapshotLease(cfg) -def _acquire_generation_lock(workspace_root: Path, *, exclusive: bool) -> int: +def _acquire_generation_lock(workspace_root: Path, *, exclusive: bool) -> _GenerationLease: root = workspace_root / ".tht-dwh" DwhPreprocessPipeline._ensure_owned_dir(root) - fd = os.open(root / "generation.lock", os.O_RDWR | os.O_CREAT | os.O_NOFOLLOW, 0o600) + root_fd = os.open(root, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW) + try: + fd = os.open( + "generation.lock", os.O_RDWR | os.O_CREAT | os.O_NOFOLLOW, + 0o600, dir_fd=root_fd, + ) + except BaseException: + os.close(root_fd) + raise try: info = os.fstat(fd) if ( @@ -226,13 +277,16 @@ def _acquire_generation_lock(workspace_root: Path, *, exclusive: bool) -> int: ): raise OSError("unsafe DWH generation lock") fcntl.flock(fd, fcntl.LOCK_EX if exclusive else fcntl.LOCK_SH) - return fd + return _GenerationLease(root_fd, fd) except BaseException: os.close(fd) + os.close(root_fd) raise -def _acquire_existing_generation_lock(workspace_root: Path, *, exclusive: bool) -> int: +def _acquire_existing_generation_lock( + workspace_root: Path, *, exclusive: bool +) -> _GenerationLease: root = workspace_root / ".tht-dwh" try: root_fd = os.open(root, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW) @@ -241,8 +295,9 @@ def _acquire_existing_generation_lock(workspace_root: Path, *, exclusive: bool) if root_info.st_uid != os.getuid() or stat.S_IMODE(root_info.st_mode) != 0o700: raise OSError("unsafe DWH workspace root") fd = os.open("generation.lock", os.O_RDWR | os.O_NOFOLLOW, dir_fd=root_fd) - finally: + except BaseException: os.close(root_fd) + raise try: info = os.fstat(fd) if ( @@ -253,9 +308,10 @@ def _acquire_existing_generation_lock(workspace_root: Path, *, exclusive: bool) ): raise OSError("unsafe DWH generation lock") fcntl.flock(fd, fcntl.LOCK_EX if exclusive else fcntl.LOCK_SH) - return fd + return _GenerationLease(root_fd, fd) except BaseException: os.close(fd) + os.close(root_fd) raise except OSError as error: raise CorruptCheckpointError( @@ -364,21 +420,22 @@ def validate_generation( def resolve_dwh_snapshot(cfg) -> DwhArtifactSnapshot: binding = config_dwh_binding(cfg) - lease_fd = _acquire_existing_generation_lock( + lease = _acquire_existing_generation_lock( cfg.paths.artifacts.parent, exclusive=False ) try: - return _resolve_dwh_snapshot_locked(cfg, binding) + return _resolve_dwh_snapshot_locked(cfg, binding, lease.root_fd) finally: - fcntl.flock(lease_fd, fcntl.LOCK_UN) - os.close(lease_fd) + lease.close() def _resolve_dwh_snapshot_locked( - cfg, binding: dict[str, str], + cfg, binding: dict[str, str], root_fd: int, ) -> DwhArtifactSnapshot: - _validate_root_binding(cfg.paths.artifacts.parent, binding) - target = _active_generation_dir_locked(cfg.paths.artifacts.parent, binding) + _validate_root_binding_fd(root_fd, binding) + target = _active_generation_dir_locked( + cfg.paths.artifacts.parent, binding, root_fd=root_fd + ) if target is None: return DwhArtifactSnapshot( None, cfg.paths.artifacts / "mschema" / "physical.yaml", cfg.paths.indexes / "lsh" @@ -390,18 +447,22 @@ def _resolve_dwh_snapshot_locked( def active_generation_dir( workspace_root: Path, expected_binding: dict[str, str] ) -> Path | None: - lease_fd = _acquire_existing_generation_lock(workspace_root, exclusive=False) + lease = _acquire_existing_generation_lock(workspace_root, exclusive=False) try: - return _active_generation_dir_locked(workspace_root, expected_binding) + return _active_generation_dir_locked( + workspace_root, expected_binding, root_fd=lease.root_fd + ) finally: - fcntl.flock(lease_fd, fcntl.LOCK_UN) - os.close(lease_fd) + lease.close() def _active_generation_dir_locked( - workspace_root: Path, expected_binding: dict[str, str] + workspace_root: Path, expected_binding: dict[str, str], *, root_fd: int | None = None ) -> Path | None: - _validate_root_binding(workspace_root, expected_binding) + if root_fd is None: + _validate_root_binding(workspace_root, expected_binding) + else: + _validate_root_binding_fd(root_fd, expected_binding) pointer = workspace_root / ".tht-dwh" / "ACTIVE" try: generation = _read_owned(pointer, readonly=False).decode("utf-8").strip() @@ -465,9 +526,11 @@ class DwhPreprocessPipeline: "input_fingerprint": self.input_fingerprint, } - def _assert_active_binding(self) -> None: - _validate_root_binding(self.workspace_root, self.binding) - active = _active_generation_dir_locked(self.workspace_root, self.binding) + def _assert_active_binding(self, root_fd: int) -> None: + _validate_root_binding_fd(root_fd, self.binding) + active = _active_generation_dir_locked( + self.workspace_root, self.binding, root_fd=root_fd + ) if active is not None: validate_generation(active, self.binding) @@ -475,18 +538,20 @@ class DwhPreprocessPipeline: self, steps: tuple[str, ...] = DWH_STAGE_IDS, *, resume_run_id: str | None = None ) -> JobReport: self._validate_steps(steps) - lease_fd = _acquire_generation_lock(self.workspace_root, exclusive=True) + lease = _acquire_generation_lock(self.workspace_root, exclusive=True) try: - self._assert_no_legacy_artifacts() - _claim_or_validate_root_binding(self.workspace_root, self.binding, lease_fd) - self._assert_active_binding() - active = _active_generation_dir_locked(self.workspace_root, self.binding) + self._assert_no_legacy_artifacts(lease.root_fd) + _claim_or_validate_root_binding(lease.root_fd, self.binding) + lease.assert_root_identity(self.workspace_root / ".tht-dwh") + self._assert_active_binding(lease.root_fd) + active = _active_generation_dir_locked( + self.workspace_root, self.binding, root_fd=lease.root_fd + ) if active is not None: self.current_physical = active / "physical.yaml" self.current_lsh_dir = active finally: - fcntl.flock(lease_fd, fcntl.LOCK_UN) - os.close(lease_fd) + lease.close() if resume_run_id is not None: self._validate_resume_publication(resume_run_id) spec = JobSpec( @@ -575,9 +640,8 @@ class DwhPreprocessPipeline: raise CorruptCheckpointError("resume artifact manifest is invalid") from error self._validate_published(target, run_dir / "artifacts", required) - def _assert_no_legacy_artifacts(self) -> None: - marker = self.workspace_root / ".tht-dwh" / OWNER_MARKER - if marker.exists() or marker.is_symlink(): + def _assert_no_legacy_artifacts(self, root_fd: int) -> None: + if OWNER_MARKER in os.listdir(root_fd): return legacy_physical = self.current_physical is not None and ( self.current_physical.exists() or self.current_physical.is_symlink() @@ -624,17 +688,16 @@ class DwhPreprocessPipeline: artifacts = self._artifacts(context) required = self._required(artifacts) seal_stage_artifacts(context, stage, required, spec) - lease_fd = _acquire_generation_lock(self.workspace_root, exclusive=True) + lease = _acquire_generation_lock(self.workspace_root, exclusive=True) try: - _validate_root_binding(self.workspace_root, self.binding) - self._assert_active_binding() + _validate_root_binding_fd(lease.root_fd, self.binding) + self._assert_active_binding(lease.root_fd) self._publish(context.run_id, artifacts, required) if self.after_publish is not None: self.after_publish(context.run_id) self._cleanup_generations() finally: - fcntl.flock(lease_fd, fcntl.LOCK_UN) - os.close(lease_fd) + lease.close() return StageArtifacts(required) def _publish(self, generation: str, artifacts: Path, required: tuple[str, ...]) -> None: