gh-160535: Keep the ZstdDict alive until the zstd context is freed

Lib/test/test_zstd.py

@@ -7,10 +7,12 @@
 import os
 import unittest
 import tempfile
+import textwrap
 import threading
 
 from test import support
 from test.support.import_helper import import_module
+from test.support.script_helper import assert_python_ok
 from test.support import threading_helper
 from test.support import _1M
 
@@ -1571,6 +1573,42 @@ def test_len(self):
         self.assertEqual(len(TRAINED_DICT), len(TRAINED_DICT.dict_content))
         self.assertIn(str(len(TRAINED_DICT)), str(TRAINED_DICT))
 
+    @unittest.skipUnless(SUPPORT_MULTITHREADING,
+                         "zstd build doesn't support multi-threaded compression")
+    @support.requires_subprocess()
+    def test_collect_cycle_with_busy_compressor(self):
+        # The dictionary must outlive the compression context when both are
+        # collected by the garbage collector while a worker thread is still
+        # using the dictionary.  Whether the 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)
+
+            class Cycle:
+                pass
+
+            job_size = 1 << 20
+            content = random.randbytes(1 << 10) * (1 << 8)
+            zd = ZstdDict(content, is_raw=True)
+            options = {CompressionParameter.compression_level: 10,
+                       CompressionParameter.nb_workers: 1,
+                       CompressionParameter.job_size: job_size}
+            c = ZstdCompressor(options=options,
+                               zstd_dict=getattr(zd, sys.argv[1]))
+            cycle = Cycle()
+            cycle.cycle = cycle
+            cycle.compressor = c
+            # Returns while the worker thread is compressing the first job.
+            c.compress(random.randbytes(1 << 10) * ((job_size >> 10) + 64))
+            del zd, c, cycle
+            gc.collect()
+        """)
+        for attr in 'as_digested_dict', 'as_prefix':
+            with self.subTest(attr=attr):
+                assert_python_ok('-c', code, attr)
+
 class FileTestCase(unittest.TestCase):
     def setUp(self):
         self.DECOMPRESSED_42 = b'a'*42

Misc/NEWS.d/next/Library/2026-10-23-22-34-55.gh-issue-160535.Zk3dGc.rst

@@ -0,0 +1,4 @@
+Fix a crash in :mod:`compression.zstd` when the garbage collector frees a
+:class:`~compression.zstd.ZstdCompressor` that uses worker threads together
+with its :class:`~compression.zstd.ZstdDict` before the compression of
+already submitted data has finished.

Modules/_zstd/compressor.c

@@ -766,17 +766,13 @@ static int
 ZstdCompressor_traverse(PyObject *ob, visitproc visit, void *arg)
 {
     ZstdCompressor *self = ZstdCompressor_CAST(ob);
+    Py_VISIT(Py_TYPE(ob));
     Py_VISIT(self->dict);
     return 0;
 }
 
-static int
-ZstdCompressor_clear(PyObject *ob)
-{
-    ZstdCompressor *self = ZstdCompressor_CAST(ob);
-    Py_CLEAR(self->dict);
-    return 0;
-}
+/* No tp_clear: the compression context (and its worker threads) uses the
+   dictionary until ZstdCompressor_dealloc() has freed the context. */
 
 static PyType_Slot zstdcompressor_slots[] = {
     {Py_tp_new, _zstd_ZstdCompressor_new},
@@ -785,7 +781,6 @@ static PyType_Slot zstdcompressor_slots[] = {
     {Py_tp_members, ZstdCompressor_members},
     {Py_tp_doc, (void *)_zstd_ZstdCompressor_new__doc__},
     {Py_tp_traverse,  ZstdCompressor_traverse},
-    {Py_tp_clear, ZstdCompressor_clear},
     {0, 0}
 };
 

Modules/_zstd/decompressor.c

@@ -689,18 +689,13 @@ static int
 ZstdDecompressor_traverse(PyObject *ob, visitproc visit, void *arg)
 {
     ZstdDecompressor *self = ZstdDecompressor_CAST(ob);
+    Py_VISIT(Py_TYPE(ob));
     Py_VISIT(self->dict);
     return 0;
 }
 
-static int
-ZstdDecompressor_clear(PyObject *ob)
-{
-    ZstdDecompressor *self = ZstdDecompressor_CAST(ob);
-    Py_CLEAR(self->dict);
-    Py_CLEAR(self->unused_data);
-    return 0;
-}
+/* No tp_clear: the decompression context references the dictionary until
+   ZstdDecompressor_dealloc() has freed the context. */
 
 static PyType_Slot ZstdDecompressor_slots[] = {
     {Py_tp_new, _zstd_ZstdDecompressor_new},
@@ -710,7 +705,6 @@ static PyType_Slot ZstdDecompressor_slots[] = {
     {Py_tp_getset, ZstdDecompressor_getset},
     {Py_tp_doc, (void *)_zstd_ZstdDecompressor_new__doc__},
     {Py_tp_traverse, ZstdDecompressor_traverse},
-    {Py_tp_clear, ZstdDecompressor_clear},
     {0, 0}
 };
 

Modules/_zstd/zstddict.c

@@ -240,16 +240,11 @@ ZstdDict_length(PyObject *ob)
 static int
 ZstdDict_traverse(PyObject *ob, visitproc visit, void *arg)
 {
-    ZstdDict *self = ZstdDict_CAST(ob);
-    Py_VISIT(self->c_dicts);
-    return 0;
-}
-
-static int
-ZstdDict_clear(PyObject *ob)
-{
-    ZstdDict *self = ZstdDict_CAST(ob);
-    Py_CLEAR(self->c_dicts);
+    /* c_dicts is not visited and there is no tp_clear: it only maps ints
+       to capsules, and the ZSTD_CDict instances they own must stay alive
+       until ZstdDict_dealloc() as a compression context being collected
+       together with this object may still use them. */
+    Py_VISIT(Py_TYPE(ob));
     return 0;
 }
 
@@ -262,7 +257,6 @@ static PyType_Slot zstddict_slots[] = {
     {Py_tp_doc, (void *)_zstd_ZstdDict_new__doc__},
     {Py_sq_length, ZstdDict_length},
     {Py_tp_traverse, ZstdDict_traverse},
-    {Py_tp_clear, ZstdDict_clear},
     {0, 0}
 };