From b4b8250b260b7cc7e663e9aba66c54b7f6231a35 Mon Sep 17 00:00:00 2001 From: "Jens W. Klein" Date: Mon, 29 Jun 2026 23:26:18 +0200 Subject: [PATCH] fix(tika-worker): stream S3 blobs to Tika instead of buffering (#189) TikaWorker materialized each blob fully in memory before sending it to Tika (peak ~2-3x blob size). A single large S3-tiered PDF/image OOM-killed a small worker pod (256Mi) into CrashLoopBackOff. _extract now streams the S3 path: get_object()['Body'].iter_chunks() is handed straight to httpx as the request content, with Content-Length from the object metadata. The worker's memory stays ~flat regardless of blob size. The PG-bytea path (small sub-threshold inline blobs) still sends bytes. Blob-location logic is split into _blob_source() so _extract streams while _fetch_blob keeps materializing for its callers. Closes #189 Co-Authored-By: Claude Opus 4.8 (1M context) --- CHANGES.md | 11 ++++ src/plone/pgcatalog/tika_worker.py | 51 +++++++++++++--- tests/test_tika_worker_extras.py | 93 +++++++++++++++++++++++++++++- 3 files changed, 146 insertions(+), 9 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index d64faba..c24c7fb 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -1,5 +1,16 @@ # Changelog +## 1.0.0b68 (unreleased) + +### Fixed + +- The Tika worker now streams S3-tiered blobs straight from S3 into the Tika + request body in chunks, instead of buffering the whole blob in memory (peak + ~2–3× blob size). A single large PDF/image could OOM-kill a small worker pod + (256Mi) into `CrashLoopBackOff`; streaming keeps the worker's memory roughly + constant regardless of blob size. The PG-bytea path (small inline blobs below + the S3 tiering threshold) is unchanged. #189 + ## 1.0.0b67 (2026-06-29) ### Changed diff --git a/src/plone/pgcatalog/tika_worker.py b/src/plone/pgcatalog/tika_worker.py index bba8293..59434fa 100644 --- a/src/plone/pgcatalog/tika_worker.py +++ b/src/plone/pgcatalog/tika_worker.py @@ -65,6 +65,9 @@ " pip install plone-pgcatalog[tika-s3] # + S3-tiered blobs" ) +# Chunk size for streaming S3 blobs into the Tika request body (#189). +_S3_STREAM_CHUNK = 1024 * 1024 # 1 MiB + class TikaWorker: """PostgreSQL-backed text extraction worker using Apache Tika.""" @@ -190,24 +193,45 @@ def _process_one(self): return True def _extract(self, conn, zoid, tid, content_type): - """Fetch blob and send to Tika, return extracted text.""" - blob_data = self._fetch_blob(conn, zoid, tid) + """Fetch blob and send to Tika, return extracted text. + + S3-tiered blobs are streamed straight from S3 into the Tika request + body in chunks, so the worker never holds the whole blob in memory + (#189). PG-bytea blobs are sub-threshold (small) and sent as bytes. + """ if httpx is None: raise RuntimeError(MISSING_EXTRA_HINT) headers = {"Accept": "text/plain"} if content_type: headers["Content-Type"] = content_type + + source = self._blob_source(conn, zoid, tid) with httpx.Client(timeout=self.http_timeout) as client: + if source["kind"] == "s3": + # Stream S3 -> Tika without buffering the whole blob. An + # explicit Content-Length (from get_object) lets us avoid + # chunked transfer-encoding. + obj = self._get_s3_client().get_object( + Bucket=self.s3_config["bucket_name"], Key=source["s3_key"] + ) + headers["Content-Length"] = str(obj["ContentLength"]) + content = obj["Body"].iter_chunks(chunk_size=_S3_STREAM_CHUNK) + else: + content = source["data"] resp = client.put( f"{self.tika_url}/tika", - content=blob_data, + content=content, headers=headers, ) resp.raise_for_status() return resp.text - def _fetch_blob(self, conn, zoid, tid): - """Fetch blob bytes from PG bytea or S3.""" + def _blob_source(self, conn, zoid, tid): + """Locate a blob without materializing S3 data. + + Returns ``{"kind": "bytes", "data": ...}`` for inline PG-bytea blobs + or ``{"kind": "s3", "s3_key": ...}`` for S3-tiered blobs. + """ with conn.cursor(row_factory=dict_row) as cur: cur.execute( "SELECT data, s3_key FROM blob_state " @@ -218,13 +242,24 @@ def _fetch_blob(self, conn, zoid, tid): if row is None: raise ValueError(f"No blob for zoid={zoid} tid={tid}") if row["data"] is not None: - return bytes(row["data"]) + return {"kind": "bytes", "data": bytes(row["data"])} if row["s3_key"]: - return self._fetch_from_s3(row["s3_key"]) + return {"kind": "s3", "s3_key": row["s3_key"]} raise ValueError(f"Blob row has neither data nor s3_key for zoid={zoid}") + def _fetch_blob(self, conn, zoid, tid): + """Fetch blob bytes from PG bytea or S3 (materializes the blob). + + Used by callers that need the bytes in hand; ``_extract`` streams S3 + instead (see #189). + """ + source = self._blob_source(conn, zoid, tid) + if source["kind"] == "bytes": + return source["data"] + return self._fetch_from_s3(source["s3_key"]) + def _fetch_from_s3(self, s3_key): - """Download blob from S3.""" + """Download a full blob from S3 into memory.""" import io client = self._get_s3_client() diff --git a/tests/test_tika_worker_extras.py b/tests/test_tika_worker_extras.py index 31034d2..7ae1422 100644 --- a/tests/test_tika_worker_extras.py +++ b/tests/test_tika_worker_extras.py @@ -7,6 +7,8 @@ it must give a clear hint pointing at the extra. See issue #171. """ +from unittest import mock + import subprocess import sys import textwrap @@ -185,7 +187,9 @@ def put(self, *a, **k): monkeypatch.setattr(tika_worker, "httpx", types.SimpleNamespace(Client=_Client)) worker = tika_worker.TikaWorker(dsn="x", tika_url="http://t", http_timeout=7.5) - monkeypatch.setattr(worker, "_fetch_blob", lambda *a, **k: b"blob") + monkeypatch.setattr( + worker, "_blob_source", lambda *a, **k: {"kind": "bytes", "data": b"blob"} + ) out = worker._extract(conn=None, zoid=1, tid=1, content_type="application/pdf") assert out == "extracted text" assert captured["timeout"] == 7.5 @@ -219,3 +223,90 @@ def run(self): tika_worker.main() assert captured["http_timeout"] == 300.0 + + +# --------------------------------------------------------------------------- +# #189: stream S3 -> Tika instead of buffering the whole blob in memory. +# --------------------------------------------------------------------------- + + +def _capture_httpx(monkeypatch): + """Patch tika_worker.httpx with a client that records the put() call.""" + from plone.pgcatalog import tika_worker + + import types + + captured = {} + + class _Resp: + text = "ok" + + def raise_for_status(self): + pass + + class _Client: + def __init__(self, *a, **k): + pass + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def put(self, url, content=None, headers=None, **k): + captured["url"] = url + captured["content"] = content + captured["headers"] = headers + return _Resp() + + monkeypatch.setattr(tika_worker, "httpx", types.SimpleNamespace(Client=_Client)) + return captured + + +def test_s3_blob_is_streamed_not_buffered(monkeypatch): + from plone.pgcatalog import tika_worker + + captured = _capture_httpx(monkeypatch) + worker = tika_worker.TikaWorker( + dsn="x", tika_url="http://t", s3_config={"bucket_name": "b"} + ) + monkeypatch.setattr( + worker, "_blob_source", lambda *a: {"kind": "s3", "s3_key": "k"} + ) + + chunks = iter([b"chunk1", b"chunk2"]) + + class _Body: + def iter_chunks(self, chunk_size): + _Body.chunk_size = chunk_size + return chunks + + fake_s3 = mock.Mock() + fake_s3.get_object.return_value = {"Body": _Body(), "ContentLength": 12} + monkeypatch.setattr(worker, "_get_s3_client", lambda: fake_s3) + + out = worker._extract(conn=None, zoid=1, tid=1, content_type="application/pdf") + assert out == "ok" + # streamed: the iterator itself is handed to httpx, not a bytes buffer + assert captured["content"] is chunks + assert not isinstance(captured["content"], (bytes, bytearray)) + # explicit Content-Length avoids chunked transfer-encoding + assert captured["headers"]["Content-Length"] == "12" + # never buffered the whole blob + fake_s3.download_fileobj.assert_not_called() + fake_s3.get_object.assert_called_once_with(Bucket="b", Key="k") + + +def test_pg_bytea_blob_sent_as_bytes(monkeypatch): + from plone.pgcatalog import tika_worker + + captured = _capture_httpx(monkeypatch) + worker = tika_worker.TikaWorker(dsn="x", tika_url="http://t") + monkeypatch.setattr( + worker, "_blob_source", lambda *a: {"kind": "bytes", "data": b"PDFBYTES"} + ) + out = worker._extract(conn=None, zoid=1, tid=1, content_type="application/pdf") + assert out == "ok" + assert captured["content"] == b"PDFBYTES" + assert "Content-Length" not in (captured["headers"] or {})