From 693cd72bea8b7ee132860f1b83e8c4c8f25bc4c4 Mon Sep 17 00:00:00 2001 From: Thomas Kowalski Date: Fri, 19 Jun 2026 14:56:58 +0200 Subject: [PATCH 1/5] Fix: make `py_zstd_compress_mt2` thread-safe --- src/python-zstd.c | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/src/python-zstd.c b/src/python-zstd.c index 828e546c7..e43163855 100644 --- a/src/python-zstd.c +++ b/src/python-zstd.c @@ -203,9 +203,9 @@ static PyObject *py_zstd_compress_mt2(PyObject* self, PyObject *args) size_t cSize; // size_t sum=0; int32_t level = ZSTD_CLEVEL_DEFAULT; - static int32_t lastLevel =0; int32_t threads = 0; int32_t strict = 0; + ZSTD_CCtx* cctx = NULL; #if PY_MAJOR_VERSION >= 3 if (!PyArg_ParseTuple(args, "y#|iii", &source, &source_size, &level, &threads, &strict)) @@ -272,14 +272,20 @@ static PyObject *py_zstd_compress_mt2(PyObject* self, PyObject *args) if (source_size >= 0) { dest = PyBytes_AS_STRING(result); - if(level != lastLevel) { - reset_cContext(threads, level); - } + cctx = ZSTD_createCCtx(); + if (cctx == NULL) { + PyErr_Format(ZstdError, "Could not create compression context"); + Py_CLEAR(result); + return NULL; + } + ZSTD_CCtx_setParameter(cctx, ZSTD_c_compressionLevel, level); + ZSTD_CCtx_setParameter(cctx, ZSTD_c_nbWorkers, threads); Py_BEGIN_ALLOW_THREADS - cSize = ZSTD_compress2(m_cctx, dest, (size_t)dest_size, source, (size_t)source_size); + cSize = ZSTD_compress2(cctx, dest, (size_t)dest_size, source, (size_t)source_size); Py_END_ALLOW_THREADS - lastLevel = level; + + ZSTD_freeCCtx(cctx); printdn("Compression result: %d\n", cSize); if (ZSTD_isError(cSize)) { From 4fdfe7d4ea127b9b5e1909465360c8fe559279ce Mon Sep 17 00:00:00 2001 From: Thomas Kowalski Date: Sat, 11 Jul 2026 14:15:04 +0200 Subject: [PATCH 2/5] Refactor: use a pool for contexts --- src/python-zstd.c | 89 ++++++++++++++++++++++++++++++++++++++++------- src/python-zstd.h | 7 ++-- 2 files changed, 78 insertions(+), 18 deletions(-) diff --git a/src/python-zstd.c b/src/python-zstd.c index e43163855..fd408a12b 100644 --- a/src/python-zstd.c +++ b/src/python-zstd.c @@ -40,6 +40,7 @@ #include #include #include +#include "pythread.h" #include "pythoncapi_compat.h" // Py_SET_SIZE() for Python 3.8 and older #include "bytesobject.h" @@ -164,22 +165,81 @@ static PyObject *py_zstd_compress_mt(PyObject* self, PyObject *args) return result; } -void init_cContext( int32_t threads, int32_t level) +/* Pool of reusable ZSTD_CCtx objects. */ +static PyThread_type_lock cctx_pool_lock = NULL; +static ZSTD_CCtx** cctx_pool = NULL; +static size_t cctx_pool_count = 0; +static size_t cctx_pool_capacity = 0; + +void init_cctx_pool(void) +{ + if (cctx_pool_lock == NULL) { + cctx_pool_lock = PyThread_allocate_lock(); + } +} + +void free_cctx_pool(void) { - m_cctx = ZSTD_createCCtx(); - ZSTD_CCtx_setParameter(m_cctx, ZSTD_c_compressionLevel, level); - ZSTD_CCtx_setParameter(m_cctx, ZSTD_c_nbWorkers, threads); + if (cctx_pool_lock != NULL) { + PyThread_acquire_lock(cctx_pool_lock, WAIT_LOCK); + } + for (size_t i = 0; i < cctx_pool_count; i++) { + ZSTD_freeCCtx(cctx_pool[i]); + } + free(cctx_pool); + cctx_pool = NULL; + cctx_pool_count = 0; + cctx_pool_capacity = 0; + if (cctx_pool_lock != NULL) { + PyThread_release_lock(cctx_pool_lock); + PyThread_free_lock(cctx_pool_lock); + cctx_pool_lock = NULL; + } } -void free_cContext(void) +static ZSTD_CCtx* cctx_pool_acquire(void) { - ZSTD_freeCCtx(m_cctx); + ZSTD_CCtx* cctx = NULL; + + if (cctx_pool_lock != NULL) { + PyThread_acquire_lock(cctx_pool_lock, WAIT_LOCK); + if (cctx_pool_count > 0) { + cctx = cctx_pool[--cctx_pool_count]; + } + PyThread_release_lock(cctx_pool_lock); + } + + if (cctx == NULL) { + cctx = ZSTD_createCCtx(); + } + return cctx; } -void reset_cContext(int32_t threads, int32_t level) +static void cctx_pool_release(ZSTD_CCtx* cctx) { - free_cContext(); - init_cContext(threads, level); + if (cctx == NULL) { + return; + } + + if (cctx_pool_lock == NULL) { + ZSTD_freeCCtx(cctx); + return; + } + + PyThread_acquire_lock(cctx_pool_lock, WAIT_LOCK); + if (cctx_pool_count == cctx_pool_capacity) { + size_t new_capacity = cctx_pool_capacity ? cctx_pool_capacity * 2 : 8; + ZSTD_CCtx** new_pool = (ZSTD_CCtx**)realloc(cctx_pool, new_capacity * sizeof(ZSTD_CCtx*)); + if (new_pool == NULL) { + PyThread_release_lock(cctx_pool_lock); + ZSTD_freeCCtx(cctx); + return; + } + cctx_pool = new_pool; + cctx_pool_capacity = new_capacity; + } + cctx_pool[cctx_pool_count++] = cctx; + PyThread_release_lock(cctx_pool_lock); } /** @@ -272,12 +332,15 @@ static PyObject *py_zstd_compress_mt2(PyObject* self, PyObject *args) if (source_size >= 0) { dest = PyBytes_AS_STRING(result); - cctx = ZSTD_createCCtx(); + cctx = cctx_pool_acquire(); if (cctx == NULL) { PyErr_Format(ZstdError, "Could not create compression context"); Py_CLEAR(result); return NULL; } + /* Context may have been used before: drop any leftover session state + * and parameters before configuring it for this call. */ + ZSTD_CCtx_reset(cctx, ZSTD_reset_session_and_parameters); ZSTD_CCtx_setParameter(cctx, ZSTD_c_compressionLevel, level); ZSTD_CCtx_setParameter(cctx, ZSTD_c_nbWorkers, threads); @@ -285,7 +348,7 @@ static PyObject *py_zstd_compress_mt2(PyObject* self, PyObject *args) cSize = ZSTD_compress2(cctx, dest, (size_t)dest_size, source, (size_t)source_size); Py_END_ALLOW_THREADS - ZSTD_freeCCtx(cctx); + cctx_pool_release(cctx); printdn("Compression result: %d\n", cSize); if (ZSTD_isError(cSize)) { @@ -775,7 +838,7 @@ static int init_py_zstd(PyObject *module) { int32_t threads = UTIL_countAvailableCores(); UNUSED(threads); - init_cContext(1, 3); + init_cctx_pool(); return 0; } @@ -810,7 +873,7 @@ static void myextension_free(void *self) { if (state != NULL) { Py_CLEAR(state->error); } - free_cContext(); + free_cctx_pool(); printdi("ZSTD module->free\n",0); return; } diff --git a/src/python-zstd.h b/src/python-zstd.h index 84ba9e7ba..214df65be 100644 --- a/src/python-zstd.h +++ b/src/python-zstd.h @@ -127,11 +127,8 @@ typedef unsigned long long U64; static PyObject *ZstdError; -static ZSTD_CCtx* m_cctx; - -void free_cContext(void); -void reset_cContext(int32_t threads, int32_t level); -void init_cContext( int32_t threads, int32_t level); +void init_cctx_pool(void); +void free_cctx_pool(void); static PyObject *py_zstd_compress_mt(PyObject* self, PyObject *args); static PyObject *py_zstd_compress_mt2(PyObject* self, PyObject *args); From 78c2e558c06d62923ab3848d2ea94d74c9d172a2 Mon Sep 17 00:00:00 2001 From: Thomas Kowalski Date: Sun, 23 Aug 2026 11:57:22 +0200 Subject: [PATCH 3/5] refactor: move pool to a different file --- setup.py | 1 + src/cctx_pool.c | 84 +++++++++++++++++++++++++++++++++++++++++++++++ src/cctx_pool.h | 35 ++++++++++++++++++++ src/python-zstd.c | 79 +------------------------------------------- src/python-zstd.h | 3 -- 5 files changed, 121 insertions(+), 81 deletions(-) create mode 100644 src/cctx_pool.c create mode 100644 src/cctx_pool.h diff --git a/setup.py b/setup.py index f2e48bfbd..42d41f95d 100644 --- a/setup.py +++ b/setup.py @@ -555,6 +555,7 @@ def build_extensions(self): # files needed always, even for external zstdFiles.append('src/debug.c') zstdFiles.append('src/util.c') +zstdFiles.append('src/cctx_pool.c') zstdFiles.append('src/python-zstd.c') diff --git a/src/cctx_pool.c b/src/cctx_pool.c new file mode 100644 index 000000000..8bc645630 --- /dev/null +++ b/src/cctx_pool.c @@ -0,0 +1,84 @@ +#include +#include + +#include +#include "pythread.h" + +#include "zstd.h" +#include "cctx_pool.h" + +static PyThread_type_lock cctx_pool_lock = NULL; +static ZSTD_CCtx** cctx_pool = NULL; +static size_t cctx_pool_count = 0; +static size_t cctx_pool_capacity = 0; + +void init_cctx_pool(void) +{ + if (cctx_pool_lock == NULL) { + cctx_pool_lock = PyThread_allocate_lock(); + } +} + +void free_cctx_pool(void) +{ + if (cctx_pool_lock != NULL) { + PyThread_acquire_lock(cctx_pool_lock, WAIT_LOCK); + } + for (size_t i = 0; i < cctx_pool_count; i++) { + ZSTD_freeCCtx(cctx_pool[i]); + } + free(cctx_pool); + cctx_pool = NULL; + cctx_pool_count = 0; + cctx_pool_capacity = 0; + if (cctx_pool_lock != NULL) { + PyThread_release_lock(cctx_pool_lock); + PyThread_free_lock(cctx_pool_lock); + cctx_pool_lock = NULL; + } +} + +ZSTD_CCtx* cctx_pool_acquire(void) +{ + ZSTD_CCtx* cctx = NULL; + + if (cctx_pool_lock != NULL) { + PyThread_acquire_lock(cctx_pool_lock, WAIT_LOCK); + if (cctx_pool_count > 0) { + cctx = cctx_pool[--cctx_pool_count]; + } + PyThread_release_lock(cctx_pool_lock); + } + + if (cctx == NULL) { + cctx = ZSTD_createCCtx(); + } + return cctx; +} + +void cctx_pool_release(ZSTD_CCtx* cctx) +{ + if (cctx == NULL) { + return; + } + + if (cctx_pool_lock == NULL) { + ZSTD_freeCCtx(cctx); + return; + } + + PyThread_acquire_lock(cctx_pool_lock, WAIT_LOCK); + if (cctx_pool_count == cctx_pool_capacity) { + size_t new_capacity = cctx_pool_capacity ? cctx_pool_capacity * 2 : 8; + ZSTD_CCtx** new_pool = (ZSTD_CCtx**)realloc(cctx_pool, new_capacity * sizeof(ZSTD_CCtx*)); + if (new_pool == NULL) { + PyThread_release_lock(cctx_pool_lock); + ZSTD_freeCCtx(cctx); + return; + } + cctx_pool = new_pool; + cctx_pool_capacity = new_capacity; + } + cctx_pool[cctx_pool_count++] = cctx; + PyThread_release_lock(cctx_pool_lock); +} diff --git a/src/cctx_pool.h b/src/cctx_pool.h new file mode 100644 index 000000000..362f30788 --- /dev/null +++ b/src/cctx_pool.h @@ -0,0 +1,35 @@ +#ifndef CCTX_POOL_H +#define CCTX_POOL_H + +#include "zstd.h" + +/* + * Pool of reusable ZSTD_CCtx objects, used by py_zstd_compress_mt2 to + * avoid re-allocating a fresh context on every call while remaining + * thread-safe. Each acquired context is owned by exactly one caller + * until it is released back into the pool. + * + * The pool has no fixed upper bound: it grows on demand up to the peak + * number of concurrent callers, and never shrinks until free_cctx_pool + * is invoked at module teardown. + */ + +/* Initialize the pool. Safe to call more than once. */ +void init_cctx_pool(void); + +/* Free every context still held by the pool and release the lock. */ +void free_cctx_pool(void); + +/* + * Return a context ready for use. If the pool is empty, a new context + * is created. Returns NULL if allocation fails. + */ +ZSTD_CCtx* cctx_pool_acquire(void); + +/* + * Return a context to the pool for later reuse. If the pool cannot + * grow to hold it, the context is freed instead. + */ +void cctx_pool_release(ZSTD_CCtx* cctx); + +#endif /* CCTX_POOL_H */ diff --git a/src/python-zstd.c b/src/python-zstd.c index fd408a12b..73a585617 100644 --- a/src/python-zstd.c +++ b/src/python-zstd.c @@ -40,13 +40,13 @@ #include #include #include -#include "pythread.h" #include "pythoncapi_compat.h" // Py_SET_SIZE() for Python 3.8 and older #include "bytesobject.h" #include "util.h" #include "debug.h" #include "python-zstd.h" +#include "cctx_pool.h" /** * @deprecated, use *mt2 instead. @@ -165,83 +165,6 @@ static PyObject *py_zstd_compress_mt(PyObject* self, PyObject *args) return result; } -/* Pool of reusable ZSTD_CCtx objects. */ -static PyThread_type_lock cctx_pool_lock = NULL; -static ZSTD_CCtx** cctx_pool = NULL; -static size_t cctx_pool_count = 0; -static size_t cctx_pool_capacity = 0; - -void init_cctx_pool(void) -{ - if (cctx_pool_lock == NULL) { - cctx_pool_lock = PyThread_allocate_lock(); - } -} - -void free_cctx_pool(void) -{ - if (cctx_pool_lock != NULL) { - PyThread_acquire_lock(cctx_pool_lock, WAIT_LOCK); - } - for (size_t i = 0; i < cctx_pool_count; i++) { - ZSTD_freeCCtx(cctx_pool[i]); - } - free(cctx_pool); - cctx_pool = NULL; - cctx_pool_count = 0; - cctx_pool_capacity = 0; - if (cctx_pool_lock != NULL) { - PyThread_release_lock(cctx_pool_lock); - PyThread_free_lock(cctx_pool_lock); - cctx_pool_lock = NULL; - } -} - -static ZSTD_CCtx* cctx_pool_acquire(void) -{ - ZSTD_CCtx* cctx = NULL; - - if (cctx_pool_lock != NULL) { - PyThread_acquire_lock(cctx_pool_lock, WAIT_LOCK); - if (cctx_pool_count > 0) { - cctx = cctx_pool[--cctx_pool_count]; - } - PyThread_release_lock(cctx_pool_lock); - } - - if (cctx == NULL) { - cctx = ZSTD_createCCtx(); - } - return cctx; -} - -static void cctx_pool_release(ZSTD_CCtx* cctx) -{ - if (cctx == NULL) { - return; - } - - if (cctx_pool_lock == NULL) { - ZSTD_freeCCtx(cctx); - return; - } - - PyThread_acquire_lock(cctx_pool_lock, WAIT_LOCK); - if (cctx_pool_count == cctx_pool_capacity) { - size_t new_capacity = cctx_pool_capacity ? cctx_pool_capacity * 2 : 8; - ZSTD_CCtx** new_pool = (ZSTD_CCtx**)realloc(cctx_pool, new_capacity * sizeof(ZSTD_CCtx*)); - if (new_pool == NULL) { - PyThread_release_lock(cctx_pool_lock); - ZSTD_freeCCtx(cctx); - return; - } - cctx_pool = new_pool; - cctx_pool_capacity = new_capacity; - } - cctx_pool[cctx_pool_count++] = cctx; - PyThread_release_lock(cctx_pool_lock); -} - /** * New function for multi-threaded compression. * Uses origin zstd header, nothing more. diff --git a/src/python-zstd.h b/src/python-zstd.h index 214df65be..d8883413d 100644 --- a/src/python-zstd.h +++ b/src/python-zstd.h @@ -127,9 +127,6 @@ typedef unsigned long long U64; static PyObject *ZstdError; -void init_cctx_pool(void); -void free_cctx_pool(void); - static PyObject *py_zstd_compress_mt(PyObject* self, PyObject *args); static PyObject *py_zstd_compress_mt2(PyObject* self, PyObject *args); static PyObject *py_zstd_uncompress(PyObject* self, PyObject *args); From 34b1aaf79046775c7abed7808b5ba8296b8b4b71 Mon Sep 17 00:00:00 2001 From: Thomas Kowalski Date: Sun, 23 Aug 2026 12:02:00 +0200 Subject: [PATCH 4/5] test: add regression test --- setup.py | 1 + tests/test_thread_safety.py | 62 +++++++++++++++++++++++++++++++++++++ 2 files changed, 63 insertions(+) create mode 100644 tests/test_thread_safety.py diff --git a/setup.py b/setup.py index 42d41f95d..e1fce6621 100644 --- a/setup.py +++ b/setup.py @@ -571,6 +571,7 @@ def my_test_suite(): test_suite.addTest(unittest.defaultTestLoader.loadTestsFromName("tests.test_decompress")) test_suite.addTest(unittest.defaultTestLoader.loadTestsFromName("tests.test_version")) test_suite.addTest(unittest.defaultTestLoader.loadTestsFromName("tests.test_speed")) + test_suite.addTest(unittest.defaultTestLoader.loadTestsFromName("tests.test_thread_safety")) return test_suite test_func_name = "setup.my_test_suite" diff --git a/tests/test_thread_safety.py b/tests/test_thread_safety.py new file mode 100644 index 000000000..9fe23a121 --- /dev/null +++ b/tests/test_thread_safety.py @@ -0,0 +1,62 @@ +# Regression test for https://github.com/sergey-dryabzhinsky/python-zstd/pull/321 +# +# Before the fix, py_zstd_compress_mt2 shared a single global ZSTD_CCtx +# across all calls. Because ZSTD_compress2 releases the GIL, two threads +# calling zstd.compress2 could use the context concurrently and corrupt +# it, leading to out-of-bounds reads inside zstd. +# +# A second race existed as well: with different `level` values the +# global context was freed and recreated while another thread was still +# compressing into it. +# +# This test runs many threads that all call zstd.compress2 with a mix +# of levels, and it must complete without raising and with every output +# round-tripping through zstd.decompress. + +import threading +import unittest + +import zstd + + +class TestZstdThreadSafety(unittest.TestCase): + # Small enough for CI, large enough that the threads actually overlap + # inside ZSTD_compress2. + DATA = b"hello world this is a test " * 40000 # ~1.1 MB + THREADS = 32 + ITERS = 100 + # Mix of levels used to trigger the reset_cContext race in older + # revisions. Kept here to exercise the pool with contexts that would + # have been reconfigured before. + LEVELS = [1, 3, 5, 9, 3, 1] + + def test_compress2_is_thread_safe(self): + errors = [] + errors_lock = threading.Lock() + + def worker(): + try: + for i in range(self.ITERS): + level = self.LEVELS[i % len(self.LEVELS)] + compressed = zstd.compress2(self.DATA, level) + self.assertEqual(self.DATA, zstd.decompress(compressed)) + except BaseException as exc: + with errors_lock: + errors.append(exc) + + threads = [ + threading.Thread(target=worker) for _ in range(self.THREADS) + ] + for t in threads: + t.start() + for t in threads: + t.join() + + self.assertEqual( + [], errors, + msg="zstd.compress2 raised under concurrent use: %r" % errors, + ) + + +if __name__ == '__main__': + unittest.main() From 12680f531d20ca9956ae14bd67ed58659996bcdc Mon Sep 17 00:00:00 2001 From: Thomas Kowalski Date: Sun, 23 Aug 2026 12:05:36 +0200 Subject: [PATCH 5/5] review: clarify growth pattern --- src/cctx_pool.c | 2 ++ src/cctx_pool.h | 7 ++++--- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/src/cctx_pool.c b/src/cctx_pool.c index 8bc645630..d62e60d5a 100644 --- a/src/cctx_pool.c +++ b/src/cctx_pool.c @@ -69,6 +69,8 @@ void cctx_pool_release(ZSTD_CCtx* cctx) PyThread_acquire_lock(cctx_pool_lock, WAIT_LOCK); if (cctx_pool_count == cctx_pool_capacity) { + /* Grow the pool by doubling capacity whenever a released context does not fit. + * The maximum size of the pool is bounded by the peak number of threads. */ size_t new_capacity = cctx_pool_capacity ? cctx_pool_capacity * 2 : 8; ZSTD_CCtx** new_pool = (ZSTD_CCtx**)realloc(cctx_pool, new_capacity * sizeof(ZSTD_CCtx*)); if (new_pool == NULL) { diff --git a/src/cctx_pool.h b/src/cctx_pool.h index 362f30788..5b7335952 100644 --- a/src/cctx_pool.h +++ b/src/cctx_pool.h @@ -9,9 +9,10 @@ * thread-safe. Each acquired context is owned by exactly one caller * until it is released back into the pool. * - * The pool has no fixed upper bound: it grows on demand up to the peak - * number of concurrent callers, and never shrinks until free_cctx_pool - * is invoked at module teardown. + * The pool grows on demand but its size is bounded by the peak number of + * threads that have ever been concurrently inside py_zstd_compress_mt2, + * one context per in-flight call and no more. + * The pool never shrinks; it is only freed at module teardown. */ /* Initialize the pool. Safe to call more than once. */