From b577e7bb710b49e3977fb9ddd92f80d8d6c7f138 Mon Sep 17 00:00:00 2001 From: dajiaohuang Date: Thu, 27 Aug 2026 20:48:09 +0800 Subject: [PATCH] perf: coalesce concurrent document normalization --- src/sagasmith_core/documents.py | 169 +++++++++++++++++++++++--------- tests/test_state_documents.py | 86 ++++++++++++++++ 2 files changed, 207 insertions(+), 48 deletions(-) diff --git a/src/sagasmith_core/documents.py b/src/sagasmith_core/documents.py index 92f65b8..234ea2a 100644 --- a/src/sagasmith_core/documents.py +++ b/src/sagasmith_core/documents.py @@ -19,6 +19,7 @@ from threading import RLock from typing import Any, Protocol from uuid import uuid4 +from weakref import WeakValueDictionary DOCUMENT_NORMALIZER_VERSION = "37" _MAX_STRUCTURAL_HEADING_CHARS = 200 @@ -30,6 +31,8 @@ _BOOKMARK_OCR_MAX_NON_WHITESPACE = 800 _RAPIDOCR_ENGINES: dict[str, tuple[Any, RLock]] = {} _RAPIDOCR_ENGINES_LOCK = RLock() +_NORMALIZATION_CACHE_LOCKS: WeakValueDictionary[Path, RLock] = WeakValueDictionary() +_NORMALIZATION_CACHE_LOCKS_GUARD = RLock() class DocumentQualityError(RuntimeError): @@ -2834,57 +2837,75 @@ def _cache_path(cache_dir: Path, checksum: str, profile: str) -> Path: return cache_dir / checksum[:2] / f"{checksum}-{profile_hash}.json" -def normalize_document( - path: str | Path, +def _normalization_cache_lock(target: Path) -> RLock: + """Return the process-local in-flight lock for one immutable cache key.""" + with _NORMALIZATION_CACHE_LOCKS_GUARD: + lock = _NORMALIZATION_CACHE_LOCKS.get(target) + if lock is None: + lock = RLock() + _NORMALIZATION_CACHE_LOCKS[target] = lock + return lock + + +def _read_normalized_document_cache( + target: Path, *, - ocr_provider: OcrProvider | None = None, - cache_dir: str | Path | None = None, - expected_checksum: str | None = None, - layout_profile: DocumentLayoutProfile = GENERIC_DOCUMENT_LAYOUT_PROFILE, + source: Path, + checksum: str, + profile: str, +) -> NormalizedDocument | None: + if not target.is_file(): + return None + try: + value = json.loads(target.read_text(encoding="utf-8")) + content = str(value["content"]) + if ( + value.get("schema") == _DOCUMENT_CACHE_SCHEMA + and value.get("checksum") == checksum + and value.get("profile") == profile + and value.get("content_checksum") + == hashlib.sha256(content.encode("utf-8")).hexdigest() + ): + return NormalizedDocument( + content=content, + media_type=str(value["media_type"]), + source_path=str(source), + checksum=checksum, + page_count=int(value.get("page_count", 1)), + bookmarks=tuple( + DocumentBookmark(str(item["title"]), int(item["page"]), int(item["depth"])) + for item in value.get("bookmarks", []) + ), + warnings=tuple(str(item) for item in value.get("warnings", [])), + metadata={ + **dict(value.get("metadata") or {}), + "normalization_cache_hit": True, + }, + ) + except (KeyError, TypeError, ValueError, json.JSONDecodeError): + pass + return None + + +def _normalize_document_with_cache( + source: Path, + *, + checksum: str, + profile: str, + target: Path | None, + ocr_provider: OcrProvider | None, + cache_dir: str | Path | None, + layout_profile: DocumentLayoutProfile, ) -> NormalizedDocument: - """Convert a document once and reuse a content-addressed normalized form.""" - source = Path(path).expanduser().resolve() - checksum = file_sha256(source) - if expected_checksum and checksum != expected_checksum: - raise DocumentQualityError( - "source_checksum_mismatch", - "managed document checksum no longer matches its staged import job", + if target is not None: + cached = _read_normalized_document_cache( + target, + source=source, + checksum=checksum, + profile=profile, ) - profile = _cache_profile(source, ocr_provider, layout_profile) - target = ( - _cache_path(Path(cache_dir).expanduser().resolve(), checksum, profile) - if cache_dir is not None - else None - ) - if target is not None and target.is_file(): - try: - value = json.loads(target.read_text(encoding="utf-8")) - content = str(value["content"]) - if ( - value.get("schema") == _DOCUMENT_CACHE_SCHEMA - and value.get("checksum") == checksum - and value.get("profile") == profile - and value.get("content_checksum") - == hashlib.sha256(content.encode("utf-8")).hexdigest() - ): - return NormalizedDocument( - content=content, - media_type=str(value["media_type"]), - source_path=str(source), - checksum=checksum, - page_count=int(value.get("page_count", 1)), - bookmarks=tuple( - DocumentBookmark(str(item["title"]), int(item["page"]), int(item["depth"])) - for item in value.get("bookmarks", []) - ), - warnings=tuple(str(item) for item in value.get("warnings", [])), - metadata={ - **dict(value.get("metadata") or {}), - "normalization_cache_hit": True, - }, - ) - except (KeyError, TypeError, ValueError, json.JSONDecodeError): - pass + if cached is not None: + return cached document = converter_for( source, @@ -2925,6 +2946,58 @@ def normalize_document( return document +def normalize_document( + path: str | Path, + *, + ocr_provider: OcrProvider | None = None, + cache_dir: str | Path | None = None, + expected_checksum: str | None = None, + layout_profile: DocumentLayoutProfile = GENERIC_DOCUMENT_LAYOUT_PROFILE, +) -> NormalizedDocument: + """Convert a document once and reuse a content-addressed normalized form.""" + source = Path(path).expanduser().resolve() + checksum = file_sha256(source) + if expected_checksum and checksum != expected_checksum: + raise DocumentQualityError( + "source_checksum_mismatch", + "managed document checksum no longer matches its staged import job", + ) + profile = _cache_profile(source, ocr_provider, layout_profile) + target = ( + _cache_path(Path(cache_dir).expanduser().resolve(), checksum, profile) + if cache_dir is not None + else None + ) + if target is None: + return _normalize_document_with_cache( + source, + checksum=checksum, + profile=profile, + target=None, + ocr_provider=ocr_provider, + cache_dir=cache_dir, + layout_profile=layout_profile, + ) + cached = _read_normalized_document_cache( + target, + source=source, + checksum=checksum, + profile=profile, + ) + if cached is not None: + return cached + with _normalization_cache_lock(target): + return _normalize_document_with_cache( + source, + checksum=checksum, + profile=profile, + target=target, + ocr_provider=ocr_provider, + cache_dir=cache_dir, + layout_profile=layout_profile, + ) + + def render_pdf_page( path: str | Path, page_number: int, diff --git a/tests/test_state_documents.py b/tests/test_state_documents.py index fad1c66..3989c05 100644 --- a/tests/test_state_documents.py +++ b/tests/test_state_documents.py @@ -3,7 +3,9 @@ import hashlib import json import zlib +from concurrent.futures import ThreadPoolExecutor from pathlib import Path +from threading import Barrier, Event, Lock from types import SimpleNamespace import pytest @@ -214,6 +216,90 @@ def fake_engine(model_name: str): documents._RAPIDOCR_ENGINES.clear() +def test_normalize_document_coalesces_concurrent_cache_misses(monkeypatch, tmp_path) -> None: + import sagasmith_core.documents as documents + + source = tmp_path / "shared.md" + source.write_text("# Shared source\n", encoding="utf-8") + cache_dir = tmp_path / "cache" + first_started = Event() + release_first = Event() + duplicate_started = Event() + calls_lock = Lock() + calls = 0 + + class BlockingConverter: + def convert(self, path, *, source_checksum=None): + nonlocal calls + with calls_lock: + calls += 1 + call_number = calls + if call_number == 1: + first_started.set() + assert release_first.wait(timeout=5) + else: + duplicate_started.set() + return NormalizedDocument( + content="# Normalized once\n", + media_type="text/markdown", + source_path=str(path), + checksum=str(source_checksum), + ) + + monkeypatch.setattr(documents, "converter_for", lambda *args, **kwargs: BlockingConverter()) + + with ThreadPoolExecutor(max_workers=2) as pool: + first = pool.submit(normalize_document, source, cache_dir=cache_dir) + assert first_started.wait(timeout=5) + second = pool.submit(normalize_document, source, cache_dir=cache_dir) + try: + assert not duplicate_started.wait(timeout=0.25) + finally: + release_first.set() + results = [first.result(timeout=5), second.result(timeout=5)] + + assert calls == 1 + assert [result.content for result in results] == [ + "# Normalized once\n", + "# Normalized once\n", + ] + assert sorted(result.metadata["normalization_cache_hit"] for result in results) == [ + False, + True, + ] + + +def test_normalize_document_keeps_distinct_cache_keys_parallel(monkeypatch, tmp_path) -> None: + import sagasmith_core.documents as documents + + sources = [tmp_path / "one.md", tmp_path / "two.md"] + for source in sources: + source.write_text(f"# {source.stem}\n", encoding="utf-8") + conversion_barrier = Barrier(len(sources)) + + class ParallelConverter: + def convert(self, path, *, source_checksum=None): + conversion_barrier.wait(timeout=5) + return NormalizedDocument( + content=f"# {Path(path).stem}\n", + media_type="text/markdown", + source_path=str(path), + checksum=str(source_checksum), + ) + + monkeypatch.setattr(documents, "converter_for", lambda *args, **kwargs: ParallelConverter()) + + with ThreadPoolExecutor(max_workers=len(sources)) as pool: + results = list( + pool.map( + lambda source: normalize_document(source, cache_dir=tmp_path / "cache"), + sources, + ) + ) + + assert [result.content for result in results] == ["# one\n", "# two\n"] + + def test_rapidocr_page_cache_survives_provider_restart_and_rejects_tampering( tmp_path, monkeypatch,