Mercurial > public > mercurial-scm > hg

diff contrib/python-zstandard/c-ext/compressionwriter.c @ 31796:e0dc40530c5a

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
zstd: vendor python-zstandard 0.8.0 Commit 81e1f5bbf1fc54808649562d3ed829730765c540 from https://github.com/indygreg/python-zstandard is imported without modifications (other than removing unwanted files). Updates relevant to Mercurial include: * Support for multi-threaded compression (we can use this for bundle and wire protocol compression). * APIs for batch compression and decompression operations using multiple threads and optimal memory allocation mechanism. (Can be useful for revlog perf improvements.) * A ``BufferWithSegments`` type that models a single memory buffer containing N discrete items of known lengths. This type can be used for very efficient 0-copy data operations. # no-check-commit
author Gregory Szorc <gregory.szorc@gmail.com>
date Sat, 01 Apr 2017 15:24:03 -0700
parents c32454d69b85
children b1fb341d8a61
line wrap: on
line diff
--- a/contrib/python-zstandard/c-ext/compressionwriter.c	Sat Apr 01 13:43:52 2017 -0700
+++ b/contrib/python-zstandard/c-ext/compressionwriter.c	Sat Apr 01 15:24:03 2017 -0700
@@ -18,11 +18,6 @@
 	Py_XDECREF(self->compressor);
 	Py_XDECREF(self->writer);
 
-	if (self->cstream) {
-		ZSTD_freeCStream(self->cstream);
-		self->cstream = NULL;
-	}
-
 	PyObject_Del(self);
 }
 
@@ -32,9 +27,15 @@
 		return NULL;
 	}
 
-	self->cstream = CStream_from_ZstdCompressor(self->compressor, self->sourceSize);
-	if (!self->cstream) {
-		return NULL;
+	if (self->compressor->mtcctx) {
+		if (init_mtcstream(self->compressor, self->sourceSize)) {
+			return NULL;
+		}
+	}
+	else {
+		if (0 != init_cstream(self->compressor, self->sourceSize)) {
+			return NULL;
+		}
 	}
 
 	self->entered = 1;
@@ -58,8 +59,8 @@
 
 	self->entered = 0;
 
-	if (self->cstream && exc_type == Py_None && exc_value == Py_None &&
-		exc_tb == Py_None) {
+	if ((self->compressor->cstream || self->compressor->mtcctx) && exc_type == Py_None
+		&& exc_value == Py_None && exc_tb == Py_None) {
 
 		output.dst = PyMem_Malloc(self->outSize);
 		if (!output.dst) {
@@ -69,7 +70,12 @@
 		output.pos = 0;
 
 		while (1) {
-			zresult = ZSTD_endStream(self->cstream, &output);
+			if (self->compressor->mtcctx) {
+				zresult = ZSTDMT_endStream(self->compressor->mtcctx, &output);
+			}
+			else {
+				zresult = ZSTD_endStream(self->compressor->cstream, &output);
+			}
 			if (ZSTD_isError(zresult)) {
 				PyErr_Format(ZstdError, "error ending compression stream: %s",
 					ZSTD_getErrorName(zresult));
@@ -95,21 +101,19 @@
 		}
 
 		PyMem_Free(output.dst);
-		ZSTD_freeCStream(self->cstream);
-		self->cstream = NULL;
 	}
 
 	Py_RETURN_FALSE;
 }
 
 static PyObject* ZstdCompressionWriter_memory_size(ZstdCompressionWriter* self) {
-	if (!self->cstream) {
+	if (!self->compressor->cstream) {
 		PyErr_SetString(ZstdError, "cannot determine size of an inactive compressor; "
 			"call when a context manager is active");
 		return NULL;
 	}
 
-	return PyLong_FromSize_t(ZSTD_sizeof_CStream(self->cstream));
+	return PyLong_FromSize_t(ZSTD_sizeof_CStream(self->compressor->cstream));
 }
 
 static PyObject* ZstdCompressionWriter_write(ZstdCompressionWriter* self, PyObject* args) {
@@ -147,7 +151,13 @@
 
 	while ((ssize_t)input.pos < sourceSize) {
 		Py_BEGIN_ALLOW_THREADS
-		zresult = ZSTD_compressStream(self->cstream, &output, &input);
+		if (self->compressor->mtcctx) {
+			zresult = ZSTDMT_compressStream(self->compressor->mtcctx,
+				&output, &input);
+		}
+		else {
+			zresult = ZSTD_compressStream(self->compressor->cstream, &output, &input);
+		}
 		Py_END_ALLOW_THREADS
 
 		if (ZSTD_isError(zresult)) {
@@ -195,7 +205,12 @@
 
 	while (1) {
 		Py_BEGIN_ALLOW_THREADS
-		zresult = ZSTD_flushStream(self->cstream, &output);
+		if (self->compressor->mtcctx) {
+			zresult = ZSTDMT_flushStream(self->compressor->mtcctx, &output);
+		}
+		else {
+			zresult = ZSTD_flushStream(self->compressor->cstream, &output);
+		}
 		Py_END_ALLOW_THREADS
 
 		if (ZSTD_isError(zresult)) {