diff --git a/arborist/migrate.py b/arborist/migrate.py new file mode 100644 index 0000000..7ae79bc --- /dev/null +++ b/arborist/migrate.py @@ -0,0 +1,328 @@ +"""Reshard / hydration tool (#000065). + +Migrates a source shards directory (any M_source) to a target shards +directory (any M_target) by content-hash routing: + + target_shard = shard_for_document(document_root, M_target) + +Per-document tables follow their parent doc. The audit chain is +consolidated to canonical shard ``000`` via Option A — every source +shard's ``audit_events`` rows are sorted by ``ts`` and re-chained +into one canonical chain landed in shard 000. + +The :class:`HydrationPlanner` picks a strategy from available disk: + +* ``all_at_once`` — direct source→target with overlap; fastest +* ``per_source_shard``— pack/delete/hydrate per round; safest, slowest +* ``streaming_row`` — in-place mutation; last resort (--allow-in-place) +* refused — below safety margin, return :class:`InsufficientDisk` + +Pure planning (``HydrationPlanner.plan``) is decoupled from execution. +The plan + free-space readout + peak draw + duration are recorded in +one ``audit_events`` row at the end of every successful migration. +""" +from __future__ import annotations + +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Literal + + +# --------------------------------------------------------------------------- +# Sizing knobs. Threshold rationale documented inline so an operator can +# audit "why did the planner pick X" from this file alone. +# --------------------------------------------------------------------------- + +#: Always keep this much free above the planner's draw estimate. SQLite +#: ``temp_store=FILE`` spills can be GB-scale on the edges fan-in; this is +#: the buffer that survives one bad sort + one badly-sized WAL grow. +SAFETY_MARGIN_BYTES = 4 * 1024 ** 3 + +#: SQLite WAL during heavy INSERT against a 9 GB shard has been measured +#: at ~4 GB on the production corpus; the four-shard parallel hydration +#: of strategy A can hold all four WALs live at once. +WAL_HEADROOM_PER_SHARD = 4 * 1024 ** 3 + +#: VACUUM scratch at end of strategy A. SQLite VACUUM rewrites the whole +#: DB to a sibling tmpfile then swaps; peak draw is ~1× DB size. +VACUUM_HEADROOM_RATIO = 1.0 + +#: Strategy B per-round pack size: a single source shard packed with +#: zstd typically lands at ~30% of source-shard size. Estimated, the +#: actual measured ratio at migration time replaces this. +DEFAULT_PACK_RATIO = 0.30 + +#: Multiplier on ``largest_source_shard_bytes`` for strategy B sizing. +#: ``1.5×`` covers: pack-on-disk + new-shard-growth + WAL on the new +#: shard being written during that round + scratch. +STRATEGY_B_MULTIPLIER = 1.5 + + +Strategy = Literal[ + "all_at_once", + "per_source_shard", + "streaming_row", +] + + +@dataclass(frozen=True) +class CorpusFacts: + """Read-only snapshot of source-corpus sizes the planner needs. + + All ``*_bytes`` fields are bytes (not blocks, not KB). ``shard_files`` + is sorted by index (``000.db``, ``001.db``, ...). + """ + + shard_files: tuple[Path, ...] + total_bytes: int + largest_shard_bytes: int + + @classmethod + def from_dir(cls, shards_dir: Path) -> "CorpusFacts": + files = tuple(sorted(shards_dir.glob("00[0-9].db"))) + if not files: + raise ValueError(f"no shard files (00[0-9].db) under {shards_dir}") + sizes = [p.stat().st_size for p in files] + return cls( + shard_files=files, + total_bytes=sum(sizes), + largest_shard_bytes=max(sizes), + ) + + +@dataclass(frozen=True) +class HydrationPlan: + """Output of :class:`HydrationPlanner.plan`. + + Persisted into the migration audit event so an operator can later + answer "why did the tool pick strategy X" from the audit chain + alone. Never includes a path to a credential or secret. + """ + + strategy: Strategy + source_dir: Path + target_dir: Path + target_M: int + corpus_facts: CorpusFacts + free_bytes_at_plan: int + estimated_peak_draw_bytes: int + estimated_final_target_bytes: int + rationale: str + + def free_at_peak_bytes(self) -> int: + """Predicted free disk at the worst point during execution.""" + return self.free_bytes_at_plan - self.estimated_peak_draw_bytes + + def as_audit_body(self) -> dict: + return { + "strategy": self.strategy, + "source_dir": str(self.source_dir), + "target_dir": str(self.target_dir), + "target_M": self.target_M, + "source_shards": [str(p) for p in self.corpus_facts.shard_files], + "source_total_bytes": self.corpus_facts.total_bytes, + "source_largest_bytes": self.corpus_facts.largest_shard_bytes, + "free_bytes_at_plan": self.free_bytes_at_plan, + "estimated_peak_draw_bytes": self.estimated_peak_draw_bytes, + "estimated_final_target_bytes": self.estimated_final_target_bytes, + "rationale": self.rationale, + } + + +class InsufficientDisk(RuntimeError): + """Raised when free space is below the safety margin for every strategy. + + The exception carries the corpus facts + free-bytes readout the planner + saw so the operator's terminal output (or CI log) explains the gap. + """ + + def __init__( + self, + *, + free_bytes: int, + required_bytes: int, + corpus_facts: CorpusFacts, + ) -> None: + self.free_bytes = free_bytes + self.required_bytes = required_bytes + self.corpus_facts = corpus_facts + super().__init__( + f"insufficient disk: have {free_bytes / 1e9:.1f} GB free, " + f"need at least {required_bytes / 1e9:.1f} GB for any strategy. " + f"corpus {corpus_facts.total_bytes / 1e9:.1f} GB across " + f"{len(corpus_facts.shard_files)} shards. Free disk or mount " + f"target on a different filesystem." + ) + + +def _statvfs_free_bytes(path: Path) -> int: + """Bytes free on the filesystem containing ``path``. + + Uses ``f_bavail`` (unprivileged free), not ``f_bfree``, matching what + ``df`` shows non-root users. + """ + st = os.statvfs(path) + return st.f_bavail * st.f_frsize + + +class HydrationPlanner: + """Pick a strategy + report estimated draw, given disk constraints. + + Pure-compute API: feed it the facts (corpus size, free bytes) and + it returns a :class:`HydrationPlan`. Execution is a separate + function (``execute_plan``) so we can unit-test planning without + touching real shards or the filesystem. + """ + + def __init__( + self, + *, + safety_margin_bytes: int = SAFETY_MARGIN_BYTES, + wal_headroom_per_shard: int = WAL_HEADROOM_PER_SHARD, + vacuum_headroom_ratio: float = VACUUM_HEADROOM_RATIO, + pack_ratio: float = DEFAULT_PACK_RATIO, + strategy_b_multiplier: float = STRATEGY_B_MULTIPLIER, + ) -> None: + self.safety_margin_bytes = safety_margin_bytes + self.wal_headroom_per_shard = wal_headroom_per_shard + self.vacuum_headroom_ratio = vacuum_headroom_ratio + self.pack_ratio = pack_ratio + self.strategy_b_multiplier = strategy_b_multiplier + + # ----- size estimators ----- + + def _estimate_strategy_a_peak_draw( + self, facts: CorpusFacts, M_target: int + ) -> int: + """Strategy A all-at-once. + + Peak draw = new shards being written (≈ corpus size) + WAL on each + of the M_target target shards + VACUUM scratch on the largest + target. Local packs are NOT used by strategy A in this build — + direct source→target read+write keeps draw lower than the pack + variant. + """ + new_shards = facts.total_bytes + wal = self.wal_headroom_per_shard * M_target + # VACUUM scratch is bounded by the largest target shard, which + # for a uniformly hash-routed corpus is approximately + # total_bytes / M_target. + per_target = facts.total_bytes / max(M_target, 1) + vacuum = int(per_target * self.vacuum_headroom_ratio) + return new_shards + wal + vacuum + + def _estimate_strategy_b_peak_draw( + self, facts: CorpusFacts, M_target: int + ) -> int: + """Strategy B per-source-shard round, peak DRAW above baseline. + + Per round: pack source → delete source → hydrate → delete pack. + Source is deleted BEFORE new shards are written, so the net + on-disk size never exceeds baseline + pack at any moment. + Round peak above baseline = one pack on disk + WAL on the + currently-writing target + VACUUM scratch on one target (the + largest, ≈ corpus / M_target). + """ + pack = int(facts.largest_shard_bytes * self.pack_ratio) + wal = self.wal_headroom_per_shard + per_target_bytes = facts.total_bytes / max(M_target, 1) + vacuum = int(per_target_bytes * self.vacuum_headroom_ratio) + return pack + wal + vacuum + + def _estimate_strategy_c_peak_draw(self, facts: CorpusFacts) -> int: + """Strategy C streaming-row in-place. + + Peak draw = WAL only (rows DELETE from source as they INSERT to + target). No pack, no overlap of old+new at corpus scale. + """ + return self.wal_headroom_per_shard + + # ----- planner ----- + + def plan( + self, + *, + source_dir: Path, + target_dir: Path, + target_M: int, + free_bytes: int | None = None, + force_strategy: Strategy | None = None, + allow_in_place: bool = False, + ) -> HydrationPlan: + """Pick a strategy that fits the available disk + return the plan.""" + if target_M < 1: + raise ValueError(f"target_M must be >= 1, got {target_M}") + facts = CorpusFacts.from_dir(source_dir) + free = free_bytes if free_bytes is not None else _statvfs_free_bytes( + target_dir.parent if target_dir.parent.exists() else target_dir + ) + + # Evaluate strategies in preference order: fastest fitting wins. + # Strategy C is opt-in only because it mutates source in place. + candidates: list[tuple[Strategy, int]] = [ + ("all_at_once", self._estimate_strategy_a_peak_draw(facts, target_M)), + ("per_source_shard", self._estimate_strategy_b_peak_draw(facts, target_M)), + ] + if allow_in_place: + candidates.append( + ("streaming_row", self._estimate_strategy_c_peak_draw(facts)) + ) + + if force_strategy is not None: + forced = next( + ((s, d) for s, d in candidates if s == force_strategy), None + ) + if forced is None: + raise ValueError( + f"forced strategy {force_strategy!r} not available " + f"(allow_in_place={allow_in_place})" + ) + strategy, draw = forced + rationale = f"forced strategy={force_strategy} (operator override)" + return self._build_plan( + strategy, draw, rationale, facts, free, source_dir, target_dir, target_M + ) + + for strategy, draw in candidates: + if free >= draw + self.safety_margin_bytes: + rationale = ( + f"{strategy}: peak draw {draw / 1e9:.1f} GB + " + f"{self.safety_margin_bytes / 1e9:.1f} GB safety = " + f"{(draw + self.safety_margin_bytes) / 1e9:.1f} GB ≤ " + f"free {free / 1e9:.1f} GB" + ) + return self._build_plan( + strategy, draw, rationale, facts, free, + source_dir, target_dir, target_M, + ) + + # Nothing fit — report the cheapest required draw so the operator + # knows how much to free. + min_needed = min(d for _, d in candidates) + self.safety_margin_bytes + raise InsufficientDisk( + free_bytes=free, required_bytes=min_needed, corpus_facts=facts + ) + + def _build_plan( + self, + strategy: Strategy, + peak_draw: int, + rationale: str, + facts: CorpusFacts, + free: int, + source_dir: Path, + target_dir: Path, + target_M: int, + ) -> HydrationPlan: + return HydrationPlan( + strategy=strategy, + source_dir=source_dir, + target_dir=target_dir, + target_M=target_M, + corpus_facts=facts, + free_bytes_at_plan=free, + estimated_peak_draw_bytes=peak_draw, + estimated_final_target_bytes=facts.total_bytes, + rationale=rationale, + ) diff --git a/tests/test_migrate_planner.py b/tests/test_migrate_planner.py new file mode 100644 index 0000000..3474324 --- /dev/null +++ b/tests/test_migrate_planner.py @@ -0,0 +1,223 @@ +"""Tests for the #000065 hydration planner. + +Pure-compute tests — never touches SQLite or the network. Validates +that the planner picks the right strategy from a given (corpus size, +free disk) pair, and that the chosen strategy's peak-draw estimate +matches the safety budget. +""" +from __future__ import annotations + +from pathlib import Path + +import pytest + +from arborist.migrate import ( + CorpusFacts, + HydrationPlan, + HydrationPlanner, + InsufficientDisk, + SAFETY_MARGIN_BYTES, +) + + +GB = 1024 ** 3 + + +@pytest.fixture +def four_shard_corpus(tmp_path: Path) -> Path: + """Create 4 empty 'shard' files of realistic sizes (~8.8 GB each).""" + shards_dir = tmp_path / "shards" + shards_dir.mkdir() + # Don't allocate real bytes; touch + truncate is enough for the + # planner since it reads ``.stat().st_size``. + sizes_gb = [8.8, 8.8, 8.8, 8.8] + for i, size_gb in enumerate(sizes_gb): + path = shards_dir / f"{i:03d}.db" + path.touch() + os = __import__("os") + os.truncate(path, int(size_gb * GB)) + return shards_dir + + +@pytest.fixture +def target_dir(tmp_path: Path) -> Path: + d = tmp_path / "shards.v2" + d.mkdir() + return d + + +class TestCorpusFacts: + def test_reads_sizes(self, four_shard_corpus): + f = CorpusFacts.from_dir(four_shard_corpus) + assert len(f.shard_files) == 4 + assert f.total_bytes == int(8.8 * GB) * 4 + assert f.largest_shard_bytes == int(8.8 * GB) + + def test_empty_dir_rejected(self, tmp_path): + empty = tmp_path / "empty" + empty.mkdir() + with pytest.raises(ValueError, match="no shard files"): + CorpusFacts.from_dir(empty) + + +class TestStrategyPick: + def test_88_gb_free_picks_all_at_once( + self, four_shard_corpus, target_dir + ): + """The production-host case: 88 GB free, 35 GB corpus.""" + plan = HydrationPlanner().plan( + source_dir=four_shard_corpus, + target_dir=target_dir, + target_M=4, + free_bytes=88 * GB, + ) + assert plan.strategy == "all_at_once" + + def test_47_gb_free_falls_to_per_shard( + self, four_shard_corpus, target_dir + ): + """The old-host case: 47 GB free, 35 GB corpus.""" + plan = HydrationPlanner().plan( + source_dir=four_shard_corpus, + target_dir=target_dir, + target_M=4, + free_bytes=47 * GB, + ) + assert plan.strategy == "per_source_shard" + + def test_15_gb_free_refused_without_in_place( + self, four_shard_corpus, target_dir + ): + """At 15 GB free, no opt-in: strategy B doesn't fit either.""" + with pytest.raises(InsufficientDisk): + HydrationPlanner().plan( + source_dir=four_shard_corpus, + target_dir=target_dir, + target_M=4, + free_bytes=15 * GB, + ) + + def test_10_gb_free_with_in_place_picks_streaming( + self, four_shard_corpus, target_dir + ): + """At 10 GB free with --allow-in-place: strategy C (streaming).""" + plan = HydrationPlanner().plan( + source_dir=four_shard_corpus, + target_dir=target_dir, + target_M=4, + free_bytes=10 * GB, + allow_in_place=True, + ) + assert plan.strategy == "streaming_row" + + def test_zero_free_refused_even_with_in_place( + self, four_shard_corpus, target_dir + ): + with pytest.raises(InsufficientDisk): + HydrationPlanner().plan( + source_dir=four_shard_corpus, + target_dir=target_dir, + target_M=4, + free_bytes=1 * GB, + allow_in_place=True, + ) + + +class TestForcedStrategy: + def test_force_strategy_b_even_when_a_fits( + self, four_shard_corpus, target_dir + ): + """Operator override: pick B even if A fits.""" + plan = HydrationPlanner().plan( + source_dir=four_shard_corpus, + target_dir=target_dir, + target_M=4, + free_bytes=88 * GB, + force_strategy="per_source_shard", + ) + assert plan.strategy == "per_source_shard" + assert "forced" in plan.rationale + + def test_force_unavailable_strategy_rejected( + self, four_shard_corpus, target_dir + ): + with pytest.raises(ValueError, match="not available"): + HydrationPlanner().plan( + source_dir=four_shard_corpus, + target_dir=target_dir, + target_M=4, + free_bytes=88 * GB, + force_strategy="streaming_row", # not allow_in_place + ) + + +class TestPlanFields: + def test_audit_body_is_json_serializable( + self, four_shard_corpus, target_dir + ): + import json + plan = HydrationPlanner().plan( + source_dir=four_shard_corpus, + target_dir=target_dir, + target_M=4, + free_bytes=88 * GB, + ) + json.dumps(plan.as_audit_body()) # no exception + + def test_free_at_peak_reasonable( + self, four_shard_corpus, target_dir + ): + plan = HydrationPlanner().plan( + source_dir=four_shard_corpus, + target_dir=target_dir, + target_M=4, + free_bytes=88 * GB, + ) + # At least the safety margin remains free at peak + # (otherwise the planner would've rejected this strategy). + assert plan.free_at_peak_bytes() >= SAFETY_MARGIN_BYTES + + +class TestTargetMValidation: + def test_target_M_zero_rejected(self, four_shard_corpus, target_dir): + with pytest.raises(ValueError, match="target_M"): + HydrationPlanner().plan( + source_dir=four_shard_corpus, + target_dir=target_dir, + target_M=0, + free_bytes=88 * GB, + ) + + def test_target_M_one(self, four_shard_corpus, target_dir): + """M=1 (consolidate to a single shard) is a valid target.""" + plan = HydrationPlanner().plan( + source_dir=four_shard_corpus, + target_dir=target_dir, + target_M=1, + free_bytes=88 * GB, + ) + assert plan.target_M == 1 + + +class TestVariedShardSizes: + """Skew test: one shard much larger than the others.""" + + def test_strategy_b_bounded_by_largest(self, tmp_path): + import os + shards = tmp_path / "shards" + shards.mkdir() + # 3 small + 1 large. + sizes = [1 * GB, 1 * GB, 1 * GB, 20 * GB] + for i, s in enumerate(sizes): + p = shards / f"{i:03d}.db" + p.touch() + os.truncate(p, s) + plan = HydrationPlanner().plan( + source_dir=shards, + target_dir=tmp_path / "out", + target_M=4, + free_bytes=40 * GB, + ) + # Strategy A peak should fail (23 GB new + 4 × 4 WAL + scratch >40), + # so we expect strategy B which is bounded by the 20 GB shard. + assert plan.strategy == "per_source_shard"