Skip to content

Commit 577956f

Browse files
committed
Merge remote-tracking branch 'origin/main' into feat/pageindex-max-concurrency
# Conflicts: # openkb/cli.py # tests/test_add_command.py
2 parents 8a8e65c + 5a2b48c commit 577956f

7 files changed

Lines changed: 794 additions & 166 deletions

File tree

openkb/add_coordinator.py

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
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

Comments
 (0)