contrib/python-zstandard/zstd.c

wireproto: in batch queries, support queries with immediate responses

listkeys and pushkey return without querying the remote when the
remote doesn't support such queries. Before this change, the batching
code didn't handle this convention, resulting in this kind of error:

  $ hg pull ssh://user@dummy/repo1 -r tip -B a
  pulling from ssh://user@dummy/repo1
  remote: ** Unknown exception encountered with possibly-broken third-party extension disable-lookup
  remote: ** which supports versions unknown of Mercurial.
  remote: ** Please disable disable-lookup and try your action again.
  remote: ** If that fixes the bug please report it to the extension author.
  remote: ** Python 2.7.15+ (default, Oct  2 2018, 22:12:08) [GCC 8.2.0]
  remote: ** Mercurial Distributed SCM (version 4.8.1+586-ef54bd33b476+20181224)
  remote: ** Extensions loaded: disable-lookup
  remote: Traceback (most recent call last):
  remote:   File "/tmp/hgtests.i66Npc/install/bin/hg", line 43, in <module>
  remote:     dispatch.run()
  remote:   File "/tmp/hgtests.i66Npc/install/lib/python/mercurial/dispatch.py", line 99, in run
  remote:     status = dispatch(req)
  remote:   File "/tmp/hgtests.i66Npc/install/lib/python/mercurial/dispatch.py", line 225, in dispatch
  remote:     ret = _runcatch(req) or 0
  remote:   File "/tmp/hgtests.i66Npc/install/lib/python/mercurial/dispatch.py", line 376, in _runcatch
  remote:     return _callcatch(ui, _runcatchfunc)
  remote:   File "/tmp/hgtests.i66Npc/install/lib/python/mercurial/dispatch.py", line 384, in _callcatch
  remote:     return scmutil.callcatch(ui, func)
  remote:   File "/tmp/hgtests.i66Npc/install/lib/python/mercurial/scmutil.py", line 166, in callcatch
  remote:     return func()
  remote:   File "/tmp/hgtests.i66Npc/install/lib/python/mercurial/dispatch.py", line 367, in _runcatchfunc
  remote:     return _dispatch(req)
  remote:   File "/tmp/hgtests.i66Npc/install/lib/python/mercurial/dispatch.py", line 1021, in _dispatch
  remote:     cmdpats, cmdoptions)
  remote:   File "/tmp/hgtests.i66Npc/install/lib/python/mercurial/dispatch.py", line 756, in runcommand
  remote:     ret = _runcommand(ui, options, cmd, d)
  remote:   File "/tmp/hgtests.i66Npc/install/lib/python/mercurial/dispatch.py", line 1030, in _runcommand
  remote:     return cmdfunc()
  remote:   File "/tmp/hgtests.i66Npc/install/lib/python/mercurial/dispatch.py", line 1018, in <lambda>
  remote:     d = lambda: util.checksignature(func)(ui, *args, **strcmdopt)
  remote:   File "/tmp/hgtests.i66Npc/install/lib/python/mercurial/util.py", line 1670, in check
  remote:     return func(*args, **kwargs)
  remote:   File "/tmp/hgtests.i66Npc/install/lib/python/mercurial/commands.py", line 5257, in serve
  remote:     s.serve_forever()
  remote:   File "/tmp/hgtests.i66Npc/install/lib/python/mercurial/wireprotoserver.py", line 797, in serve_forever
  remote:     self.serveuntil(threading.Event())
  remote:   File "/tmp/hgtests.i66Npc/install/lib/python/mercurial/wireprotoserver.py", line 804, in serveuntil
  remote:     _runsshserver(self._ui, self._repo, self._fin, self._fout, ev)
  remote:   File "/tmp/hgtests.i66Npc/install/lib/python/mercurial/wireprotoserver.py", line 656, in _runsshserver
  remote:     rsp = wireprotov1server.dispatch(repo, proto, request)
  remote:   File "/tmp/hgtests.i66Npc/install/lib/python/mercurial/wireprotov1server.py", line 74, in dispatch
  remote:     return func(repo, proto, *args)
  remote:   File "/tmp/hgtests.i66Npc/install/lib/python/mercurial/wireprotov1server.py", line 195, in batch
  remote:     data[k] = vals[k]
  remote: KeyError: 'namespace'
  abort: unexpected response: empty string
  [255]

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

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

/* A Python C extension for Zstandard. */

#if defined(_WIN32)
#define WIN32_LEAN_AND_MEAN
#include <Windows.h>
#elif defined(__APPLE__) || defined(__OpenBSD__) || defined(__FreeBSD__) || defined(__NetBSD__) || defined(__DragonFly__)
#include <sys/types.h>
#include <sys/sysctl.h>
#endif

#include "python-zstandard.h"

PyObject *ZstdError;

PyDoc_STRVAR(estimate_decompression_context_size__doc__,
"estimate_decompression_context_size()\n"
"\n"
"Estimate the amount of memory allocated to a decompression context.\n"
);

static PyObject* estimate_decompression_context_size(PyObject* self) {
	return PyLong_FromSize_t(ZSTD_estimateDCtxSize());
}

PyDoc_STRVAR(frame_content_size__doc__,
"frame_content_size(data)\n"
"\n"
"Obtain the decompressed size of a frame."
);

static PyObject* frame_content_size(PyObject* self, PyObject* args, PyObject* kwargs) {
	static char* kwlist[] = {
		"source",
		NULL
	};

	Py_buffer source;
	PyObject* result = NULL;
	unsigned long long size;

#if PY_MAJOR_VERSION >= 3
	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*:frame_content_size",
#else
	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s*:frame_content_size",
#endif
		kwlist, &source)) {
		return NULL;
	}

	if (!PyBuffer_IsContiguous(&source, 'C') || source.ndim > 1) {
		PyErr_SetString(PyExc_ValueError,
			"data buffer should be contiguous and have at most one dimension");
		goto finally;
	}

	size = ZSTD_getFrameContentSize(source.buf, source.len);

	if (size == ZSTD_CONTENTSIZE_ERROR) {
		PyErr_SetString(ZstdError, "error when determining content size");
	}
	else if (size == ZSTD_CONTENTSIZE_UNKNOWN) {
		result = PyLong_FromLong(-1);
	}
	else {
		result = PyLong_FromUnsignedLongLong(size);
	}

finally:
	PyBuffer_Release(&source);

	return result;
}

PyDoc_STRVAR(frame_header_size__doc__,
"frame_header_size(data)\n"
"\n"
"Obtain the size of a frame header.\n"
);

static PyObject* frame_header_size(PyObject* self, PyObject* args, PyObject* kwargs) {
	static char* kwlist[] = {
		"source",
		NULL
	};

	Py_buffer source;
	PyObject* result = NULL;
	size_t zresult;

#if PY_MAJOR_VERSION >= 3
	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*:frame_header_size",
#else
	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s*:frame_header_size",
#endif
		kwlist, &source)) {
		return NULL;
	}

	if (!PyBuffer_IsContiguous(&source, 'C') || source.ndim > 1) {
		PyErr_SetString(PyExc_ValueError,
			"data buffer should be contiguous and have at most one dimension");
		goto finally;
	}

	zresult = ZSTD_frameHeaderSize(source.buf, source.len);
	if (ZSTD_isError(zresult)) {
		PyErr_Format(ZstdError, "could not determine frame header size: %s",
			ZSTD_getErrorName(zresult));
	}
	else {
		result = PyLong_FromSize_t(zresult);
	}

finally:

	PyBuffer_Release(&source);

	return result;
}

PyDoc_STRVAR(get_frame_parameters__doc__,
"get_frame_parameters(data)\n"
"\n"
"Obtains a ``FrameParameters`` instance by parsing data.\n");

PyDoc_STRVAR(train_dictionary__doc__,
"train_dictionary(dict_size, samples, k=None, d=None, steps=None,\n"
"                 threads=None,notifications=0, dict_id=0, level=0)\n"
"\n"
"Train a dictionary from sample data using the COVER algorithm.\n"
"\n"
"A compression dictionary of size ``dict_size`` will be created from the\n"
"iterable of ``samples``. The raw dictionary bytes will be returned.\n"
"\n"
"The COVER algorithm has 2 parameters: ``k`` and ``d``. These control the\n"
"*segment size* and *dmer size*. A reasonable range for ``k`` is\n"
"``[16, 2048+]``. A reasonable range for ``d`` is ``[6, 16]``.\n"
"``d`` must be less than or equal to ``k``.\n"
"\n"
"``steps`` can be specified to control the number of steps through potential\n"
"values of ``k`` and ``d`` to try. ``k`` and ``d`` will only be varied if\n"
"those arguments are not defined. i.e. if ``d`` is ``8``, then only ``k``\n"
"will be varied in this mode.\n"
"\n"
"``threads`` can specify how many threads to use to test various ``k`` and\n"
"``d`` values. ``-1`` will use as many threads as available CPUs. By default,\n"
"a single thread is used.\n"
"\n"
"When ``k`` and ``d`` are not defined, default values are used and the\n"
"algorithm will perform multiple iterations - or steps - to try to find\n"
"ideal parameters. If both ``k`` and ``d`` are specified, then those values\n"
"will be used. ``steps`` or ``threads`` triggers optimization mode to test\n"
"multiple ``k`` and ``d`` variations.\n"
);

static char zstd_doc[] = "Interface to zstandard";

static PyMethodDef zstd_methods[] = {
	{ "estimate_decompression_context_size", (PyCFunction)estimate_decompression_context_size,
	METH_NOARGS, estimate_decompression_context_size__doc__ },
	{ "frame_content_size", (PyCFunction)frame_content_size,
	METH_VARARGS | METH_KEYWORDS, frame_content_size__doc__ },
	{ "frame_header_size", (PyCFunction)frame_header_size,
	METH_VARARGS | METH_KEYWORDS, frame_header_size__doc__ },
	{ "get_frame_parameters", (PyCFunction)get_frame_parameters,
	METH_VARARGS | METH_KEYWORDS, get_frame_parameters__doc__ },
	{ "train_dictionary", (PyCFunction)train_dictionary,
	METH_VARARGS | METH_KEYWORDS, train_dictionary__doc__ },
	{ NULL, NULL }
};

void bufferutil_module_init(PyObject* mod);
void compressobj_module_init(PyObject* mod);
void compressor_module_init(PyObject* mod);
void compressionparams_module_init(PyObject* mod);
void constants_module_init(PyObject* mod);
void compressionchunker_module_init(PyObject* mod);
void compressiondict_module_init(PyObject* mod);
void compressionreader_module_init(PyObject* mod);
void compressionwriter_module_init(PyObject* mod);
void compressoriterator_module_init(PyObject* mod);
void decompressor_module_init(PyObject* mod);
void decompressobj_module_init(PyObject* mod);
void decompressionreader_module_init(PyObject *mod);
void decompressionwriter_module_init(PyObject* mod);
void decompressoriterator_module_init(PyObject* mod);
void frameparams_module_init(PyObject* mod);

void zstd_module_init(PyObject* m) {
	/* python-zstandard relies on unstable zstd C API features. This means
	   that changes in zstd may break expectations in python-zstandard.

	   python-zstandard is distributed with a copy of the zstd sources.
	   python-zstandard is only guaranteed to work with the bundled version
	   of zstd.

	   However, downstream redistributors or packagers may unbundle zstd
	   from python-zstandard. This can result in a mismatch between zstd
	   versions and API semantics. This essentially "voids the warranty"
	   of python-zstandard and may cause undefined behavior.

	   We detect this mismatch here and refuse to load the module if this
	   scenario is detected.
	*/
	if (ZSTD_VERSION_NUMBER != 10306 || ZSTD_versionNumber() != 10306) {
		PyErr_SetString(PyExc_ImportError, "zstd C API mismatch; Python bindings not compiled against expected zstd version");
		return;
	}

	bufferutil_module_init(m);
	compressionparams_module_init(m);
	compressiondict_module_init(m);
	compressobj_module_init(m);
	compressor_module_init(m);
	compressionchunker_module_init(m);
	compressionreader_module_init(m);
	compressionwriter_module_init(m);
	compressoriterator_module_init(m);
	constants_module_init(m);
	decompressor_module_init(m);
	decompressobj_module_init(m);
	decompressionreader_module_init(m);
	decompressionwriter_module_init(m);
	decompressoriterator_module_init(m);
	frameparams_module_init(m);
}

#if defined(__GNUC__) && (__GNUC__ >= 4)
#  define PYTHON_ZSTD_VISIBILITY __attribute__ ((visibility ("default")))
#else
#  define PYTHON_ZSTD_VISIBILITY
#endif

#if PY_MAJOR_VERSION >= 3
static struct PyModuleDef zstd_module = {
	PyModuleDef_HEAD_INIT,
	"zstd",
	zstd_doc,
	-1,
	zstd_methods
};

PYTHON_ZSTD_VISIBILITY PyMODINIT_FUNC PyInit_zstd(void) {
	PyObject *m = PyModule_Create(&zstd_module);
	if (m) {
		zstd_module_init(m);
		if (PyErr_Occurred()) {
			Py_DECREF(m);
			m = NULL;
		}
	}
	return m;
}
#else
PYTHON_ZSTD_VISIBILITY PyMODINIT_FUNC initzstd(void) {
	PyObject *m = Py_InitModule3("zstd", zstd_methods, zstd_doc);
	if (m) {
		zstd_module_init(m);
	}
}
#endif

/* Attempt to resolve the number of CPUs in the system. */
int cpu_count() {
	int count = 0;

#if defined(_WIN32)
	SYSTEM_INFO si;
	si.dwNumberOfProcessors = 0;
	GetSystemInfo(&si);
	count = si.dwNumberOfProcessors;
#elif defined(__APPLE__)
	int num;
	size_t size = sizeof(int);

	if (0 == sysctlbyname("hw.logicalcpu", &num, &size, NULL, 0)) {
		count = num;
	}
#elif defined(__linux__)
	count = sysconf(_SC_NPROCESSORS_ONLN);
#elif defined(__OpenBSD__) || defined(__FreeBSD__) || defined(__NetBSD__) || defined(__DragonFly__)
	int mib[2];
	size_t len = sizeof(count);
	mib[0] = CTL_HW;
	mib[1] = HW_NCPU;
	if (0 != sysctl(mib, 2, &count, &len, NULL, 0)) {
		count = 0;
	}
#elif defined(__hpux)
	count = mpctl(MPC_GETNUMSPUS, NULL, NULL);
#endif

	return count;
}

size_t roundpow2(size_t i) {
	i--;
	i |= i >> 1;
	i |= i >> 2;
	i |= i >> 4;
	i |= i >> 8;
	i |= i >> 16;
	i++;

	return i;
}

/* Safer version of _PyBytes_Resize().
 *
 * _PyBytes_Resize() only works if the refcount is 1. In some scenarios,
 * we can get an object with a refcount > 1, even if it was just created
 * with PyBytes_FromStringAndSize()! That's because (at least) CPython
 * pre-allocates PyBytes instances of size 1 for every possible byte value.
 *
 * If non-0 is returned, obj may or may not be NULL.
 */
int safe_pybytes_resize(PyObject** obj, Py_ssize_t size) {
	PyObject* tmp;

	if ((*obj)->ob_refcnt == 1) {
		return _PyBytes_Resize(obj, size);
	}

	tmp = PyBytes_FromStringAndSize(NULL, size);
	if (!tmp) {
		return -1;
	}

	memcpy(PyBytes_AS_STRING(tmp), PyBytes_AS_STRING(*obj),
		PyBytes_GET_SIZE(*obj));

	Py_DECREF(*obj);
	*obj = tmp;

	return 0;
}