Introduce getKeysResult for getKeysFromCommand.

Avoid using a static buffer for short key index responses, and make it
caller's responsibility to stack-allocate a result type. Responses that
don't fit are still allocated on the heap.

(cherry picked from commit bf5beab64a196214c3c741d9ef67d0446c6480c3)
This commit is contained in:
Yossi Gottlieb 2020-10-05 17:03:17 +03:00 committed by Oran Agra
parent ec9b1cca59
commit ad1ed7dcd0
7 changed files with 170 additions and 121 deletions

View File

@ -1115,8 +1115,9 @@ int ACLCheckCommandPerm(client *c, int *keyidxptr) {
if (!(c->user->flags & USER_FLAG_ALLKEYS) && if (!(c->user->flags & USER_FLAG_ALLKEYS) &&
(c->cmd->getkeys_proc || c->cmd->firstkey)) (c->cmd->getkeys_proc || c->cmd->firstkey))
{ {
int numkeys; getKeysResult result = GETKEYS_RESULT_INIT;
int *keyidx = getKeysFromCommand(c->cmd,c->argv,c->argc,&numkeys); int numkeys = getKeysFromCommand(c->cmd,c->argv,c->argc,&result);
int *keyidx = result.keys;
for (int j = 0; j < numkeys; j++) { for (int j = 0; j < numkeys; j++) {
listIter li; listIter li;
listNode *ln; listNode *ln;
@ -1137,11 +1138,11 @@ int ACLCheckCommandPerm(client *c, int *keyidxptr) {
} }
if (!match) { if (!match) {
if (keyidxptr) *keyidxptr = keyidx[j]; if (keyidxptr) *keyidxptr = keyidx[j];
getKeysFreeResult(keyidx); getKeysFreeResult(&result);
return ACL_DENIED_KEY; return ACL_DENIED_KEY;
} }
} }
getKeysFreeResult(keyidx); getKeysFreeResult(&result);
} }
/* If we survived all the above checks, the user can execute the /* If we survived all the above checks, the user can execute the

View File

@ -5640,7 +5640,10 @@ clusterNode *getNodeByQuery(client *c, struct redisCommand *cmd, robj **argv, in
margc = ms->commands[i].argc; margc = ms->commands[i].argc;
margv = ms->commands[i].argv; margv = ms->commands[i].argv;
keyindex = getKeysFromCommand(mcmd,margv,margc,&numkeys); getKeysResult result = GETKEYS_RESULT_INIT;
numkeys = getKeysFromCommand(mcmd,margv,margc,&result);
keyindex = result.keys;
for (j = 0; j < numkeys; j++) { for (j = 0; j < numkeys; j++) {
robj *thiskey = margv[keyindex[j]]; robj *thiskey = margv[keyindex[j]];
int thisslot = keyHashSlot((char*)thiskey->ptr, int thisslot = keyHashSlot((char*)thiskey->ptr,
@ -5658,7 +5661,7 @@ clusterNode *getNodeByQuery(client *c, struct redisCommand *cmd, robj **argv, in
* not trapped earlier in processCommand(). Report the same * not trapped earlier in processCommand(). Report the same
* error to the client. */ * error to the client. */
if (n == NULL) { if (n == NULL) {
getKeysFreeResult(keyindex); getKeysFreeResult(&result);
if (error_code) if (error_code)
*error_code = CLUSTER_REDIR_DOWN_UNBOUND; *error_code = CLUSTER_REDIR_DOWN_UNBOUND;
return NULL; return NULL;
@ -5682,7 +5685,7 @@ clusterNode *getNodeByQuery(client *c, struct redisCommand *cmd, robj **argv, in
if (!equalStringObjects(firstkey,thiskey)) { if (!equalStringObjects(firstkey,thiskey)) {
if (slot != thisslot) { if (slot != thisslot) {
/* Error: multiple keys from different slots. */ /* Error: multiple keys from different slots. */
getKeysFreeResult(keyindex); getKeysFreeResult(&result);
if (error_code) if (error_code)
*error_code = CLUSTER_REDIR_CROSS_SLOT; *error_code = CLUSTER_REDIR_CROSS_SLOT;
return NULL; return NULL;
@ -5701,7 +5704,7 @@ clusterNode *getNodeByQuery(client *c, struct redisCommand *cmd, robj **argv, in
missing_keys++; missing_keys++;
} }
} }
getKeysFreeResult(keyindex); getKeysFreeResult(&result);
} }
/* No key at all in command? then we can serve the request /* No key at all in command? then we can serve the request

171
src/db.c
View File

@ -1322,27 +1322,54 @@ int expireIfNeeded(redisDb *db, robj *key) {
/* ----------------------------------------------------------------------------- /* -----------------------------------------------------------------------------
* API to get key arguments from commands * API to get key arguments from commands
* ---------------------------------------------------------------------------*/ * ---------------------------------------------------------------------------*/
#define MAX_KEYS_BUFFER 256
static int getKeysTempBuffer[MAX_KEYS_BUFFER]; /* Prepare the getKeysResult struct to hold numkeys, either by using the
* pre-allocated keysbuf or by allocating a new array on the heap.
*
* This function must be called at least once before starting to populate
* the result, and can be called repeatedly to enlarge the result array.
*/
int *getKeysPrepareResult(getKeysResult *result, int numkeys) {
/* GETKEYS_RESULT_INIT initializes keys to NULL, point it to the pre-allocated stack
* buffer here. */
if (!result->keys) {
serverAssert(!result->numkeys);
result->keys = result->keysbuf;
}
/* Resize if necessary */
if (numkeys > result->size) {
if (result->keys != result->keysbuf) {
/* We're not using a static buffer, just (re)alloc */
result->keys = zrealloc(result->keys, numkeys * sizeof(int));
} else {
/* We are using a static buffer, copy its contents */
result->keys = zmalloc(numkeys * sizeof(int));
if (result->numkeys)
memcpy(result->keys, result->keysbuf, result->numkeys * sizeof(int));
}
result->size = numkeys;
}
return result->keys;
}
/* The base case is to use the keys position as given in the command table /* The base case is to use the keys position as given in the command table
* (firstkey, lastkey, step). */ * (firstkey, lastkey, step). */
int *getKeysUsingCommandTable(struct redisCommand *cmd,robj **argv, int argc, int *numkeys) { int getKeysUsingCommandTable(struct redisCommand *cmd,robj **argv, int argc, getKeysResult *result) {
int j, i = 0, last, *keys; int j, i = 0, last, *keys;
UNUSED(argv); UNUSED(argv);
if (cmd->firstkey == 0) { if (cmd->firstkey == 0) {
*numkeys = 0; result->numkeys = 0;
return NULL; return 0;
} }
last = cmd->lastkey; last = cmd->lastkey;
if (last < 0) last = argc+last; if (last < 0) last = argc+last;
int count = ((last - cmd->firstkey)+1); int count = ((last - cmd->firstkey)+1);
keys = getKeysTempBuffer; keys = getKeysPrepareResult(result, count);
if (count > MAX_KEYS_BUFFER)
keys = zmalloc(sizeof(int)*count);
for (j = cmd->firstkey; j <= last; j += cmd->keystep) { for (j = cmd->firstkey; j <= last; j += cmd->keystep) {
if (j >= argc) { if (j >= argc) {
@ -1353,17 +1380,17 @@ int *getKeysUsingCommandTable(struct redisCommand *cmd,robj **argv, int argc, in
* return no keys and expect the command implementation to report * return no keys and expect the command implementation to report
* an arity or syntax error. */ * an arity or syntax error. */
if (cmd->flags & CMD_MODULE || cmd->arity < 0) { if (cmd->flags & CMD_MODULE || cmd->arity < 0) {
getKeysFreeResult(keys); getKeysFreeResult(result);
*numkeys = 0; result->numkeys = 0;
return NULL; return 0;
} else { } else {
serverPanic("Redis built-in command declared keys positions not matching the arity requirements."); serverPanic("Redis built-in command declared keys positions not matching the arity requirements.");
} }
} }
keys[i++] = j; keys[i++] = j;
} }
*numkeys = i; result->numkeys = i;
return keys; return i;
} }
/* Return all the arguments that are keys in the command passed via argc / argv. /* Return all the arguments that are keys in the command passed via argc / argv.
@ -1377,26 +1404,26 @@ int *getKeysUsingCommandTable(struct redisCommand *cmd,robj **argv, int argc, in
* *
* This function uses the command table if a command-specific helper function * This function uses the command table if a command-specific helper function
* is not required, otherwise it calls the command-specific function. */ * is not required, otherwise it calls the command-specific function. */
int *getKeysFromCommand(struct redisCommand *cmd, robj **argv, int argc, int *numkeys) { int getKeysFromCommand(struct redisCommand *cmd, robj **argv, int argc, getKeysResult *result) {
if (cmd->flags & CMD_MODULE_GETKEYS) { if (cmd->flags & CMD_MODULE_GETKEYS) {
return moduleGetCommandKeysViaAPI(cmd,argv,argc,numkeys); return moduleGetCommandKeysViaAPI(cmd,argv,argc,result);
} else if (!(cmd->flags & CMD_MODULE) && cmd->getkeys_proc) { } else if (!(cmd->flags & CMD_MODULE) && cmd->getkeys_proc) {
return cmd->getkeys_proc(cmd,argv,argc,numkeys); return cmd->getkeys_proc(cmd,argv,argc,result);
} else { } else {
return getKeysUsingCommandTable(cmd,argv,argc,numkeys); return getKeysUsingCommandTable(cmd,argv,argc,result);
} }
} }
/* Free the result of getKeysFromCommand. */ /* Free the result of getKeysFromCommand. */
void getKeysFreeResult(int *result) { void getKeysFreeResult(getKeysResult *result) {
if (result != getKeysTempBuffer) if (result && result->keys != result->keysbuf)
zfree(result); zfree(result->keys);
} }
/* Helper function to extract keys from following commands: /* Helper function to extract keys from following commands:
* ZUNIONSTORE <destkey> <num-keys> <key> <key> ... <key> <options> * ZUNIONSTORE <destkey> <num-keys> <key> <key> ... <key> <options>
* ZINTERSTORE <destkey> <num-keys> <key> <key> ... <key> <options> */ * ZINTERSTORE <destkey> <num-keys> <key> <key> ... <key> <options> */
int *zunionInterGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkeys) { int zunionInterGetKeys(struct redisCommand *cmd, robj **argv, int argc, getKeysResult *result) {
int i, num, *keys; int i, num, *keys;
UNUSED(cmd); UNUSED(cmd);
@ -1404,30 +1431,30 @@ int *zunionInterGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *nu
/* Sanity check. Don't return any key if the command is going to /* Sanity check. Don't return any key if the command is going to
* reply with syntax error. */ * reply with syntax error. */
if (num < 1 || num > (argc-3)) { if (num < 1 || num > (argc-3)) {
*numkeys = 0; result->numkeys = 0;
return NULL; return 0;
} }
/* Keys in z{union,inter}store come from two places: /* Keys in z{union,inter}store come from two places:
* argv[1] = storage key, * argv[1] = storage key,
* argv[3...n] = keys to intersect */ * argv[3...n] = keys to intersect */
keys = getKeysTempBuffer; /* Total keys = {union,inter} keys + storage key */
if (num+1>MAX_KEYS_BUFFER) keys = getKeysPrepareResult(result, num+1);
keys = zmalloc(sizeof(int)*(num+1)); result->numkeys = num+1;
/* Add all key positions for argv[3...n] to keys[] */ /* Add all key positions for argv[3...n] to keys[] */
for (i = 0; i < num; i++) keys[i] = 3+i; for (i = 0; i < num; i++) keys[i] = 3+i;
/* Finally add the argv[1] key position (the storage key target). */ /* Finally add the argv[1] key position (the storage key target). */
keys[num] = 1; keys[num] = 1;
*numkeys = num+1; /* Total keys = {union,inter} keys + storage key */
return keys; return result->numkeys;
} }
/* Helper function to extract keys from the following commands: /* Helper function to extract keys from the following commands:
* EVAL <script> <num-keys> <key> <key> ... <key> [more stuff] * EVAL <script> <num-keys> <key> <key> ... <key> [more stuff]
* EVALSHA <script> <num-keys> <key> <key> ... <key> [more stuff] */ * EVALSHA <script> <num-keys> <key> <key> ... <key> [more stuff] */
int *evalGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkeys) { int evalGetKeys(struct redisCommand *cmd, robj **argv, int argc, getKeysResult *result) {
int i, num, *keys; int i, num, *keys;
UNUSED(cmd); UNUSED(cmd);
@ -1435,20 +1462,17 @@ int *evalGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkeys)
/* Sanity check. Don't return any key if the command is going to /* Sanity check. Don't return any key if the command is going to
* reply with syntax error. */ * reply with syntax error. */
if (num <= 0 || num > (argc-3)) { if (num <= 0 || num > (argc-3)) {
*numkeys = 0; result->numkeys = 0;
return NULL; return 0;
} }
keys = getKeysTempBuffer; keys = getKeysPrepareResult(result, num);
if (num>MAX_KEYS_BUFFER) result->numkeys = num;
keys = zmalloc(sizeof(int)*num);
*numkeys = num;
/* Add all key positions for argv[3...n] to keys[] */ /* Add all key positions for argv[3...n] to keys[] */
for (i = 0; i < num; i++) keys[i] = 3+i; for (i = 0; i < num; i++) keys[i] = 3+i;
return keys; return result->numkeys;
} }
/* Helper function to extract keys from the SORT command. /* Helper function to extract keys from the SORT command.
@ -1458,13 +1482,12 @@ int *evalGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkeys)
* The first argument of SORT is always a key, however a list of options * The first argument of SORT is always a key, however a list of options
* follow in SQL-alike style. Here we parse just the minimum in order to * follow in SQL-alike style. Here we parse just the minimum in order to
* correctly identify keys in the "STORE" option. */ * correctly identify keys in the "STORE" option. */
int *sortGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkeys) { int sortGetKeys(struct redisCommand *cmd, robj **argv, int argc, getKeysResult *result) {
int i, j, num, *keys, found_store = 0; int i, j, num, *keys, found_store = 0;
UNUSED(cmd); UNUSED(cmd);
num = 0; num = 0;
keys = getKeysTempBuffer; /* Alloc 2 places for the worst case. */ keys = getKeysPrepareResult(result, 2); /* Alloc 2 places for the worst case. */
keys[num++] = 1; /* <sort-key> is always present. */ keys[num++] = 1; /* <sort-key> is always present. */
/* Search for STORE option. By default we consider options to don't /* Search for STORE option. By default we consider options to don't
@ -1496,11 +1519,11 @@ int *sortGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkeys)
} }
} }
} }
*numkeys = num + found_store; result->numkeys = num + found_store;
return keys; return result->numkeys;
} }
int *migrateGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkeys) { int migrateGetKeys(struct redisCommand *cmd, robj **argv, int argc, getKeysResult *result) {
int i, num, first, *keys; int i, num, first, *keys;
UNUSED(cmd); UNUSED(cmd);
@ -1521,20 +1544,17 @@ int *migrateGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkey
} }
} }
keys = getKeysTempBuffer; keys = getKeysPrepareResult(result, num);
if (num>MAX_KEYS_BUFFER)
keys = zmalloc(sizeof(int)*num);
for (i = 0; i < num; i++) keys[i] = first+i; for (i = 0; i < num; i++) keys[i] = first+i;
*numkeys = num; result->numkeys = num;
return keys; return num;
} }
/* Helper function to extract keys from following commands: /* Helper function to extract keys from following commands:
* GEORADIUS key x y radius unit [WITHDIST] [WITHHASH] [WITHCOORD] [ASC|DESC] * GEORADIUS key x y radius unit [WITHDIST] [WITHHASH] [WITHCOORD] [ASC|DESC]
* [COUNT count] [STORE key] [STOREDIST key] * [COUNT count] [STORE key] [STOREDIST key]
* GEORADIUSBYMEMBER key member radius unit ... options ... */ * GEORADIUSBYMEMBER key member radius unit ... options ... */
int *georadiusGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkeys) { int georadiusGetKeys(struct redisCommand *cmd, robj **argv, int argc, getKeysResult *result) {
int i, num, *keys; int i, num, *keys;
UNUSED(cmd); UNUSED(cmd);
@ -1557,24 +1577,21 @@ int *georadiusGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numk
* argv[1] = key, * argv[1] = key,
* argv[5...n] = stored key if present * argv[5...n] = stored key if present
*/ */
keys = getKeysTempBuffer; keys = getKeysPrepareResult(result, num);
if (num>MAX_KEYS_BUFFER)
keys = zmalloc(sizeof(int) * num);
/* Add all key positions to keys[] */ /* Add all key positions to keys[] */
keys[0] = 1; keys[0] = 1;
if(num > 1) { if(num > 1) {
keys[1] = stored_key; keys[1] = stored_key;
} }
*numkeys = num; result->numkeys = num;
return keys; return num;
} }
/* LCS ... [KEYS <key1> <key2>] ... */ /* LCS ... [KEYS <key1> <key2>] ... */
int *lcsGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkeys) int lcsGetKeys(struct redisCommand *cmd, robj **argv, int argc, getKeysResult *result) {
{
int i; int i;
int *keys = getKeysTempBuffer; int *keys = getKeysPrepareResult(result, 2);
UNUSED(cmd); UNUSED(cmd);
/* We need to parse the options of the command in order to check for the /* We need to parse the options of the command in order to check for the
@ -1588,33 +1605,32 @@ int *lcsGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkeys)
} else if (!strcasecmp(arg, "keys") && moreargs >= 2) { } else if (!strcasecmp(arg, "keys") && moreargs >= 2) {
keys[0] = i+1; keys[0] = i+1;
keys[1] = i+2; keys[1] = i+2;
*numkeys = 2; result->numkeys = 2;
return keys; return result->numkeys;
} }
} }
*numkeys = 0; result->numkeys = 0;
return keys; return result->numkeys;
} }
/* Helper function to extract keys from memory command. /* Helper function to extract keys from memory command.
* MEMORY USAGE <key> */ * MEMORY USAGE <key> */
int *memoryGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkeys) { int memoryGetKeys(struct redisCommand *cmd, robj **argv, int argc, getKeysResult *result) {
int *keys;
UNUSED(cmd); UNUSED(cmd);
getKeysPrepareResult(result, 1);
if (argc >= 3 && !strcasecmp(argv[1]->ptr,"usage")) { if (argc >= 3 && !strcasecmp(argv[1]->ptr,"usage")) {
keys = getKeysTempBuffer; result->keys[0] = 2;
keys[0] = 2; result->numkeys = 1;
*numkeys = 1; return result->numkeys;
return keys;
} }
*numkeys = 0; result->numkeys = 0;
return NULL; return 0;
} }
/* XREAD [BLOCK <milliseconds>] [COUNT <count>] [GROUP <groupname> <ttl>] /* XREAD [BLOCK <milliseconds>] [COUNT <count>] [GROUP <groupname> <ttl>]
* STREAMS key_1 key_2 ... key_N ID_1 ID_2 ... ID_N */ * STREAMS key_1 key_2 ... key_N ID_1 ID_2 ... ID_N */
int *xreadGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkeys) { int xreadGetKeys(struct redisCommand *cmd, robj **argv, int argc, getKeysResult *result) {
int i, num = 0, *keys; int i, num = 0, *keys;
UNUSED(cmd); UNUSED(cmd);
@ -1644,19 +1660,16 @@ int *xreadGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkeys)
/* Syntax error. */ /* Syntax error. */
if (streams_pos == -1 || num == 0 || num % 2 != 0) { if (streams_pos == -1 || num == 0 || num % 2 != 0) {
*numkeys = 0; result->numkeys = 0;
return NULL; return 0;
} }
num /= 2; /* We have half the keys as there are arguments because num /= 2; /* We have half the keys as there are arguments because
there are also the IDs, one per key. */ there are also the IDs, one per key. */
keys = getKeysTempBuffer; keys = getKeysPrepareResult(result, num);
if (num>MAX_KEYS_BUFFER)
keys = zmalloc(sizeof(int) * num);
for (i = streams_pos+1; i < argc-num; i++) keys[i-streams_pos-1] = i; for (i = streams_pos+1; i < argc-num; i++) keys[i-streams_pos-1] = i;
*numkeys = num; result->numkeys = num;
return keys; return num;
} }
/* Slot to Key API. This is used by Redis Cluster in order to obtain in /* Slot to Key API. This is used by Redis Cluster in order to obtain in

View File

@ -147,8 +147,7 @@ struct RedisModuleCtx {
on keys. */ on keys. */
/* Used if there is the REDISMODULE_CTX_KEYS_POS_REQUEST flag set. */ /* Used if there is the REDISMODULE_CTX_KEYS_POS_REQUEST flag set. */
int *keys_pos; getKeysResult *keys_result;
int keys_count;
struct RedisModulePoolAllocBlock *pa_head; struct RedisModulePoolAllocBlock *pa_head;
redisOpArray saved_oparray; /* When propagating commands in a callback redisOpArray saved_oparray; /* When propagating commands in a callback
@ -158,7 +157,7 @@ struct RedisModuleCtx {
}; };
typedef struct RedisModuleCtx RedisModuleCtx; typedef struct RedisModuleCtx RedisModuleCtx;
#define REDISMODULE_CTX_INIT {(void*)(unsigned long)&RM_GetApi, NULL, NULL, NULL, NULL, 0, 0, 0, NULL, 0, NULL, NULL, NULL, 0, NULL, {0}} #define REDISMODULE_CTX_INIT {(void*)(unsigned long)&RM_GetApi, NULL, NULL, NULL, NULL, 0, 0, 0, NULL, 0, NULL, NULL, NULL, NULL, {0}}
#define REDISMODULE_CTX_MULTI_EMITTED (1<<0) #define REDISMODULE_CTX_MULTI_EMITTED (1<<0)
#define REDISMODULE_CTX_AUTO_MEMORY (1<<1) #define REDISMODULE_CTX_AUTO_MEMORY (1<<1)
#define REDISMODULE_CTX_KEYS_POS_REQUEST (1<<2) #define REDISMODULE_CTX_KEYS_POS_REQUEST (1<<2)
@ -670,18 +669,24 @@ void RedisModuleCommandDispatcher(client *c) {
* In order to accomplish its work, the module command is called, flagging * In order to accomplish its work, the module command is called, flagging
* the context in a way that the command can recognize this is a special * the context in a way that the command can recognize this is a special
* "get keys" call by calling RedisModule_IsKeysPositionRequest(ctx). */ * "get keys" call by calling RedisModule_IsKeysPositionRequest(ctx). */
int *moduleGetCommandKeysViaAPI(struct redisCommand *cmd, robj **argv, int argc, int *numkeys) { int moduleGetCommandKeysViaAPI(struct redisCommand *cmd, robj **argv, int argc, getKeysResult *result) {
RedisModuleCommandProxy *cp = (void*)(unsigned long)cmd->getkeys_proc; RedisModuleCommandProxy *cp = (void*)(unsigned long)cmd->getkeys_proc;
RedisModuleCtx ctx = REDISMODULE_CTX_INIT; RedisModuleCtx ctx = REDISMODULE_CTX_INIT;
ctx.module = cp->module; ctx.module = cp->module;
ctx.client = NULL; ctx.client = NULL;
ctx.flags |= REDISMODULE_CTX_KEYS_POS_REQUEST; ctx.flags |= REDISMODULE_CTX_KEYS_POS_REQUEST;
/* Initialize getKeysResult */
getKeysPrepareResult(result, MAX_KEYS_BUFFER);
ctx.keys_result = result;
cp->func(&ctx,(void**)argv,argc); cp->func(&ctx,(void**)argv,argc);
int *res = ctx.keys_pos; /* We currently always use the array allocated by RM_KeyAtPos() and don't try
if (numkeys) *numkeys = ctx.keys_count; * to optimize for the pre-allocated buffer.
*/
moduleFreeContext(&ctx); moduleFreeContext(&ctx);
return res; return result->numkeys;
} }
/* Return non-zero if a module command, that was declared with the /* Return non-zero if a module command, that was declared with the
@ -706,10 +711,18 @@ int RM_IsKeysPositionRequest(RedisModuleCtx *ctx) {
* keys are at fixed positions. This interface is only used for commands * keys are at fixed positions. This interface is only used for commands
* with a more complex structure. */ * with a more complex structure. */
void RM_KeyAtPos(RedisModuleCtx *ctx, int pos) { void RM_KeyAtPos(RedisModuleCtx *ctx, int pos) {
if (!(ctx->flags & REDISMODULE_CTX_KEYS_POS_REQUEST)) return; if (!(ctx->flags & REDISMODULE_CTX_KEYS_POS_REQUEST) || !ctx->keys_result) return;
if (pos <= 0) return; if (pos <= 0) return;
ctx->keys_pos = zrealloc(ctx->keys_pos,sizeof(int)*(ctx->keys_count+1));
ctx->keys_pos[ctx->keys_count++] = pos; getKeysResult *res = ctx->keys_result;
/* Check overflow */
if (res->numkeys == res->size) {
int newsize = res->size + (res->size > 8192 ? 8192 : res->size);
getKeysPrepareResult(res, newsize);
}
res->keys[res->numkeys++] = pos;
} }
/* Helper for RM_CreateCommand(). Turns a string representing command /* Helper for RM_CreateCommand(). Turns a string representing command

View File

@ -3177,7 +3177,7 @@ struct redisCommand *lookupCommand(sds name) {
return dictFetchValue(server.commands, name); return dictFetchValue(server.commands, name);
} }
struct redisCommand *lookupCommandByCString(char *s) { struct redisCommand *lookupCommandByCString(const char *s) {
struct redisCommand *cmd; struct redisCommand *cmd;
sds name = sdsnew(s); sds name = sdsnew(s);
@ -4041,7 +4041,8 @@ NULL
addReplyLongLong(c, dictSize(server.commands)); addReplyLongLong(c, dictSize(server.commands));
} else if (!strcasecmp(c->argv[1]->ptr,"getkeys") && c->argc >= 3) { } else if (!strcasecmp(c->argv[1]->ptr,"getkeys") && c->argc >= 3) {
struct redisCommand *cmd = lookupCommand(c->argv[2]->ptr); struct redisCommand *cmd = lookupCommand(c->argv[2]->ptr);
int *keys, numkeys, j; getKeysResult result = GETKEYS_RESULT_INIT;
int j;
if (!cmd) { if (!cmd) {
addReplyError(c,"Invalid command specified"); addReplyError(c,"Invalid command specified");
@ -4056,14 +4057,13 @@ NULL
return; return;
} }
keys = getKeysFromCommand(cmd,c->argv+2,c->argc-2,&numkeys); if (!getKeysFromCommand(cmd,c->argv+2,c->argc-2,&result)) {
if (!keys) {
addReplyError(c,"Invalid arguments specified for command"); addReplyError(c,"Invalid arguments specified for command");
} else { } else {
addReplyArrayLen(c,numkeys); addReplyArrayLen(c,result.numkeys);
for (j = 0; j < numkeys; j++) addReplyBulk(c,c->argv[keys[j]+2]); for (j = 0; j < result.numkeys; j++) addReplyBulk(c,c->argv[result.keys[j]+2]);
getKeysFreeResult(keys);
} }
getKeysFreeResult(&result);
} else { } else {
addReplySubcommandSyntaxError(c); addReplySubcommandSyntaxError(c);
} }

View File

@ -1482,8 +1482,21 @@ typedef struct pubsubPattern {
robj *pattern; robj *pattern;
} pubsubPattern; } pubsubPattern;
#define MAX_KEYS_BUFFER 256
/* A result structure for the various getkeys function calls. It lists the
* keys as indices to the provided argv.
*/
typedef struct {
int keysbuf[MAX_KEYS_BUFFER]; /* Pre-allocated buffer, to save heap allocations */
int *keys; /* Key indices array, points to keysbuf or heap */
int numkeys; /* Number of key indices return */
int size; /* Available array size */
} getKeysResult;
#define GETKEYS_RESULT_INIT { {0}, NULL, 0, MAX_KEYS_BUFFER }
typedef void redisCommandProc(client *c); typedef void redisCommandProc(client *c);
typedef int *redisGetKeysProc(struct redisCommand *cmd, robj **argv, int argc, int *numkeys); typedef int redisGetKeysProc(struct redisCommand *cmd, robj **argv, int argc, getKeysResult *result);
struct redisCommand { struct redisCommand {
char *name; char *name;
redisCommandProc *proc; redisCommandProc *proc;
@ -1592,7 +1605,7 @@ extern dictType modulesDictType;
void moduleInitModulesSystem(void); void moduleInitModulesSystem(void);
int moduleLoad(const char *path, void **argv, int argc); int moduleLoad(const char *path, void **argv, int argc);
void moduleLoadFromQueue(void); void moduleLoadFromQueue(void);
int *moduleGetCommandKeysViaAPI(struct redisCommand *cmd, robj **argv, int argc, int *numkeys); int moduleGetCommandKeysViaAPI(struct redisCommand *cmd, robj **argv, int argc, getKeysResult *result);
moduleType *moduleTypeLookupModuleByID(uint64_t id); moduleType *moduleTypeLookupModuleByID(uint64_t id);
void moduleTypeNameByID(char *name, uint64_t moduleid); void moduleTypeNameByID(char *name, uint64_t moduleid);
void moduleFreeContext(struct RedisModuleCtx *ctx); void moduleFreeContext(struct RedisModuleCtx *ctx);
@ -1998,7 +2011,7 @@ int freeMemoryIfNeededAndSafe(void);
int processCommand(client *c); int processCommand(client *c);
void setupSignalHandlers(void); void setupSignalHandlers(void);
struct redisCommand *lookupCommand(sds name); struct redisCommand *lookupCommand(sds name);
struct redisCommand *lookupCommandByCString(char *s); struct redisCommand *lookupCommandByCString(const char *s);
struct redisCommand *lookupCommandOrOriginal(sds name); struct redisCommand *lookupCommandOrOriginal(sds name);
void call(client *c, int flags); void call(client *c, int flags);
void propagate(struct redisCommand *cmd, int dbid, robj **argv, int argc, int flags); void propagate(struct redisCommand *cmd, int dbid, robj **argv, int argc, int flags);
@ -2159,16 +2172,17 @@ size_t lazyfreeGetPendingObjectsCount(void);
void freeObjAsync(robj *o); void freeObjAsync(robj *o);
/* API to get key arguments from commands */ /* API to get key arguments from commands */
int *getKeysFromCommand(struct redisCommand *cmd, robj **argv, int argc, int *numkeys); int *getKeysPrepareResult(getKeysResult *result, int numkeys);
void getKeysFreeResult(int *result); int getKeysFromCommand(struct redisCommand *cmd, robj **argv, int argc, getKeysResult *result);
int *zunionInterGetKeys(struct redisCommand *cmd,robj **argv, int argc, int *numkeys); void getKeysFreeResult(getKeysResult *result);
int *evalGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkeys); int zunionInterGetKeys(struct redisCommand *cmd,robj **argv, int argc, getKeysResult *result);
int *sortGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkeys); int evalGetKeys(struct redisCommand *cmd, robj **argv, int argc, getKeysResult *result);
int *migrateGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkeys); int sortGetKeys(struct redisCommand *cmd, robj **argv, int argc, getKeysResult *result);
int *georadiusGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkeys); int migrateGetKeys(struct redisCommand *cmd, robj **argv, int argc, getKeysResult *result);
int *xreadGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkeys); int georadiusGetKeys(struct redisCommand *cmd, robj **argv, int argc, getKeysResult *result);
int *memoryGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkeys); int xreadGetKeys(struct redisCommand *cmd, robj **argv, int argc, getKeysResult *result);
int *lcsGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkeys); int memoryGetKeys(struct redisCommand *cmd, robj **argv, int argc, getKeysResult *result);
int lcsGetKeys(struct redisCommand *cmd, robj **argv, int argc, getKeysResult *result);
/* Cluster */ /* Cluster */
void clusterInit(void); void clusterInit(void);

View File

@ -171,9 +171,14 @@ void trackingRememberKeys(client *c) {
uint64_t caching_given = c->flags & CLIENT_TRACKING_CACHING; uint64_t caching_given = c->flags & CLIENT_TRACKING_CACHING;
if ((optin && !caching_given) || (optout && caching_given)) return; if ((optin && !caching_given) || (optout && caching_given)) return;
int numkeys; getKeysResult result = GETKEYS_RESULT_INIT;
int *keys = getKeysFromCommand(c->cmd,c->argv,c->argc,&numkeys); int numkeys = getKeysFromCommand(c->cmd,c->argv,c->argc,&result);
if (keys == NULL) return; if (!numkeys) {
getKeysFreeResult(&result);
return;
}
int *keys = result.keys;
for(int j = 0; j < numkeys; j++) { for(int j = 0; j < numkeys; j++) {
int idx = keys[j]; int idx = keys[j];
@ -188,7 +193,7 @@ void trackingRememberKeys(client *c) {
if (raxTryInsert(ids,(unsigned char*)&c->id,sizeof(c->id),NULL,NULL)) if (raxTryInsert(ids,(unsigned char*)&c->id,sizeof(c->id),NULL,NULL))
TrackingTableTotalItems++; TrackingTableTotalItems++;
} }
getKeysFreeResult(keys); getKeysFreeResult(&result);
} }
/* Given a key name, this function sends an invalidation message in the /* Given a key name, this function sends an invalidation message in the