From 70c9c2778fc2626b41ade930609c61b848162df5 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sun, 14 Jun 2026 10:50:42 +0300 Subject: [PATCH] Serialize all in-process pdfium access under a shared lock --- CHANGELOG.md | 1 + haiku_rag_slim/haiku/rag/client/documents.py | 93 ++++++++++++------- .../haiku/rag/converters/pdf_split.py | 23 ++--- tests/test_pdf_concurrency.py | 62 +++++++++++++ 4 files changed, 133 insertions(+), 46 deletions(-) create mode 100644 tests/test_pdf_concurrency.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 549bee08..296dd7a2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ ### Fixed +- All in-process pdfium access (page slicing and embedded-attachment scanning) is serialized under a single shared lock. Concurrent ingester workers no longer corrupt libpdfium's global state, which previously failed valid PDFs with "Data format error". - Embedded PDF attachment extension is derived from the attachment filename, not the parent's synthetic `...#attachment=` URI; non-PDF attachments (e.g. `.joboptions`) are no longer misrouted to docling's PDF backend, and unsupported extensions are skipped. ## [0.57.0] - 2026-06-11 diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py index 67def6d9..fd868546 100644 --- a/haiku_rag_slim/haiku/rag/client/documents.py +++ b/haiku_rag_slim/haiku/rag/client/documents.py @@ -438,6 +438,62 @@ async def _ingest_fetch_result( return created +def _extract_pdf_attachments( + parent_body: bytes, parent_uri: str, *, depth: int +) -> dict[str, tuple[str, bytes, str, str]] | None: + """Open the parent PDF and return its embedded attachments keyed by child + URI. Returns ``None`` when the PDF can't be opened or the recursion depth + cap is reached — in both cases the caller skips reconciliation entirely. + + Every pdfium call is held under ``PDFIUM_LOCK`` (shared with page slicing) + because libpdfium's global C state is not thread-safe; concurrent access + from another worker corrupts it and then fails valid PDFs with "Data format + error" until the process restarts. + """ + import pypdfium2 as pdfium + + from haiku.rag.converters.pdf_split import PDFIUM_LOCK + + with PDFIUM_LOCK: + try: + pdf = pdfium.PdfDocument(parent_body) + except pdfium.PdfiumError as exc: + logger.warning( + "Cannot scan %s for embedded attachments: %s", parent_uri, exc + ) + return None + try: + attachment_count = pdf.count_attachments() + + if depth + 1 >= MAX_ATTACHMENT_DEPTH: + if attachment_count > 0: + logger.warning( + "Attachment depth cap (%d) reached at %s; skipping %d nested " + "attachment(s).", + MAX_ATTACHMENT_DEPTH, + parent_uri, + attachment_count, + ) + return None + + new_attachments: dict[str, tuple[str, bytes, str, str]] = {} + for i in range(attachment_count): + att = pdf.get_attachment(i) + name = att.get_name() + if not name: + continue + data = bytes(att.get_data()) + child_uri = f"{parent_uri}#attachment={quote(name, safe='')}" + content_type = ( + mimetypes.guess_type(name)[0] or "application/octet-stream" + ) + content_hash = hashlib.md5(data, usedforsecurity=False).hexdigest() + new_attachments[child_uri] = (name, data, content_type, content_hash) + return new_attachments + finally: + pdf.close() + + async def _reconcile_pdf_attachments( client: "HaikuRAG", parent_doc: Document, @@ -460,42 +516,9 @@ async def _reconcile_pdf_attachments( if (parent_doc.metadata or {}).get("content_type") != "application/pdf": return - import pypdfium2 as pdfium - - try: - pdf = pdfium.PdfDocument(parent_body) - except pdfium.PdfiumError as exc: - logger.warning( - "Cannot scan %s for embedded attachments: %s", parent_doc.uri, exc - ) + new_attachments = _extract_pdf_attachments(parent_body, parent_doc.uri, depth=depth) + if new_attachments is None: return - try: - attachment_count = pdf.count_attachments() - - if depth + 1 >= MAX_ATTACHMENT_DEPTH: - if attachment_count > 0: - logger.warning( - "Attachment depth cap (%d) reached at %s; skipping %d nested " - "attachment(s).", - MAX_ATTACHMENT_DEPTH, - parent_doc.uri, - attachment_count, - ) - return - - new_attachments: dict[str, tuple[str, bytes, str, str]] = {} - for i in range(attachment_count): - att = pdf.get_attachment(i) - name = att.get_name() - if not name: - continue - data = bytes(att.get_data()) - child_uri = f"{parent_doc.uri}#attachment={quote(name, safe='')}" - content_type = mimetypes.guess_type(name)[0] or "application/octet-stream" - content_hash = hashlib.md5(data, usedforsecurity=False).hexdigest() - new_attachments[child_uri] = (name, data, content_type, content_hash) - finally: - pdf.close() existing = await client.list_documents(filter=parent_uri_filter(parent_doc.uri)) existing_by_uri: dict[str, Document] = {d.uri: d for d in existing if d.uri} diff --git a/haiku_rag_slim/haiku/rag/converters/pdf_split.py b/haiku_rag_slim/haiku/rag/converters/pdf_split.py index f5a6a7f4..4057b60e 100644 --- a/haiku_rag_slim/haiku/rag/converters/pdf_split.py +++ b/haiku_rag_slim/haiku/rag/converters/pdf_split.py @@ -12,14 +12,15 @@ from haiku.rag.client.exceptions import UnsupportedSourceError from haiku.rag.telemetry import logfire # pypdfium2 wraps libpdfium, which has global C state and is not thread-safe. -# Multiple workers calling iter_pdf_slices concurrently race on that state -# and corrupt it — first error surfaces as e.g. "Failed to import pages", -# and after that every subsequent PDF load fails with "Failed to load -# document" until the process restarts. Hold this lock around every pdfium -# call so only one worker's slicing runs at a time. Crucially we release it -# between slices so other workers' slicing can interleave; the heavy work -# (docling convert) happens between yields with the lock free. -_PDFIUM_LOCK = threading.Lock() +# Two workers calling into pdfium concurrently race on that state and corrupt +# it — the first error surfaces as e.g. "Failed to import pages", and after +# that every subsequent PDF load fails with "Data format error" until the +# process restarts. This is the single process-wide lock around *all* in-process +# pdfium access (page slicing here and embedded-attachment scanning in +# client.documents); every pdfium call must hold it so only one runs at a time. +# Slicing releases it between slices so other callers can interleave; the heavy +# work (docling convert) happens between yields with the lock free. +PDFIUM_LOCK = threading.Lock() if TYPE_CHECKING: from docling_core.types.doc.document import DoclingDocument @@ -44,7 +45,7 @@ def iter_pdf_slices( """ if slice_size <= 0: raise ValueError(f"slice_size must be >= 1, got {slice_size}") - with _PDFIUM_LOCK: + with PDFIUM_LOCK: try: src = pdfium.PdfDocument(str(path)) except pdfium.PdfiumError as exc: @@ -55,7 +56,7 @@ def iter_pdf_slices( try: for start in range(0, total, slice_size): end = min(start + slice_size, total) - with _PDFIUM_LOCK: + with PDFIUM_LOCK: dst = pdfium.PdfDocument.new() try: dst.import_pages(src, list(range(start, end))) @@ -66,7 +67,7 @@ def iter_pdf_slices( dst.close() yield (start + 1, end, slice_bytes) finally: - with _PDFIUM_LOCK: + with PDFIUM_LOCK: src.close() diff --git a/tests/test_pdf_concurrency.py b/tests/test_pdf_concurrency.py new file mode 100644 index 00000000..b9a24e32 --- /dev/null +++ b/tests/test_pdf_concurrency.py @@ -0,0 +1,62 @@ +import io +import threading + +import pypdfium2 as pdfium + +from haiku.rag.client.documents import _extract_pdf_attachments +from haiku.rag.converters.pdf_split import iter_pdf_slices + + +def _make_pdf(pages: int, attachment: tuple[str, bytes] | None) -> bytes: + pdf = pdfium.PdfDocument.new() + for _ in range(pages): + pdf.new_page(200, 200) + if attachment is not None: + name, data = attachment + pdf.new_attachment(name).set_data(data) + buf = io.BytesIO() + pdf.save(buf) + return buf.getvalue() + + +def test_concurrent_pdfium_access_does_not_corrupt_global_state(tmp_path): + """libpdfium has global, non-thread-safe C state. The attachment scan and + the page slicer both call into it; without a single shared lock across both + sites, concurrent workers corrupt that state and then fail otherwise-valid + PDFs with "Data format error". Both paths operate on valid PDFs here, so any + failure means the global state was corrupted by a concurrent caller.""" + body = _make_pdf(pages=6, attachment=("notes.txt", b"payload")) + path = tmp_path / "doc.pdf" + path.write_bytes(body) + + scan_results: list[dict | None] = [] + slice_errors: list[str] = [] + lock = threading.Lock() + + def scan() -> None: + for _ in range(40): + r = _extract_pdf_attachments(body, "file:///doc.pdf", depth=0) + with lock: + scan_results.append(r) + + def slice_() -> None: + for _ in range(40): + try: + slices = list(iter_pdf_slices(path, 2)) + assert len(slices) == 3 + except Exception as exc: # noqa: BLE001 + with lock: + slice_errors.append(repr(exc)) + + threads = [threading.Thread(target=scan) for _ in range(4)] + threads += [threading.Thread(target=slice_) for _ in range(4)] + for t in threads: + t.start() + for t in threads: + t.join() + + failed_scans = sum(r is None or len(r) != 1 for r in scan_results) + assert failed_scans == 0, ( + f"{failed_scans}/{len(scan_results)} attachment scans failed" + ) + assert slice_errors == [], slice_errors