mptyl/ThothII
1
0
Fork 0
Code Issues 25 Pull Requests Actions Packages Projects Releases 1 Wiki Activity

fix(preprocess): retain DWH root fd through lease

Browse Source
This commit is contained in:
marcopan
2026-07-12 07:12:28 +02:00
parent 4cf0adbdd2
commit 24e61d5713
2 changed files with 128 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files
harness/tests/test_dwh_preprocess_job.py
+2 -2
View File
@@ -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):
harness/tht/jobs/dwh_pipeline.py
+126 -63
View File
@@ -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: