Implement zlib compression.

This commit is contained in:
lericson 2009-11-30 09:43:05 +01:00
parent 9fec8f59f6
commit c5590258fb
5 changed files with 220 additions and 10 deletions

View File

@ -32,6 +32,12 @@
*/
#include "_pylibmcmodule.h"
#ifdef USE_ZLIB
# include <zlib.h>
# define ZLIB_BUFSZ (1 << 14)
# define _ZLIB_ERR(s, rc) \
PyErr_Format(PylibMCExc_MemcachedError, "zlib error %d in " s, rc);
#endif
/* {{{ _pylibmc.client implementation */
@ -155,12 +161,160 @@ error:
return -1;
}
/* {{{ Compression helpers */
#ifdef USE_ZLIB
static PyObject *_PylibMC_Deflate(PyObject *value) {
int rc;
char *out;
z_stream strm;
Py_ssize_t length, out_sz;
if (!PyString_Check(value)) {
return NULL;
}
/* Don't ask me about this one. Got it from zlibmodule.c in Python 2.6. */
length = PyString_GET_SIZE(value);
out_sz = length + length / 1000 + 12 + 1;
if ((out = PyMem_New(char, out_sz)) == NULL) {
return NULL;
}
strm.avail_in = length;
strm.avail_out = out_sz;
strm.next_in = (Byte *)PyString_AS_STRING(value);
strm.next_out = (Byte *)out;
strm.zalloc = (alloc_func)NULL;
strm.zfree = (free_func)Z_NULL;
/* TODO Expose compression level somehow. */
if ((rc = deflateInit((z_streamp)&strm, Z_BEST_SPEED)) != Z_OK) {
_ZLIB_ERR("deflateInit", rc);
goto error;
}
if ((rc = deflate((z_streamp)&strm, Z_FINISH)) != Z_STREAM_END) {
_ZLIB_ERR("deflate", rc);
goto error;
}
if ((rc = deflateEnd((z_streamp)&strm)) != Z_OK) {
_ZLIB_ERR("deflateEnd", rc);
goto error;
}
return PyString_FromStringAndSize(out, strm.total_out);
error:
PyMem_Del(out);
return NULL;
}
static PyObject *_PylibMC_Inflate(char *value, size_t size) {
int rc;
char *out;
PyObject *out_obj;
Py_ssize_t rvalsz;
z_stream strm;
/* Output buffer */
rvalsz = ZLIB_BUFSZ;
out_obj = PyString_FromStringAndSize(NULL, rvalsz);
if (out_obj == NULL) {
return NULL;
}
out = PyString_AS_STRING(out_obj);
/* Set up zlib stream. */
strm.avail_in = size;
strm.avail_out = (uInt)rvalsz;
strm.next_in = (Byte *)value;
strm.next_out = (Byte *)out;
strm.zalloc = (alloc_func)NULL;
strm.zfree = (free_func)Z_NULL;
/* TODO Add controlling of windowBits with inflateInit2? */
if ((rc = inflateInit((z_streamp)&strm)) != Z_OK) {
_ZLIB_ERR("inflateInit", rc);
goto error;
}
do {
rc = inflate((z_streamp)&strm, Z_FINISH);
printf("(loop) rc = %d\n", rc);
switch (rc) {
case Z_STREAM_END:
printf("Z_STREAM_END\n");
break;
/* When a Z_BUF_ERROR occurs, we should be out of memory.
* This is also true for Z_OK, hence the fall-through. */
case Z_BUF_ERROR:
if (strm.avail_out) {
_ZLIB_ERR("inflate", rc);
inflateEnd(&strm);
goto error;
}
/* Fall-through */
case Z_OK:
printf("Z_OK\n");
if (_PyString_Resize(&out_obj, (Py_ssize_t)(rvalsz << 1)) < 0) {
inflateEnd(&strm);
goto error;
}
/* Wind forward */
out = PyString_AS_STRING(out_obj);
strm.next_out = (Byte *)(out + rvalsz);
strm.avail_out = rvalsz;
rvalsz = rvalsz << 1;
break;
default:
inflateEnd(&strm);
goto error;
}
} while (rc != Z_STREAM_END);
if ((rc = inflateEnd(&strm)) != Z_OK) {
_ZLIB_ERR("inflateEnd", rc);
goto error;
}
printf("size %d -> size %d\n", (int)size, (int)strm.total_out);
_PyString_Resize(&out_obj, strm.total_out);
return out_obj;
error:
Py_DECREF(out_obj);
return NULL;
}
#endif
static PyObject *_PylibMC_parse_memcached_value(char *value, size_t size,
uint32_t flags) {
PyObject *retval, *tmp;
#if USE_ZLIB
PyObject *inflated = NULL;
/* Decompress value if necessary. */
if (flags & PYLIBMC_FLAG_ZLIB) {
inflated = _PylibMC_Inflate(value, size);
value = PyString_AS_STRING(inflated);
size = PyString_GET_SIZE(inflated);
}
#else
if (flags & PYLIBMC_FLAG_ZLIB) {
PyErr_SetString(PylibMCExc_MemcachedError,
"value for key compressed, unable to inflate");
return NULL;
}
#endif
retval = NULL;
switch (flags) {
switch (flags & PYLIBMC_FLAG_TYPES) {
case PYLIBMC_FLAG_PICKLE:
retval = _PylibMC_Unpickle(value, size);
break;
@ -183,6 +337,12 @@ static PyObject *_PylibMC_parse_memcached_value(char *value, size_t size,
"unknown memcached key flags %u", flags);
}
#if USE_ZLIB
if (inflated != NULL) {
Py_DECREF(inflated);
}
#endif
return retval;
}
@ -229,18 +389,26 @@ static PyObject *_PylibMC_RunSetCommand(PylibMC_Client *self,
PyObject *retval = NULL;
PyObject *store_val = NULL;
unsigned int time = 0;
unsigned int min_compress = 0;
uint32_t store_flags = 0;
static char *kws[] = { "key", "val", "time", NULL };
static char *kws[] = { "key", "val", "time", "min_compress_len", NULL };
if (!PyArg_ParseTupleAndKeywords(args, kwds, "s#O|I", kws,
&key, &key_sz, &val, &time)) {
if (!PyArg_ParseTupleAndKeywords(args, kwds, "s#O|II", kws,
&key, &key_sz, &val, &time, &min_compress)) {
return NULL;
}
if (!_PylibMC_CheckKeyStringAndSize(key, key_sz)) {
return NULL;
}
#ifndef USE_ZLIB
if (min_compress) {
PyErr_SetString(PyExc_TypeError, "min_compress_len without zlib");
return NULL;
}
#endif
/* Adapt val to a str. */
if (PyString_Check(val)) {
store_val = val;
@ -270,6 +438,22 @@ static PyObject *_PylibMC_RunSetCommand(PylibMC_Client *self,
return NULL;
}
#ifdef USE_ZLIB
if (min_compress && PyString_GET_SIZE(store_val) > min_compress) {
PyObject *deflated = _PylibMC_Deflate(val);
if (deflated != NULL) {
Py_DECREF(store_val);
store_val = deflated;
store_flags |= PYLIBMC_FLAG_ZLIB;
} else {
/* XXX What to do here, really? */
PyErr_Print();
PyErr_Clear();
}
}
#endif
rc = f(self->mc, key, key_sz,
PyString_AS_STRING(store_val), PyString_GET_SIZE(store_val),
time, store_flags);
@ -1029,6 +1213,14 @@ Oh, and: plankton.\n");
PyModule_AddStringConstant(module, "__version__", PYLIBMC_VERSION);
#ifdef USE_ZLIB
Py_INCREF(Py_True);
PyModule_AddObject(module, "support_compression", Py_True);
#else
Py_INCREF(Py_False);
PyModule_AddObject(module, "support_compression", Py_False);
#endif
PylibMCExc_MemcachedError = PyErr_NewException(
"_pylibmc.MemcachedError", NULL, NULL);
PyModule_AddObject(module, "MemcachedError",

View File

@ -53,12 +53,20 @@ typedef ssize_t Py_ssize_t;
#define PYLIBMC_SERVER_UDP (1 << 1)
#define PYLIBMC_SERVER_UNIX (1 << 2)
/* Key flags from python-memcache. */
/* {{{ Key flags from python-memcached
* Some flags (like the compression one, ZLIB) are combined with others.
*/
#define PYLIBMC_FLAG_NONE 0
#define PYLIBMC_FLAG_PICKLE (1 << 0)
#define PYLIBMC_FLAG_INTEGER (1 << 1)
#define PYLIBMC_FLAG_LONG (1 << 2)
#define PYLIBMC_FLAG_BOOL (1 << 3)
/* Note: this is an addition! python-memcached doesn't handle bools. */
#define PYLIBMC_FLAG_BOOL (1 << 4)
#define PYLIBMC_FLAG_TYPES (PYLIBMC_FLAG_PICKLE | PYLIBMC_FLAG_INTEGER | \
PYLIBMC_FLAG_LONG | PYLIBMC_FLAG_BOOL)
/* Modifier flags */
#define PYLIBMC_FLAG_ZLIB (1 << 3)
/* }}} */
#define PYLIBMC_INC (1 << 0)
#define PYLIBMC_DEC (1 << 1)

View File

@ -26,6 +26,7 @@ import _pylibmc
__all__ = ["hashers", "distributions", "Client"]
__version__ = _pylibmc.__version__
support_compression = _pylibmc.support_compression
errors = tuple(e for (n, e) in _pylibmc.exceptions)
# *Cough* Uhm, not the prettiest of things but this unpacks all exception

View File

@ -1,6 +1,10 @@
import os
from distutils.core import setup, Extension
use_zlib = True # TODO Configurable somehow?
libs = ["memcached"]
defs = []
incdirs = []
libdirs = []
@ -9,12 +13,16 @@ if "LIBMEMCACHED_DIR" in os.environ:
incdirs.append(os.path.join(libdir, "include"))
libdirs.append(os.path.join(libdir, "lib"))
readme_text = open("README.rst", "U").read()
version = open("pylibmc-version.h", "U").read().strip().split("\"")[1]
if use_zlib:
libs.append("z")
defs.append(("USE_ZLIB", None))
pylibmc_ext = Extension("_pylibmc", ["_pylibmcmodule.c"],
libraries=["memcached"],
include_dirs=incdirs, library_dirs=libdirs)
libraries=libs, include_dirs=incdirs,
library_dirs=libdirs, define_macros=defs)
readme_text = open("README.rst", "U").read()
version = open("pylibmc-version.h", "U").read().strip().split("\"")[1]
setup(name="pylibmc", version=version,
url="http://lericson.blogg.se/code/category/pylibmc.html",

View File

@ -342,6 +342,7 @@ if __name__ == "__main__":
print "Starting tests with _pylibmc at", _pylibmc.__file__
print "Reported libmemcached version:", _pylibmc.libmemcached_version
print "Reported pylibmc version:", _pylibmc.__version__
print "Support compression:", _pylibmc.support_compression
if not is_alive(test_server):
raise SystemExit("Test server (%r) not alive." % (test_server,))