diff --git a/_pylibmcmodule.c b/_pylibmcmodule.c index 85d1571..57cec50 100644 --- a/_pylibmcmodule.c +++ b/_pylibmcmodule.c @@ -32,6 +32,12 @@ */ #include "_pylibmcmodule.h" +#ifdef USE_ZLIB +# include +# 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", diff --git a/_pylibmcmodule.h b/_pylibmcmodule.h index 252700c..590f3fd 100644 --- a/_pylibmcmodule.h +++ b/_pylibmcmodule.h @@ -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) diff --git a/pylibmc.py b/pylibmc.py index e5113cb..f98d1f3 100644 --- a/pylibmc.py +++ b/pylibmc.py @@ -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 diff --git a/setup.py b/setup.py index 638968e..3bea5ff 100644 --- a/setup.py +++ b/setup.py @@ -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", diff --git a/tests.py b/tests.py index 0f11b70..0465957 100644 --- a/tests.py +++ b/tests.py @@ -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,))