Make py_zstd_compress_mt2 thread-safe

src/python-zstd.c

@@ -40,6 +40,7 @@
 #include <stdio.h>
 #include <string.h>
 #include <Python.h>
+#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)
 {
-    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) {
+        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;
+    }
 }
 
-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);
 }
 
 /**
@@ -203,9 +263,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 +332,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 +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;
 }
 
@@ -804,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;
 }

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