diff --git a/setup.py b/setup.py index f2e48bfbd..e1fce6621 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') @@ -570,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/src/cctx_pool.c b/src/cctx_pool.c new file mode 100644 index 000000000..d62e60d5a --- /dev/null +++ b/src/cctx_pool.c @@ -0,0 +1,86 @@ +#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) { + /* 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) { + 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..5b7335952 --- /dev/null +++ b/src/cctx_pool.h @@ -0,0 +1,36 @@ +#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 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. */ +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 828e546c7..73a585617 100644 --- a/src/python-zstd.c +++ b/src/python-zstd.c @@ -46,6 +46,7 @@ #include "util.h" #include "debug.h" #include "python-zstd.h" +#include "cctx_pool.h" /** * @deprecated, use *mt2 instead. @@ -164,24 +165,6 @@ static PyObject *py_zstd_compress_mt(PyObject* self, PyObject *args) return result; } -void init_cContext( int32_t threads, int32_t level) -{ - m_cctx = ZSTD_createCCtx(); - ZSTD_CCtx_setParameter(m_cctx, ZSTD_c_compressionLevel, level); - ZSTD_CCtx_setParameter(m_cctx, ZSTD_c_nbWorkers, threads); -} - -void free_cContext(void) -{ - ZSTD_freeCCtx(m_cctx); -} - -void reset_cContext(int32_t threads, int32_t level) -{ - free_cContext(); - init_cContext(threads, level); -} - /** * New function for multi-threaded compression. * Uses origin zstd header, nothing more. @@ -203,9 +186,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 +255,23 @@ 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 = 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); 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; + + cctx_pool_release(cctx); printdn("Compression result: %d\n", cSize); if (ZSTD_isError(cSize)) { @@ -769,7 +761,7 @@ static int init_py_zstd(PyObject *module) { int32_t threads = UTIL_countAvailableCores(); UNUSED(threads); - init_cContext(1, 3); + init_cctx_pool(); return 0; } @@ -804,7 +796,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..d8883413d 100644 --- a/src/python-zstd.h +++ b/src/python-zstd.h @@ -127,12 +127,6 @@ 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); - 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); 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()