Add 'count' option to the 'RANDOMKEY' command

src/commands.def

@@ -2475,6 +2475,11 @@ const char *RANDOMKEY_Tips[] = {
 #define RANDOMKEY_Keyspecs NULL
 #endif
 
+/* RANDOMKEY argument table */
+struct COMMAND_ARG RANDOMKEY_Args[] = {
+{MAKE_ARG("count",ARG_TYPE_INTEGER,-1,"COUNT",NULL,NULL,CMD_ARG_OPTIONAL,0,NULL)},
+};
+
 /********** RENAME ********************/
 
 #ifndef SKIP_CMD_HISTORY_TABLE
@@ -11088,7 +11093,7 @@ struct COMMAND_STRUCT serverCommandTable[] = {
 {MAKE_CMD("pexpireat","Sets the expiration time of a key to a Unix milliseconds timestamp.","O(1)","2.6.0",CMD_DOC_NONE,NULL,NULL,"generic",COMMAND_GROUP_GENERIC,PEXPIREAT_History,1,PEXPIREAT_Tips,0,pexpireatCommand,-3,CMD_WRITE|CMD_FAST,ACL_CATEGORY_KEYSPACE,PEXPIREAT_Keyspecs,1,NULL,3),.args=PEXPIREAT_Args},
 {MAKE_CMD("pexpiretime","Returns the expiration time of a key as a Unix milliseconds timestamp.","O(1)","7.0.0",CMD_DOC_NONE,NULL,NULL,"generic",COMMAND_GROUP_GENERIC,PEXPIRETIME_History,0,PEXPIRETIME_Tips,0,pexpiretimeCommand,2,CMD_READONLY|CMD_FAST,ACL_CATEGORY_KEYSPACE,PEXPIRETIME_Keyspecs,1,NULL,1),.args=PEXPIRETIME_Args},
 {MAKE_CMD("pttl","Returns the expiration time in milliseconds of a key.","O(1)","2.6.0",CMD_DOC_NONE,NULL,NULL,"generic",COMMAND_GROUP_GENERIC,PTTL_History,1,PTTL_Tips,1,pttlCommand,2,CMD_READONLY|CMD_FAST,ACL_CATEGORY_KEYSPACE,PTTL_Keyspecs,1,NULL,1),.args=PTTL_Args},
-{MAKE_CMD("randomkey","Returns a random key name from the database.","O(1)","1.0.0",CMD_DOC_NONE,NULL,NULL,"generic",COMMAND_GROUP_GENERIC,RANDOMKEY_History,0,RANDOMKEY_Tips,3,randomkeyCommand,1,CMD_READONLY|CMD_TOUCHES_ARBITRARY_KEYS,ACL_CATEGORY_KEYSPACE,RANDOMKEY_Keyspecs,0,NULL,0)},
+{MAKE_CMD("randomkey","Returns a random key name or keys names from the database.","O(1) for count being 1 and O(N) with N being the number of keys in the database, under the assumption that the key names in the database and the given pattern have limited length.","1.0.0",CMD_DOC_NONE,NULL,NULL,"generic",COMMAND_GROUP_GENERIC,RANDOMKEY_History,0,RANDOMKEY_Tips,3,randomkeyCommand,-1,CMD_READONLY|CMD_TOUCHES_ARBITRARY_KEYS,ACL_CATEGORY_KEYSPACE,RANDOMKEY_Keyspecs,0,NULL,1),.args=RANDOMKEY_Args},
 {MAKE_CMD("rename","Renames a key and overwrites the destination.","O(1)","1.0.0",CMD_DOC_NONE,NULL,NULL,"generic",COMMAND_GROUP_GENERIC,RENAME_History,0,RENAME_Tips,0,renameCommand,3,CMD_WRITE,ACL_CATEGORY_KEYSPACE,RENAME_Keyspecs,2,NULL,2),.args=RENAME_Args},
 {MAKE_CMD("renamenx","Renames a key only when the target key name doesn't exist.","O(1)","1.0.0",CMD_DOC_NONE,NULL,NULL,"generic",COMMAND_GROUP_GENERIC,RENAMENX_History,1,RENAMENX_Tips,0,renamenxCommand,3,CMD_WRITE|CMD_FAST,ACL_CATEGORY_KEYSPACE,RENAMENX_Keyspecs,2,NULL,2),.args=RENAMENX_Args},
 {MAKE_CMD("restore","Creates a key from the serialized representation of a value.","O(1) to create the new key and additional O(N*M) to reconstruct the serialized value, where N is the number of objects composing the value and M their average size. For small string values the time complexity is thus O(1)+O(1*M) where M is small, so simply O(1). However for sorted set values the complexity is O(N*M*log(N)) because inserting values into sorted sets is O(log(N)).","2.6.0",CMD_DOC_NONE,NULL,NULL,"generic",COMMAND_GROUP_GENERIC,RESTORE_History,3,RESTORE_Tips,0,restoreCommand,-4,CMD_WRITE|CMD_DENYOOM,ACL_CATEGORY_KEYSPACE|ACL_CATEGORY_DANGEROUS,RESTORE_Keyspecs,1,NULL,7),.args=RESTORE_Args},

src/commands/randomkey.json

@@ -1,10 +1,10 @@
 {
     "RANDOMKEY": {
-        "summary": "Returns a random key name from the database.",
-        "complexity": "O(1)",
+        "summary": "Returns a random key name or keys names from the database.",
+        "complexity": "O(1) for count being 1 and O(N) with N being the number of keys in the database, under the assumption that the key names in the database and the given pattern have limited length.",
         "group": "generic",
         "since": "1.0.0",
-        "arity": 1,
+        "arity": -1,
         "function": "randomkeyCommand",
         "command_flags": [
             "READONLY",
@@ -21,14 +21,25 @@
         "reply_schema": {
             "oneOf": [
                 {
-                    "description": "When the database is empty.",
+                    "description": "when the database is empty or pattern not found",
                     "type": "null"
                 },
                 {
-                    "description": "Random key in db.",
-                    "type": "string"
+                    "description": "random key or keys in db",
+                    "type": "array",
+                    "items": {
+                        "type": "string"
+                    }
                 }
             ]
-        }
+        },
+        "arguments": [
+            {
+                "token": "COUNT",
+                "name": "count",
+                "type": "integer",
+                "optional": true
+            }
+        ]
     }
 }

src/db.c

@@ -871,7 +871,74 @@ void selectCommand(client *c) {
     }
 }
 
+void randomWithCountCommand(client *c) {
+    long l;
+    unsigned long count;
+    unsigned long numkeys = 0;
+    int canDuplicated = 0;
+    dict *section_dict;
+
+    if (getRangeLongFromObjectOrReply(c, c->argv[2], -LONG_MAX, LONG_MAX, &l, NULL) != C_OK) return;
+    if (l >= 0) {
+        count = (unsigned long)l;
+    } else {
+        /* A negative count means: return the same elements multiple times */
+        count = -l;
+        canDuplicated = 1;
+    }
+
+    /* If count is zero, serve it ASAP to avoid special cases later. */
+    if (count == 0) {
+        addReply(c, shared.emptyarray);
+        return;
+    }
+
+    if (canDuplicated) {
+        section_dict = dictCreate(&stringSetDictType);
+        dictExpand(section_dict, count);
+    }
+
+    int maxtries = 100 * count;
+    void *replylen = addReplyDeferredLen(c);
+    while (maxtries-- > 0 && count > 0) {
+        void *entry;
+        int randomDictIndex = kvstoreGetFairRandomHashtableIndex(c->db->keys);
+        int ok = kvstoreHashtableFairRandomEntry(c->db->keys, randomDictIndex, &entry);
+        if (!ok) {
+            continue;
+        }
+        robj *valkey = entry;
+        sds key = objectGetKey(valkey);
+        if (canDuplicated) {
+            sds dictKey = sdsdup(key);
+            if (dictAdd(section_dict, dictKey, NULL) != DICT_OK) {
+                sdsfree(dictKey);
+                continue;
+            }
+        }
+        addReplyBulkCBuffer(c, key, sdslen(key));
+        numkeys++;
+        count--;
+    }
+    setDeferredArrayLen(c, replylen, numkeys);
+    if (canDuplicated) {
+        dictRelease(section_dict);
+    }
+}
+
+/************************************************************
+RANDOMKEY                     numOfArgs = 1
+RANDOMKEY [COUNT <count> ]           numOfArgs = 3
+******************************************************************/
 void randomkeyCommand(client *c) {
+    if (c->argc == 3) {
+        randomWithCountCommand(c);
+        return;
+    } else if (c->argc > 3) {
+        addReplyErrorObject(c, shared.syntaxerr);
+        return;
+    }
+
     robj *key;
 
     if ((key = dbRandomKey(c->db)) == NULL) {

tests/unit/keyspace.tcl

@@ -492,6 +492,12 @@ foreach {type large} [array get largevalue] {
         r randomkey
     } {}
 
+    test {RANDOMKEY count} {
+        r flushdb
+        r set y 10
+        r randomkey count 1
+    } {y}
+
     test {RANDOMKEY regression 1} {
         r flushdb
         r set x 10