adjust zstd handling to avoid possible core dumps with concurrency
This commit is contained in:
parent
65f1339aa0
commit
2ca12c0096
2 changed files with 30 additions and 5 deletions
|
|
@ -10,15 +10,18 @@ try: # pragma: no cover
|
||||||
except ImportError:
|
except ImportError:
|
||||||
from zstandard import ZstdCompressor, ZstdDecompressor, get_frame_parameters
|
from zstandard import ZstdCompressor, ZstdDecompressor, get_frame_parameters
|
||||||
|
|
||||||
_zstd_compressor = ZstdCompressor()
|
# ZstdCompressor/ZstdDecompressor are not thread-safe: each wraps a single
|
||||||
_zstd_decompressor = ZstdDecompressor()
|
# 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:
|
def _zstd_compress(data: bytes) -> bytes:
|
||||||
return _zstd_compressor.compress(data)
|
return ZstdCompressor(threads=-1).compress(data)
|
||||||
|
|
||||||
def _zstd_decompress(data: bytes) -> bytes:
|
def _zstd_decompress(data: bytes) -> bytes:
|
||||||
content_size = get_frame_parameters(data).content_size
|
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:
|
def compress_json(json_str: str) -> bytes:
|
||||||
|
|
|
||||||
|
|
@ -1,4 +1,5 @@
|
||||||
import json
|
import json
|
||||||
|
from concurrent.futures import ThreadPoolExecutor
|
||||||
|
|
||||||
from haiku.rag.store.compression import (
|
from haiku.rag.store.compression import (
|
||||||
compress_docling_split,
|
compress_docling_split,
|
||||||
|
|
@ -24,6 +25,27 @@ class TestJsonCompression:
|
||||||
compressed = compress_json('{"test": true}')
|
compressed = compress_json('{"test": true}')
|
||||||
assert compressed[:4] == b"\x28\xb5\x2f\xfd"
|
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:
|
class TestDoclingCompressionSplit:
|
||||||
def test_split_with_pages(self):
|
def test_split_with_pages(self):
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue