Resuming an Interrupted OSM Import Jump to heading#

#!/usr/bin/env python3
"""A resumable import: atomic units, markers written after the work, safe replay."""
from __future__ import annotations

import json
import logging
import os
import signal
import tempfile
from collections.abc import Callable, Iterable
from dataclasses import dataclass
from pathlib import Path

logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
logger = logging.getLogger(__name__)


class Markers:
    """Durable record of completed units. One line per unit, fsync'd on write."""

    def __init__(self, path: Path) -> None:
        self.path = path
        self.done: set[str] = set()
        if path.exists():
            for line in path.read_text().splitlines():
                if line.strip():
                    self.done.add(json.loads(line)["unit"])
            logger.info("resuming: %d unit(s) already complete", len(self.done))

    def record(self, unit: str, rows: int) -> None:
        # Append + flush + fsync: a marker that is not on disk is not a marker.
        with self.path.open("a") as handle:
            handle.write(json.dumps({"unit": unit, "rows": rows}) + "\n")
            handle.flush()
            os.fsync(handle.fileno())
        self.done.add(unit)

    def __contains__(self, unit: str) -> bool:
        return unit in self.done


def write_atomically(target: Path, write: Callable[[Path], int]) -> int:
    """Write via a temp file in the same directory, then rename.

    Same directory matters: rename is only atomic within one filesystem, and a
    temp file in /tmp is frequently on a different one.
    """
    target.parent.mkdir(parents=True, exist_ok=True)
    fd, tmp_name = tempfile.mkstemp(dir=target.parent, suffix=".tmp")
    os.close(fd)
    tmp = Path(tmp_name)
    try:
        rows = write(tmp)
        # Durability before visibility: the data must be on disk before the rename.
        with tmp.open("rb") as handle:
            os.fsync(handle.fileno())
        tmp.replace(target)                     # atomic within the filesystem
        return rows
    except BaseException:
        tmp.unlink(missing_ok=True)             # never leave a half-written .tmp
        raise


@dataclass
class Interrupted(Exception):
    """Raised on SIGTERM so the current unit unwinds cleanly rather than being killed."""
    signum: int


def _install_signal_handlers() -> None:
    def handler(signum, _frame):
        raise Interrupted(signum)
    for sig in (signal.SIGINT, signal.SIGTERM):
        signal.signal(sig, handler)


def run_import(units: Iterable[str],
               process: Callable[[str, Path], int],
               out_dir: Path,
               marker_path: Path) -> None:
    """Process each unit exactly once across any number of runs."""
    _install_signal_handlers()
    markers = Markers(marker_path)
    total_rows = skipped = 0

    for unit in units:
        if unit in markers:
            skipped += 1
            continue
        target = out_dir / f"{unit}.parquet"
        try:
            rows = write_atomically(target, lambda tmp, u=unit: process(u, tmp))
        except Interrupted as stop:
            logger.warning("interrupted by signal %d during unit %s — "
                           "that unit will be replayed on the next run", stop.signum, unit)
            raise SystemExit(130)
        markers.record(unit, rows)              # AFTER the rename, never before
        total_rows += rows
        logger.info("unit %s: %d row(s)", unit, rows)

    logger.info("import complete: %d unit(s) skipped, %d row(s) written this run",
                skipped, total_rows)


def reconcile(units: Iterable[str], out_dir: Path, marker_path: Path) -> list[str]:
    """Find units whose marker and output disagree — the crash-window casualties."""
    markers = Markers(marker_path)
    problems: list[str] = []
    for unit in units:
        exists = (out_dir / f"{unit}.parquet").exists()
        recorded = unit in markers
        if exists and not recorded:
            problems.append(f"{unit}: written but not marked (crash before marker)")
        if recorded and not exists:
            problems.append(f"{unit}: marked but missing (output deleted?)")
    for problem in problems:
        logger.warning("%s", problem)
    return problems

# Start the import, kill it partway, then run it again.
timeout 60 python3 import.py; echo "exit $?"
python3 import.py            # should log "resuming: N unit(s) already complete"

python3 import.py            # run to completion
find out -name '*.parquet' -printf '%s %p\n' | sort > /tmp/first
rm -f markers.jsonl          # force a full replay
python3 import.py
find out -name '*.parquet' -printf '%s %p\n' | sort > /tmp/second
diff /tmp/first /tmp/second && echo "idempotent"

find out -name '*.tmp' -print | head       # expect nothing
python3 -c "from import_ import reconcile; ..."   # expect no problems