diff --git a/src/quant_data_kit/capture_v2/epoch.py b/src/quant_data_kit/capture_v2/epoch.py index fb5a5d2..1a7bc32 100644 --- a/src/quant_data_kit/capture_v2/epoch.py +++ b/src/quant_data_kit/capture_v2/epoch.py @@ -30,6 +30,7 @@ from quant_data_kit.data_lake import ( RawObjectReference, StoragePolicy, + _capacity_tree_lock, validate_raw_reference, write_normalized_batches, ) @@ -1548,16 +1549,17 @@ def _seal_part(self) -> None: final_name = f"part-{index:08d}-sha256-{digest}.ndjson" final_path = self.root / final_name checked_final = _validate_safe_path(self.hot_root, final_path, allow_missing=True) - try: - os.link(self._open_path, checked_final) - except FileExistsError as exc: - raise ValidationError( - f"Normalized journal part already exists: {checked_final}" - ) from exc + with _capacity_tree_lock(self.hot_root): + try: + os.link(self._open_path, checked_final) + except FileExistsError as exc: + raise ValidationError( + f"Normalized journal part already exists: {checked_final}" + ) from exc + self._open_path.unlink() _validate_safe_path(self.hot_root, checked_final, allow_missing=False) if _sha256_file(checked_final, trusted_root=self.hot_root) != digest: raise ValidationError(f"Normalized sealed journal part hash changed: {checked_final}") - self._open_path.unlink() _fsync_directory(self.root) self._parts.append( EpochPart( diff --git a/src/quant_data_kit/capture_v2/storage.py b/src/quant_data_kit/capture_v2/storage.py index 65756ff..5a367ec 100644 --- a/src/quant_data_kit/capture_v2/storage.py +++ b/src/quant_data_kit/capture_v2/storage.py @@ -30,6 +30,8 @@ from quant_data_kit.data_lake import ( RawObjectManifest, StoragePolicy, + _capacity_tree_lock, + _unlink_tree_entry, evaluate_capacity, load_raw_object, write_raw_bytes, @@ -758,6 +760,21 @@ def _atomic_immutable_write(path: Path, body: bytes, *, root: Path | None = None temporary = parent / f".{path.name}.{uuid.uuid4().hex}.tmp" _validate_safe_path(trusted_root, temporary, allow_missing=True) expected_hash = hashlib.sha256(body).hexdigest() + relative = checked_path.relative_to(trusted_root) + coordinated = bool( + relative.parts + and relative.parts[0] in {"capture", "curated", "normalized", "quarantine", "raw"} + ) + + def publish_link() -> None: + try: + os.link(temporary, checked_path) + _fsync_directory(parent) + except FileExistsError: + pass + if coordinated and temporary.exists(): + temporary.unlink() + try: with temporary.open("xb") as stream: stream.write(body) @@ -766,11 +783,11 @@ def _atomic_immutable_write(path: Path, body: bytes, *, root: Path | None = None if _sha256_file(temporary) != expected_hash: raise ValidationError(f"immutable staging hash mismatch: {temporary}") _validate_safe_path(trusted_root, parent, allow_missing=False) - try: - os.link(temporary, checked_path) - _fsync_directory(parent) - except FileExistsError: - pass + if coordinated: + with _capacity_tree_lock(trusted_root): + publish_link() + else: + publish_link() _validate_safe_path(trusted_root, checked_path, allow_missing=False) if not checked_path.is_file() or _sha256_file(checked_path) != expected_hash: raise ValidationError( @@ -778,7 +795,10 @@ def _atomic_immutable_write(path: Path, body: bytes, *, root: Path | None = None ) finally: if temporary.exists(): - temporary.unlink() + if coordinated: + _unlink_tree_entry(trusted_root, temporary) + else: + temporary.unlink() class DurableAuditStore: diff --git a/src/quant_data_kit/curated.py b/src/quant_data_kit/curated.py index 8f4b9a4..d67f802 100644 --- a/src/quant_data_kit/curated.py +++ b/src/quant_data_kit/curated.py @@ -4,7 +4,6 @@ import hashlib import json -import os import re from collections import defaultdict from collections.abc import Iterable, Mapping @@ -22,6 +21,7 @@ StoragePolicy, _atomic_write_bytes, _mkdir_in_lake, + _publish_tree_entry, _resolved_lake_root, _stable_staging_directory, _validate_lake_path, @@ -388,8 +388,7 @@ def _publish_curated_snapshot( if existing != snapshot: raise ValidationError(f"Curated snapshot collision: {snapshot_dir}") else: - require_collection_capacity(lake_root, projected_write_bytes=0, policy=policy) - os.replace(stage, snapshot_dir) + _publish_tree_entry(lake_root, stage, snapshot_dir, policy=policy) if not revision_path.exists(): _atomic_write_bytes( lake_root, @@ -429,11 +428,6 @@ def _write_curated_bars( groups[(str(record["trading_day"]), str(record["instrument_id"]))].append(record) estimated_bytes = sum(len(_canonical(_json_value(item))) for item in records) - require_collection_capacity( - lake_root, - projected_write_bytes=estimated_bytes, - policy=policy, - ) curated_root = _mkdir_in_lake(lake_root, lake_root / "curated" / dataset) staging_root = curated_root / "staging" partition_items: list[CuratedPartition] = [] @@ -444,6 +438,11 @@ def _write_curated_bars( namespace="curated-revision", identity=revision_identity, ) as stage: + require_collection_capacity( + lake_root, + projected_write_bytes=estimated_bytes, + policy=policy, + ) for (trading_date, instrument_id), group in sorted(groups.items()): ordered = sorted(group, key=lambda row: (row["event_time"], row["event_id"])) table = pa.Table.from_pylist( diff --git a/src/quant_data_kit/data_lake.py b/src/quant_data_kit/data_lake.py index e93c67b..c57477e 100644 --- a/src/quant_data_kit/data_lake.py +++ b/src/quant_data_kit/data_lake.py @@ -9,6 +9,7 @@ import re import shutil import stat +import threading import uuid from collections import defaultdict from collections.abc import Iterable, Iterator, Mapping, Sequence @@ -66,6 +67,16 @@ "trade": TRADE_EVENT_SCHEMA_ID, } _GIB = 1024**3 +_CAPACITY_TREE_LOCK_STATE = threading.local() +_TRANSIENT_LAKE_RELATIVE_DIRECTORIES = { + (".locks",), + ("normalized", ".stage-owners"), + ("normalized", "event-claim-index-v3", ".legacy-staging"), + ("normalized", "event-claim-index-v3", ".staging"), + ("normalized", "staging"), + ("quarantine", ".staging"), + ("raw", ".staging"), +} class CollectionStoppedError(ValidationError): @@ -456,7 +467,7 @@ def _atomic_write_bytes(root: Path, target: Path, body: bytes) -> None: checked_stale = _validate_lake_path(root, stale, allow_missing=False) if not checked_stale.is_file(): raise ValidationError(f"Atomic staging entry is not a file: {checked_stale}") - checked_stale.unlink() + _unlink_tree_entry(root, checked_stale) temporary = parent / f"{temporary_prefix}{uuid.uuid4().hex}.tmp" _validate_lake_path(root, temporary, allow_missing=True) try: @@ -466,10 +477,10 @@ def _atomic_write_bytes(root: Path, target: Path, body: bytes) -> None: os.fsync(stream.fileno()) if _sha256_file(temporary) != _sha256_bytes(body): raise ValidationError(f"Atomic staging verification failed: {temporary}") - os.replace(temporary, checked_target) + _replace_tree_entry(root, temporary, checked_target) finally: if temporary.exists(): - temporary.unlink() + _unlink_tree_entry(root, temporary) @contextmanager @@ -482,6 +493,71 @@ def _lake_lock(root: Path, namespace: str, identity: Mapping[str, Any]) -> Itera yield +@contextmanager +def _capacity_tree_lock(root: Path) -> Iterable[None]: + """Serialize lake-wide capacity scans with topology-removing mutations.""" + key = str(Path(root).absolute()) + depths = getattr(_CAPACITY_TREE_LOCK_STATE, "depths", None) + if depths is None: + depths = {} + _CAPACITY_TREE_LOCK_STATE.depths = depths + if depths.get(key, 0): + depths[key] += 1 + try: + yield + finally: + depths[key] -= 1 + return + with _lake_lock(root, "capacity-tree", {"scope": "lake-wide"}): + depths[key] = 1 + try: + yield + finally: + depths.pop(key, None) + + +def _replace_tree_entry(root: Path, source: Path, target: Path) -> None: + """Atomically replace one lake entry without racing a capacity tree scan.""" + with _capacity_tree_lock(root): + os.replace(source, target) + + +def _publish_tree_entry( + root: Path, + source: Path, + target: Path, + *, + policy: StoragePolicy, +) -> CapacityDecision: + """Capacity-check and publish one staged tree as a single lake-wide transaction.""" + with _capacity_tree_lock(root): + decision = require_collection_capacity( + root, + projected_write_bytes=_tree_size(source), + policy=policy, + ) + os.replace(source, target) + return decision + + +def _remove_tree(root: Path, target: Path) -> None: + """Remove one lake subtree without racing a capacity tree scan.""" + with _capacity_tree_lock(root): + shutil.rmtree(target) + + +def _remove_empty_tree(root: Path, target: Path) -> None: + """Remove one empty lake directory without racing a capacity tree scan.""" + with _capacity_tree_lock(root): + target.rmdir() + + +def _unlink_tree_entry(root: Path, target: Path, *, missing_ok: bool = False) -> None: + """Remove one lake file without racing a capacity tree scan.""" + with _capacity_tree_lock(root): + target.unlink(missing_ok=missing_ok) + + @contextmanager def _stable_staging_directory( root: Path, @@ -499,7 +575,7 @@ def _stable_staging_directory( checked_stale = _validate_lake_path(root, stale, allow_missing=False) if not checked_stale.is_dir(): raise ValidationError(f"Stable staging entry is not a directory: {checked_stale}") - shutil.rmtree(checked_stale) + _remove_tree(root, checked_stale) stage = checked_root / f"{prefix}{uuid.uuid4().hex}" _validate_lake_path(root, stage, allow_missing=True) stage.mkdir(exist_ok=False) @@ -507,13 +583,32 @@ def _stable_staging_directory( yield stage finally: if stage.exists() and stage.parent == checked_root: - shutil.rmtree(stage) + _remove_tree(root, stage) def _tree_size(root: Path) -> int: if not root.exists(): return 0 - return sum(path.stat().st_size for path in root.rglob("*") if path.is_file()) + total = 0 + for directory, child_directories, filenames in os.walk(root): + relative_directory = Path(directory).relative_to(root) + child_directories[:] = [ + name + for name in child_directories + if not _is_transient_lake_directory(relative_directory / name) + ] + for filename in filenames: + if filename.startswith(".atomic-") and filename.endswith(".tmp"): + continue + total += (Path(directory) / filename).stat().st_size + return total + + +def _is_transient_lake_directory(relative_path: Path) -> bool: + parts = relative_path.parts + return parts in _TRANSIENT_LAKE_RELATIVE_DIRECTORIES or ( + len(parts) == 3 and parts[0] == "curated" and parts[2] == "staging" + ) def _disk_probe_path(root: Path) -> Path: @@ -535,7 +630,11 @@ def evaluate_capacity( """Evaluate the 150GB hot quota and max(20% volume, 100GB) free-space gate.""" if projected_write_bytes < 0: raise ValidationError("projected_write_bytes must be non-negative") - hot_bytes = _tree_size(Path(root)) if current_hot_bytes is None else current_hot_bytes + if current_hot_bytes is None: + with _capacity_tree_lock(Path(root)): + hot_bytes = _tree_size(Path(root)) + else: + hot_bytes = current_hot_bytes if hot_bytes < 0: raise ValidationError("current_hot_bytes must be non-negative") if disk_total_bytes is None or disk_free_bytes is None: @@ -797,7 +896,7 @@ def _relocate_invalid_raw(root: Path, object_dir: Path, *, reason: str) -> Path: quarantine_root = _mkdir_in_lake(root, Path(root) / "quarantine" / "raw-unpublished") target = quarantine_root / f"{uuid.uuid4().hex}-{object_dir.name}" _validate_lake_path(root, target, allow_missing=True) - os.replace(object_dir, target) + _replace_tree_entry(root, object_dir, target) evidence = { "schema_version": "2.0.0", "layer": "quarantine", @@ -817,13 +916,15 @@ def _remove_staging_directory(root: Path, stage: Path) -> None: staging_root = _validate_lake_path(root, Path(root) / "raw" / ".staging", allow_missing=False) if checked.parent != staging_root: raise ValidationError("Refused to remove a non-staging Raw path") - shutil.rmtree(checked) + _remove_tree(root, checked) def _recover_raw_staging( root: Path, reference: RawObjectReference, manifest: RawObjectManifest, + *, + policy: StoragePolicy = _DEFAULT_STORAGE_POLICY, ) -> RawObjectManifest | None: staging_root = _mkdir_in_lake(root, Path(root) / "raw" / ".staging") object_dir = _raw_object_dir(root, reference) @@ -850,7 +951,7 @@ def _recover_raw_staging( _remove_staging_directory(root, stage) recovered = existing continue - os.replace(stage, object_dir) + _publish_tree_entry(root, stage, object_dir, policy=policy) recovered = _load_raw_from_dir(root, object_dir, expected=reference) return recovered @@ -930,7 +1031,7 @@ def write_raw_bytes( if list(key_dir.glob("deleting=*")): raise ValidationError(f"Raw idempotency key cleanup is in progress: {resolved_key}") object_dir = _raw_object_dir(lake_root, reference) - recovered = _recover_raw_staging(lake_root, reference, manifest) + recovered = _recover_raw_staging(lake_root, reference, manifest, policy=policy) if recovered is not None: return recovered for existing_dir in key_dir.glob("object=*"): @@ -969,7 +1070,7 @@ def write_raw_bytes( enforce_directory_identity=False, ) _validate_lake_path(lake_root, object_dir, allow_missing=True) - os.replace(stage, object_dir) + _publish_tree_entry(lake_root, stage, object_dir, policy=policy) return _load_raw_from_dir(lake_root, object_dir, expected=reference) finally: if stage.exists(): @@ -1016,7 +1117,7 @@ def _verify_archive_restore(root: Path, archive_path: Path) -> tuple[str, str]: return archive_hash, _sha256_file(restore_path) finally: if restore_path.exists(): - restore_path.unlink() + _unlink_tree_entry(root, restore_path) def _read_cleanup_audit(root: Path, reference: RawObjectReference) -> dict[str, Any]: @@ -1098,7 +1199,7 @@ def _finalize_raw_deleting( raise ValidationError("Raw deleting payload length changed") if _sha256_file(payload_path) != manifest.content_sha256: raise ValidationError("Raw deleting payload hash changed") - payload_path.unlink() + _unlink_tree_entry(root, payload_path) manifest_path = deleting_dir / "manifest.json" if manifest_path.exists(): manifest_path = _validate_lake_path(root, manifest_path, allow_missing=False) @@ -1110,8 +1211,8 @@ def _finalize_raw_deleting( raise ValidationError("Raw deleting manifest is unreadable or malformed") from exc if remaining_manifest != manifest: raise ValidationError("Raw deleting manifest changed") - manifest_path.unlink() - deleting_dir.rmdir() + _unlink_tree_entry(root, manifest_path) + _remove_empty_tree(root, deleting_dir) def validate_raw_reference( @@ -1188,7 +1289,7 @@ def cleanup_archived_raw_object( if object_dir.exists(): if deleting_dir.exists(): raise ValidationError("Raw cleanup has both live and deleting states") - os.replace(object_dir, deleting_dir) + _replace_tree_entry(lake_root, object_dir, deleting_dir) if deleting_dir.exists(): _finalize_raw_deleting(lake_root, deleting_dir, manifest) return tombstone @@ -1214,7 +1315,7 @@ def cleanup_archived_raw_object( current_time=current_time, ) if should_rename: - os.replace(object_dir, deleting_dir) + _replace_tree_entry(lake_root, object_dir, deleting_dir) audit = { "schema_version": "2.0.0", "action": "verified_local_archive_cleanup", @@ -1495,7 +1596,7 @@ def _write_quarantine( for stale in staging_root.glob(f"{batch_id}-*"): if stale.is_dir(): checked = _validate_lake_path(lake_root, stale, allow_missing=False) - shutil.rmtree(checked) + _remove_tree(lake_root, checked) require_collection_capacity( lake_root, projected_write_bytes=len(body) + len(manifest_bytes), @@ -1515,11 +1616,11 @@ def _write_quarantine( stream.flush() os.fsync(stream.fileno()) _validate_quarantine_batch(lake_root, stage, manifest) - os.replace(stage, batch_dir) + _publish_tree_entry(lake_root, stage, batch_dir, policy=policy) return _validate_quarantine_batch(lake_root, batch_dir, manifest) finally: if stage.exists(): - shutil.rmtree(stage) + _remove_tree(lake_root, stage) def _write_normalized_events_legacy( @@ -1754,8 +1855,7 @@ def _write_normalized_events_legacy( quarantined_rows=len(quarantined), quarantine_manifest=quarantine_manifest, ) - require_collection_capacity(lake_root, projected_write_bytes=0, policy=policy) - os.replace(stage, snapshot_dir) + _publish_tree_entry(lake_root, stage, snapshot_dir, policy=policy) stage = snapshot_dir verified = load_normalized_snapshot(root, snapshot_id) return NormalizationResult( @@ -1831,6 +1931,7 @@ def _load_normalized_snapshot( snapshot_id: str, *, verify_event_claim_files: bool, + recovery_policy: StoragePolicy | None = None, ) -> NormalizedSnapshot: lake_root = _resolved_lake_root(root, create=False) snapshot_id = _segment(snapshot_id, "snapshot_id") @@ -1854,6 +1955,7 @@ def _load_normalized_snapshot( lake_root, snapshot_id, payload=payload, + recovery_policy=recovery_policy, ) payload["upstream_raw_references"] = tuple( RawObjectReference(**item) for item in payload["upstream_raw_references"] @@ -1936,8 +2038,18 @@ def _load_normalized_snapshot( return snapshot -def load_normalized_snapshot(root: Path, snapshot_id: str) -> NormalizedSnapshot: - return _load_normalized_snapshot(root, snapshot_id, verify_event_claim_files=True) +def load_normalized_snapshot( + root: Path, + snapshot_id: str, + *, + recovery_policy: StoragePolicy | None = None, +) -> NormalizedSnapshot: + return _load_normalized_snapshot( + root, + snapshot_id, + verify_event_claim_files=True, + recovery_policy=recovery_policy, + ) def read_normalized_events( diff --git a/src/quant_data_kit/normalized_v3.py b/src/quant_data_kit/normalized_v3.py index f462a44..f7c3797 100644 --- a/src/quant_data_kit/normalized_v3.py +++ b/src/quant_data_kit/normalized_v3.py @@ -45,6 +45,8 @@ _lake_lock, _mkdir_in_lake, _partition_segment, + _publish_tree_entry, + _replace_tree_entry, _resolved_lake_root, _safe_snapshot_partition, _segment, @@ -866,7 +868,7 @@ def _publish_quarantine_file( (stage / "manifest.json").write_bytes(manifest_bytes) _validate_quarantine_batch(root, stage, manifest) _mkdir_in_lake(root, batch_dir.parent) - os.replace(stage, batch_dir) + _publish_tree_entry(root, stage, batch_dir, policy=policy) return _validate_quarantine_batch(root, batch_dir, manifest) finally: if stage.exists(): @@ -1339,6 +1341,7 @@ def _publish_or_validate_claim_index( snapshot_id: str, snapshot_logical_sha256: str, expected: EventClaimIndexManifest, + policy: StoragePolicy | None = None, ) -> tuple[Path, ...]: final = _claim_index_root(root, snapshot_id) with _lake_lock(root, "normalized-claim-index", {"snapshot_id": snapshot_id}): @@ -1351,13 +1354,17 @@ def _publish_or_validate_claim_index( expected, ) except _ClaimIndexMissingError: + if policy is None: + raise ValidationError( + "Normalized claim-index recovery requires an explicit StoragePolicy" + ) evidence_root = _mkdir_in_lake( root, root / "normalized" / "event-claim-index-v3" / "recovery-evidence", ) evidence = evidence_root / f"{snapshot_id}-{uuid.uuid4().hex}" - os.replace(final, evidence) - os.replace(staged_index, final) + _replace_tree_entry(root, final, evidence) + _publish_tree_entry(root, staged_index, final, policy=policy) return _validate_claim_index( root, snapshot_id, @@ -1366,7 +1373,11 @@ def _publish_or_validate_claim_index( ) _mkdir_in_lake(root, final.parent) _validate_lake_path(root, final, allow_missing=True) - os.replace(staged_index, final) + if policy is None: + raise ValidationError( + "Normalized claim-index recovery requires an explicit StoragePolicy" + ) + _publish_tree_entry(root, staged_index, final, policy=policy) return _validate_claim_index( root, snapshot_id, @@ -1409,6 +1420,7 @@ def _index_for_snapshot( snapshot_id: str, snapshot_logical_sha256: str, expected: EventClaimIndexManifest, + policy: StoragePolicy | None = None, ) -> tuple[Path, ...]: final = _claim_index_root(root, snapshot_id) verification_root = _mkdir_in_lake( @@ -1455,6 +1467,7 @@ def _index_for_snapshot( snapshot_id=snapshot_id, snapshot_logical_sha256=snapshot_logical_sha256, expected=expected, + policy=policy, ) finally: if connection is not None: @@ -1500,7 +1513,12 @@ def iter_event_claims_v3( connection.close() -def _historical_index_paths(root: Path, candidate_snapshot_id: str) -> tuple[Path, ...]: +def _historical_index_paths( + root: Path, + candidate_snapshot_id: str, + *, + policy: StoragePolicy | None = None, +) -> tuple[Path, ...]: snapshots_root = root / "normalized" / "snapshots" paths: list[Path] = [] if not snapshots_root.exists(): @@ -1514,6 +1532,7 @@ def _historical_index_paths(root: Path, candidate_snapshot_id: str) -> tuple[Pat root, snapshot_dir.name, verify_event_claim_files=True, + recovery_policy=policy, ) if snapshot.layout_version != LAYOUT_VERSION or snapshot.event_claim_index is None: legacy_root = _mkdir_in_lake( @@ -1553,6 +1572,7 @@ def _historical_index_paths(root: Path, candidate_snapshot_id: str) -> tuple[Pat snapshot_id=snapshot.snapshot_id, snapshot_logical_sha256=snapshot.logical_sha256, expected=expected, + policy=policy, ) ) finally: @@ -1577,8 +1597,13 @@ def _assert_lake_wide_claims( candidate_paths: tuple[Path, ...], *, candidate_snapshot_id: str, + policy: StoragePolicy | None = None, ) -> None: - historical_paths = _historical_index_paths(root, candidate_snapshot_id) + historical_paths = _historical_index_paths( + root, + candidate_snapshot_id, + policy=policy, + ) if not historical_paths: return connection = duckdb.connect(database=":memory:") @@ -1604,6 +1629,7 @@ def load_normalized_snapshot_v3( snapshot_id: str, *, payload: Mapping[str, Any] | None = None, + recovery_policy: StoragePolicy | None = None, _trusted_publish: bool = False, ) -> NormalizedSnapshot: lake_root = _resolved_lake_root(root, create=False) @@ -1750,6 +1776,7 @@ def load_normalized_snapshot_v3( snapshot_id=snapshot_id, snapshot_logical_sha256=logical_sha256, expected=claim_index, + policy=recovery_policy, ) claims = EventClaimSequence(lake_root, snapshot_id, claim_index) return NormalizedSnapshot( @@ -2508,6 +2535,7 @@ def _finalize_strict_batch_snapshot( lake_root, candidate_paths, candidate_snapshot_id=snapshot_id, + policy=policy, ) require_collection_capacity(lake_root, projected_write_bytes=0, policy=policy) _mkdir_in_lake(lake_root, snapshot_dir.parent) @@ -2520,12 +2548,13 @@ def _finalize_strict_batch_snapshot( snapshot_id=snapshot_id, snapshot_logical_sha256=logical_sha256, expected=claim_index, + policy=policy, ) try: - os.replace(snapshot_stage, snapshot_dir) + _publish_tree_entry(lake_root, snapshot_stage, snapshot_dir, policy=policy) except OSError: if not index_preexisted and index_root.exists(): - os.replace(index_root, staged_index) + _replace_tree_entry(lake_root, index_root, staged_index) raise verified = load_normalized_snapshot_v3( lake_root, @@ -2883,6 +2912,7 @@ def write_normalized_events_v3( lake_root, candidate_paths, candidate_snapshot_id=snapshot_id, + policy=policy, ) require_collection_capacity(lake_root, projected_write_bytes=0, policy=policy) _mkdir_in_lake(lake_root, snapshot_dir.parent) @@ -2895,12 +2925,13 @@ def write_normalized_events_v3( snapshot_id=snapshot_id, snapshot_logical_sha256=logical_sha256, expected=claim_index, + policy=policy, ) try: - os.replace(snapshot_stage, snapshot_dir) + _publish_tree_entry(lake_root, snapshot_stage, snapshot_dir, policy=policy) except OSError: if not index_preexisted and index_root.exists(): - os.replace(index_root, staged_index) + _replace_tree_entry(lake_root, index_root, staged_index) raise verified = load_normalized_snapshot_v3( lake_root, diff --git a/tests/test_capture_v2_remediation.py b/tests/test_capture_v2_remediation.py index 2fb9e3e..3841354 100644 --- a/tests/test_capture_v2_remediation.py +++ b/tests/test_capture_v2_remediation.py @@ -16,6 +16,7 @@ import quant_data_kit.capture_v2.collector as collector_module import quant_data_kit.capture_v2.epoch as epoch_module import quant_data_kit.capture_v2.storage as storage_module +import quant_data_kit.data_lake as lake_module from quant_data_kit.capture_v2.collector import CaptureStreamRunner, CryptoL2CaptureCoordinator from quant_data_kit.capture_v2.epoch import NormalizedEpochJournal from quant_data_kit.capture_v2.models import ( @@ -230,6 +231,61 @@ def counted(descriptor: int) -> None: assert calls > before_append + 3 +def test_epoch_seal_waits_for_lake_capacity_scan( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + hot, _archive, guard = epoch_fixtures.storage(tmp_path) + config = epoch_fixtures.stream(Provider.BINANCE, epoch_fixtures.MarketKind.SPOT) + journal = NormalizedEpochJournal( + hot, + epoch_id="capacity-scan-seal", + stream_id=config.stream_id, + provider=config.provider.value, + venue=config.venue, + storage_guard=guard, + policy=epoch_fixtures.POLICY, + max_part_rows=100, + ) + journal.append(({"value": 1},)) + open_path = journal._open_path + assert open_path.is_file() + scan_ready = threading.Event() + scan_release = threading.Event() + seal_started = threading.Event() + real_tree_size = lake_module._tree_size + + def held_tree_size(root: Path) -> int: + scan_ready.set() + if not scan_release.wait(20): + raise TimeoutError("capacity scan was not released") + return real_tree_size(root) + + def seal() -> None: + seal_started.set() + journal._seal_part() + + monkeypatch.setattr(lake_module, "_tree_size", held_tree_size) + with ThreadPoolExecutor(max_workers=2) as executor: + scan = executor.submit( + lake_module.evaluate_capacity, + hot, + projected_write_bytes=0, + policy=epoch_fixtures.POLICY, + ) + assert scan_ready.wait(20) + sealed = executor.submit(seal) + assert seal_started.wait(20) + assert not sealed.done() + scan_release.set() + assert scan.result(timeout=20).allowed + sealed.result(timeout=20) + assert journal._open_path == open_path + assert journal._open_path.is_file() and journal._open_path.stat().st_size == 0 + assert len(tuple(journal.root.glob("part-[0-9]*-sha256-*.ndjson"))) == 1 + journal.abort_visible("test-complete") + + def _journal_with_lineage(tmp_path: Path, epoch_id: str) -> NormalizedEpochJournal: tmp_path.mkdir(parents=True, exist_ok=True) hot, _archive, guard = epoch_fixtures.storage(tmp_path) diff --git a/tests/test_m2_process_integrity.py b/tests/test_m2_process_integrity.py index 54d3d6e..0c74e3d 100644 --- a/tests/test_m2_process_integrity.py +++ b/tests/test_m2_process_integrity.py @@ -2,6 +2,7 @@ import multiprocessing import os +from contextlib import contextmanager from copy import deepcopy from dataclasses import asdict from datetime import datetime, timedelta, timezone @@ -197,14 +198,14 @@ def _curated_process( def _curated_hard_exit_process(root: str, normalized_snapshot_id: str) -> None: - real_replace = curated_module.os.replace + real_replace = lake_module.os.replace def crash_before_publish(source: Path, destination: Path) -> None: if Path(destination).parent.name == "snapshots": os._exit(73) real_replace(source, destination) - curated_module.os.replace = crash_before_publish + lake_module.os.replace = crash_before_publish curate_trade_bars_from_snapshot( Path(root), normalized_snapshot_id=normalized_snapshot_id, @@ -217,6 +218,90 @@ def crash_before_publish(source: Path, destination: Path) -> None: ) +def _curated_capacity_window_process( + root: str, + normalized_snapshot_id: str, + role: str, + publisher_ready: Any, + publisher_release: Any, + contender_transition: Any, + results: Any, +) -> None: + real_capacity_check = curated_module.require_collection_capacity + real_publish = curated_module._publish_curated_snapshot + real_staging_directory = curated_module._stable_staging_directory + + def controlled_capacity_check(*args: Any, **kwargs: Any) -> Any: + if role == "contender": + contender_transition.put("capacity-check") + return real_capacity_check(*args, **kwargs) + + def controlled_publish(*args: Any, **kwargs: Any) -> Any: + if role == "publisher": + publisher_ready.set() + if not publisher_release.wait(20): + raise TimeoutError("publisher capacity window was not released") + return real_publish(*args, **kwargs) + + @contextmanager + def tracked_staging_directory(*args: Any, **kwargs: Any): + if role == "contender": + contender_transition.put("revision-stage-lock") + with real_staging_directory(*args, **kwargs) as stage: + yield stage + + curated_module.require_collection_capacity = controlled_capacity_check + curated_module._publish_curated_snapshot = controlled_publish + curated_module._stable_staging_directory = tracked_staging_directory + try: + snapshot = _curated_process_once(Path(root), normalized_snapshot_id) + results.put(("ok", snapshot.snapshot_id)) + except WORKER_ERRORS as exc: + results.put(("error", f"{type(exc).__name__}: {exc}")) + + +def _held_capacity_scan_process( + root: str, + scan_ready: Any, + scan_release: Any, + results: Any, +) -> None: + real_tree_size = lake_module._tree_size + + def held_tree_size(path: Path) -> int: + scan_ready.set() + if not scan_release.wait(20): + raise TimeoutError("capacity scan was not released") + return real_tree_size(path) + + lake_module._tree_size = held_tree_size + try: + decision = lake_module.require_collection_capacity( + Path(root), + projected_write_bytes=0, + policy=TEST_POLICY, + ) + results.put(("scan", str(decision.hot_bytes))) + except WORKER_ERRORS as exc: + results.put(("error", f"{type(exc).__name__}: {exc}")) + + +def _published_tree_remove_process( + root: str, + relative_path: str, + remove_attempted: Any, + remove_completed: Any, + results: Any, +) -> None: + remove_attempted.set() + try: + lake_module._remove_tree(Path(root), Path(root) / relative_path) + remove_completed.set() + results.put(("remove", relative_path)) + except WORKER_ERRORS as exc: + results.put(("error", f"{type(exc).__name__}: {exc}")) + + def _hold_staging_process( root: str, staging_relative: str, @@ -775,6 +860,144 @@ def _curated_process_once(root: Path, normalized_snapshot_id: str): ) +def test_curated_capacity_scan_waits_for_active_revision_stage(tmp_path: Path) -> None: + root = tmp_path / "curated-capacity-window" + first_snapshot = _normalized( + root, + key="capacity-window-1", + record=trade("capacity-window-1", timestamp="2026-01-02T00:00:01Z"), + ) + second_snapshot = _normalized( + root, + key="capacity-window-2", + record=trade("capacity-window-2", timestamp="2026-01-02T00:00:02Z"), + ) + context = multiprocessing.get_context("spawn") + publisher_ready = context.Event() + publisher_release = context.Event() + contender_transition = context.Queue() + results = context.Queue() + publisher = context.Process( + target=_curated_capacity_window_process, + args=( + str(root), + first_snapshot.snapshot_id, + "publisher", + publisher_ready, + publisher_release, + contender_transition, + results, + ), + ) + contender = context.Process( + target=_curated_capacity_window_process, + args=( + str(root), + second_snapshot.snapshot_id, + "contender", + publisher_ready, + publisher_release, + contender_transition, + results, + ), + ) + try: + publisher.start() + assert publisher_ready.wait(20) + active_stages = list( + (root / "curated" / "concurrent-bars" / "staging").glob("curated-revision-*-*") + ) + assert len(active_stages) == 1 and active_stages[0].is_dir() + contender.start() + assert contender_transition.get(timeout=20) == "revision-stage-lock" + finally: + publisher_release.set() + for process in (publisher, contender): + if process.pid is None: + continue + process.join(30) + if process.is_alive(): + process.terminate() + process.join(5) + assert publisher.exitcode == 0 + assert contender.exitcode == 0 + received = [results.get(timeout=5), results.get(timeout=5)] + assert len([1 for status, _ in received if status == "ok"]) == 1 + conflicts = [value for status, value in received if status == "error"] + assert len(conflicts) == 1 and "maps to different content" in conflicts[0] + + +def test_lake_wide_capacity_scan_serializes_published_tree_removal(tmp_path: Path) -> None: + root = tmp_path / "capacity-tree-lock" + published = root / "curated" / "dataset-a" / "snapshots" / "identity-a" + published.mkdir(parents=True) + (published / "data.parquet").write_bytes(b"published") + relative = published.relative_to(root).as_posix() + context = multiprocessing.get_context("spawn") + scan_ready = context.Event() + scan_release = context.Event() + remove_attempted = context.Event() + remove_completed = context.Event() + results = context.Queue() + scanner = context.Process( + target=_held_capacity_scan_process, + args=(str(root), scan_ready, scan_release, results), + ) + remover = context.Process( + target=_published_tree_remove_process, + args=( + str(root), + relative, + remove_attempted, + remove_completed, + results, + ), + ) + try: + scanner.start() + assert scan_ready.wait(20) + remover.start() + assert remove_attempted.wait(20) + assert not remove_completed.wait(0.5) + finally: + scan_release.set() + for process in (scanner, remover): + if process.pid is None: + continue + process.join(30) + if process.is_alive(): + process.terminate() + process.join(5) + assert scanner.exitcode == 0 + assert remover.exitcode == 0 + received = dict(results.get(timeout=5) for _ in range(2)) + assert received["remove"] == relative + assert int(received["scan"]) >= len(b"published") + assert not published.exists() + + +def test_capacity_scan_counts_published_dataset_named_staging(tmp_path: Path) -> None: + root = tmp_path / "published-staging-dataset" + published = root / "curated" / "staging" / "snapshots" / "sha256-demo" + published.mkdir(parents=True) + payload = b"published-content" + (published / "data.parquet").write_bytes(payload) + decision = lake_module.evaluate_capacity( + root, + projected_write_bytes=0, + policy=StoragePolicy( + hot_quota_bytes=len(payload) - 1, + minimum_free_bytes=1, + minimum_free_fraction=0.000001, + ), + disk_total_bytes=10**9, + disk_free_bytes=10**9, + ) + assert decision.hot_bytes >= len(payload) + assert not decision.allowed + assert any("hot quota exceeded" in reason for reason in decision.reasons) + + def test_atomic_normalized_and_curated_staging_recover_after_hard_exit(tmp_path: Path) -> None: root = tmp_path / "staging-hard-exit" root.mkdir() diff --git a/tests/test_normalized_v3.py b/tests/test_normalized_v3.py index ef0c539..ff9ed2d 100644 --- a/tests/test_normalized_v3.py +++ b/tests/test_normalized_v3.py @@ -211,7 +211,13 @@ def test_missing_claim_index_recovers_but_tampering_fails_closed(tmp_path: Path) tmp_path / "normalized" / "event-claim-index-v3" / f"snapshot={result.snapshot.snapshot_id}" ) next(index_root.rglob("*.parquet")).unlink() - recovered = load_normalized_snapshot(tmp_path, result.snapshot.snapshot_id) + with pytest.raises(ValidationError, match="explicit StoragePolicy"): + load_normalized_snapshot(tmp_path, result.snapshot.snapshot_id) + recovered = load_normalized_snapshot( + tmp_path, + result.snapshot.snapshot_id, + recovery_policy=TEST_POLICY, + ) assert recovered.snapshot_id == result.snapshot.snapshot_id assert list(index_root.rglob("*.parquet")) assert list((tmp_path / "normalized" / "event-claim-index-v3" / "recovery-evidence").iterdir()) @@ -222,7 +228,13 @@ def test_missing_claim_index_recovers_but_tampering_fails_closed(tmp_path: Path) load_normalized_snapshot(tmp_path, result.snapshot.snapshot_id) shutil.rmtree(index_root) - recovered_again = load_normalized_snapshot(tmp_path, result.snapshot.snapshot_id) + with pytest.raises(ValidationError, match="explicit StoragePolicy"): + load_normalized_snapshot(tmp_path, result.snapshot.snapshot_id) + recovered_again = load_normalized_snapshot( + tmp_path, + result.snapshot.snapshot_id, + recovery_policy=TEST_POLICY, + ) assert recovered_again.snapshot_id == result.snapshot.snapshot_id diff --git a/tests/test_normalized_v3_failures.py b/tests/test_normalized_v3_failures.py index 1d7c643..b6da245 100644 --- a/tests/test_normalized_v3_failures.py +++ b/tests/test_normalized_v3_failures.py @@ -252,8 +252,15 @@ def test_claim_index_schema_logical_shape_and_missing_manifest_paths( / f"snapshot={recovery.snapshot.snapshot_id}" ) (recovery_index / "manifest.json").unlink() + with pytest.raises(ValidationError, match="explicit StoragePolicy"): + load_normalized_snapshot(recovery_root, recovery.snapshot.snapshot_id) assert ( - load_normalized_snapshot(recovery_root, recovery.snapshot.snapshot_id) == recovery.snapshot + load_normalized_snapshot( + recovery_root, + recovery.snapshot.snapshot_id, + recovery_policy=TEST_POLICY, + ) + == recovery.snapshot ) assert (recovery_index / "manifest.json").is_file()