Added cache for COMMAND

Signed-off-by: Evgeny Barskiy <[email protected]>

Author: ebarskiy

src/module.c

@@ -1385,6 +1385,8 @@ int VM_CreateCommand(ValkeyModuleCtx *ctx,
     serverAssert(hashtableAdd(server.commands, cp->serverCmd));
     serverAssert(hashtableAdd(server.orig_commands, cp->serverCmd));
     cp->serverCmd->id = ACLGetCommandID(declared_name); /* ID used for ACL. */
+    /* Invalidate COMMAND response cache since we added a new command */
+    invalidateCommandCache();
     return VALKEYMODULE_OK;
 }
 
@@ -12530,6 +12532,14 @@ int moduleFreeCommand(struct ValkeyModule *module, struct serverCommand *cmd) {
         hdr_close(cmd->latency_histogram);
         cmd->latency_histogram = NULL;
     }
+    if (cmd->info_cache_resp2) {
+        sdsfree(cmd->info_cache_resp2);
+        cmd->info_cache_resp2 = NULL;
+    }
+    if (cmd->info_cache_resp3) {
+        sdsfree(cmd->info_cache_resp3);
+        cmd->info_cache_resp3 = NULL;
+    }
     moduleFreeArgs(cmd->args, cmd->num_args);
     zfree(cp);
 
@@ -12571,6 +12581,8 @@ void moduleUnregisterCommands(struct ValkeyModule *module) {
         zfree(cmd);
     }
     hashtableResetIterator(&iter);
+    /* Invalidate COMMAND response cache since we removed commands */
+    invalidateCommandCache();
 }
 
 /* We parse argv to add sds "NAME VALUE" pairs to the server.module_configs_queue list of configs.

src/server.c

@@ -2342,6 +2342,8 @@ void initServerConfig(void) {
      * valkey.conf using the rename-command directive. */
     server.commands = hashtableCreate(&commandSetType);
     server.orig_commands = hashtableCreate(&originalCommandSetType);
+    server.command_response_cache_resp2 = NULL;
+    server.command_response_cache_resp3 = NULL;
     populateCommandTable();
 
     /* Debugging */
@@ -3282,6 +3284,10 @@ int populateCommandStructure(struct serverCommand *c) {
      * has been issued for the first time */
     c->latency_histogram = NULL;
 
+    /* Initialize command info cache */
+    c->info_cache_resp2 = NULL;
+    c->info_cache_resp3 = NULL;
+
     /* Handle the legacy range spec and the "movablekeys" flag (must be done after populating all key specs). */
     populateCommandLegacyRangeSpec(c);
 
@@ -5222,30 +5228,77 @@ void addReplyCommandSubCommands(client *c,
     hashtableResetIterator(&iter);
 }
 
+/* Forward declaration */
+void addReplyCommandInfo(client *c, struct serverCommand *cmd);
+
+/* Collect all output from a caching client (both buffer and reply list) */
+static sds collectCachedResponse(client *c) {
+    sds response = sdsempty();
+    
+    /* First, collect from the fixed buffer if any */
+    if (c->bufpos > 0) {
+        response = sdscatlen(response, c->buf, c->bufpos);
+    }
+    
+    /* Then, collect from the reply list */
+    listIter li;
+    listNode *ln;
+    clientReplyBlock *val_block;
+    listRewind(c->reply, &li);
+    while ((ln = listNext(&li)) != NULL) {
+        val_block = (clientReplyBlock *)listNodeValue(ln);
+        response = sdscatlen(response, val_block->buf, val_block->used);
+    }
+    
+    return response;
+}
+
+/* Generate and cache the command info response for a given protocol version */
+static void cacheCommandInfo(struct serverCommand *cmd, int resp) {
+    client *caching_client = createCachedResponseClient(resp);
+    
+    int firstkey = 0, lastkey = 0, keystep = 0;
+    if (cmd->legacy_range_key_spec.begin_search_type != KSPEC_BS_INVALID) {
+        firstkey = cmd->legacy_range_key_spec.bs.index.pos;
+        lastkey = cmd->legacy_range_key_spec.fk.range.lastkey;
+        if (lastkey >= 0) lastkey += firstkey;
+        keystep = cmd->legacy_range_key_spec.fk.range.keystep;
+    }
+
+    addReplyArrayLen(caching_client, 10);
+    addReplyBulkCBuffer(caching_client, cmd->fullname, sdslen(cmd->fullname));
+    addReplyLongLong(caching_client, cmd->arity);
+    addReplyFlagsForCommand(caching_client, cmd);
+    addReplyLongLong(caching_client, firstkey);
+    addReplyLongLong(caching_client, lastkey);
+    addReplyLongLong(caching_client, keystep);
+    addReplyCommandCategories(caching_client, cmd);
+    addReplyCommandTips(caching_client, cmd);
+    addReplyCommandKeySpecs(caching_client, cmd);
+    addReplyCommandSubCommands(caching_client, cmd, addReplyCommandInfo, 0);
+
+    if (resp == 2) {
+        cmd->info_cache_resp2 = collectCachedResponse(caching_client);
+    } else {
+        cmd->info_cache_resp3 = collectCachedResponse(caching_client);
+    }
+    
+    deleteCachedResponseClient(caching_client);
+}
+
 /* Output the representation of a server command. Used by the COMMAND command and COMMAND INFO. */
 void addReplyCommandInfo(client *c, struct serverCommand *cmd) {
     if (!cmd) {
         addReplyNull(c);
     } else {
-        int firstkey = 0, lastkey = 0, keystep = 0;
-        if (cmd->legacy_range_key_spec.begin_search_type != KSPEC_BS_INVALID) {
-            firstkey = cmd->legacy_range_key_spec.bs.index.pos;
-            lastkey = cmd->legacy_range_key_spec.fk.range.lastkey;
-            if (lastkey >= 0) lastkey += firstkey;
-            keystep = cmd->legacy_range_key_spec.fk.range.keystep;
+        /* Use cached response if available for the client's protocol version */
+        sds *cache = (c->resp == 2) ? &cmd->info_cache_resp2 : &cmd->info_cache_resp3;
+        
+        if (*cache == NULL) {
+            cacheCommandInfo(cmd, c->resp);
         }
-
-        addReplyArrayLen(c, 10);
-        addReplyBulkCBuffer(c, cmd->fullname, sdslen(cmd->fullname));
-        addReplyLongLong(c, cmd->arity);
-        addReplyFlagsForCommand(c, cmd);
-        addReplyLongLong(c, firstkey);
-        addReplyLongLong(c, lastkey);
-        addReplyLongLong(c, keystep);
-        addReplyCommandCategories(c, cmd);
-        addReplyCommandTips(c, cmd);
-        addReplyCommandKeySpecs(c, cmd);
-        addReplyCommandSubCommands(c, cmd, addReplyCommandInfo, 0);
+        
+        addReplyProto(c, *cache, sdslen(*cache));
     }
 }
 
@@ -5371,16 +5424,50 @@ void getKeysSubcommand(client *c) {
 }
 
 /* COMMAND (no args) */
-void commandCommand(client *c) {
+/* Invalidate the cached COMMAND response when command table changes */
+void invalidateCommandCache(void) {
+    if (server.command_response_cache_resp2) {
+        sdsfree(server.command_response_cache_resp2);
+        server.command_response_cache_resp2 = NULL;
+    }
+    if (server.command_response_cache_resp3) {
+        sdsfree(server.command_response_cache_resp3);
+        server.command_response_cache_resp3 = NULL;
+    }
+}
+
+/* Generate and cache the full COMMAND response */
+static void cacheCommandResponse(int resp) {
+    client *caching_client = createCachedResponseClient(resp);
+    
     hashtableIterator iter;
     void *next;
-    addReplyArrayLen(c, hashtableSize(server.commands));
+    addReplyArrayLen(caching_client, hashtableSize(server.commands));
     hashtableInitIterator(&iter, server.commands, 0);
     while (hashtableNext(&iter, &next)) {
         struct serverCommand *cmd = next;
-        addReplyCommandInfo(c, cmd);
+        addReplyCommandInfo(caching_client, cmd);
     }
     hashtableResetIterator(&iter);
+    
+    if (resp == 2) {
+        server.command_response_cache_resp2 = collectCachedResponse(caching_client);
+    } else {
+        server.command_response_cache_resp3 = collectCachedResponse(caching_client);
+    }
+    
+    deleteCachedResponseClient(caching_client);
+}
+
+void commandCommand(client *c) {
+    /* Use cached response if available for the client's protocol version */
+    sds *cache = (c->resp == 2) ? &server.command_response_cache_resp2 : &server.command_response_cache_resp3;
+    
+    if (*cache == NULL) {
+        cacheCommandResponse(c->resp);
+    }
+    
+    addReplyProto(c, *cache, sdslen(*cache));
 }
 
 /* COMMAND COUNT */

src/server.h

@@ -1677,6 +1677,8 @@ struct valkeyServer {
     serverDb **db;            /* each db created when it's first used */
     hashtable *commands;      /* Command table */
     hashtable *orig_commands; /* Command table before command renaming. */
+    sds command_response_cache_resp2; /* Cached COMMAND response for RESP2 */
+    sds command_response_cache_resp3; /* Cached COMMAND response for RESP3 */
     aeEventLoop *el;
     _Atomic AeIoState io_poll_state;     /* Indicates the state of the IO polling. */
     int io_ae_fired_events;              /* Number of poll events received by the IO thread. */
@@ -2626,6 +2628,8 @@ struct serverCommand {
                                     * (not the fullname), and the value is the serverCommand structure pointer. */
     struct serverCommand *parent;
     struct ValkeyModuleCommand *module_cmd; /* A pointer to the module command data (NULL if native command) */
+    sds info_cache_resp2;                   /* Cached COMMAND INFO response for RESP2 */
+    sds info_cache_resp3;                   /* Cached COMMAND INFO response for RESP3 */
 };
 
 struct serverError {
@@ -3783,6 +3787,7 @@ void commandCommand(client *c);
 void commandCountCommand(client *c);
 void commandListCommand(client *c);
 void commandInfoCommand(client *c);
+void invalidateCommandCache(void);
 void commandGetKeysCommand(client *c);
 void commandGetKeysAndFlagsCommand(client *c);
 void commandHelpCommand(client *c);