Fix zuiFind crash / RM_ScanKey hang on SET object listpack encoding (redis#11581)

In redis#11290, we added listpack encoding for SET object.
But forgot to support it in zuiFind, causes ZINTER, ZINTERSTORE,
ZINTERCARD, ZIDFF, ZDIFFSTORE to crash.
And forgot to support it in RM_ScanKey, causes it hang.

This PR add support SET listpack in zuiFind, and in RM_ScanKey.
And add tests for related commands to cover this case.

Other changes:
- There is no reason for zuiFind to go into the internals of the SET.
  It can simply use setTypeIsMember and don't care about encoding.
- Remove the `#include "intset.h"` from server.h reduce the chance of
  accidental intset API use.
- Move setTypeAddAux, setTypeRemoveAux and setTypeIsMemberAux
  interfaces to the header.
- In scanGenericCommand, use setTypeInitIterator and setTypeNext
  to handle OBJ_SET scan.
- In RM_ScanKey, improve hash scan mode, use lpGetValue like zset,
  they can share code and better performance.

The zuiFind part fixes redis#11578

Co-authored-by: Oran Agra <[email protected]>
Co-authored-by: Viktor Söderqvist <[email protected]>

Author: 3 people

src/db.c

@@ -946,14 +946,21 @@ void scanGenericCommand(client *c, robj *o, unsigned long cursor) {
         } while (cursor &&
               maxiterations-- &&
               listLength(keys) < (unsigned long)count);
-    } else if (o->type == OBJ_SET && o->encoding == OBJ_ENCODING_INTSET) {
-        int pos = 0;
-        int64_t ll;
-
-        while(intsetGet(o->ptr,pos++,&ll))
-            listAddNodeTail(keys,createStringObjectFromLongLong(ll));
+    } else if (o->type == OBJ_SET) {
+        char *str;
+        size_t len;
+        int64_t llele;
+        setTypeIterator *si = setTypeInitIterator(o);
+        while (setTypeNext(si, &str, &len, &llele) != -1) {
+            if (str == NULL) {
+                listAddNodeTail(keys, createStringObjectFromLongLong(llele));
+            } else {
+                listAddNodeTail(keys, createStringObject(str, len));
+            }
+        }
+        setTypeReleaseIterator(si);
         cursor = 0;
-    } else if ((o->type == OBJ_HASH || o->type == OBJ_ZSET || o->type == OBJ_SET) &&
+    } else if ((o->type == OBJ_HASH || o->type == OBJ_ZSET) &&
                o->encoding == OBJ_ENCODING_LISTPACK)
     {
         unsigned char *p = lpFirst(o->ptr);

src/module.c

@@ -10425,18 +10425,19 @@ int RM_ScanKey(RedisModuleKey *key, RedisModuleScanCursor *cursor, RedisModuleSc
             cursor->done = 1;
             ret = 0;
         }
-    } else if (o->type == OBJ_SET && o->encoding == OBJ_ENCODING_INTSET) {
-        int pos = 0;
-        int64_t ll;
-        while(intsetGet(o->ptr,pos++,&ll)) {
-            robj *field = createObject(OBJ_STRING,sdsfromlonglong(ll));
+    } else if (o->type == OBJ_SET) {
+        setTypeIterator *si = setTypeInitIterator(o);
+        sds sdsele;
+        while ((sdsele = setTypeNextObject(si)) != NULL) {
+            robj *field = createObject(OBJ_STRING, sdsele);
             fn(key, field, NULL, privdata);
             decrRefCount(field);
         }
+        setTypeReleaseIterator(si);
         cursor->cursor = 1;
         cursor->done = 1;
         ret = 0;
-    } else if (o->type == OBJ_ZSET) {
+    } else if (o->type == OBJ_ZSET || o->type == OBJ_HASH) {
         unsigned char *p = lpSeek(o->ptr,0);
         unsigned char *vstr;
         unsigned int vlen;
@@ -10459,25 +10460,6 @@ int RM_ScanKey(RedisModuleKey *key, RedisModuleScanCursor *cursor, RedisModuleSc
         cursor->cursor = 1;
         cursor->done = 1;
         ret = 0;
-    } else if (o->type == OBJ_HASH) {
-        unsigned char *p = lpFirst(o->ptr);
-        unsigned char *vstr;
-        int64_t vlen;
-        unsigned char intbuf[LP_INTBUF_SIZE];
-        while(p) {
-            vstr = lpGet(p,&vlen,intbuf);
-            robj *field = createStringObject((char*)vstr,vlen);
-            p = lpNext(o->ptr,p);
-            vstr = lpGet(p,&vlen,intbuf);
-            robj *value = createStringObject((char*)vstr,vlen);
-            fn(key, field, value, privdata);
-            p = lpNext(o->ptr,p);
-            decrRefCount(field);
-            decrRefCount(value);
-        }
-        cursor->cursor = 1;
-        cursor->done = 1;
-        ret = 0;
     }
     errno = 0;
     return ret;

src/object.c

@@ -30,6 +30,7 @@
 
 #include "server.h"
 #include "functions.h"
+#include "intset.h"  /* Compact integer set structure */
 #include <math.h>
 #include <ctype.h>
 

src/rdb.c

@@ -34,6 +34,7 @@
 #include "fpconv_dtoa.h"
 #include "stream.h"
 #include "functions.h"
+#include "intset.h"  /* Compact integer set structure */
 
 #include <math.h>
 #include <fcntl.h>

src/server.c

@@ -6856,6 +6856,7 @@ int iAmMaster(void) {
 
 #ifdef REDIS_TEST
 #include "testhelp.h"
+#include "intset.h"  /* Compact integer set structure */
 
 int __failed_tests = 0;
 int __test_num = 0;

src/server.h

@@ -71,7 +71,6 @@ typedef long long ustime_t; /* microsecond time type. */
 #include "adlist.h"  /* Linked lists */
 #include "zmalloc.h" /* total memory usage aware version of malloc/free */
 #include "anet.h"    /* Networking the easy way */
-#include "intset.h"  /* Compact integer set structure */
 #include "version.h" /* Version macro */
 #include "util.h"    /* Misc functions useful in many places */
 #include "latency.h" /* Latency monitor API */
@@ -3012,8 +3011,11 @@ int restartServer(int flags, mstime_t delay);
 /* Set data type */
 robj *setTypeCreate(sds value);
 int setTypeAdd(robj *subject, sds value);
+int setTypeAddAux(robj *set, char *str, size_t len, int64_t llval, int str_is_sds);
 int setTypeRemove(robj *subject, sds value);
+int setTypeRemoveAux(robj *set, char *str, size_t len, int64_t llval, int str_is_sds);
 int setTypeIsMember(robj *subject, sds value);
+int setTypeIsMemberAux(robj *set, char *str, size_t len, int64_t llval, int str_is_sds);
 setTypeIterator *setTypeInitIterator(robj *subject);
 void setTypeReleaseIterator(setTypeIterator *si);
 int setTypeNext(setTypeIterator *si, char **str, size_t *len, int64_t *llele);

src/t_set.c

@@ -28,12 +28,7 @@
  */
 
 #include "server.h"
-
-/* Internal prototypes */
-
-int setTypeAddAux(robj *set, char *str, size_t len, int64_t llval, int str_is_sds);
-int setTypeRemoveAux(robj *set, char *str, size_t len, int64_t llval, int str_is_sds);
-int setTypeIsMemberAux(robj *set, char *str, size_t len, int64_t llval, int str_is_sds);
+#include "intset.h"  /* Compact integer set structure */
 
 /*-----------------------------------------------------------------------------
  * Set Commands

src/t_zset.c

@@ -57,6 +57,7 @@
  * from tail to head, useful for ZREVRANGE. */
 
 #include "server.h"
+#include "intset.h"  /* Compact integer set structure */
 #include <math.h>
 
 /*-----------------------------------------------------------------------------
@@ -2257,26 +2258,13 @@ int zuiFind(zsetopsrc *op, zsetopval *val, double *score) {
         return 0;
 
     if (op->type == OBJ_SET) {
-        if (op->encoding == OBJ_ENCODING_INTSET) {
-            if (zuiLongLongFromValue(val) &&
-                intsetFind(op->subject->ptr,val->ell))
-            {
-                *score = 1.0;
-                return 1;
-            } else {
-                return 0;
-            }
-        } else if (op->encoding == OBJ_ENCODING_HT) {
-            dict *ht = op->subject->ptr;
-            zuiSdsFromValue(val);
-            if (dictFind(ht,val->ele) != NULL) {
-                *score = 1.0;
-                return 1;
-            } else {
-                return 0;
-            }
+        char *str = val->ele ? val->ele : (char *)val->estr;
+        size_t len = val->ele ? sdslen(val->ele) : val->elen;
+        if (setTypeIsMemberAux(op->subject, str, len, val->ell, val->ele != NULL)) {
+            *score = 1.0;
+            return 1;
         } else {
-            serverPanic("Unknown set encoding");
+            return 0;
         }
     } else if (op->type == OBJ_ZSET) {
         zuiSdsFromValue(val);

tests/unit/moduleapi/scan.tcl

@@ -12,45 +12,58 @@ start_server {tags {"modules"}} {
         lsort [r scan.scan_strings]
     } {{x 1} {y 2} {z 3}}
 
-    test {Module scan hash ziplist} {
+    test {Module scan hash listpack} {
         r hmset hh f1 v1 f2 v2
+        assert_encoding listpack hh
         lsort [r scan.scan_key hh]
     } {{f1 v1} {f2 v2}}
-    
-    test {Module scan hash dict with int value} {
-        r hmset hh1 f1 1 
+
+    test {Module scan hash listpack with int value} {
+        r hmset hh1 f1 1
+        assert_encoding listpack hh1
         lsort [r scan.scan_key hh1]
     } {{f1 1}}
 
     test {Module scan hash dict} {
         r config set hash-max-ziplist-entries 2
         r hmset hh f3 v3
+        assert_encoding hashtable hh
         lsort [r scan.scan_key hh]
     } {{f1 v1} {f2 v2} {f3 v3}}
 
-    test {Module scan zset ziplist} {
+    test {Module scan zset listpack} {
         r zadd zz 1 f1 2 f2
+        assert_encoding listpack zz
         lsort [r scan.scan_key zz]
     } {{f1 1} {f2 2}}
 
-    test {Module scan zset dict} {
+    test {Module scan zset skiplist} {
         r config set zset-max-ziplist-entries 2
         r zadd zz 3 f3
+        assert_encoding skiplist zz
         lsort [r scan.scan_key zz]
     } {{f1 1} {f2 2} {f3 3}}
 
     test {Module scan set intset} {
         r sadd ss 1 2
+        assert_encoding intset ss
         lsort [r scan.scan_key ss]
     } {{1 {}} {2 {}}}
 
     test {Module scan set dict} {
         r config set set-max-intset-entries 2
         r sadd ss 3
+        assert_encoding hashtable ss
         lsort [r scan.scan_key ss]
     } {{1 {}} {2 {}} {3 {}}}
 
+    test {Module scan set listpack} {
+        r sadd ss1 a b c
+        assert_encoding listpack ss1
+        lsort [r scan.scan_key ss1]
+    } {{a {}} {b {}} {c {}}}
+
     test "Unload the module - scan" {
         assert_equal {OK} [r module unload scan]
     }
-}
+}

tests/unit/type/zset.tcl

@@ -2480,4 +2480,83 @@ start_server {tags {"zset"}} {
         r zscore zz dblmax
     } {1.7976931348623157e+308}
 
+    test {zunionInterDiffGenericCommand acts on SET and ZSET} {
+        r del set_small{t} set_big{t} zset_small{t} zset_big{t} zset_dest{t}
+
+        foreach set_type {intset listpack hashtable} {
+            # Restore all default configurations before each round of testing.
+            r config set set-max-intset-entries 512
+            r config set set-max-listpack-entries 128
+            r config set zset-max-listpack-entries 128
+
+            r del set_small{t} set_big{t}
+
+            if {$set_type == "intset"} {
+                r sadd set_small{t} 1 2 3
+                r sadd set_big{t} 1 2 3 4 5
+                assert_encoding intset set_small{t}
+                assert_encoding intset set_big{t}
+            } elseif {$set_type == "listpack"} {
+                # Add an "a" and then remove it, make sure the set is listpack encoding.
+                r sadd set_small{t} a 1 2 3
+                r sadd set_big{t} a 1 2 3 4 5
+                r srem set_small{t} a
+                r srem set_big{t} a
+                assert_encoding listpack set_small{t}
+                assert_encoding listpack set_big{t}
+            } elseif {$set_type == "hashtable"} {
+                r config set set-max-intset-entries 0
+                r config set set-max-listpack-entries 0
+                r sadd set_small{t} 1 2 3
+                r sadd set_big{t} 1 2 3 4 5
+                assert_encoding hashtable set_small{t}
+                assert_encoding hashtable set_big{t}
+            }
+
+            foreach zset_type {listpack skiplist} {
+                r del zset_small{t} zset_big{t}
+
+                if {$zset_type == "listpack"} {
+                    r zadd zset_small{t} 1 1 2 2 3 3
+                    r zadd zset_big{t} 1 1 2 2 3 3 4 4 5 5
+                    assert_encoding listpack zset_small{t}
+                    assert_encoding listpack zset_big{t}
+                } elseif {$zset_type == "skiplist"} {
+                    r config set zset-max-listpack-entries 0
+                    r zadd zset_small{t} 1 1 2 2 3 3
+                    r zadd zset_big{t} 1 1 2 2 3 3 4 4 5 5
+                    assert_encoding skiplist zset_small{t}
+                    assert_encoding skiplist zset_big{t}
+                }
+
+                # Test one key is big and one key is small separately.
+                # The reason for this is because we will sort the sets from smallest to largest.
+                # So set one big key and one small key, then the test can cover more code paths.
+                foreach {small_or_big set_key zset_key} {
+                    small set_small{t} zset_big{t}
+                    big set_big{t} zset_small{t}
+                } {
+                    # The result of these commands are not related to the order of the keys.
+                    assert_equal {1 2 3 4 5} [lsort [r zunion 2 $set_key $zset_key]]
+                    assert_equal {5} [r zunionstore zset_dest{t} 2 $set_key $zset_key]
+                    assert_equal {1 2 3} [lsort [r zinter 2 $set_key $zset_key]]
+                    assert_equal {3} [r zinterstore zset_dest{t} 2 $set_key $zset_key]
+                    assert_equal {3} [r zintercard 2 $set_key $zset_key]
+
+                    # The result of sdiff is related to the order of the keys.
+                    if {$small_or_big == "small"} {
+                        assert_equal {} [r zdiff 2 $set_key $zset_key]
+                        assert_equal {0} [r zdiffstore zset_dest{t} 2 $set_key $zset_key]
+                    } else {
+                        assert_equal {4 5} [lsort [r zdiff 2 $set_key $zset_key]]
+                        assert_equal {2} [r zdiffstore zset_dest{t} 2 $set_key $zset_key]
+                    }
+                }
+            }
+        }
+
+        r config set set-max-intset-entries 512
+        r config set set-max-listpack-entries 128
+        r config set zset-max-listpack-entries 128
+    }
 }