perf: coalesce concurrent document normalization

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,

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,