Skip to content

Commit fa1a212

Browse files
[3.15] gh-160535: Keep the ZstdDict alive until the zstd context is freed (GH-160536) (#160635)
gh-160535: Keep the ZstdDict alive until the zstd context is freed (GH-160536) The garbage collector could release a ZstdCompressor's dictionary, or the digested dictionaries cached by the ZstdDict, while worker threads of the compression context were still reading them. They are now only released on deallocation, after the context has been freed. ANT-2026-CT3P1CSJ + Visit the type in the compressor and decompressor traverse; move NEWS to Library (cherry picked from commit b949033) Co-authored-by: Gregory P. Smith <68491+gpshead@users.noreply.github.com>
1 parent 3255d5d commit fa1a212

5 files changed

Lines changed: 54 additions & 28 deletions

File tree

‎Lib/test/test_zstd.py‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,12 @@
77
import os
88
import unittest
99
import tempfile
10+
import textwrap
1011
import threading
1112

13+
from test import support
1214
from test.support.import_helper import import_module
15+
from test.support.script_helper import assert_python_ok
1316
from test.support import threading_helper
1417
from test.support import _1M
1518

@@ -1548,6 +1551,42 @@ def test_len(self):
15481551
self.assertEqual(len(TRAINED_DICT), len(TRAINED_DICT.dict_content))
15491552
self.assertIn(str(len(TRAINED_DICT)), str(TRAINED_DICT))
15501553

1554+
@unittest.skipUnless(SUPPORT_MULTITHREADING,
1555+
"zstd build doesn't support multi-threaded compression")
1556+
@support.requires_subprocess()
1557+
def test_collect_cycle_with_busy_compressor(self):
1558+
# The dictionary must outlive the compression context when both are
1559+
# collected by the garbage collector while a worker thread is still
1560+
# using the dictionary. Whether the failure is detected depends on
1561+
# the memory allocator, so use a fresh process for each case.
1562+
code = textwrap.dedent("""
1563+
import gc, random, sys
1564+
from compression.zstd import (CompressionParameter,
1565+
ZstdCompressor, ZstdDict)
1566+
1567+
class Cycle:
1568+
pass
1569+
1570+
job_size = 1 << 20
1571+
content = random.randbytes(1 << 10) * (1 << 8)
1572+
zd = ZstdDict(content, is_raw=True)
1573+
options = {CompressionParameter.compression_level: 10,
1574+
CompressionParameter.nb_workers: 1,
1575+
CompressionParameter.job_size: job_size}
1576+
c = ZstdCompressor(options=options,
1577+
zstd_dict=getattr(zd, sys.argv[1]))
1578+
cycle = Cycle()
1579+
cycle.cycle = cycle
1580+
cycle.compressor = c
1581+
# Returns while the worker thread is compressing the first job.
1582+
c.compress(random.randbytes(1 << 10) * ((job_size >> 10) + 64))
1583+
del zd, c, cycle
1584+
gc.collect()
1585+
""")
1586+
for attr in 'as_digested_dict', 'as_prefix':
1587+
with self.subTest(attr=attr):
1588+
assert_python_ok('-c', code, attr)
1589+
15511590
class FileTestCase(unittest.TestCase):
15521591
def setUp(self):
15531592
self.DECOMPRESSED_42 = b'a'*42
Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
Fix a crash in :mod:`compression.zstd` when the garbage collector frees a
2+
:class:`~compression.zstd.ZstdCompressor` that uses worker threads together
3+
with its :class:`~compression.zstd.ZstdDict` before the compression of
4+
already submitted data has finished.

‎Modules/_zstd/compressor.c‎

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -766,17 +766,13 @@ static int
766766
ZstdCompressor_traverse(PyObject *ob, visitproc visit, void *arg)
767767
{
768768
ZstdCompressor *self = ZstdCompressor_CAST(ob);
769+
Py_VISIT(Py_TYPE(ob));
769770
Py_VISIT(self->dict);
770771
return 0;
771772
}
772773

773-
static int
774-
ZstdCompressor_clear(PyObject *ob)
775-
{
776-
ZstdCompressor *self = ZstdCompressor_CAST(ob);
777-
Py_CLEAR(self->dict);
778-
return 0;
779-
}
774+
/* No tp_clear: the compression context (and its worker threads) uses the
775+
dictionary until ZstdCompressor_dealloc() has freed the context. */
780776

781777
static PyType_Slot zstdcompressor_slots[] = {
782778
{Py_tp_new, _zstd_ZstdCompressor_new},
@@ -785,7 +781,6 @@ static PyType_Slot zstdcompressor_slots[] = {
785781
{Py_tp_members, ZstdCompressor_members},
786782
{Py_tp_doc, (void *)_zstd_ZstdCompressor_new__doc__},
787783
{Py_tp_traverse, ZstdCompressor_traverse},
788-
{Py_tp_clear, ZstdCompressor_clear},
789784
{0, 0}
790785
};
791786

‎Modules/_zstd/decompressor.c‎

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -689,18 +689,13 @@ static int
689689
ZstdDecompressor_traverse(PyObject *ob, visitproc visit, void *arg)
690690
{
691691
ZstdDecompressor *self = ZstdDecompressor_CAST(ob);
692+
Py_VISIT(Py_TYPE(ob));
692693
Py_VISIT(self->dict);
693694
return 0;
694695
}
695696

696-
static int
697-
ZstdDecompressor_clear(PyObject *ob)
698-
{
699-
ZstdDecompressor *self = ZstdDecompressor_CAST(ob);
700-
Py_CLEAR(self->dict);
701-
Py_CLEAR(self->unused_data);
702-
return 0;
703-
}
697+
/* No tp_clear: the decompression context references the dictionary until
698+
ZstdDecompressor_dealloc() has freed the context. */
704699

705700
static PyType_Slot ZstdDecompressor_slots[] = {
706701
{Py_tp_new, _zstd_ZstdDecompressor_new},
@@ -710,7 +705,6 @@ static PyType_Slot ZstdDecompressor_slots[] = {
710705
{Py_tp_getset, ZstdDecompressor_getset},
711706
{Py_tp_doc, (void *)_zstd_ZstdDecompressor_new__doc__},
712707
{Py_tp_traverse, ZstdDecompressor_traverse},
713-
{Py_tp_clear, ZstdDecompressor_clear},
714708
{0, 0}
715709
};
716710

‎Modules/_zstd/zstddict.c‎

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -240,16 +240,11 @@ ZstdDict_length(PyObject *ob)
240240
static int
241241
ZstdDict_traverse(PyObject *ob, visitproc visit, void *arg)
242242
{
243-
ZstdDict *self = ZstdDict_CAST(ob);
244-
Py_VISIT(self->c_dicts);
245-
return 0;
246-
}
247-
248-
static int
249-
ZstdDict_clear(PyObject *ob)
250-
{
251-
ZstdDict *self = ZstdDict_CAST(ob);
252-
Py_CLEAR(self->c_dicts);
243+
/* c_dicts is not visited and there is no tp_clear: it only maps ints
244+
to capsules, and the ZSTD_CDict instances they own must stay alive
245+
until ZstdDict_dealloc() as a compression context being collected
246+
together with this object may still use them. */
247+
Py_VISIT(Py_TYPE(ob));
253248
return 0;
254249
}
255250

@@ -262,7 +257,6 @@ static PyType_Slot zstddict_slots[] = {
262257
{Py_tp_doc, (void *)_zstd_ZstdDict_new__doc__},
263258
{Py_sq_length, ZstdDict_length},
264259
{Py_tp_traverse, ZstdDict_traverse},
265-
{Py_tp_clear, ZstdDict_clear},
266260
{0, 0}
267261
};
268262

0 commit comments

Comments
 (0)