[3.14] gh-160549: Fix crash when freeing a busy multi-threaded ZstdCompressor (GH-160550) (#160638)

gh-160549: Fix crash when freeing a busy multi-threaded ZstdCompressor (GH-160550)

ZSTD_freeCCtx() frees a dictionary loaded with ZSTD_CCtx_loadDictionary()
before it stops the worker threads.  Make libzstd wait for unfinished
jobs before freeing the compression context of a ZstdCompressor that
uses worker threads and a ZstdDict.

Works around facebook/zstd#4862
(cherry picked from commit 23a5ba1)

Co-authored-by: Gregory P. Smith <68491+gpshead@users.noreply.github.com>

Author: miss-islington

Lib/test/test_zstd.py

@@ -350,6 +350,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
 

Misc/NEWS.d/next/Library/2026-10-23-23-21-21.gh-issue-160549.q7ZdMt.rst

@@ -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.

Modules/_zstd/compressor.c

@@ -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;
 
@@ -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};
@@ -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);
     }
 
@@ -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;
@@ -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;