mostlyright-data 0.9.0__py3-none-any.whl

mostlyright/data_harness/ingestion/spool.py

@@ -0,0 +1,130 @@
+"""Bounded-memory streaming byte spool and digest primitives."""
+
+from __future__ import annotations
+
+import fcntl
+import hashlib
+import os
+import secrets
+import stat
+from collections.abc import Callable, Iterable
+from dataclasses import dataclass
+from pathlib import Path
+
+from mostlyright.data_harness.ingestion.contracts import MAX_CHUNK_BYTES, IngestionError
+from mostlyright.data_harness.ingestion.faults import DeterministicFaultInjector
+
+
+@dataclass(frozen=True)
+class SpooledChunk:
+    """A verified local temporary file whose content can become one immutable chunk object."""
+
+    path: Path
+    content_sha256: str
+    byte_length: int
+    device: int
+    inode: int
+
+
+def spool_hash_stream(
+    directory: Path,
+    parts: Iterable[bytes | bytearray | memoryview],
+    *,
+    maximum_bytes: int = MAX_CHUNK_BYTES,
+    fault_injector: DeterministicFaultInjector | None = None,
+    reserve_bytes: Callable[[int, int], None] | None = None,
+) -> SpooledChunk:
+    """Spool an iterable once while computing its exact SHA-256 without buffering it all.
+
+    The returned file is fsynced before it is observable.  A caller may safely retry after
+    a failure: temporary names are unguessable and are never used as durable references.
+    """
+
+    if not directory.is_absolute() or directory.is_symlink() or not directory.is_dir():
+        raise IngestionError(
+            "INGESTION_SPOOL_ROOT", "spool.directory", "must be an absolute directory"
+        )
+    if isinstance(maximum_bytes, bool) or not isinstance(maximum_bytes, int) or maximum_bytes < 0:
+        raise IngestionError(
+            "INGESTION_SPOOL_LIMIT", "spool.maximum_bytes", "must be a non-negative integer"
+        )
+    name = f".spool.{os.getpid()}.{secrets.token_hex(16)}.tmp"
+    path = directory / name
+    descriptor = -1
+    published = False
+    digest = hashlib.sha256()
+    total = 0
+    try:
+        descriptor = os.open(
+            path,
+            os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0),
+            0o600,
+        )
+        fcntl.flock(descriptor, fcntl.LOCK_EX)
+        for part in parts:
+            if not isinstance(part, (bytes, bytearray, memoryview)):
+                raise IngestionError(
+                    "INGESTION_SPOOL_PART", "spool.parts", "must yield bytes-like values"
+                )
+            view = memoryview(part)
+            if not view.c_contiguous:
+                raise IngestionError(
+                    "INGESTION_SPOOL_PART", "spool.parts", "must yield contiguous bytes-like values"
+                )
+            if view.format != "B" or view.ndim != 1:
+                view = view.cast("B")
+            if view.nbytes > maximum_bytes - total:
+                raise IngestionError(
+                    "INGESTION_CHUNK_BYTES", "spool.parts", "exceeds the configured byte limit"
+                )
+            if reserve_bytes is not None:
+                # The local store reserves the exact next logical file length while holding its
+                # short global quota lock.  The caller iterable is deliberately *outside* that
+                # lock: a connector may block on the network or re-enter the store while
+                # constructing its next part.  ``ftruncate`` makes an unfilled tail visible to
+                # every competing quota check, so a crash leaves a conservative reservation that
+                # stale-spool cleanup can reclaim.
+                reserve_bytes(descriptor, total + view.nbytes)
+                if fault_injector is not None:
+                    fault_injector.checkpoint("spool.after_reservation")
+            _write_all(descriptor, view)
+            digest.update(view)
+            total += view.nbytes
+            if fault_injector is not None:
+                fault_injector.checkpoint("spool.after_part")
+        if fault_injector is not None:
+            fault_injector.checkpoint("spool.before_fsync")
+        os.fsync(descriptor)
+        if fault_injector is not None:
+            fault_injector.checkpoint("spool.after_fsync")
+        info = os.fstat(descriptor)
+        if not stat.S_ISREG(info.st_mode) or info.st_nlink != 1:
+            raise IngestionError(
+                "INGESTION_SPOOL_PATH", "spool.path", "is not a private regular file"
+            )
+        published = True
+        return SpooledChunk(
+            path=path,
+            content_sha256=digest.hexdigest(),
+            byte_length=total,
+            device=info.st_dev,
+            inode=info.st_ino,
+        )
+    except OSError as error:
+        raise IngestionError("INGESTION_SPOOL_IO", "spool.path", "stream spool failed") from error
+    finally:
+        if descriptor >= 0:
+            os.close(descriptor)
+        if not published:
+            try:
+                path.unlink(missing_ok=True)
+            except OSError:
+                pass
+
+
+def _write_all(descriptor: int, view: memoryview) -> None:
+    while view:
+        written = os.write(descriptor, view)
+        if written <= 0:
+            raise OSError("short spool write")
+        view = view[written:]