From f20d338d1de43d6676275630737aa89530ab8721 Mon Sep 17 00:00:00 2001
From: Binbin <binloveplay1314@qq.com>
Date: Tue, 27 Aug 2024 12:04:27 +0800
Subject: [PATCH] Make KEYS to be an exact match if there is no pattern (#792)

Although KEYS is a dangerous command and we recommend people
to avoid using it, some people who are not familiar with it
still using it, and even use KEYS with no pattern at all.

Once KEYS is using with no pattern, we can convert it to an
exact match to avoid iterating over all data.

Signed-off-by: Binbin <binloveplay1314@qq.com>
Signed-off-by: Ping Xie <pingxie@google.com>
---
 src/db.c                | 35 ++++++++++++++++++++++++++
 tests/unit/keyspace.tcl | 55 ++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 89 insertions(+), 1 deletion(-)

diff --git a/src/db.c b/src/db.c
index ac9b5d3d8ec..040dd0dbe14 100644
--- a/src/db.c
+++ b/src/db.c
@@ -808,6 +808,21 @@ void randomkeyCommand(client *c) {
     decrRefCount(key);
 }
 
+/* Returns 1 if the pattern can be an exact match in KEYS context. */
+int patternExactMatch(const char *pattern, int length) {
+    for (int i = 0; i < length; i++) {
+        if (pattern[i] == '*' || pattern[i] == '?' || pattern[i] == '[') {
+            /* Wildcard or character class found. Keys can be in anywhere. */
+            return 0;
+        } else if (pattern[i] == '\\') {
+            /* Escaped character. Computing the key name in this case is not
+             * implemented. We would need a temp buffer. */
+            return 0;
+        }
+    }
+    return 1;
+}
+
 void keysCommand(client *c) {
     dictEntry *de;
     sds pattern = c->argv[1]->ptr;
@@ -817,7 +832,27 @@ void keysCommand(client *c) {
     allkeys = (pattern[0] == '*' && plen == 1);
     if (server.cluster_enabled && !allkeys) {
         pslot = patternHashSlot(pattern, plen);
+    } else if (!server.cluster_enabled) {
+        pslot = 0;
     }
+
+    /* Once the pattern can do an exact match, we can convert
+     * it to a kvstoreDictFind to avoid iterating over all data. */
+    if (patternExactMatch(pattern, plen) && pslot != -1) {
+        de = kvstoreDictFind(c->db->keys, pslot, pattern);
+        if (de) {
+            robj keyobj;
+            sds key = dictGetKey(de);
+            initStaticStringObject(keyobj, key);
+            if (!keyIsExpired(c->db, &keyobj)) {
+                addReplyBulkCBuffer(c, key, sdslen(key));
+                numkeys++;
+            }
+        }
+        setDeferredArrayLen(c, replylen, numkeys);
+        return;
+    }
+
     kvstoreDictIterator *kvs_di = NULL;
     kvstoreIterator *kvs_it = NULL;
     if (pslot != -1) {
diff --git a/tests/unit/keyspace.tcl b/tests/unit/keyspace.tcl
index c8bdcac9847..fd2a622c8ec 100644
--- a/tests/unit/keyspace.tcl
+++ b/tests/unit/keyspace.tcl
@@ -40,7 +40,25 @@ start_server {tags {"keyspace"}} {
         }
         assert_equal [lsort [r keys "{a}*"]] [list "{a}x" "{a}y" "{a}z"]
         assert_equal [lsort [r keys "*{b}*"]] [list "{b}a" "{b}b" "{b}c"]
-    } 
+    }
+
+    test {KEYS with no pattern} {
+        r debug populate 1000
+        r set foo bar
+        r set foo{t} bar
+        r set foo\\ bar
+        r del non-exist
+        assert_equal [r keys foo] {foo}
+        assert_equal [r keys foo{t}] {foo{t}}
+        assert_equal [r keys foo\\] {foo\\}
+        assert_equal [r keys non-exist] {}
+
+        # Make sure the return value is consistent with the * one
+        assert_equal [r keys foo] [r keys *foo]
+        assert_equal [r keys foo{t}] [r keys *foo{t}]
+        assert_equal [r keys foo\\] [r keys *foo\\]
+        assert_equal [r keys non-exist] [r keys *non-exist]
+    } {} {needs:debug}
 
     test {DEL all keys} {
         foreach key [r keys *] {r del $key}
@@ -548,3 +566,38 @@ foreach {type large} [array get largevalue] {
        r flushall
     } {OK} {singledb:skip}
 }
+
+start_cluster 1 0 {tags {"keyspace external:skip cluster"}} {
+    # SET keys for random slots, for random noise.
+    set num_keys 0
+    while {$num_keys < 1000} {
+        set random_key [randomInt 163840]
+        r SET $random_key VALUE
+        incr num_keys 1
+    }
+
+    test {KEYS with hashtag in cluster mode} {
+        foreach key {"{a}x" "{a}y" "{a}z" "{b}a" "{b}b" "{b}c"} {
+            r set $key hello
+        }
+        assert_equal [lsort [r keys "{a}*"]] [list "{a}x" "{a}y" "{a}z"]
+        assert_equal [lsort [r keys "*{b}*"]] [list "{b}a" "{b}b" "{b}c"]
+    }
+
+    test {KEYS with no pattern in cluster mode} {
+        r set foo bar
+        r set foo{t} bar
+        r set foo\\ bar
+        r del non-exist
+        assert_equal [r keys foo] {foo}
+        assert_equal [r keys foo{t}] {foo{t}}
+        assert_equal [r keys foo\\] {foo\\}
+        assert_equal [r keys non-exist] {}
+
+        # Make sure the return value is consistent with the * one
+        assert_equal [r keys foo] [r keys *foo]
+        assert_equal [r keys foo{t}] [r keys *foo{t}]
+        assert_equal [r keys foo\\] [r keys *foo\\]
+        assert_equal [r keys non-exist] [r keys *non-exist]
+    }
+}