#000065 step 2: hydration planner (pure compute, no execution yet)

Disk-aware strategy picker for the corpus reshard. Pure-compute API:
feed it (corpus facts, free bytes) and it returns a HydrationPlan
naming one of three strategies + a peak-draw estimate.

Strategies (preference order):
  all_at_once       - direct source→target with overlap; fastest;
                      peak draw = corpus + WAL × M_target + VACUUM
  per_source_shard  - pack→delete→hydrate per round; safest;
                      peak draw = pack + WAL + 1 target VACUUM
  streaming_row     - in-place mutation, --allow-in-place opt-in;
                      peak draw = WAL only

Safety budget: 4 GB margin above the strategy's peak draw — the
buffer that survives one bad sort + one badly-sized WAL grow.

Production-host preview at 95 GB free / 38 GB corpus:
  strategy: all_at_once
  peak draw est: 64.2 GB
  free at peak: 31.2 GB (well above safety)
  rationale: all_at_once: peak draw 64.2 GB + 4.3 GB safety ≤ free 95.4 GB

The plan's full readout (strategy, free bytes, peak draw, rationale)
is JSON-serializable so it goes into the migration audit event as a
single forensic record.

14 tests cover: strategy pick by free-disk level, --force_strategy
override, --allow-in-place gating, skewed shard sizes, json-
serializable audit body, target_M validation.

No execution yet — this is just the planner. Next: the strategy A
executor (direct source→target read+write).

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,
+        )

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"