Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 52 additions & 0 deletions Lib/test/test_zstd.py
Original file line number Diff line number Diff line change
Expand Up @@ -349,6 +349,58 @@ def test_zstd_multithread_compress(self):
with ZstdFile(io.BytesIO(), 'w', options=options) as f:
f.write(b)

@unittest.skipIf(not SUPPORT_MULTITHREADING,
"zstd build doesn't support multi-threaded compression")
def test_free_busy_multithread_compressor_with_dict(self):
# A compressor can be freed while a worker thread is still
# compressing with a dictionary that the compression context owns.
# Whether a failure is detected depends on the memory allocator,
# so use a fresh process for each case.
code = textwrap.dedent("""
import gc, random, sys
from compression.zstd import (CompressionParameter,
ZstdCompressor, ZstdDict)

job_size = 1 << 20
zd = ZstdDict(random.randbytes(1 << 10) * (1 << 10), is_raw=True)
options = {CompressionParameter.compression_level: 10,
CompressionParameter.nb_workers: 1,
CompressionParameter.job_size: job_size}
c = ZstdCompressor(options=options, zstd_dict=zd)
# Returns while the worker thread is compressing the first job.
c.compress(random.randbytes(1 << 10) * ((job_size >> 10) + 64))
if sys.argv[1] == 'error':
# A failed call does not wait for the worker thread either.
import _testcapi
_testcapi.set_nomemory(0)
try:
c.compress(b'x')
except MemoryError:
pass
finally:
_testcapi.remove_mem_hooks()
if sys.argv[1] == 'gc':
# The garbage collector clears the compressor before
# deallocating it.
cycle = [c]
cycle.append(cycle)
del c, cycle
gc.collect()
elif sys.argv[1] != 'exit':
del c
# Otherwise the compressor is freed at interpreter exit.
""")
cases = ['del', 'gc', 'exit']
try:
import _testcapi # noqa: F401
except ImportError:
pass
else:
cases.append('error')
for case in cases:
with self.subTest(case=case):
assert_python_ok('-c', code, case)

def test_compress_flushblock(self):
point = len(THIS_FILE_BYTES) // 2

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
Fix a crash when a :class:`compression.zstd.ZstdCompressor` that uses worker
threads and a :class:`~compression.zstd.ZstdDict` is deallocated, or the
interpreter exits, while data passed to
:meth:`~compression.zstd.ZstdCompressor.compress` is still being compressed.
21 changes: 21 additions & 0 deletions Modules/_zstd/compressor.c
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,10 @@ typedef struct {
/* (nbWorker >= 1) ? 1 : 0 */
int use_multithread;

/* Nonzero if worker threads may still be compressing submitted data:
set by any call that did not complete a frame, including a failed one */
int mt_maybe_busy;

/* Compression level */
int compression_level;

Expand Down Expand Up @@ -347,6 +351,7 @@ _zstd_ZstdCompressor_new_impl(PyTypeObject *type, PyObject *level,
}

self->use_multithread = 0;
self->mt_maybe_busy = 0;
self->compression_level = ZSTD_CLEVEL_DEFAULT;
self->dict = NULL;
self->lock = (PyMutex){0};
Expand Down Expand Up @@ -427,6 +432,18 @@ ZstdCompressor_dealloc(PyObject *ob)

/* Free compression context */
if (self->cctx) {
if (self->mt_maybe_busy) {
/* ZSTD_freeCCtx() (as of libzstd 1.5.7) frees the context's own
copy of a dictionary loaded with ZSTD_CCtx_loadDictionary()
before it stops the worker threads, which may still be
reading it. Starting a new frame makes libzstd wait for the
jobs of the previous one, and with no input it starts none. */
ZSTD_inBuffer in = {NULL, 0, 0};
ZSTD_outBuffer out = {NULL, 0, 0};
(void)ZSTD_CCtx_reset(self->cctx, ZSTD_reset_session_only);
(void)ZSTD_compressStream2(self->cctx, &out, &in,
ZSTD_e_continue);
}
ZSTD_freeCCtx(self->cctx);
}

Expand Down Expand Up @@ -637,6 +654,8 @@ _zstd_ZstdCompressor_compress_impl(ZstdCompressor *self, Py_buffer *data,
/* Resetting cctx's session never fail */
ZSTD_CCtx_reset(self->cctx, ZSTD_reset_session_only);
}
self->mt_maybe_busy = (self->use_multithread
&& (ret == NULL || mode != ZSTD_e_end));
PyMutex_Unlock(&self->lock);

return ret;
Expand Down Expand Up @@ -685,6 +704,8 @@ _zstd_ZstdCompressor_flush_impl(ZstdCompressor *self, int mode)
/* Resetting cctx's session never fail */
ZSTD_CCtx_reset(self->cctx, ZSTD_reset_session_only);
}
self->mt_maybe_busy = (self->use_multithread
&& (ret == NULL || mode != ZSTD_e_end));
PyMutex_Unlock(&self->lock);

return ret;
Expand Down
Loading