Make py_zstd_compress_mt2 thread-safe

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"

src/cctx_pool.c

@@ -0,0 +1,86 @@
+#include <stdlib.h>
+#include <stddef.h>
+
+#include <Python.h>
+#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);
+}

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 */

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

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

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()