Rework flow of multi_get.
This should fix some corner-case bugs, and makes the code read a lot better.
This commit is contained in:
lericson
parent
c3d46783d3
commit
c023d18a20
233
_pylibmcmodule.c
233
_pylibmcmodule.c
@ -352,132 +352,143 @@ static PyObject *PylibMC_Client_get_multi(PylibMC_Client *self, PyObject *args,
|
|||||||
PyObject *key_seq, **key_objs, *retval = NULL;
|
PyObject *key_seq, **key_objs, *retval = NULL;
|
||||||
char **keys, *prefix = NULL;
|
char **keys, *prefix = NULL;
|
||||||
unsigned int prefix_len = 0;
|
unsigned int prefix_len = 0;
|
||||||
|
Py_ssize_t i;
|
||||||
|
PyObject *key_it, *ckey;
|
||||||
size_t *key_lens;
|
size_t *key_lens;
|
||||||
size_t nkeys;
|
size_t nkeys;
|
||||||
unsigned int valid_keys_len = 0;
|
|
||||||
memcached_return rc;
|
memcached_return rc;
|
||||||
|
|
||||||
char curr_key[MEMCACHED_MAX_KEY];
|
char curr_key[MEMCACHED_MAX_KEY];
|
||||||
size_t curr_key_len, curr_val_len;
|
size_t curr_key_len, curr_val_len;
|
||||||
uint32_t curr_flags;
|
uint32_t curr_flags;
|
||||||
|
char *curr_val;
|
||||||
|
|
||||||
static char *kws[] = { "keys", "key_prefix", NULL };
|
static char *kws[] = { "keys", "key_prefix", NULL };
|
||||||
|
|
||||||
if (PyArg_ParseTupleAndKeywords(args, kwds, "O|s#", kws,
|
if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|s#", kws,
|
||||||
&key_seq, &prefix, &prefix_len)) {
|
&key_seq, &prefix, &prefix_len)) {
|
||||||
PyObject *key_it, *ckey;
|
return NULL;
|
||||||
Py_ssize_t i;
|
|
||||||
|
|
||||||
/* First clear any potential earlier mishap because we rely on it in
|
|
||||||
* our iteration over keys. */
|
|
||||||
PyErr_Clear();
|
|
||||||
/* Populate keys and key_lens. */
|
|
||||||
nkeys = PySequence_Length(key_seq);
|
|
||||||
keys = malloc(sizeof(char *) * nkeys);
|
|
||||||
key_lens = malloc(sizeof(size_t) * nkeys);
|
|
||||||
key_objs = malloc(sizeof(PyObject *) * nkeys);
|
|
||||||
if (keys == NULL || key_lens == NULL || key_objs == NULL) {
|
|
||||||
free(keys);
|
|
||||||
free(key_lens);
|
|
||||||
free(key_objs);
|
|
||||||
return PyErr_NoMemory();
|
|
||||||
}
|
|
||||||
/* Iterate through all keys and set lengths etc. */
|
|
||||||
i = 0;
|
|
||||||
key_it = PyObject_GetIter(key_seq);
|
|
||||||
while (!PyErr_Occurred()
|
|
||||||
&& i < nkeys
|
|
||||||
&& (ckey = PyIter_Next(key_it)) != NULL) {
|
|
||||||
PyObject *rkey;
|
|
||||||
|
|
||||||
if (!_PylibMC_CheckKey(ckey)) {
|
|
||||||
break;
|
|
||||||
} else {
|
|
||||||
key_lens[i] = PyString_GET_SIZE(ckey) + prefix_len;
|
|
||||||
if (prefix != NULL) {
|
|
||||||
rkey = PyString_FromFormat("%s%s",
|
|
||||||
prefix, PyString_AS_STRING(ckey));
|
|
||||||
Py_DECREF(ckey);
|
|
||||||
} else {
|
|
||||||
rkey = ckey;
|
|
||||||
}
|
|
||||||
keys[i] = PyString_AS_STRING(rkey);
|
|
||||||
key_objs[i++] = rkey;
|
|
||||||
valid_keys_len++;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Py_DECREF(key_it);
|
|
||||||
|
|
||||||
if (PyErr_Occurred()) {
|
|
||||||
free(key_lens);
|
|
||||||
free(keys);
|
|
||||||
nkeys = i;
|
|
||||||
for (i = 0; i < nkeys; i++)
|
|
||||||
Py_DECREF(key_objs[i]);
|
|
||||||
free(key_objs);
|
|
||||||
return NULL;
|
|
||||||
}
|
|
||||||
|
|
||||||
/* TODO Change this interface or at least provide an alternative that
|
|
||||||
* returns some kind of iterable which fetches keys sequentially
|
|
||||||
* instead of doing it all at once. The current way is grossly
|
|
||||||
* inefficient for larger datasets, as a dict has to be allocated that
|
|
||||||
* is large enough to hold all the data at once.
|
|
||||||
*/
|
|
||||||
retval = PyDict_New();
|
|
||||||
|
|
||||||
if(valid_keys_len > 0) {
|
|
||||||
if ((rc = memcached_mget(self->mc, keys, key_lens, nkeys))
|
|
||||||
== MEMCACHED_SUCCESS) {
|
|
||||||
char *curr_val;
|
|
||||||
|
|
||||||
while ((curr_val = memcached_fetch(
|
|
||||||
self->mc, curr_key, &curr_key_len, &curr_val_len,
|
|
||||||
&curr_flags, &rc)) != NULL
|
|
||||||
&& !PyErr_Occurred()) {
|
|
||||||
if (curr_val == NULL && rc == MEMCACHED_END) {
|
|
||||||
break;
|
|
||||||
} else if (rc == MEMCACHED_BAD_KEY_PROVIDED
|
|
||||||
|| rc == MEMCACHED_NO_KEY_PROVIDED) {
|
|
||||||
/* Do nothing at all. :-) */
|
|
||||||
} else if (rc != MEMCACHED_SUCCESS) {
|
|
||||||
Py_DECREF(retval);
|
|
||||||
retval = PylibMC_ErrFromMemcached(
|
|
||||||
self, "memcached_fetch", rc);
|
|
||||||
} else {
|
|
||||||
PyObject *val;
|
|
||||||
|
|
||||||
/* This is safe because libmemcached's max key length
|
|
||||||
* includes space for a NUL-byte. */
|
|
||||||
curr_key[curr_key_len] = 0;
|
|
||||||
val = _PylibMC_parse_memcached_value(
|
|
||||||
curr_val, curr_val_len, curr_flags);
|
|
||||||
PyDict_SetItemString(retval, curr_key + prefix_len, val);
|
|
||||||
Py_DECREF(val);
|
|
||||||
}
|
|
||||||
free(curr_val);
|
|
||||||
}
|
|
||||||
/* Need to cleanup. */
|
|
||||||
if (PyErr_Occurred()) {
|
|
||||||
/* Not checking rc because an exception already occured, and
|
|
||||||
* we wouldn't want to mask it. */
|
|
||||||
memcached_quit(self->mc);
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
retval = PylibMC_ErrFromMemcached(self, "memcached_mget", rc);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
free(key_lens);
|
|
||||||
free(keys);
|
|
||||||
for (i = 0; i < nkeys; i++)
|
|
||||||
Py_DECREF(key_objs[i]);
|
|
||||||
free(key_objs);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if ((nkeys = PySequence_Length(key_seq)) == -1) {
|
||||||
|
return NULL;
|
||||||
|
}
|
||||||
|
|
||||||
|
/* Populate keys and key_lens. */
|
||||||
|
keys = PyMem_New(char *, nkeys);
|
||||||
|
key_lens = PyMem_New(size_t, nkeys);
|
||||||
|
key_objs = PyMem_New(PyObject *, nkeys);
|
||||||
|
if (keys == NULL || key_lens == NULL || key_objs == NULL) {
|
||||||
|
PyMem_Free(keys);
|
||||||
|
PyMem_Free(key_lens);
|
||||||
|
PyMem_Free(key_objs);
|
||||||
|
return PyErr_NoMemory();
|
||||||
|
}
|
||||||
|
|
||||||
|
/* Clear potential previous exception, because we explicitly check for
|
||||||
|
* exceptions as a loop predicate. */
|
||||||
|
PyErr_Clear();
|
||||||
|
|
||||||
|
/* Iterate through all keys and set lengths etc. */
|
||||||
|
i = 0;
|
||||||
|
key_it = PyObject_GetIter(key_seq);
|
||||||
|
while (!PyErr_Occurred()
|
||||||
|
&& i < nkeys
|
||||||
|
&& (ckey = PyIter_Next(key_it)) != NULL) {
|
||||||
|
PyObject *rkey;
|
||||||
|
|
||||||
|
if (!_PylibMC_CheckKey(ckey)) {
|
||||||
|
break;
|
||||||
|
} else {
|
||||||
|
key_lens[i] = PyString_GET_SIZE(ckey) + prefix_len;
|
||||||
|
if (prefix != NULL) {
|
||||||
|
rkey = PyString_FromFormat("%s%s",
|
||||||
|
prefix, PyString_AS_STRING(ckey));
|
||||||
|
Py_DECREF(ckey);
|
||||||
|
} else {
|
||||||
|
rkey = ckey;
|
||||||
|
}
|
||||||
|
keys[i] = PyString_AS_STRING(rkey);
|
||||||
|
key_objs[i++] = rkey;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Py_XDECREF(key_it);
|
||||||
|
|
||||||
|
if (i == 0) {
|
||||||
|
/* No usable keys to fetch. */
|
||||||
|
nkeys = 0;
|
||||||
|
goto cleanup;
|
||||||
|
} else if (PyErr_Occurred()) {
|
||||||
|
nkeys--;
|
||||||
|
goto cleanup;
|
||||||
|
}
|
||||||
|
|
||||||
|
/* TODO Make an iterator interface for getting each key separately.
|
||||||
|
*
|
||||||
|
* This would help large retrievals, as a single dictionary containing all
|
||||||
|
* the data at once isn't needed. (Should probably look into if it's even
|
||||||
|
* worth it.)
|
||||||
|
*/
|
||||||
|
retval = PyDict_New();
|
||||||
|
|
||||||
|
if ((rc = memcached_mget(self->mc, keys, key_lens, nkeys))
|
||||||
|
!= MEMCACHED_SUCCESS) {
|
||||||
|
PylibMC_ErrFromMemcached(self, "memcached_mget", rc);
|
||||||
|
goto cleanup;
|
||||||
|
}
|
||||||
|
|
||||||
|
while ((curr_val = memcached_fetch(
|
||||||
|
self->mc, curr_key, &curr_key_len, &curr_val_len,
|
||||||
|
&curr_flags, &rc)) != NULL
|
||||||
|
&& !PyErr_Occurred()) {
|
||||||
|
if (curr_val == NULL && rc == MEMCACHED_END) {
|
||||||
|
break;
|
||||||
|
} else if (rc == MEMCACHED_BAD_KEY_PROVIDED
|
||||||
|
|| rc == MEMCACHED_NO_KEY_PROVIDED) {
|
||||||
|
/* Do nothing at all. :-) */
|
||||||
|
} else if (rc != MEMCACHED_SUCCESS) {
|
||||||
|
PylibMC_ErrFromMemcached(self, "memcached_fetch", rc);
|
||||||
|
memcached_quit(self->mc);
|
||||||
|
goto cleanup;
|
||||||
|
} else {
|
||||||
|
PyObject *val;
|
||||||
|
|
||||||
|
/* This is safe because libmemcached's max key length
|
||||||
|
* includes space for a NUL-byte. */
|
||||||
|
curr_key[curr_key_len] = 0;
|
||||||
|
val = _PylibMC_parse_memcached_value(
|
||||||
|
curr_val, curr_val_len, curr_flags);
|
||||||
|
if (val == NULL) {
|
||||||
|
memcached_quit(self->mc);
|
||||||
|
goto cleanup;
|
||||||
|
}
|
||||||
|
PyDict_SetItemString(retval, curr_key + prefix_len, val);
|
||||||
|
Py_DECREF(val);
|
||||||
|
}
|
||||||
|
/* Although Python prohibits you from using the libc memory allocation
|
||||||
|
* interface, we have to since libmemcached goes around doing
|
||||||
|
* malloc()... */
|
||||||
|
free(curr_val);
|
||||||
|
}
|
||||||
|
|
||||||
|
PyMem_Free(key_lens);
|
||||||
|
PyMem_Free(keys);
|
||||||
|
for (i = 0; i < nkeys; i++) {
|
||||||
|
Py_DECREF(key_objs[i]);
|
||||||
|
}
|
||||||
|
PyMem_Free(key_objs);
|
||||||
|
|
||||||
/* Not INCREFing because the only two outcomes are NULL and a new dict.
|
/* Not INCREFing because the only two outcomes are NULL and a new dict.
|
||||||
* We're the owner of that dict already, so. */
|
* We're the owner of that dict already, so. */
|
||||||
return retval;
|
return retval;
|
||||||
|
cleanup:
|
||||||
|
Py_XDECREF(retval);
|
||||||
|
PyMem_Free(key_lens);
|
||||||
|
PyMem_Free(keys);
|
||||||
|
for (i = 0; i < nkeys; i++)
|
||||||
|
Py_DECREF(key_objs[i]);
|
||||||
|
PyMem_Free(key_objs);
|
||||||
|
return NULL;
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
|
|||||||
Reference in New Issue
Block a user