From 991aff1c0fe818627b0f8a7d9a6cc7ff720d3682 Mon Sep 17 00:00:00 2001
From: Chen Tianjie <chentianjie.ctj@alibaba-inc.com>
Date: Tue, 5 Dec 2023 22:21:50 +0800
Subject: [PATCH] Optimize KEYS when pattern includes hashtag and implies a
 single slot. (#12754)

in #12536 we made a similar optimization for SCAN, now that hashtags in
patterns. When we can make sure all keys matching the pettern will be in
the same slot, we can limit the iteration to run only one one.
---
 src/db.c                | 21 +++++++++++++++++----
 tests/unit/keyspace.tcl |  8 ++++++++
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/src/db.c b/src/db.c
index aded076b7..3d0a5ebca 100644
--- a/src/db.c
+++ b/src/db.c
@@ -33,6 +33,7 @@
 #include "latency.h"
 #include "script.h"
 #include "functions.h"
+#include "cluster.h"
 
 #include <signal.h>
 #include <ctype.h>
@@ -993,13 +994,22 @@ void randomkeyCommand(client *c) {
 void keysCommand(client *c) {
     dictEntry *de;
     sds pattern = c->argv[1]->ptr;
-    int plen = sdslen(pattern), allkeys;
+    int plen = sdslen(pattern), allkeys, pslot = -1;
     long numkeys = 0;
     void *replylen = addReplyDeferredLen(c);
-    dbIterator *dbit = dbIteratorInit(c->db, DB_MAIN);
     allkeys = (pattern[0] == '*' && plen == 1);
+    if (!allkeys) {
+        pslot = patternHashSlot(pattern, plen);
+    }
+    dictIterator *di = NULL;
+    dbIterator *dbit = NULL;
+    if (server.cluster_enabled && !allkeys && pslot != -1) {
+        di = dictGetSafeIterator(c->db->dict[pslot]);
+    } else {
+        dbit = dbIteratorInit(c->db, DB_MAIN);
+    }
     robj keyobj;
-    while ((de = dbIteratorNext(dbit)) != NULL) {
+    while ((de = di ? dictNext(di) : dbIteratorNext(dbit)) != NULL) {
         sds key = dictGetKey(de);
 
         if (allkeys || stringmatchlen(pattern,plen,key,sdslen(key),0)) {
@@ -1012,7 +1022,10 @@ void keysCommand(client *c) {
         if (c->flags & CLIENT_CLOSE_ASAP)
             break;
     }
-    dbReleaseIterator(dbit);
+    if (di)
+        dictReleaseIterator(di);
+    if (dbit)
+        dbReleaseIterator(dbit);
     setDeferredArrayLen(c,replylen,numkeys);
 }
 
diff --git a/tests/unit/keyspace.tcl b/tests/unit/keyspace.tcl
index 0769dbf3e..31130e4c6 100644
--- a/tests/unit/keyspace.tcl
+++ b/tests/unit/keyspace.tcl
@@ -34,6 +34,14 @@ start_server {tags {"keyspace"}} {
         r dbsize
     } {6}
 
+    test {KEYS with hashtag} {
+        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 {DEL all keys} {
         foreach key [r keys *] {r del $key}
         r dbsize