view contrib/python-zstandard/c-ext/compressoriterator.c @ 40132:e67522413ca8

wireprotov2: define and use stream encoders

Now that we have basic support for defining stream encoding, it is
time to start doing something with it.

We define various classes implementing stream encoders/decoders for
the defined encoding profiles. This is relatively straightforward.

We teach the inputstream and outputstream classes how to encode,
decode, and flush data.

We then teach the clientreactor how to filter received data through
the inputstream decoder.

One of the features of the framing format is that streams can span
requests. This is a differentiating feature from say HTTP/2, which
associates streams with requests. By allowing streams to span requests,
we can reuse compression context data across requests/responses. But
in order to do this, we need a mechanism to "flush" the encoder at
logical boundaries so that receivers receive all data where it is
expected. And a "flush" event is distinct from a "finish" event from
the perspective of certain compressors because a "flush" will retain
compression context state whereas a "finish" operation will not. This
is why encoders have both a flush() and a finish() and each uses
specific flushing semantics on the underlying compressor.

The added tests verify various behavior of decoders via clientreactor.
These tests do test some compression behavior via use of outputstream.
But for all intents and purposes, server reactor support for encoding
is not yet implemented.

Differential Revision: https://phab.mercurial-scm.org/D4921

/**
* Copyright (c) 2016-present, Gregory Szorc
* All rights reserved.
*
* This software may be modified and distributed under the terms
* of the BSD license. See the LICENSE file for details.
*/

#include "python-zstandard.h"

#define min(a, b) (((a) < (b)) ? (a) : (b))

extern PyObject* ZstdError;

PyDoc_STRVAR(ZstdCompressorIterator__doc__,
"Represents an iterator of compressed data.\n"
);

static void ZstdCompressorIterator_dealloc(ZstdCompressorIterator* self) {
	Py_XDECREF(self->readResult);
	Py_XDECREF(self->compressor);
	Py_XDECREF(self->reader);

	if (self->buffer.buf) {
		PyBuffer_Release(&self->buffer);
		memset(&self->buffer, 0, sizeof(self->buffer));
	}

	if (self->output.dst) {
		PyMem_Free(self->output.dst);
		self->output.dst = NULL;
	}

	PyObject_Del(self);
}

static PyObject* ZstdCompressorIterator_iter(PyObject* self) {
	Py_INCREF(self);
	return self;
}

static PyObject* ZstdCompressorIterator_iternext(ZstdCompressorIterator* self) {
	size_t zresult;
	PyObject* readResult = NULL;
	PyObject* chunk;
	char* readBuffer;
	Py_ssize_t readSize = 0;
	Py_ssize_t bufferRemaining;

	if (self->finishedOutput) {
		PyErr_SetString(PyExc_StopIteration, "output flushed");
		return NULL;
	}

feedcompressor:

	/* If we have data left in the input, consume it. */
	if (self->input.pos < self->input.size) {
		Py_BEGIN_ALLOW_THREADS
		zresult = ZSTD_compress_generic(self->compressor->cctx, &self->output,
			&self->input, ZSTD_e_continue);
		Py_END_ALLOW_THREADS

		/* Release the Python object holding the input buffer. */
		if (self->input.pos == self->input.size) {
			self->input.src = NULL;
			self->input.pos = 0;
			self->input.size = 0;
			Py_DECREF(self->readResult);
			self->readResult = NULL;
		}

		if (ZSTD_isError(zresult)) {
			PyErr_Format(ZstdError, "zstd compress error: %s", ZSTD_getErrorName(zresult));
			return NULL;
		}

		/* If it produced output data, emit it. */
		if (self->output.pos) {
			chunk = PyBytes_FromStringAndSize(self->output.dst, self->output.pos);
			self->output.pos = 0;
			return chunk;
		}
	}

	/* We should never have output data sitting around after a previous call. */
	assert(self->output.pos == 0);

	/* The code above should have either emitted a chunk and returned or consumed
	the entire input buffer. So the state of the input buffer is not
	relevant. */
	if (!self->finishedInput) {
		if (self->reader) {
			readResult = PyObject_CallMethod(self->reader, "read", "I", self->inSize);
			if (!readResult) {
				PyErr_SetString(ZstdError, "could not read() from source");
				return NULL;
			}

			PyBytes_AsStringAndSize(readResult, &readBuffer, &readSize);
		}
		else {
			assert(self->buffer.buf);

			/* Only support contiguous C arrays. */
			assert(self->buffer.strides == NULL && self->buffer.suboffsets == NULL);
			assert(self->buffer.itemsize == 1);

			readBuffer = (char*)self->buffer.buf + self->bufferOffset;
			bufferRemaining = self->buffer.len - self->bufferOffset;
			readSize = min(bufferRemaining, (Py_ssize_t)self->inSize);
			self->bufferOffset += readSize;
		}

		if (0 == readSize) {
			Py_XDECREF(readResult);
			self->finishedInput = 1;
		}
		else {
			self->readResult = readResult;
		}
	}

	/* EOF */
	if (0 == readSize) {
		self->input.src = NULL;
		self->input.size = 0;
		self->input.pos = 0;

		zresult = ZSTD_compress_generic(self->compressor->cctx, &self->output,
			&self->input, ZSTD_e_end);
		if (ZSTD_isError(zresult)) {
			PyErr_Format(ZstdError, "error ending compression stream: %s",
				ZSTD_getErrorName(zresult));
			return NULL;
		}

		assert(self->output.pos);

		if (0 == zresult) {
			self->finishedOutput = 1;
		}

		chunk = PyBytes_FromStringAndSize(self->output.dst, self->output.pos);
		self->output.pos = 0;
		return chunk;
	}

	/* New data from reader. Feed into compressor. */
	self->input.src = readBuffer;
	self->input.size = readSize;
	self->input.pos = 0;

	Py_BEGIN_ALLOW_THREADS
	zresult = ZSTD_compress_generic(self->compressor->cctx, &self->output,
		&self->input, ZSTD_e_continue);
	Py_END_ALLOW_THREADS

	/* The input buffer currently points to memory managed by Python
	(readBuffer). This object was allocated by this function. If it wasn't
	fully consumed, we need to release it in a subsequent function call.
	If it is fully consumed, do that now.
	*/
	if (self->input.pos == self->input.size) {
		self->input.src = NULL;
		self->input.pos = 0;
		self->input.size = 0;
		Py_XDECREF(self->readResult);
		self->readResult = NULL;
	}

	if (ZSTD_isError(zresult)) {
		PyErr_Format(ZstdError, "zstd compress error: %s", ZSTD_getErrorName(zresult));
		return NULL;
	}

	assert(self->input.pos <= self->input.size);

	/* If we didn't write anything, start the process over. */
	if (0 == self->output.pos) {
		goto feedcompressor;
	}

	chunk = PyBytes_FromStringAndSize(self->output.dst, self->output.pos);
	self->output.pos = 0;
	return chunk;
}

PyTypeObject ZstdCompressorIteratorType = {
	PyVarObject_HEAD_INIT(NULL, 0)
	"zstd.ZstdCompressorIterator",   /* tp_name */
	sizeof(ZstdCompressorIterator),  /* tp_basicsize */
	0,                               /* tp_itemsize */
	(destructor)ZstdCompressorIterator_dealloc, /* tp_dealloc */
	0,                               /* tp_print */
	0,                               /* tp_getattr */
	0,                               /* tp_setattr */
	0,                               /* tp_compare */
	0,                               /* tp_repr */
	0,                               /* tp_as_number */
	0,                               /* tp_as_sequence */
	0,                               /* tp_as_mapping */
	0,                               /* tp_hash */
	0,                               /* tp_call */
	0,                               /* tp_str */
	0,                               /* tp_getattro */
	0,                               /* tp_setattro */
	0,                               /* tp_as_buffer */
	Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /* tp_flags */
	ZstdCompressorIterator__doc__,   /* tp_doc */
	0,                               /* tp_traverse */
	0,                               /* tp_clear */
	0,                               /* tp_richcompare */
	0,                               /* tp_weaklistoffset */
	ZstdCompressorIterator_iter,     /* tp_iter */
	(iternextfunc)ZstdCompressorIterator_iternext, /* tp_iternext */
	0,                               /* tp_methods */
	0,                               /* tp_members */
	0,                               /* tp_getset */
	0,                               /* tp_base */
	0,                               /* tp_dict */
	0,                               /* tp_descr_get */
	0,                               /* tp_descr_set */
	0,                               /* tp_dictoffset */
	0,                               /* tp_init */
	0,                               /* tp_alloc */
	PyType_GenericNew,              /* tp_new */
};

void compressoriterator_module_init(PyObject* mod) {
	Py_TYPE(&ZstdCompressorIteratorType) = &PyType_Type;
	if (PyType_Ready(&ZstdCompressorIteratorType) < 0) {
		return;
	}
}