First take on initialize compress2 context only once

Author: sergey-dryabzhinsky

src/python-zstd.c

@@ -159,6 +159,136 @@ 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()
+{
+    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.
+ * Simple version: not for streaming, no dict support, full block compression.
+ * Uses new API with context object, initialize only once.
+ */
+static PyObject *py_zstd_compress_mt2(PyObject* self, PyObject *args)
+{
+    UNUSED(self);
+
+    PyObject *result;
+//    PyObject *resultT;
+    const char *source;
+    Py_ssize_t source_size;
+//    Py_ssize_t chunk_size;
+    char *dest;
+//    char *destT;
+    Py_ssize_t dest_size;
+    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))
+        return NULL;
+#else
+    if (!PyArg_ParseTuple(args, "s#|iii", &source, &source_size, &level, &threads, &strict))
+        return NULL;
+#endif
+
+    printdn("got Compression level:%d\n",level);
+    if (0 == level) level=ZSTD_defaultCLevel();
+    /* Fast levels (zstd >= 1.3.4) - [-1..-100] */
+    /* Usual levels                - [ 1..22] */
+    /* If level less than -100 or 1 - raise Error, level 0 handled before. */
+    printdn("Compression min level:%d\n",ZSTD_MIN_CLEVEL);
+    if (level < ZSTD_MIN_CLEVEL) {
+        printd2e("Bad compression level - less than %d: %d\n", ZSTD_MIN_CLEVEL, level);
+	if (strict) {
+        PyErr_Format(ZstdError, "Bad compression level - less than %d: %d", ZSTD_MIN_CLEVEL, level);
+        return NULL;
+	} else {
+	    level = ZSTD_MIN_CLEVEL;
+	}
+    }
+    /* If level more than 22 - raise Error. */
+    printdn("Compression max level:%d\n",ZSTD_maxCLevel());
+    if (level > ZSTD_maxCLevel()) {
+        printd2e("Bad compression level - more than %d: %d\n", ZSTD_maxCLevel(), level);
+	if (strict) {
+        PyErr_Format(ZstdError, "Bad compression level - more than %d: %d", ZSTD_MAX_CLEVEL, level);
+        return NULL;
+	} else {
+	    level = ZSTD_maxCLevel();
+	}
+    }
+    printdn("Compression level will be:%d\n",level);
+
+    printdn("got Compression threads:%d\n",threads);
+    if (threads < 0) {
+        printd2e("Bad threads count - less than %d: %d\n", 0, threads);
+	if (strict) {
+        PyErr_Format(ZstdError, "Bad threads count - less than %d: %d", 0, threads);
+        return NULL;
+	} else threads = 1;
+    }
+    if (0 == threads) threads = UTIL_countAvailableCores();
+    printdn("got CPU cores:%d\n",threads);
+    /* If threads more than 200 - raise Error. */
+    if (threads > ZSTDMT_NBWORKERS_MAX) {
+        printd2e("Bad threads count - more than %d: %d\n", ZSTDMT_NBWORKERS_MAX, threads);
+        threads = ZSTDMT_NBWORKERS_MAX;
+        // do not fail here, due auto thread counter
+        //PyErr_Format(ZstdError, "Bad threads count - more than %d: %d", ZSTDMT_NBWORKERS_MAX, threads);
+        //return NULL;
+    }
+    printdn("Compression will use:%d threads\n",threads);
+
+    dest_size = (Py_ssize_t)ZSTD_compressBound(source_size);
+    result = PyBytes_FromStringAndSize(NULL, dest_size);
+    if (result == NULL) {
+        return NULL;
+    }
+
+    if (source_size >= 0) {
+        dest = PyBytes_AS_STRING(result);
+
+		if(level != lastLevel) {
+			reset_cContext(threads, level);
+		}
+
+        Py_BEGIN_ALLOW_THREADS
+        cSize = ZSTD_compress2(m_cctx, dest, (size_t)dest_size, source, (size_t)source_size);
+        Py_END_ALLOW_THREADS
+		lastLevel = level;
+
+        printdn("Compression result: %d\n", cSize);
+        if (ZSTD_isError(cSize)) {
+            printdes("debug INFO: Compression error: %s", ZSTD_getErrorName(cSize));
+            PyErr_Format(ZstdError, "Compression error: %s", ZSTD_getErrorName(cSize));
+            Py_CLEAR(result);
+            return NULL;
+        }
+        Py_SET_SIZE(result, cSize);
+    }
+    return result;
+}
+
 
 /**
  * New more interoperable function
@@ -538,6 +668,7 @@ static PyMethodDef ZstdMethods[] = {
     {"verify",  py_zstd_check, METH_VARARGS, CHECK_DOCSTRING},
     {"compress",  py_zstd_compress_mt, METH_VARARGS, COMPRESS_DOCSTRING},
     {"compress_real_mt",  py_zstd_compress_mt, METH_VARARGS, COMPRESS_DOCSTRING},
+    {"compress2",  py_zstd_compress_mt2, METH_VARARGS, COMPRESS_DOCSTRING},
     {"uncompress",  py_zstd_uncompress, METH_VARARGS, UNCOMPRESS_DOCSTRING},
     {"encode",  py_zstd_compress_mt, METH_VARARGS, COMPRESS_DOCSTRING},
     {"decode",  py_zstd_uncompress, METH_VARARGS, UNCOMPRESS_DOCSTRING},
@@ -586,6 +717,7 @@ static int init_py_zstd(PyObject *module) {
     Py_INCREF(ZstdError);
     PyModule_AddObject(module, "Error", ZstdError);
 
+	init_cContext(1, 3);
     return 0;
 }
 
@@ -611,6 +743,7 @@ static int myextension_clear(PyObject *self) {
 
 static void myextension_free(void *self) {
     Py_CLEAR(GETSTATE((PyObject *)self)->error);
+	free_cContext();
     printdi("ZSTD module->free\n",0);
     return;
 }

src/python-zstd.h

@@ -92,7 +92,10 @@
 
 static PyObject *ZstdError;
 
+static ZSTD_CCtx* m_cctx;
+
 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);
 static PyObject *py_zstd_check(PyObject* self, PyObject *args);
 static PyObject *py_zstd_module_version(PyObject* self, PyObject *args);

tests/test_compress.py

@@ -54,17 +54,22 @@ def test_compression_MT_and_not(self):
         cdata2 = zstd.compress_real_mt(tDATA, 14)
         self.assertEqual(cdata1, cdata2)
 
+    def test_compression_oldMT_and_new2(self):
+        cdata1 = zstd.compress(tDATA, 14)
+        cdata2 = zstd.compress_real_mt(tDATA, 14)
+        self.assertEqual(cdata1, cdata2)
+
     def test_compression_MT_and_raw(self):
         cdata2 = zstd.compress_real_mt(tDATA, 14)
         data2 = zstd.decompress(cdata2)
         self.assertEqual(tDATA, data2)
 
-    def test_compression_equal_low_level(self):
+    def test_compression_equal_too_low_level(self):
         cdata1 = zstd.compress(tDATA, -100)
         cdata2 = zstd.compress(tDATA, -1000)
         self.assertEqual(cdata1, cdata2)
 
-    def test_compression_equal_high_threads(self):
+    def test_compression_equal_too_high_threads(self):
         cdata1 = zstd.compress(tDATA, 22, 256)
         cdata2 = zstd.compress(tDATA, 22, 2560)
         self.assertEqual(cdata1, cdata2)

tests/test_speed.py

@@ -30,6 +30,25 @@ def test_compression_speed(self):
         log.info("Compression speed average = %6.2f Mb/sec" % (1.0*sum/1024/1024/wait,))
         log.info("diff Compression memory usage = %6.2f kb" % (1.0*(endMemoryUsage-beginMemoryUsage)/1024,))
 
+    def test_compression2_speed_init_context_once(self):
+        wait = 10
+        if "ZSTD_FULLTIME_TESTS" in os.environ:
+            wait = 30
+        log.info("\nWait %d seconds..." % wait)
+        sum = 0
+        l=len(tDATA)
+        tbegin = time()
+        beginMemoryUsage=get_real_memory_usage()
+        log.info("begin Compression2 memory usage = %6.2f kb" % (1.0*beginMemoryUsage/1024,))
+        while time()-tbegin<wait:
+            cdata = zstd.compress2(tDATA,3,0)
+            sum+=l
+
+        endMemoryUsage=get_real_memory_usage()
+        log.info("end Compression2 memory usage = %6.2f kb" % (1.0*endMemoryUsage/1024,))
+        log.info("Compression2 speed average = %6.2f Mb/sec" % (1.0*sum/1024/1024/wait,))
+        log.info("diff Compression2 memory usage = %6.2f kb" % (1.0*(endMemoryUsage-beginMemoryUsage)/1024,))
+
     def test_compression_speed_no_cpu_cores_cache(self):
         wait = 10
         if "ZSTD_FULLTIME_TESTS" in os.environ: