fix(dwh): bind snapshots to validated bytes

harness/tht/jobs/dwh_pipeline.py

@@ -366,7 +366,7 @@ def _open_generation_fd(generations_fd: int, generation: str) -> int:
 
 
 def _active_generation_fd(
-    root_fd: int, expected_binding: dict[str, str]
+    root_fd: int, expected_binding: dict[str, str], *, validate: bool = True
 ) -> tuple[str, int] | None:
     _validate_root_binding_fd(root_fd, expected_binding)
     try:
@@ -394,7 +394,8 @@ def _active_generation_fd(
     finally:
         os.close(generations_fd)
     try:
-        validate_generation_fd(generation_fd, generation, expected_binding)
+        if validate:
+            validate_generation_fd(generation_fd, generation, expected_binding)
         return generation, generation_fd
     except BaseException:
         os.close(generation_fd)
@@ -408,10 +409,10 @@ def _materialize_generation_fd(
     _SNAPSHOT_DIRS.add(holder)
     snapshot_root = holder
     try:
-        manifest = validate_generation_fd(generation_fd, generation, binding)
-        for name in manifest["files"]:
+        _, payloads = _read_validated_generation_fd(generation_fd, generation, binding)
+        for name, payload in payloads.items():
             destination = snapshot_root / name
-            destination.write_bytes(_read_owned_at(generation_fd, name, readonly=True))
+            destination.write_bytes(payload)
             destination.chmod(0o400)
         return holder, snapshot_root
     except BaseException:
@@ -423,6 +424,15 @@ def _materialize_generation_fd(
 def validate_generation_fd(
     directory_fd: int, generation: str, expected_binding: dict[str, str] | None = None,
 ) -> dict:
+    manifest, _ = _read_validated_generation_fd(
+        directory_fd, generation, expected_binding
+    )
+    return manifest
+
+
+def _read_validated_generation_fd(
+    directory_fd: int, generation: str, expected_binding: dict[str, str] | None = None,
+) -> tuple[dict, dict[str, bytes]]:
     try:
         directory_info = os.fstat(directory_fd)
         if (
@@ -453,13 +463,15 @@ def validate_generation_fd(
             )
         if set(os.listdir(directory_fd)) != set(files) | {GENERATION_MANIFEST}:
             raise ValueError
+        payloads = {}
         for name, expected in files.items():
             if name != "physical.yaml" and not _SAFE_FILE.fullmatch(name):
                 raise ValueError
             payload = _read_owned_at(directory_fd, name, readonly=True)
             if hashlib.sha256(payload).hexdigest() != expected:
                 raise ValueError
-        return manifest
+            payloads[name] = payload
+        return manifest, payloads
     except (OSError, KeyError, TypeError, ValueError, UnicodeDecodeError, json.JSONDecodeError) as error:
         raise CorruptCheckpointError("published DWH generation is invalid") from error
 
@@ -478,7 +490,7 @@ def resolve_dwh_snapshot(cfg) -> DwhArtifactSnapshot:
 def _resolve_dwh_snapshot_locked(
     cfg, binding: dict[str, str], root_fd: int,
 ) -> DwhArtifactSnapshot:
-    active = _active_generation_fd(root_fd, binding)
+    active = _active_generation_fd(root_fd, binding, validate=False)
     if active is None:
         return DwhArtifactSnapshot(
             None, cfg.paths.artifacts / "mschema" / "physical.yaml", cfg.paths.indexes / "lsh"
@@ -576,7 +588,7 @@ class DwhPreprocessPipeline:
             _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_fd(lease.root_fd, self.binding)
+            active = _active_generation_fd(lease.root_fd, self.binding, validate=False)
             if active is not None:
                 generation, generation_fd = active
                 try:
@@ -633,10 +645,22 @@ class DwhPreprocessPipeline:
             return self._publish_stage(context, "lsh", spec)
 
         implementations = {"introspect": introspect_stage, "lsh": lsh_stage}
-        return run_job(
-            spec, tuple(implementations[step] for step in steps),
-            reconcile_effects=self._reconcile_effects,
-        )
+        try:
+            return run_job(
+                spec, tuple(implementations[step] for step in steps),
+                reconcile_effects=self._reconcile_effects,
+            )
+        finally:
+            holder, self._snapshot_holder = self._snapshot_holder, None
+            if isinstance(holder, Path):
+                shutil.rmtree(holder, ignore_errors=True)
+                _SNAPSHOT_DIRS.discard(holder)
+                if self.current_physical is not None and holder in self.current_physical.parents:
+                    self.current_physical = None
+                if self.current_lsh_dir is not None and (
+                    self.current_lsh_dir == holder or holder in self.current_lsh_dir.parents
+                ):
+                    self.current_lsh_dir = None
 
     def _reconcile_effects(self, source, run_dir: Path) -> set[str]:
         running = next(
@@ -654,9 +678,12 @@ class DwhPreprocessPipeline:
                     run_dir / "artifacts", running.artifact_files,
                 )
                 active = _active_generation_fd(lease.root_fd, self.binding)
-                if active is None or active[0] != source.run_id:
-                    raise CorruptCheckpointError("sealed DWH publication is not ACTIVE")
-                os.close(active[1])
+                try:
+                    if active is None or active[0] != source.run_id:
+                        raise CorruptCheckpointError("sealed DWH publication is not ACTIVE")
+                finally:
+                    if active is not None:
+                        os.close(active[1])
             finally:
                 os.close(generations_fd)
         finally: