|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +import logging |
| 4 | +import shutil |
| 5 | +from collections.abc import Callable, Sequence |
| 6 | +from dataclasses import dataclass, field |
| 7 | +from pathlib import Path |
| 8 | + |
| 9 | +import click |
| 10 | + |
| 11 | +from openkb.locks import kb_ingest_lock_held |
| 12 | +from openkb.mutation import MutationSnapshot, snapshot_paths |
| 13 | + |
| 14 | +logger = logging.getLogger(__name__) |
| 15 | + |
| 16 | +MutationBody = Callable[[MutationSnapshot], None] |
| 17 | +PostCommitHook = Callable[[], None] |
| 18 | + |
| 19 | + |
| 20 | +class DirtyRollbackError(RuntimeError): |
| 21 | + """A mutation's rollback failed, leaving an active journal on disk. |
| 22 | +
|
| 23 | + The KB may be in a partially-applied state that the retained journal will |
| 24 | + attempt to roll back on the next exclusive-lock acquisition. Batch owners |
| 25 | + (the parallel/serial ``add`` loops) MUST stop committing further mutations |
| 26 | + on top of this dirty state instead of continuing — otherwise the next |
| 27 | + recovery rolls this journal back over the shared paths it recorded |
| 28 | + (``hashes.json``, ``index.md``, ``concepts/``, ``entities/``) and silently |
| 29 | + clobbers the later commits. Single-mutation callers should let it propagate |
| 30 | + so the command fails loudly; rerunning recovers via the drain. |
| 31 | + """ |
| 32 | + |
| 33 | + def __init__(self, operation: str, journal_path: Path) -> None: |
| 34 | + super().__init__( |
| 35 | + f"Dirty rollback for {operation}; journal retained at {journal_path}. " |
| 36 | + f"Rerun the command to recover." |
| 37 | + ) |
| 38 | + self.operation = operation |
| 39 | + self.journal_path = journal_path |
| 40 | + |
| 41 | + |
| 42 | +@dataclass(slots=True) |
| 43 | +class AddMutationPlan: |
| 44 | + operation: str |
| 45 | + details: dict |
| 46 | + touched_paths: Sequence[Path] |
| 47 | + body: MutationBody |
| 48 | + post_commit_hooks: Sequence[PostCommitHook] = field(default_factory=tuple) |
| 49 | + hardlink_dirs: set[Path] = field(default_factory=set) |
| 50 | + staging_dirs: Sequence[Path | None] = field(default_factory=tuple) |
| 51 | + |
| 52 | + |
| 53 | +def _cleanup_staging_dirs(staging_dirs: Sequence[Path | None]) -> None: |
| 54 | + for staging_dir in staging_dirs: |
| 55 | + if staging_dir is not None: |
| 56 | + shutil.rmtree(staging_dir, ignore_errors=True) |
| 57 | + |
| 58 | + |
| 59 | +def _rollback_snapshot(plan: AddMutationPlan, snapshot) -> Path | None: |
| 60 | + """Best-effort rollback; returns the retained journal path on dirty failure. |
| 61 | +
|
| 62 | + Returns ``snapshot.journal_path`` when the snapshot existed but rollback |
| 63 | + FAILED (the active journal is retained for next-run recovery), otherwise |
| 64 | + ``None`` — covering both "snapshot is None" (nothing was applied; the |
| 65 | + failure happened during snapshot setup before the body ran) and a clean |
| 66 | + rollback that discarded its journal. |
| 67 | + """ |
| 68 | + if snapshot is None: |
| 69 | + _cleanup_staging_dirs(plan.staging_dirs) |
| 70 | + return None |
| 71 | + rollback_error = snapshot.rollback_best_effort() |
| 72 | + if rollback_error is None: |
| 73 | + snapshot.discard_best_effort() |
| 74 | + else: |
| 75 | + click.echo( |
| 76 | + " [ERROR] Rollback failed; mutation journal retained for recovery: " |
| 77 | + f"{snapshot.journal_path}" |
| 78 | + ) |
| 79 | + _cleanup_staging_dirs(plan.staging_dirs) |
| 80 | + return snapshot.journal_path if rollback_error is not None else None |
| 81 | + |
| 82 | + |
| 83 | +def _failure_target(details: dict) -> str: |
| 84 | + for key in ("name", "doc_name", "doc_id"): |
| 85 | + value = details.get(key) |
| 86 | + if value: |
| 87 | + return f" for {value}" |
| 88 | + return "" |
| 89 | + |
| 90 | + |
| 91 | +def run_add_mutation(kb_dir: Path, plan: AddMutationPlan) -> bool: |
| 92 | + if not kb_ingest_lock_held(kb_dir / ".openkb"): |
| 93 | + raise RuntimeError("run_add_mutation requires the caller to hold kb_ingest_lock") |
| 94 | + snapshot = None |
| 95 | + try: |
| 96 | + snapshot = snapshot_paths( |
| 97 | + kb_dir, |
| 98 | + list(plan.touched_paths), |
| 99 | + operation=plan.operation, |
| 100 | + details=plan.details, |
| 101 | + hardlink_dirs=plan.hardlink_dirs, |
| 102 | + ) |
| 103 | + plan.body(snapshot) |
| 104 | + snapshot.mark_committed() |
| 105 | + except Exception as exc: |
| 106 | + dirty_journal = _rollback_snapshot(plan, snapshot) |
| 107 | + if dirty_journal is not None: |
| 108 | + # Rollback failed and left an active journal. Stop the batch rather |
| 109 | + # than committing more docs on top of dirty state that the next |
| 110 | + # recovery would roll back over. |
| 111 | + raise DirtyRollbackError(plan.operation, dirty_journal) |
| 112 | + click.echo(f" [ERROR] {plan.operation} failed{_failure_target(plan.details)}: {exc}") |
| 113 | + logger.debug("%s mutation failed:", plan.operation, exc_info=True) |
| 114 | + return False |
| 115 | + except BaseException: |
| 116 | + # Interrupt (KeyboardInterrupt / SystemExit): best-effort rollback for |
| 117 | + # its side-effects only. Do NOT raise DirtyRollbackError — propagate the |
| 118 | + # interrupt so the user's abort is honored. Any retained journal or |
| 119 | + # orphaned staging is recovered next run by the drain + reaper. |
| 120 | + _rollback_snapshot(plan, snapshot) |
| 121 | + raise |
| 122 | + finally: |
| 123 | + _cleanup_staging_dirs(plan.staging_dirs) |
| 124 | + |
| 125 | + for hook in plan.post_commit_hooks: |
| 126 | + try: |
| 127 | + hook() |
| 128 | + except Exception as exc: |
| 129 | + logger.warning("Post-commit hook failed for %s: %s", plan.operation, exc) |
| 130 | + |
| 131 | + cleanup_error = snapshot.discard_best_effort() |
| 132 | + if cleanup_error is not None: |
| 133 | + click.echo(f" [WARN] mutation journal cleanup failed: {cleanup_error}") |
| 134 | + return True |
0 commit comments