Modules: add RM_GetCommandKeys().

This is essentially the same as calling COMMAND GETKEYS but provides a
more efficient interface that can be used in every context (i.e. not a
Redis command).

(cherry picked from commit 7d117d7591)

runtest-moduleapi

@@ -27,4 +27,5 @@ $TCLSH tests/test_helper.tcl \
 --single unit/moduleapi/auth \
 --single unit/moduleapi/keyspace_events \
 --single unit/moduleapi/blockedclient \
+--single unit/moduleapi/getkeys \
 "${@}"

src/module.c

@@ -7879,6 +7879,69 @@ int RM_ModuleTypeReplaceValue(RedisModuleKey *key, moduleType *mt, void *new_val
     return REDISMODULE_OK;
 }
 
+/* For a specified command, parse its arguments and return an array that
+ * contains the indexes of all key name arguments. This function is
+ * essnetially a more efficient way to do COMMAND GETKEYS.
+ *
+ * A NULL return value indicates the specified command has no keys, or
+ * an error condition. Error conditions are indicated by setting errno
+ * as folllows:
+ *
+ *  ENOENT: Specified command does not exist.
+ *  EINVAL: Invalid command arity specified.
+ *
+ * NOTE: The returned array is not a Redis Module object so it does not
+ * get automatically freed even when auto-memory is used. The caller
+ * must explicitly call RM_Free() to free it.
+ */
+int *RM_GetCommandKeys(RedisModuleCtx *ctx, const char *cmdname, RedisModuleString **argv, int argc, int *num_keys) {
+    UNUSED(ctx);
+    struct redisCommand *cmd;
+    int *res = NULL;
+
+    /* Find command */
+    if ((cmd = lookupCommandByCString(cmdname)) == NULL) {
+        errno = ENOENT;
+        return NULL;
+    }
+
+    /* Bail out if command has no keys */
+    if (cmd->getkeys_proc == NULL && cmd->firstkey == 0) {
+        errno = 0;
+        return NULL;
+    }
+
+    if ((cmd->arity > 0 && cmd->arity != argc) || (argc < -cmd->arity)) {
+        errno = EINVAL;
+        return NULL;
+    }
+
+    getKeysResult result = GETKEYS_RESULT_INIT;
+    getKeysFromCommand(cmd, argv, argc, &result);
+
+    *num_keys = result.numkeys;
+    if (!result.numkeys) {
+        errno = 0;
+        getKeysFreeResult(&result);
+        return NULL;
+    }
+
+    if (result.keys == result.keysbuf) {
+        /* If the result is using a stack based array, copy it. */
+        unsigned long int size = sizeof(int) * result.numkeys;
+        res = zmalloc(size);
+        memcpy(res, result.keys, size);
+    } else {
+        /* We return the heap based array and intentionally avoid calling
+         * getKeysFreeResult() here, as it is the caller's responsibility
+         * to free this array.
+         */
+        res = result.keys;
+    }
+
+    return res;
+}
+
 /* Register all the APIs we export. Keep this function at the end of the
  * file so that's easy to seek it to add new entries. */
 void moduleRegisterCoreAPI(void) {
@@ -8114,4 +8177,5 @@ void moduleRegisterCoreAPI(void) {
     REGISTER_API(DeauthenticateAndCloseClient);
     REGISTER_API(AuthenticateClientWithACLUser);
     REGISTER_API(AuthenticateClientWithUser);
+    REGISTER_API(GetCommandKeys);
 }

src/redismodule.h

@@ -724,6 +724,7 @@ REDISMODULE_API int (*RedisModule_SetModuleUserACL)(RedisModuleUser *user, const
 REDISMODULE_API int (*RedisModule_AuthenticateClientWithACLUser)(RedisModuleCtx *ctx, const char *name, size_t len, RedisModuleUserChangedFunc callback, void *privdata, uint64_t *client_id) REDISMODULE_ATTR;
 REDISMODULE_API int (*RedisModule_AuthenticateClientWithUser)(RedisModuleCtx *ctx, RedisModuleUser *user, RedisModuleUserChangedFunc callback, void *privdata, uint64_t *client_id) REDISMODULE_ATTR;
 REDISMODULE_API int (*RedisModule_DeauthenticateAndCloseClient)(RedisModuleCtx *ctx, uint64_t client_id) REDISMODULE_ATTR;
+REDISMODULE_API int *(*RedisModule_GetCommandKeys)(RedisModuleCtx *ctx, const char *cmdname, RedisModuleString **argv, int argc, int *num_keys) REDISMODULE_ATTR;
 #endif
 
 #define RedisModule_IsAOFClient(id) ((id) == CLIENT_ID_AOF)
@@ -965,6 +966,7 @@ static int RedisModule_Init(RedisModuleCtx *ctx, const char *name, int ver, int
     REDISMODULE_GET_API(DeauthenticateAndCloseClient);
     REDISMODULE_GET_API(AuthenticateClientWithACLUser);
     REDISMODULE_GET_API(AuthenticateClientWithUser);
+    REDISMODULE_GET_API(GetCommandKeys);
 #endif
 
     if (RedisModule_IsModuleNameBusy && RedisModule_IsModuleNameBusy(name)) return REDISMODULE_ERR;

tests/modules/Makefile

@@ -24,7 +24,8 @@ TEST_MODULES = \
     datatype.so \
     auth.so \
     keyspace_events.so \
-    blockedclient.so
+    blockedclient.so \
+    getkeys.so
 
 
 .PHONY: all

tests/modules/getkeys.c

@@ -0,0 +1,125 @@
+#define REDISMODULE_EXPERIMENTAL_API
+
+#include "redismodule.h"
+#include <strings.h>
+#include <assert.h>
+#include <unistd.h>
+#include <errno.h>
+
+#define UNUSED(V) ((void) V)
+
+/* A sample movable keys command that returns a list of all
+ * arguments that follow a KEY argument, i.e.
+ */
+int getkeys_command(RedisModuleCtx *ctx, RedisModuleString **argv, int argc)
+{
+    int i;
+    int count = 0;
+
+    /* Handle getkeys-api introspection */
+    if (RedisModule_IsKeysPositionRequest(ctx)) {
+        for (i = 0; i < argc; i++) {
+            size_t len;
+            const char *str = RedisModule_StringPtrLen(argv[i], &len);
+
+            if (len == 3 && !strncasecmp(str, "key", 3) && i + 1 < argc)
+                RedisModule_KeyAtPos(ctx, i + 1);
+        }
+
+        return REDISMODULE_OK;
+    }
+
+    /* Handle real command invocation */
+    RedisModule_ReplyWithArray(ctx, REDISMODULE_POSTPONED_ARRAY_LEN);
+    for (i = 0; i < argc; i++) {
+        size_t len;
+        const char *str = RedisModule_StringPtrLen(argv[i], &len);
+
+        if (len == 3 && !strncasecmp(str, "key", 3) && i + 1 < argc) {
+            RedisModule_ReplyWithString(ctx, argv[i+1]);
+            count++;
+        }
+    }
+    RedisModule_ReplySetArrayLength(ctx, count);
+
+    return REDISMODULE_OK;
+}
+
+int getkeys_fixed(RedisModuleCtx *ctx, RedisModuleString **argv, int argc)
+{
+    int i;
+
+    RedisModule_ReplyWithArray(ctx, argc - 1);
+    for (i = 1; i < argc; i++) {
+        RedisModule_ReplyWithString(ctx, argv[i]);
+    }
+    return REDISMODULE_OK;
+}
+
+/* Introspect a command using RM_GetCommandKeys() and returns the list
+ * of keys. Essentially this is COMMAND GETKEYS implemented in a module.
+ */
+int getkeys_introspect(RedisModuleCtx *ctx, RedisModuleString **argv, int argc)
+{
+    UNUSED(argv);
+    UNUSED(argc);
+
+    if (argc < 3) {
+        RedisModule_WrongArity(ctx);
+        return REDISMODULE_OK;
+    }
+
+    size_t cmd_len;
+    const char *cmd = RedisModule_StringPtrLen(argv[1], &cmd_len);
+
+    int num_keys;
+    int *keyidx = RedisModule_GetCommandKeys(ctx, cmd, &argv[1], argc - 1, &num_keys);
+
+    if (!keyidx) {
+        if (!errno)
+            RedisModule_ReplyWithEmptyArray(ctx);
+        else {
+            char err[100];
+            switch (errno) {
+                case ENOENT:
+                    RedisModule_ReplyWithError(ctx, "ERR ENOENT");
+                    break;
+                case EINVAL:
+                    RedisModule_ReplyWithError(ctx, "ERR EINVAL");
+                    break;
+                default:
+                    snprintf(err, sizeof(err) - 1, "ERR errno=%d", errno);
+                    RedisModule_ReplyWithError(ctx, err);
+                    break;
+            }
+        }
+    } else {
+        int i;
+
+        RedisModule_ReplyWithArray(ctx, num_keys);
+        for (i = 0; i < num_keys; i++)
+            RedisModule_ReplyWithString(ctx, argv[1 + keyidx[i]]);
+
+        RedisModule_Free(keyidx);
+    }
+
+    return REDISMODULE_OK;
+}
+
+int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
+    UNUSED(argv);
+    UNUSED(argc);
+    if (RedisModule_Init(ctx,"getkeys",1,REDISMODULE_APIVER_1)== REDISMODULE_ERR)
+        return REDISMODULE_ERR;
+
+    if (RedisModule_CreateCommand(ctx,"getkeys.command", getkeys_command,"getkeys-api",0,0,0) == REDISMODULE_ERR)
+        return REDISMODULE_ERR;
+
+    if (RedisModule_CreateCommand(ctx,"getkeys.fixed", getkeys_fixed,"",2,4,1) == REDISMODULE_ERR)
+        return REDISMODULE_ERR;
+
+    if (RedisModule_CreateCommand(ctx,"getkeys.introspect", getkeys_introspect,"",0,0,0) == REDISMODULE_ERR)
+        return REDISMODULE_ERR;
+
+    return REDISMODULE_OK;
+}

tests/unit/moduleapi/getkeys.tcl

@@ -0,0 +1,44 @@
+set testmodule [file normalize tests/modules/getkeys.so]
+
+start_server {tags {"modules"}} {
+    r module load $testmodule
+
+    test {COMMAND INFO correctly reports a movable keys module command} {
+        set info [lindex [r command info getkeys.command] 0]
+
+        assert_equal {movablekeys} [lindex $info 2]
+        assert_equal {0} [lindex $info 3]
+        assert_equal {0} [lindex $info 4]
+        assert_equal {0} [lindex $info 5]
+    }
+
+    test {COMMAND GETKEYS correctly reports a movable keys module command} {
+        r command getkeys getkeys.command arg1 arg2 key key1 arg3 key key2 key key3
+    } {key1 key2 key3}
+
+    test {RM_GetCommandKeys on non-existing command} {
+        catch {r getkeys.introspect non-command key1 key2} e
+        set _ $e
+    } {*ENOENT*}
+
+    test {RM_GetCommandKeys on built-in fixed keys command} {
+        r getkeys.introspect set key1 value1
+    } {key1}
+
+    test {RM_GetCommandKeys on EVAL} {
+        r getkeys.introspect eval "" 4 key1 key2 key3 key4 arg1 arg2
+    } {key1 key2 key3 key4}
+
+    test {RM_GetCommandKeys on a movable keys module command} {
+        r getkeys.introspect getkeys.command arg1 arg2 key key1 arg3 key key2 key key3
+    } {key1 key2 key3}
+
+    test {RM_GetCommandKeys on a non-movable module command} {
+        r getkeys.introspect getkeys.fixed arg1 key1 key2 key3 arg2
+    } {key1 key2 key3}
+
+    test {RM_GetCommandKeys with bad arity} {
+        catch {r getkeys.introspect set key} e
+        set _ $e
+    } {*EINVAL*}
+}