Skip to content

Commit ee22914

Browse files
committed
refactor(mutation): split lock-free prepare from serial commit for directory ingest
Stage 4 of the parallel-add roadmap (#151). `openkb add <dir>` now routes through a worker-safe prepare / serial commit split: prepare converts into private .openkb/staging/prepare output without the KB mutation lock and without touching official raw/, wiki/, or .openkb/ state; the serial owner commits under kb_ingest_lock, resolving the final name and publishing. `--jobs` (Stage 5) is not included. Batch coordinator (openkb/cli.py): - add_directory_serial runs a serial prepare -> commit loop. `add` already holds kb_ingest_lock for its whole body via @_with_kb_lock, so prepare runs under that one outer lock; the reaper (first lock acquisition, before this batch's staging exists) cannot collide with a live prepare. Per-file failure continues the batch; DirtyRollbackError stops it. Prepare (openkb/add_prepare.py, openkb/converter.py): - convert_document_for_prepare: lock-free conversion into private staging under a placeholder doc_name (sanitized stem); returns ConvertResult without registering the hash or resolving the final name. - prepare_document owns the staging-dir lifecycle (rmtree on interrupt). prepare/commit are coordinator-internal, called only by add_directory_serial. Serial commit (openkb/cli.py): - commit_prepared_document requires kb_ingest_lock held (reentrant acquire). - The prepared branch of _add_single_file_locked re-validates under the lock because prepare ran without it: re-decides skip from live registry state, re-hashes the source, and re-converts when the source changed or prepare had short-circuited with no artifacts (the stale-prepare contract). - _retarget_prepared_document_artifacts renames staged raw/source/images from the placeholder name to the owner-resolved final name. Reaper (openkb/locks.py): - _reap_prepare_staging reclaims orphaned prepare staging at first exclusive acquisition; skips symlinks, unlinks stray files, and logs INFO on success / WARNING on failure. No per-prepare marker is needed: directory add holds the lock across its whole batch, so a live batch's staging is never visible to another reaper. Tests: prepare writes only private staging and takes no lock; commit resolves the final name under the owner and requires its lock; the reaper reaps orphans and skips symlinks; stale skip and source-changed (TOCTOU) re-prepare at commit; `add <dir>` end-to-end lands every file via prepare/commit; a prepare failure isolates the file while the batch continues.
1 parent 5a2b48c commit ee22914

8 files changed

Lines changed: 797 additions & 129 deletions

File tree

openkb/add_prepare.py

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
from __future__ import annotations
2+
3+
import shutil
4+
import uuid
5+
from dataclasses import dataclass
6+
from pathlib import Path
7+
8+
from openkb.converter import ConvertResult, _sanitize_stem, convert_document_for_prepare
9+
10+
11+
@dataclass(slots=True)
12+
class PreparedDocument:
13+
input_index: int
14+
source_path: Path
15+
staging_dir: Path
16+
result: ConvertResult
17+
18+
@property
19+
def doc_name_candidate(self) -> str | None:
20+
return self.result.doc_name
21+
22+
23+
def _prepare_staging_dir(kb_dir: Path, input_index: int, source: Path) -> Path:
24+
safe = _sanitize_stem(source.stem)
25+
path = (
26+
kb_dir
27+
/ ".openkb"
28+
/ "staging"
29+
/ "prepare"
30+
/ f"{input_index:06d}-{safe}-{uuid.uuid4().hex[:8]}"
31+
)
32+
path.mkdir(parents=True, exist_ok=False)
33+
return path
34+
35+
36+
def prepare_document(source: Path, kb_dir: Path, *, input_index: int) -> PreparedDocument:
37+
"""Convert ``source`` into private staging without the KB mutation lock.
38+
39+
Coordinator-internal: callers must run this under the serial batch owner's
40+
held ``kb_ingest_lock`` (the ``add`` command acquires it via
41+
``@_with_kb_lock``). The reaper reclaims any staging present at the owner's
42+
first lock acquisition, so once this runs under the held lock the staging
43+
tree is private to this batch.
44+
"""
45+
staging_dir = _prepare_staging_dir(kb_dir, input_index, source)
46+
try:
47+
result = convert_document_for_prepare(source, kb_dir, staging_dir=staging_dir)
48+
return PreparedDocument(
49+
input_index=input_index,
50+
source_path=source,
51+
staging_dir=staging_dir,
52+
result=result,
53+
)
54+
except BaseException:
55+
shutil.rmtree(staging_dir, ignore_errors=True)
56+
raise
57+
58+
59+
def _clear_staging_artifacts(staging_dir: Path) -> None:
60+
"""Drop convertible artifacts (raw/, wiki/) from a prepare staging dir.
61+
62+
Used before re-converting into the same staging dir at commit time so stale
63+
images/raw from the prior prepare don't leak into the re-convert.
64+
"""
65+
for sub in ("raw", "wiki"):
66+
shutil.rmtree(staging_dir / sub, ignore_errors=True)

openkb/cli.py

Lines changed: 164 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -62,17 +62,25 @@ def filter(self, record: logging.LogRecord) -> bool:
6262
)
6363
from openkb.add_coordinator import _cleanup_staging_dirs
6464
from openkb.converter import (
65+
ConvertResult,
6566
_registry_path,
6667
_sanitize_stem,
6768
convert_document,
69+
resolve_doc_name,
6870
resolve_doc_name_from_key,
6971
)
7072
from openkb.indexer import (
7173
_cloud_display_stem,
7274
_write_long_doc_artifacts,
7375
prepare_cloud_import,
7476
)
75-
from openkb.locks import atomic_write_json, atomic_write_text, kb_ingest_lock, kb_read_lock
77+
from openkb.locks import (
78+
atomic_write_json,
79+
atomic_write_text,
80+
kb_ingest_lock,
81+
kb_ingest_lock_held,
82+
kb_read_lock,
83+
)
7684
from openkb.log import append_log
7785
from openkb.mutation import publish_staged_tree
7886
from openkb.schema import AGENTS_MD, INDEX_SEED, PAGE_CONTENT_DIRS
@@ -452,8 +460,113 @@ def add_single_file(
452460
return _add_single_file_locked(file_path, kb_dir, stage=stage)
453461

454462

463+
def commit_prepared_document(prepared, kb_dir: Path) -> Literal["added", "skipped", "failed"]:
464+
"""Commit a prepared document under the serial KB mutation owner.
465+
466+
Coordinator-internal: called only by the serial batch owner
467+
(:func:`add_directory_serial`), which already holds ``kb_ingest_lock`` for
468+
the whole prepare+commit batch. The reentrant acquire below is only safe
469+
because of that. Propagates :class:`DirtyRollbackError` if the mutation's
470+
rollback failed, so the batch owner can stop on dirty state.
471+
"""
472+
openkb_dir = kb_dir / ".openkb"
473+
if not kb_ingest_lock_held(openkb_dir):
474+
raise RuntimeError("commit_prepared_document requires the caller to hold kb_ingest_lock")
475+
with kb_ingest_lock(openkb_dir):
476+
return _add_single_file_locked(
477+
prepared.source_path,
478+
kb_dir,
479+
stage=True,
480+
prepared=prepared,
481+
)
482+
483+
484+
def add_directory_serial(files: list[Path], kb_dir: Path) -> None:
485+
"""Serially prepare and commit each file in input order.
486+
487+
Stage 4 batch coordinator for ``openkb add <dir>``. Must be called with
488+
``kb_ingest_lock`` already held — the ``add`` command's ``@_with_kb_lock``
489+
decorator holds it for the whole batch. Prepares are lock-free and run under
490+
that held lock, so the reaper (which fires once at the owner's first lock
491+
acquisition, before this batch's staging exists) cannot collide with a live
492+
prepare. Stage 5 will add a parallel-prepare variant alongside this one.
493+
494+
A per-file prepare/commit failure is logged and the batch continues; a
495+
:class:`DirtyRollbackError` propagates so the caller can stop the batch on
496+
dirty state.
497+
"""
498+
from openkb.add_prepare import prepare_document
499+
500+
total = len(files)
501+
for i, f in enumerate(files, 1):
502+
click.echo(f"\n[{i}/{total}] ", nl=False)
503+
try:
504+
prepared = prepare_document(f, kb_dir, input_index=i - 1)
505+
except Exception as exc:
506+
click.echo(f" [ERROR] Prepare failed: {exc}")
507+
logger.debug("Prepare traceback:", exc_info=True)
508+
continue
509+
commit_prepared_document(prepared, kb_dir)
510+
511+
512+
def _retarget_prepared_document_artifacts(prepared, doc_name: str) -> None:
513+
result = prepared.result
514+
old_doc_name = prepared.doc_name_candidate or prepared.source_path.stem
515+
if old_doc_name == doc_name:
516+
return
517+
518+
suffix = prepared.source_path.suffix.lower()
519+
raw_dir = prepared.staging_dir / "raw"
520+
old_raw = raw_dir / f"{old_doc_name}{suffix}"
521+
new_raw = raw_dir / f"{doc_name}{suffix}"
522+
if old_raw.exists():
523+
old_raw.rename(new_raw)
524+
result.raw_path = new_raw
525+
# Defensive: if prepare ever writes raw_path off the <doc_name><suffix>
526+
# convention, retarget via the recorded path rather than silently no-op.
527+
elif result.raw_path is not None and result.raw_path.exists():
528+
new_raw = result.raw_path.with_name(f"{doc_name}{result.raw_path.suffix}")
529+
result.raw_path.rename(new_raw)
530+
result.raw_path = new_raw
531+
532+
sources_dir = prepared.staging_dir / "wiki" / "sources"
533+
old_source = sources_dir / f"{old_doc_name}.md"
534+
new_source = sources_dir / f"{doc_name}.md"
535+
if old_source.exists():
536+
text = old_source.read_text(encoding="utf-8")
537+
text = text.replace(f"sources/images/{old_doc_name}/", f"sources/images/{doc_name}/")
538+
old_source.write_text(text, encoding="utf-8")
539+
old_source.rename(new_source)
540+
result.source_path = new_source
541+
# Defensive: if prepare ever writes source_path off the <doc_name>.md
542+
# convention, retarget via the recorded path rather than silently no-op.
543+
elif result.source_path is not None and result.source_path.exists():
544+
new_source = result.source_path.with_name(f"{doc_name}{result.source_path.suffix}")
545+
result.source_path.rename(new_source)
546+
result.source_path = new_source
547+
548+
images_dir = sources_dir / "images"
549+
old_images = images_dir / old_doc_name
550+
new_images = images_dir / doc_name
551+
if old_images.exists() and old_images != new_images:
552+
old_images.rename(new_images)
553+
554+
result.doc_name = doc_name
555+
556+
557+
def _convert_or_fail(src: Path, kb_dir: Path, staging_dir: Path | None) -> ConvertResult | None:
558+
"""Convert ``src``; on failure log, clean staging, and return ``None``."""
559+
try:
560+
return convert_document(src, kb_dir, staging_dir=staging_dir)
561+
except Exception as exc:
562+
click.echo(f" [ERROR] Conversion failed: {exc}")
563+
logger.debug("Conversion traceback:", exc_info=True)
564+
_cleanup_staging_dirs([staging_dir])
565+
return None
566+
567+
455568
def _add_single_file_locked(
456-
file_path: Path, kb_dir: Path, *, stage: bool = True
569+
file_path: Path, kb_dir: Path, *, stage: bool = True, prepared=None
457570
) -> Literal["added", "skipped", "failed"]:
458571
"""Convert, index, and compile a single document into the knowledge base.
459572
@@ -478,18 +591,50 @@ def _add_single_file_locked(
478591
config = load_config(openkb_dir / "config.yaml")
479592
_setup_llm_key(kb_dir)
480593
model: str = config.get("model", DEFAULT_CONFIG["model"])
594+
# One registry instance covers the prepared-branch re-validation and the
595+
# commit-time registration; the held kb_ingest_lock means nothing else
596+
# mutates hashes.json in between, so reloading would just re-parse JSON.
597+
registry = HashRegistry(openkb_dir / "hashes.json")
481598

482-
staging_dir = _staging_dir_for(kb_dir, file_path) if stage else None
599+
staging_dir = (
600+
prepared.staging_dir
601+
if prepared is not None
602+
else (_staging_dir_for(kb_dir, file_path) if stage else None)
603+
)
483604

484605
# 2. Convert document into staging when possible.
485606
click.echo(f"Adding: {file_path.name}")
486-
try:
487-
result = convert_document(file_path, kb_dir, staging_dir=staging_dir)
488-
except Exception as exc:
489-
click.echo(f" [ERROR] Conversion failed: {exc}")
490-
logger.debug("Conversion traceback:", exc_info=True)
491-
_cleanup_staging_dirs([staging_dir])
492-
return "failed"
607+
result: ConvertResult
608+
if prepared is None:
609+
converted = _convert_or_fail(file_path, kb_dir, staging_dir)
610+
if converted is None:
611+
return "failed"
612+
result = converted
613+
else:
614+
result = prepared.result
615+
# Prepare ran without the lock: re-decide skip from live registry state
616+
# (not the cached prepare-time decision) and re-validate the source. A
617+
# hash removed between prepare and commit, or a source edited in that
618+
# window, leaves the staged artifacts missing/stale — re-convert.
619+
result.skipped = bool(result.file_hash and registry.is_known(result.file_hash))
620+
if not result.skipped:
621+
source_intact = (
622+
result.file_hash is not None
623+
and file_path.exists()
624+
and HashRegistry.hash_file(file_path) == result.file_hash
625+
and result.raw_path is not None
626+
)
627+
if source_intact:
628+
doc_name = resolve_doc_name(file_path, kb_dir, registry, persist_legacy=False)
629+
_retarget_prepared_document_artifacts(prepared, doc_name)
630+
else:
631+
from openkb.add_prepare import _clear_staging_artifacts
632+
633+
_clear_staging_artifacts(prepared.staging_dir)
634+
converted = _convert_or_fail(file_path, kb_dir, staging_dir)
635+
if converted is None:
636+
return "failed"
637+
result = converted
493638

494639
if result.skipped:
495640
click.echo(f" [SKIP] Already in knowledge base: {file_path.name}")
@@ -569,7 +714,6 @@ def commit_body(snapshot) -> None:
569714

570715
# Register hash only after successful compilation.
571716
if result.file_hash:
572-
registry = HashRegistry(openkb_dir / "hashes.json")
573717
doc_type = "long_pdf" if result.is_long_doc else file_path.suffix.lstrip(".")
574718
meta = {
575719
"name": file_path.name,
@@ -1056,9 +1200,15 @@ def add(ctx, path, from_pageindex_cloud):
10561200
return
10571201
total = len(files)
10581202
click.echo(f"Found {total} supported file(s) in {path}.")
1059-
for i, f in enumerate(files, 1):
1060-
click.echo(f"\n[{i}/{total}] ", nl=False)
1061-
add_single_file(f, kb_dir)
1203+
from openkb.add_coordinator import DirtyRollbackError
1204+
1205+
try:
1206+
add_directory_serial(files, kb_dir)
1207+
except DirtyRollbackError as exc:
1208+
# A normal "failed"/"skipped" outcome returns and the batch
1209+
# continues; a dirty rollback is fatal and stops the batch.
1210+
click.echo(f" [ERROR] {exc}")
1211+
ctx.exit(1)
10621212
else:
10631213
if target.suffix.lower() not in SUPPORTED_EXTENSIONS:
10641214
click.echo(

0 commit comments

Comments
 (0)