DynamisLab/src/CCD_analysis/acquisition/artifacts.py

"""Immutable, fsync-backed acquisition artifact transactions."""
from __future__ import annotations
import ctypes
import errno
from hashlib import sha256
import json, os, shutil, uuid
from pathlib import Path
from typing import Mapping, Any
import numpy as np
from .contracts import ARTIFACT_SCHEMA_ID, canonical_json
from .validation import array_sha256, validate_acquisition_semantics


def file_sha256(path: Path) -> str:
    digest = sha256()
    with path.open("rb") as stream:
        for block in iter(lambda: stream.read(8 * 1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def rename_noreplace(source: Path, destination: Path) -> None:
    """Atomically publish a directory without replacing any existing inode."""
    libc = ctypes.CDLL(None, use_errno=True)
    renameat2 = getattr(libc, "renameat2", None)
    if renameat2 is None:
        raise RuntimeError("atomic no-replace publication unavailable: libc renameat2 is missing")
    renameat2.argtypes = [ctypes.c_int, ctypes.c_char_p, ctypes.c_int, ctypes.c_char_p, ctypes.c_uint]
    renameat2.restype = ctypes.c_int
    result = renameat2(-100, os.fsencode(source), -100, os.fsencode(destination), 1)
    if result == 0:
        return
    code = ctypes.get_errno()
    if code in (errno.EEXIST, errno.ENOTEMPTY):
        raise FileExistsError(destination)
    if code in (errno.ENOSYS, errno.EINVAL, errno.ENOTSUP):
        raise RuntimeError("atomic no-replace publication unavailable; refusing unsafe fallback") from OSError(code, os.strerror(code))
    raise OSError(code, os.strerror(code), destination)


def _hash_array(value: np.ndarray) -> str:
    return sha256(np.ascontiguousarray(value).tobytes()).hexdigest()


def _fsync_file(path: Path) -> None:
    with path.open("rb") as stream:
        os.fsync(stream.fileno())


def _validate_sha(value: np.ndarray, key: str) -> None:
    if value.ndim != 0 or value.dtype.kind not in "SU": raise ValueError(f"{key} must be a scalar string")
    text = str(value.item())
    if len(text) != 64: raise ValueError(f"{key} must be SHA256")
    int(text, 16)


class ArtifactTransaction:
    def __init__(self, destination: str | Path):
        self.destination = Path(destination)
        self.partial = self.destination.with_name(f".{self.destination.name}.partial.{os.getpid()}.{uuid.uuid4().hex}")
        self.active = False

    def __enter__(self):
        if self.destination.exists(): raise FileExistsError(self.destination)
        self.destination.parent.mkdir(parents=True, exist_ok=True)
        self.partial.mkdir(); self.active = True
        return self

    def write(self, *, arrays: Mapping[str, Any], config: dict, state: Mapping[str, Any]):
        if not self.active: raise RuntimeError("transaction inactive")
        data, state_arrays = validate_acquisition_semantics(arrays=arrays, config=config, state=state)
        np.savez_compressed(self.partial/"fields.npz", **data)
        (self.partial/"config.json").write_bytes(canonical_json(config))
        np.savez_compressed(self.partial/"controller_state.npz", **state_arrays)
        for path in self.partial.iterdir():
            if path.is_file(): _fsync_file(path)
        files = {path.name:file_sha256(path) for path in sorted(self.partial.iterdir()) if path.is_file()}
        manifest = {"schema_id":ARTIFACT_SCHEMA_ID,"complete":True,"files":files,"state_array_sha256":{key:array_sha256(value) for key,value in state_arrays.items()},"config_sha256":sha256(canonical_json(config)).hexdigest(),"field_count":data["ux"].shape[0]}
        validate_acquisition_semantics(arrays=data, config=config, state=state_arrays, manifest=manifest)
        (self.partial/"manifest.json").write_bytes(canonical_json(manifest)); _fsync_file(self.partial/"manifest.json")
        return manifest

    def publish(self):
        manifest = json.loads((self.partial/"manifest.json").read_text())
        if not manifest.get("complete") or manifest.get("schema_id") != ARTIFACT_SCHEMA_ID: raise RuntimeError("partial artifact is not complete")
        for name, digest in manifest["files"].items():
            if file_sha256(self.partial/name) != digest: raise RuntimeError("artifact hash validation failed")
        with np.load(self.partial/"fields.npz", allow_pickle=False) as fields, np.load(self.partial/"controller_state.npz", allow_pickle=False) as state:
            config=json.loads((self.partial/"config.json").read_text())
            validate_acquisition_semantics(arrays={key:fields[key] for key in fields.files},config=config,state={key:state[key] for key in state.files},manifest=manifest)
        rename_noreplace(self.partial, self.destination)
        directory_fd = os.open(self.destination.parent, os.O_RDONLY)
        try: os.fsync(directory_fd)
        finally: os.close(directory_fd)
        self.active = False
        return self.destination

    def __exit__(self, typ, value, tb):
        if self.active: shutil.rmtree(self.partial, ignore_errors=True); self.active=False