adjust zstd handling to avoid possible core dumps with concurrency

This commit is contained in:
bryan davis 2026-06-23 15:26:33 -05:00
parent 65f1339aa0
commit 2ca12c0096
No known key found for this signature in database
GPG key ID: D11B4A4C0C731E5E
2 changed files with 30 additions and 5 deletions

View file

@ -10,15 +10,18 @@ try: # pragma: no cover
except ImportError:
from zstandard import ZstdCompressor, ZstdDecompressor, get_frame_parameters
_zstd_compressor = ZstdCompressor()
_zstd_decompressor = ZstdDecompressor()
# ZstdCompressor/ZstdDecompressor are not thread-safe: each wraps a single
# reused ZSTD_CCtx/ZSTD_DCtx, and concurrent .compress()/.decompress() calls
# corrupt that context and segfault in the C backend. Ingestion drives this
# path from multiple worker threads (asyncio.to_thread in
# _prepare_document_from_docling), so construct a fresh instance per call
# rather than sharing a module-level singleton.
def _zstd_compress(data: bytes) -> bytes:
return _zstd_compressor.compress(data)
return ZstdCompressor(threads=-1).compress(data)
def _zstd_decompress(data: bytes) -> bytes:
content_size = get_frame_parameters(data).content_size
return _zstd_decompressor.decompress(data, max_output_size=content_size)
return ZstdDecompressor().decompress(data, max_output_size=content_size)
def compress_json(json_str: str) -> bytes:

View file

@ -1,4 +1,5 @@
import json
from concurrent.futures import ThreadPoolExecutor
from haiku.rag.store.compression import (
compress_docling_split,
@ -24,6 +25,27 @@ class TestJsonCompression:
compressed = compress_json('{"test": true}')
assert compressed[:4] == b"\x28\xb5\x2f\xfd"
def test_concurrent_compress_decompress_is_safe(self):
"""Compression must be safe under concurrent threads.
Ingestion offloads compression to worker threads via
asyncio.to_thread; sharing a single zstandard compressor/decompressor
across threads corrupts its internal C context and segfaults the
process. Hammer both paths from many threads to guard against a
regression to module-level singletons.
"""
payloads = [
json.dumps({"i": i, "text": f"document body {i} " * 200}) for i in range(64)
]
def roundtrip(json_str: str) -> str:
return decompress_json(compress_json(json_str))
with ThreadPoolExecutor(max_workers=16) as pool:
results = list(pool.map(roundtrip, payloads * 8))
assert results == (payloads * 8)
class TestDoclingCompressionSplit:
def test_split_with_pages(self):