Implement custom command

proxy.conf

@@ -164,4 +164,28 @@ log-level error
 #
 # dump-queues no
 
-
+# Define custom commands for modules, rename-commands or new commands of redis,
+# especially, you can use redis modules easily by defining custom commands.
+#
+# The custom command is composed of the following fields:
+#   name:            A string representing the command name
+#   arity:           Number of arguments, it is possible to use -N to say >= N
+#   first_key_index: First argument that is a key
+#   last_key_index:  Last argument that is a key
+#   key_step:        Step to get all the keys from first to last argument.
+#                    For instance in redis command 'MSET' the step is two
+#                    since arguments are key,val,key,val,...
+#
+# Cross-slot queries is unsupported currently becasue they are different
+# for various commands to hanle their replie. You can use hash tag to make
+# all the keys in queries belong to the same slot.
+#
+# If rename command 'mset' to 'redis.mset' in redis server, we can define
+# the following command, but you should use hash tag as mentioned above.
+# custom-command "redis.mset -3 1 -1 2"
+#
+# For modules, we can define the following commands if use helloworld module
+# in redis repository https://github.com/antirez/redis/tree/6.0/src/modules
+#
+# custom-command "hello.push.call 3 1 1 1"
+# custom-command "hello.list.sum.LEN 2 1 1 1"

src/commands.c

@@ -17,6 +17,7 @@
 
 
 #include "commands.h"
+#include "zmalloc.h"
 
 /* Command Handlers */
 int proxyCommand(void *req);
@@ -274,3 +275,26 @@ struct redisCommandDef redisCommandTable[203] = {
     /* Custom Commands */
     {"proxy", -2, 0, 0, 0, 0, 0, NULL, proxyCommand, NULL}
 };
+
+/* Create a new custom proxy command, that is used for module-commands,
+ * rename-commands and new commands of redis.
+ *
+ * Cross-slot queries are unsupported currently becasue they are different
+ * for various commands to hanle their replies. */
+redisCommandDef *createProxyCustomCommand(char *name, int arity,
+                      int first_key, int last_key, int key_step)
+{
+    redisCommandDef *cmd = zcalloc(sizeof(redisCommandDef));
+    cmd->name = zstrdup(name);
+    cmd->arity = arity;
+    cmd->first_key = first_key;
+    cmd->last_key = last_key;
+    cmd->key_step = key_step;
+    cmd->proxy_flags = CMDFLAG_MULTISLOT_UNSUPPORTED;
+    cmd->unsupported = 0;
+    cmd->get_keys = NULL;
+    cmd->handle = NULL;
+    cmd->handleReply = NULL;
+
+    return cmd;
+}

src/commands.h

@@ -61,6 +61,8 @@ typedef struct redisCommandDef {
     redisClusterProxyReplyHandler   *handleReply;
 } redisCommandDef;
 
+redisCommandDef *createProxyCustomCommand(char *name, int arity,
+                      int first_key, int last_key, int key_step);
 
 extern struct redisCommandDef redisCommandTable[203];
 

src/config.c

@@ -43,6 +43,7 @@ void initConfig(void) {
     config.dump_queues = 0;
     config.auth = NULL;
     config.auth_user = NULL;
+    config.custom_commands = raxNew();
     config.cross_slot_enabled = 0;
     config.bindaddr_count = 0;
     config.pidfile = NULL;

src/config.h

@@ -18,6 +18,7 @@
 #ifndef __REDIS_CLUSTER_PROXY_CONFIG_H__
 #define __REDIS_CLUSTER_PROXY_CONFIG_H__
 
+#include "rax.h"
 #include "redis_config.h"
 
 #define CFG_DISABLE_MULTIPLEXING_AUTO       1
@@ -65,6 +66,7 @@ typedef struct {
     int dump_queues;
     char *auth;
     char *auth_user;
+    rax  *custom_commands;
     int disable_multiplexing;
     int cross_slot_enabled;
     int bindaddr_count;

src/proxy.c

@@ -1741,6 +1741,20 @@ int parseOptions(int argc, char **argv) {
                 exit(1);
             }
             config.bindaddr[config.bindaddr_count++] = zstrdup(argv[++i]);
+        } else if (!strcmp(arg,"--custom-command") && !lastarg) {
+            int number = 0;
+            sds *fields = sdssplitargs(argv[++i], &number);
+            if (fields == NULL || number < 5) {
+                fprintf(stderr, "Custom-command '%s' format error\n", argv[i]);
+                exit(1);
+            }
+            sdstolower(fields[0]);
+            redisCommandDef *cmd = createProxyCustomCommand(fields[0],
+                atoi(fields[1]), atoi(fields[2]),
+                atoi(fields[3]), atoi(fields[4]));
+            raxInsert(config.custom_commands, (unsigned char*)cmd->name,
+                strlen(cmd->name), cmd, NULL);
+            sdsfreesplitres(fields, number);
         } else if (!strcmp("-c", arg) && !lastarg) {
             char *cfgfile = argv[++i];
             if (!parseOptionsFromFile(cfgfile)) exit(1);
@@ -1921,6 +1935,20 @@ static void initProxy(void) {
         if (strcasecmp("auth", cmd->name) == 0) authCommandDef = cmd;
         else if (strcasecmp("scan", cmd->name) == 0) scanCommandDef = cmd;
     }
+    /* Populate custom commands. */
+    raxIterator iter;
+    raxStart(&iter, config.custom_commands);
+    if (!raxSeek(&iter, "^", NULL, 0)) {
+        raxStop(&iter);
+        exit(1);
+    }
+    while (raxNext(&iter)) {
+       redisCommandDef *cmd = (redisCommandDef *)iter.data;
+       raxInsert(proxy.commands, (unsigned char*)cmd->name,
+                  strlen(cmd->name), cmd, NULL);
+    }
+    raxStop(&iter);
+
     proxy.main_loop = aeCreateEventLoop(proxy.min_reserved_fds);
     proxy.threads = zmalloc(config.num_threads *
                             sizeof(proxyThread *));