pylibmc/_pylibmcmodule.c

/**
 * _pylibmc: hand-made libmemcached bindings for Python
 *
 * Copyright (c) 2008, Ludvig Ericson
 * All rights reserved.
 * 
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 * 
 *  - Redistributions of source code must retain the above copyright notice,
 *  this list of conditions and the following disclaimer.
 * 
 *  - Redistributions in binary form must reproduce the above copyright notice,
 *  this list of conditions and the following disclaimer in the documentation
 *  and/or other materials provided with the distribution.
 * 
 *  - Neither the name of the author nor the names of the contributors may be
 *  used to endorse or promote products derived from this software without
 *  specific prior written permission.
 * 
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

#include "_pylibmcmodule.h"
#ifdef USE_ZLIB
#  include <zlib.h>
#  define ZLIB_BUFSZ (1 << 14)
/* only callable while holding the GIL */
#  define _ZLIB_ERR(s, rc) \
  PyErr_Format(PylibMCExc_MemcachedError, "zlib error %d in " s, rc);
#endif


/* {{{ Type methods */
static PylibMC_Client *PylibMC_ClientType_new(PyTypeObject *type,
        PyObject *args, PyObject *kwds) {
    PylibMC_Client *self;

    /* GenericNew calls GenericAlloc (via the indirection type->tp_alloc) which
     * adds GC tracking if flagged for, and also calls PyObject_Init. */
    self = (PylibMC_Client *)PyType_GenericNew(type, args, kwds);

    if (self != NULL) {
        self->mc = memcached_create(NULL);
    }

    return self;
}

static void PylibMC_ClientType_dealloc(PylibMC_Client *self) {
    if (self->mc != NULL) {
        memcached_free(self->mc);
    }

    self->ob_type->tp_free(self);
}
/* }}} */

static int PylibMC_Client_init(PylibMC_Client *self, PyObject *args,
        PyObject *kwds) {
    PyObject *srvs, *srvs_it, *c_srv;
    unsigned char set_stype = 0, bin = 0, got_server = 0;

    static char *kws[] = { "servers", "binary", NULL };

    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|b", kws, &srvs, &bin)) {
        return -1;
    } else if ((srvs_it = PyObject_GetIter(srvs)) == NULL) {
        return -1;
    }

    memcached_behavior_set(self->mc, MEMCACHED_BEHAVIOR_BINARY_PROTOCOL, bin);

    while ((c_srv = PyIter_Next(srvs_it)) != NULL) {
        unsigned char stype;
        char *hostname;
        unsigned short int port;
        memcached_return rc;

        got_server |= 1;
        port = 0;
        if (PyString_Check(c_srv)) {
            memcached_server_st *list;

            list = memcached_servers_parse(PyString_AS_STRING(c_srv));
            if (list == NULL) {
                PyErr_SetString(PylibMCExc_MemcachedError,
                        "memcached_servers_parse returned NULL");
                goto it_error;
            }

            rc = memcached_server_push(self->mc, list);
            free(list);
            if (rc != MEMCACHED_SUCCESS) {
                PylibMC_ErrFromMemcached(self, "memcached_server_push", rc);
                goto it_error;
            }
        } else if (PyArg_ParseTuple(c_srv, "Bs|H", &stype, &hostname, &port)) {
            if (set_stype && set_stype != stype) {
                PyErr_SetString(PyExc_ValueError, "can't mix transport types");
                goto it_error;
            } else {
                set_stype = stype;
                if (stype == PYLIBMC_SERVER_UDP) {
                    memcached_behavior_set(self->mc,
                        MEMCACHED_BEHAVIOR_USE_UDP, 1);
                }
            }

            switch (stype) {
                case PYLIBMC_SERVER_TCP:
                    rc = memcached_server_add(self->mc, hostname, port);
                    break;
                case PYLIBMC_SERVER_UDP:
                    rc = memcached_server_add_udp(self->mc, hostname, port);
                    break;
                case PYLIBMC_SERVER_UNIX:
                    if (port) {
                        PyErr_SetString(PyExc_ValueError,
                                "can't set port on unix sockets");
                        goto it_error;
                    }
                    rc = memcached_server_add_unix_socket(self->mc, hostname);
                    break;
                default:
                    PyErr_Format(PyExc_ValueError, "bad type: %u", stype);
                    goto it_error;
            }
            if (rc != MEMCACHED_SUCCESS) {
                PylibMC_ErrFromMemcached(self, "memcached_server_add_*", rc);
                goto it_error;
            }
        }
        Py_DECREF(c_srv);
        continue;

it_error:
        Py_DECREF(c_srv);
        goto error;
    }

    if (!got_server) {
        PyErr_SetString(PylibMCExc_MemcachedError, "empty server list");
        goto error;
    }

    Py_DECREF(srvs_it);
    return 0;
error:
    Py_DECREF(srvs_it);
    return -1;
}

/* {{{ Compression helpers */
#ifdef USE_ZLIB
static int _PylibMC_Deflate(char* value, size_t value_len,
                    char** result, size_t *result_len) {
    int rc;
    /* todo: failures in here are entirely silent. this should probably
       be fixed */

    z_stream strm;
    *result = NULL;
    *result_len = 0;

    /* Don't ask me about this one. Got it from zlibmodule.c in Python 2.6. */
    size_t out_sz = value_len + value_len / 1000 + 12 + 1;

    if ((*result = malloc(sizeof(char) * out_sz)) == NULL) {
      goto error;
    }

    strm.avail_in = value_len;
    strm.avail_out = out_sz;
    strm.next_in = (Bytef*)value;
    strm.next_out = (Bytef*)*result;

    /* we just pre-allocated all of it up front */
    strm.zalloc = (alloc_func)NULL;
    strm.zfree = (free_func)Z_NULL;

    /* TODO Expose compression level somehow. */
    if (deflateInit((z_streamp)&strm, Z_BEST_SPEED) != Z_OK) {
        goto error;
    }

    Py_BEGIN_ALLOW_THREADS;
    rc = deflate((z_streamp)&strm, Z_FINISH);
    Py_END_ALLOW_THREADS;

    if (rc != Z_STREAM_END) {
        _ZLIB_ERR("deflate", rc);
        goto error;
    }

    if (deflateEnd((z_streamp)&strm) != Z_OK) {
        goto error;
    }

    if(strm.total_out >= value_len) {
      /* if we didn't actually save anything, don't bother storing it
         compressed */
      goto error;
    }

    /* *result should already be populated since that's the address we
       passed into the z_stream */
    *result_len = strm.total_out;

    return 1;
error:
    /* if any error occurred, we'll just use the original value
       instead of trying to compress it */
    if(*result != NULL) {
      free(*result);
      *result = NULL;
    }
    return 0;
}

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 {
        Py_BEGIN_ALLOW_THREADS;
        rc = inflate((z_streamp)&strm, Z_FINISH);
        Py_END_ALLOW_THREADS;

        switch (rc) {
        case Z_STREAM_END:
            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:
            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;
    }

    _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 = NULL;
    PyObject *tmp = NULL;
    uint32_t dtype = flags & PYLIBMC_FLAG_TYPES;

#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

    switch (dtype) {
        case PYLIBMC_FLAG_PICKLE:
            retval = _PylibMC_Unpickle(value, size);
            break;
        case PYLIBMC_FLAG_INTEGER:
        case PYLIBMC_FLAG_LONG:
        case PYLIBMC_FLAG_BOOL:
            /* PyInt_FromString doesn't take a length param and we're
               not NULL-terminated, so we'll have to make an
               intermediate Python string out of it */
            tmp = PyString_FromStringAndSize(value, size);
            if(tmp == NULL) {
              goto cleanup;
            }
            retval = PyInt_FromString(PyString_AS_STRING(tmp), NULL, 10);
            if(retval != NULL && dtype == PYLIBMC_FLAG_BOOL) {
              Py_DECREF(tmp);
              tmp = retval;
              retval = PyBool_FromLong(PyInt_AS_LONG(tmp));
            }
            break;
        case PYLIBMC_FLAG_NONE:
            retval = PyString_FromStringAndSize(value, (Py_ssize_t)size);
            break;
        default:
            PyErr_Format(PylibMCExc_MemcachedError,
                    "unknown memcached key flags %u", flags);
    }

cleanup:

#if USE_ZLIB
    Py_XDECREF(inflated);
#endif

    Py_XDECREF(tmp);

    return retval;
}

static PyObject *PylibMC_Client_get(PylibMC_Client *self, PyObject *arg) {
    char *mc_val;
    size_t val_size;
    uint32_t flags;
    memcached_return error;

    if (!_PylibMC_CheckKey(arg)) {
        return NULL;
    } else if (!PySequence_Length(arg) ) {
        /* Others do this, so... */
        Py_RETURN_NONE;
    }

    Py_BEGIN_ALLOW_THREADS;
    mc_val = memcached_get(self->mc,
            PyString_AS_STRING(arg), PyString_GET_SIZE(arg),
            &val_size, &flags, &error);
    Py_END_ALLOW_THREADS;

    if (mc_val != NULL) {
        PyObject *r = _PylibMC_parse_memcached_value(mc_val, val_size, flags);
        free(mc_val);
        return r;
    } else if (error == MEMCACHED_SUCCESS) {
        /* This happens for empty values, and so we fake an empty string. */
        return PyString_FromStringAndSize("", 0);
    } else if (error == MEMCACHED_NOTFOUND) {
        /* Since python-memcache returns None when the key doesn't exist,
         * so shall we. */
        Py_RETURN_NONE;
    }

    return PylibMC_ErrFromMemcached(self, "memcached_get", error);
}

static PyObject *PylibMC_Client_gets(PylibMC_Client *self, PyObject *arg) {
    const char* keys[2];
    size_t keylengths[2];
    memcached_result_st results_obj;
    memcached_result_st *results = NULL;
    memcached_return_t rc;
    PyObject* ret = NULL;

    if (!_PylibMC_CheckKey(arg)) {
        return NULL;
    } else if (!PySequence_Length(arg)) {
        /* Others do this, so... */
        Py_RETURN_NONE;
    } else if (!memcached_behavior_get(self->mc, MEMCACHED_BEHAVIOR_SUPPORT_CAS)) {
        PyErr_SetString(PyExc_ValueError, "gets without cas behavior");
        return NULL;
    }

    /* Use an mget to fetch the key.
     * mget is the only function that returns a memcached_result_st,
     * which is the only way to get at the returned cas value. */
    *keys = PyString_AS_STRING(arg);
    *keylengths = (size_t)PyString_GET_SIZE(arg);

    Py_BEGIN_ALLOW_THREADS;

    rc = memcached_mget(self->mc, keys, keylengths, 1);
    if (rc == MEMCACHED_SUCCESS) {
        memcached_result_create(self->mc, &results_obj);
        /* this will be NULL if the key wasn't found, or
           memcached_result_st if it was */
        results = memcached_fetch_result(self->mc, &results_obj, &rc);
    }

    Py_END_ALLOW_THREADS;

    if (rc == MEMCACHED_SUCCESS && results != NULL) {
        const char *mc_val = memcached_result_value(results);
        size_t val_size = memcached_result_length(results);
        uint32_t flags = memcached_result_flags(results);
        uint64_t cas = memcached_result_cas(results);

        ret = _PylibMC_parse_memcached_value((char *)mc_val, val_size, flags);
        ret = Py_BuildValue("(NL)", ret, cas);

        /* we have to fetch the last result from the mget cursor */
        memcached_fetch_result(self->mc, &results_obj, &rc);
        if (rc != MEMCACHED_END) {
            Py_DECREF(ret);
            ret = NULL;
            PyErr_SetString(PyExc_RuntimeError, "fetch not done");
        }
    } else if (rc == MEMCACHED_END) {
        /* Key not found => (None, None) */
        ret = Py_BuildValue("(OO)", Py_None, Py_None);
    } else {
        ret = PylibMC_ErrFromMemcached(self, "memcached_gets", rc);
    }

    /* deallocate any indirect buffers, even though the struct itself
       is on our stack */
    memcached_result_free(&results_obj);

    return ret;
}

/* {{{ Set commands (set, replace, add, prepend, append) */
static PyObject *_PylibMC_RunSetCommandSingle(PylibMC_Client *self,
        _PylibMC_SetCommand f, char *fname, PyObject *args,
        PyObject *kwds) {
    /* function called by the set/add/etc commands */
    static char *kws[] = { "key", "val", "time", "min_compress_len", "custom_flag", NULL };
    PyObject *key;
    PyObject *value;
    unsigned int time = 0; /* this will be turned into a time_t */
    unsigned int min_compress = 0;
    unsigned int custom_flag = 0;
    bool success = false;

    if (!PyArg_ParseTupleAndKeywords(args, kwds, "SO|III", kws,
                                     &key, &value, &time,
                                     &min_compress, &custom_flag)) {
      return NULL;
    }

#ifndef USE_ZLIB
    if (min_compress) {
      PyErr_SetString(PyExc_TypeError, "min_compress_len without zlib");
      return NULL;
    }
#endif

    pylibmc_mset serialized = { NULL };

    success = _PylibMC_SerializeValue(key, NULL, value, time, &serialized);

    if (!success)
        goto cleanup;

    success = _PylibMC_RunSetCommand(self, f, fname,
                                     &serialized, 1,
                                     min_compress,
                                     custom_flag);

cleanup:
    _PylibMC_FreeMset(&serialized);

    if(PyErr_Occurred() != NULL) {
        return NULL;
    } else if(success) {
        Py_RETURN_TRUE;
    } else {
        Py_RETURN_FALSE;
    }
}

static PyObject *_PylibMC_RunSetCommandMulti(PylibMC_Client* self,
        _PylibMC_SetCommand f, char *fname, PyObject* args,
        PyObject* kwds) {
    /* function called by the set/add/incr/etc commands */
    PyObject* keys = NULL;
    PyObject* key_prefix = NULL;
    unsigned int time = 0;
    unsigned int min_compress = 0;
    unsigned int custom_flag = 0;
    PyObject * retval = NULL;
    size_t idx = 0;

    static char *kws[] = { "keys", "time", "key_prefix", "min_compress_len", "custom_flag", NULL };

    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O!|ISII", kws,
                                     &PyDict_Type, &keys,
                                     &time, &key_prefix,
                                     &min_compress, &custom_flag)) {
        return NULL;
    }

#ifndef USE_ZLIB
    if (min_compress) {
        PyErr_SetString(PyExc_TypeError, "min_compress_len without zlib");
        return NULL;
    }
#endif

    PyObject *curr_key, *curr_value;
    size_t nkeys = (size_t)PyDict_Size(keys);

    pylibmc_mset* serialized = PyMem_New(pylibmc_mset, nkeys);
    if (serialized == NULL) {
        goto cleanup;
    }

    /**
     * We're pointing into existing Python memory with the 'key' members of
     * pylibmc_mset (extracted using PyDict_Next) and during
     * _PylibMC_RunSetCommand (which uses those same 'key' params, and
     * potentially points into value string objects too), so we don't want to
     * go around decrementing any references that risk destroying the pointed
     * objects until we're done, especially since we're going to release the
     * GIL while we do the I/O that accesses that memory. We're assuming that
     * this is safe because Python strings are immutable
     */

    Py_ssize_t pos = 0; /* PyDict_Next's 'pos' isn't an incrementing index */

    for (idx = 0; PyDict_Next(keys, &pos, &curr_key, &curr_value); idx++) {
        int success = _PylibMC_SerializeValue(curr_key, key_prefix,
                                              curr_value, time,
                                              &serialized[idx]);

        if (!success || PyErr_Occurred() != NULL) {
            /* exception should already be on the stack */
            goto cleanup;
        }
    }

    if (PyErr_Occurred() != NULL) {
        /* an iteration error of some sort */
        goto cleanup;
    }

    bool allsuccess = _PylibMC_RunSetCommand(self, f, fname,
                                             serialized, nkeys,
                                             min_compress, custom_flag);

    if (PyErr_Occurred() != NULL) {
        goto cleanup;
    }

    /* Return value for set_multi, which is a list
       of keys which failed to be set */
    if ((retval = PyList_New(0)) == NULL)
        return PyErr_NoMemory();

    for (idx = 0; !allsuccess && idx < nkeys; idx++) {
        if (serialized[idx].success)
            continue;

        if (PyList_Append(retval, serialized[idx].key_obj) != 0) {
            /* Ugh */
            Py_DECREF(retval);
            retval = PyErr_NoMemory();
            goto cleanup;
        }
    }

cleanup:
    if (serialized != NULL) {
        for (pos = 0; pos < nkeys; pos++) {
            _PylibMC_FreeMset(&serialized[pos]);
        }
        PyMem_Free(serialized);
    }

    return retval;
}

static PyObject *_PylibMC_RunCasCommand(PylibMC_Client *self,
        PyObject *args, PyObject *kwds) {
    /* function called by the set/add/etc commands */
    static char *kws[] = { "key", "val", "cas", "time", NULL };
    PyObject *ret = NULL;
    PyObject *key;
    PyObject *value;
    uint64_t cas = 0;
    unsigned int time = 0; /* this will be turned into a time_t */
    bool success = false;
    memcached_return rc;

    if (!PyArg_ParseTupleAndKeywords(args, kwds, "SOL|I", kws,
                                     &key, &value, &cas,
                                     &time)) {
      return NULL;
    }

    if (!memcached_behavior_get(self->mc, MEMCACHED_BEHAVIOR_SUPPORT_CAS)) {
      PyErr_SetString(PyExc_ValueError, "cas without cas behavior");
      return NULL;
    }

    pylibmc_mset mset = { NULL };

    /* TODO: because it's RunSetCommand that does the zlib
       compression, cas can't currently use compressed values. */
    success = _PylibMC_SerializeValue(key, NULL, value, time, &mset);

    if (!success || PyErr_Occurred() != NULL) {
        goto cleanup;
    }

    Py_BEGIN_ALLOW_THREADS;
    rc = memcached_cas(self->mc,
                       mset.key, mset.key_len,
                       mset.value, mset.value_len,
                       mset.time, mset.flags, cas);
    Py_END_ALLOW_THREADS;

    switch(rc) {
        case MEMCACHED_SUCCESS:
            Py_INCREF(Py_True);
            ret = Py_True;
            break;
        case MEMCACHED_DATA_EXISTS:
            Py_INCREF(Py_False);
            ret = Py_False;
            break;
        default:
            PylibMC_ErrFromMemcached(self, "memcached_cas", rc);
    }

cleanup:
    _PylibMC_FreeMset(&mset);

    return ret;
}

static void _PylibMC_FreeMset(pylibmc_mset *mset) {
      Py_XDECREF(mset->key_obj);
      mset->key_obj = NULL;

      Py_XDECREF(mset->prefixed_key_obj);
      mset->prefixed_key_obj = NULL;

      /* Either a ref we own, or a ref passed to us which we borrowed. */
      Py_XDECREF(mset->value_obj);
      mset->value_obj = NULL;
}

static int _PylibMC_SerializeValue(PyObject* key_obj,
                                   PyObject* key_prefix,
                                   PyObject* value_obj,
                                   time_t time,
                                   pylibmc_mset* serialized) {
    /* first zero the whole structure out */
    memset((void *)serialized, 0x0, sizeof(pylibmc_mset));

    serialized->time = time;
    serialized->success = false;
    serialized->flags = PYLIBMC_FLAG_NONE;

    if(!_PylibMC_CheckKey(key_obj)
       || PyString_AsStringAndSize(key_obj, &serialized->key,
                                   &serialized->key_len) == -1) {
        return false;
    }

    /* We need to incr our reference here so that it's guaranteed to
       exist while we release the GIL. Even if we fail after this it
       should be decremeneted by pylib_mset_free */
    Py_INCREF(key_obj);
    serialized->key_obj = key_obj;

    /* Check the key_prefix */
    if (key_prefix != NULL) {
        if (!_PylibMC_CheckKey(key_prefix)) {
            return false;
        }

        /* Ignore empty prefixes */
        if (!PyString_Size(key_prefix)) {
            key_prefix = NULL;
        }
    }

    /* Make the prefixed key if appropriate */
    if (key_prefix != NULL) {
        PyObject* prefixed_key_obj = NULL; /* freed by _PylibMC_FreeMset */

        prefixed_key_obj = PyString_FromFormat("%s%s",
                PyString_AS_STRING(key_prefix),
                PyString_AS_STRING(key_obj));

        if(prefixed_key_obj == NULL) {
            return false;
        }

        /* check the key and overwrite the C string */
        if(!_PylibMC_CheckKey(prefixed_key_obj)
           || PyString_AsStringAndSize(prefixed_key_obj,
                                       &serialized->key,
                                       &serialized->key_len) == -1) {
            Py_DECREF(prefixed_key_obj);
            return false;
        }

        serialized->prefixed_key_obj = prefixed_key_obj;
    }

    /* Key/key_size should be harmonized, now onto the value */

    PyObject* store_val = NULL;

    /* First build store_val, a Python String object, out of the object
       we were passed */
    if (PyString_Check(value_obj)) {
        store_val = value_obj;
        Py_INCREF(store_val); /* because we'll be decring it again in
                                 pylibmc_mset_free*/
    } else if (PyBool_Check(value_obj)) {
        serialized->flags |= PYLIBMC_FLAG_BOOL;
        PyObject* tmp = PyNumber_Int(value_obj);
        store_val = PyObject_Str(tmp);
        Py_DECREF(tmp);
    } else if (PyInt_Check(value_obj)) {
        serialized->flags |= PYLIBMC_FLAG_INTEGER;
        PyObject* tmp = PyNumber_Int(value_obj);
        store_val = PyObject_Str(tmp);
        Py_DECREF(tmp);
    } else if (PyLong_Check(value_obj)) {
        serialized->flags |= PYLIBMC_FLAG_LONG;
        PyObject* tmp = PyNumber_Long(value_obj);
        store_val = PyObject_Str(tmp);
        Py_DECREF(tmp);
    } else if(value_obj != NULL) {
        /* we have no idea what it is, so we'll store it pickled */
        Py_INCREF(value_obj);
        serialized->flags |= PYLIBMC_FLAG_PICKLE;
        store_val = _PylibMC_Pickle(value_obj);
        Py_DECREF(value_obj);
    }

    if (store_val == NULL) {
        return false;
    }

    if (PyString_AsStringAndSize(store_val, &serialized->value,
                                 &serialized->value_len) == -1) {
        if (serialized->flags == PYLIBMC_FLAG_NONE) {
            /* For some reason we weren't able to extract the value/size
               from a string that we were explicitly passed, that we
               INCREF'd above */
            Py_DECREF(store_val);
        }
        return false;
    }

    /* So now we have a reference to a string that we may have
       created. we need that to keep existing while we release the GIL,
       so we need to hold the reference, but we need to free it up when
       we're done */
    serialized->value_obj = store_val;

    return true;
}

/* {{{ Set commands (set, replace, add, prepend, append) */
static bool _PylibMC_RunSetCommand(PylibMC_Client* self,
                                   _PylibMC_SetCommand f, char *fname,
                                   pylibmc_mset* msets, size_t nkeys,
                                   size_t min_compress, size_t custom_flag) {
    memcached_st* mc = self->mc;
    memcached_return rc = MEMCACHED_SUCCESS;
    int pos;
    bool error = false;
    bool allsuccess = true;

    Py_BEGIN_ALLOW_THREADS;

    for (pos=0; pos < nkeys && !error; pos++) {
        pylibmc_mset* mset = &msets[pos];

        char* value = mset->value;
        size_t value_len = mset->value_len;
        uint32_t flags = mset->flags;
        
        if (custom_flag) {
            flags |= (1 << custom_flag);
        }

#ifdef USE_ZLIB
        char* compressed_value = NULL;
        size_t compressed_len = 0;

        if (min_compress && value_len >= min_compress) {
            Py_BLOCK_THREADS;
            _PylibMC_Deflate(value, value_len,
                             &compressed_value, &compressed_len);
            Py_UNBLOCK_THREADS;
        }

        if (compressed_value != NULL) {
            /* Will want to change this if this function 
             * needs to get back at the old *value at some point */
            value = compressed_value;