src/msgboard/replicate.py

"""Snapshot replication of a board to an object store (an HF bucket, or any directory).

The live database stays on local disk. A background thread periodically takes a
consistent SQLite snapshot (backup API) and uploads it, together with any new
attachment files, then moves the LATEST pointer. An upload call returns only
when the object is stored, so the pointer never names a snapshot that is not
there. On startup an empty data directory is restored from LATEST.

Writes acknowledged since the last completed snapshot are lost if the machine
disappears; a graceful shutdown (SIGTERM) pushes a final snapshot first.
"""

from __future__ import annotations

import hashlib
import json
import logging
import os
import shutil
import sqlite3
import tempfile
import threading
import time
import uuid
from pathlib import Path

log = logging.getLogger("msgboard.replicate")

LATEST = "LATEST.json"
LEASE = "LEASE.json"


# ------------------------------------------------------------------ object stores


class DirStore:
    """Objects as files under a directory (tests, or a mounted volume)."""

    def __init__(self, root: str | Path):
        self.root = Path(root)

    def put(self, items: list[tuple[str | Path | bytes, str]]) -> None:
        for source, dest in items:
            target = self.root / dest
            target.parent.mkdir(parents=True, exist_ok=True)
            tmp = target.with_name(target.name + f".tmp-{uuid.uuid4().hex[:8]}")
            if isinstance(source, bytes):
                tmp.write_bytes(source)
            else:
                shutil.copyfile(source, tmp)
            os.replace(tmp, target)

    def get(self, remote: str, local: str | Path) -> bool:
        src = self.root / remote
        if not src.exists():
            return False
        shutil.copyfile(src, local)
        return True

    def list(self, prefix: str) -> list[str]:
        base = self.root / prefix
        if not base.exists():
            return []
        return sorted(str(p.relative_to(self.root)) for p in base.rglob("*") if p.is_file() and ".tmp-" not in p.name)

    def delete(self, paths: list[str]) -> None:
        for p in paths:
            (self.root / p).unlink(missing_ok=True)


class HfBucketStore:
    """Objects in a Hugging Face bucket through the Hub API (not a FUSE mount)."""

    def __init__(self, bucket_id: str, prefix: str = "", token: str | None = None):
        import huggingface_hub  # optional dependency

        self.hub = huggingface_hub
        self.bucket_id = bucket_id
        self.prefix = prefix.strip("/")
        self.token = token

    def _p(self, path: str) -> str:
        return f"{self.prefix}/{path}" if self.prefix else path

    def put(self, items):
        self.hub.batch_bucket_files(self.bucket_id, add=[(s, self._p(d)) for s, d in items], token=self.token)

    def get(self, remote: str, local) -> bool:
        local = Path(local)
        local.unlink(missing_ok=True)
        self.hub.download_bucket_files(self.bucket_id, [(self._p(remote), str(local))], token=self.token)
        return local.exists()

    def list(self, prefix: str) -> list[str]:
        full = self._p(prefix)
        try:
            entries = self.hub.list_bucket_tree(self.bucket_id, prefix=full, recursive=True, token=self.token)
            names = [e.path for e in entries if getattr(e, "type", "file") != "directory"
                     and type(e).__name__ != "BucketFolder"]
        except Exception as exc:  # missing prefix is not an error
            if "404" in str(exc) or "not found" in str(exc).lower():
                return []
            raise
        cut = len(self.prefix) + 1 if self.prefix else 0
        return sorted(n[cut:] for n in names)

    def delete(self, paths: list[str]) -> None:
        if paths:
            self.hub.batch_bucket_files(self.bucket_id, delete=[self._p(p) for p in paths], token=self.token)


def store_from_url(url: str, token: str | None = None):
    """`hf://buckets/<namespace>/<bucket>[/<prefix>]` or a plain directory path."""
    if url.startswith("hf://buckets/"):
        parts = url[len("hf://buckets/"):].strip("/").split("/")
        if len(parts) < 2:
            raise ValueError("expected hf://buckets/<namespace>/<bucket>[/<prefix>]")
        return HfBucketStore("/".join(parts[:2]), "/".join(parts[2:]),
                             token or os.environ.get("HF_TOKEN"))
    return DirStore(url)


# -------------------------------------------------------------------- replicator


def _sha256(path: Path) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def _read_json(store, remote: str) -> dict | None:
    with tempfile.TemporaryDirectory() as d:
        local = Path(d) / "x.json"
        try:
            if not store.get(remote, local):
                return None
        except Exception as exc:
            if "404" in str(exc) or "not found" in str(exc).lower():
                return None
            raise
        return json.loads(local.read_text())


def restore(store, data_dir: Path) -> dict | None:
    """Restore board.db and attachment files from the latest snapshot. Returns the manifest or None."""
    data_dir = Path(data_dir)
    if (data_dir / "board.db").exists():
        raise FileExistsError(f"{data_dir / 'board.db'} exists; refusing to overwrite it with a restore")
    latest = _read_json(store, LATEST)
    if latest is None:
        return None
    data_dir.mkdir(parents=True, exist_ok=True, mode=0o700)
    tmp = data_dir / "board.db.restoring"
    tmp.unlink(missing_ok=True)
    if not store.get(latest["snapshot"], tmp):
        raise RuntimeError(f"LATEST names {latest['snapshot']}, which is missing")
    if _sha256(tmp) != latest["sha256"]:
        tmp.unlink()
        raise RuntimeError("snapshot checksum mismatch")
    conn = sqlite3.connect(tmp)
    try:
        if conn.execute("PRAGMA integrity_check").fetchone()[0] != "ok":
            raise RuntimeError("snapshot failed integrity_check")
        keys = [r[0] for r in conn.execute("SELECT storage_key FROM attachments WHERE state = 'ready'")]
    finally:
        conn.close()
    objects = data_dir / "blobs" / "objects"
    for key in keys:
        dest = objects / key
        dest.parent.mkdir(parents=True, exist_ok=True)
        if not store.get(f"blobs/{key}", dest):
            log.warning("attachment %s missing from replica", key)
    os.replace(tmp, data_dir / "board.db")
    log.info("restored snapshot %s (%s bytes)", latest["snapshot"], latest["size"])
    return latest


def wait_for_release(store, stale_after: float = 120.0, poll: float = 1.0, log_every: float = 10.0,
                     stop: threading.Event | None = None) -> str:
    """Before restoring, wait for the previous owner to push its final snapshot and release the lease.

    Platforms such as HF Spaces start the replacement container before stopping
    the old one, so restoring immediately would miss the old one's last writes.
    Returns why waiting ended: "free", "released", "stale" (previous owner
    stopped heartbeating, e.g. it crashed) or "stopped" (`stop` was set).
    """
    last_log = 0.0
    while stop is None or not stop.is_set():
        lease = _read_json(store, LEASE)
        if lease is None:
            return "free"
        if lease.get("released"):
            return "released"
        if time.time() - lease.get("heartbeat", lease.get("since", 0)) > stale_after:
            return "stale"
        if time.time() - last_log > log_every:
            log.info("waiting for instance %s to hand over the replica", lease.get("instance"))
            last_log = time.time()
        time.sleep(poll)
    return "stopped"


class Replicator:
    """Background snapshot uploader owned by a running Board."""

    def __init__(self, store, data_dir: Path, interval: float = 10.0, keep: int = 5, heartbeat: float = 30.0):
        self.store = store
        self.data_dir = Path(data_dir)
        self.interval = interval
        self.keep = keep
        self.heartbeat = heartbeat
        self._last_heartbeat = 0.0
        self.instance = uuid.uuid4().hex[:12]
        self.generation_fn = lambda: 0
        self._pushed_generation = None
        self._uploaded_blobs: set[str] | None = None
        self._stop = threading.Event()
        self._thread: threading.Thread | None = None
        self._lock = threading.Lock()
        self.fenced = False
        self.events: list[tuple[float, str]] = []
        self.status = {"instance": self.instance, "last_push_at": None, "last_push_seconds": None,
                       "last_snapshot": None, "last_error": None, "pushes": 0, "fenced": False}

    # Lease: the newest instance owns the replica. An older instance that is
    # still running (overlapping restart) stops pushing instead of overwriting.
    def take_lease(self, released: bool = False) -> None:
        now = time.time()
        self.store.put([(json.dumps({"instance": self.instance, "heartbeat": now, "released": released,
                                     "snapshot": self.status["last_snapshot"]}).encode(), LEASE)])
        self._last_heartbeat = now

    def _still_owner(self) -> bool:
        lease = _read_json(self.store, LEASE)
        return lease is None or lease.get("instance") == self.instance

    def record(self, event: str, wait: bool = False) -> None:
        """Append a lifecycle event to events/<instance>.json in the replica (diagnostics)."""
        self.events.append((time.time(), event))
        body = json.dumps({"instance": self.instance, "events": self.events}).encode()

        def upload():
            try:
                self.store.put([(body, f"events/{self.instance}.json")])
            except Exception as exc:
                log.warning("could not record event %s: %s", event, exc)

        if wait:
            upload()
        else:
            threading.Thread(target=upload, daemon=True).start()

    def start(self, generation_fn) -> None:
        self.generation_fn = generation_fn
        self.take_lease()
        self.record("started")
        self._thread = threading.Thread(target=self._run, name="msgboard-replicator", daemon=True)
        self._thread.start()

    def _run(self) -> None:
        while not self._stop.wait(self.interval):
            try:
                self.push_if_changed()
                if not self.fenced and time.time() - self._last_heartbeat > self.heartbeat and self._still_owner():
                    self.take_lease()
            except Exception as exc:  # keep serving; report through metrics
                self.status["last_error"] = f"{type(exc).__name__}: {exc}"[:300]
                log.warning("replication push failed: %s", exc)

    def stop(self, final_push: bool = True) -> None:
        self._stop.set()
        if self._thread is not None:
            self._thread.join(timeout=60)
        if final_push:
            self.record("final-push-start")
            try:
                pushed = self.push_if_changed()
                if not self.fenced and self._still_owner():
                    self.take_lease(released=True)  # the next instance may restore now
                self.record("final-push-done" if pushed else
                            ("final-push-fenced" if self.fenced else "final-push-unchanged"), wait=True)
            except Exception as exc:
                log.error("final snapshot push failed: %s", exc)
                self.record(f"final-push-failed: {type(exc).__name__}", wait=True)

    def push_if_changed(self) -> bool:
        with self._lock:
            generation = self.generation_fn()
            if self.fenced or generation == self._pushed_generation:
                return False
            if not self._still_owner():
                self.fenced = self.status["fenced"] = True
                log.error("another instance owns the replica; this instance stops pushing")
                return False
            started = time.time()
            self._push(generation)
            self._pushed_generation = generation
            self.status.update(last_push_at=time.time(), last_push_seconds=round(time.time() - started, 3),
                               last_error=None, pushes=self.status["pushes"] + 1)
            return True

    def _push(self, generation: int) -> None:
        with tempfile.TemporaryDirectory(dir=self.data_dir) as d:
            snap = Path(d) / "snapshot.db"
            src = sqlite3.connect(self.data_dir / "board.db")
            dst = sqlite3.connect(snap)
            try:
                src.backup(dst)  # one step: a consistent point-in-time copy
                keys = [r[0] for r in dst.execute("SELECT storage_key FROM attachments WHERE state = 'ready'")]
            finally:
                dst.close()
                src.close()
            # 1. attachment files referenced by the snapshot
            if self._uploaded_blobs is None:
                self._uploaded_blobs = {p[len("blobs/"):] for p in self.store.list("blobs/")}
            new = [k for k in keys if k not in self._uploaded_blobs]
            objects = self.data_dir / "blobs" / "objects"
            for i in range(0, len(new), 50):
                batch = [(objects / k, f"blobs/{k}") for k in new[i:i + 50] if (objects / k).exists()]
                if batch:
                    self.store.put(batch)
                self._uploaded_blobs.update(new[i:i + 50])
            # 2. the snapshot, 3. then the pointer to it
            stamp = time.strftime("%Y%m%dT%H%M%S", time.gmtime())
            name = f"snapshots/{stamp}-{self.instance}-{generation}.db"
            manifest = {"snapshot": name, "sha256": _sha256(snap), "size": snap.stat().st_size,
                        "created_at": time.time(), "instance": self.instance, "generation": generation}
            self.store.put([(snap, name)])
            self.store.put([(json.dumps(manifest).encode(), LATEST)])
            self.status["last_snapshot"] = name
        self._prune()

    def _prune(self) -> None:
        snaps = [p for p in self.store.list("snapshots/") if p.endswith(".db")]
        if len(snaps) > self.keep:
            self.store.delete(snaps[: len(snaps) - self.keep])