From b9a0500f16d0cd016398133cc7ac256ad927b679 Mon Sep 17 00:00:00 2001 From: Yang Bodong Date: Fri, 29 Jan 2021 16:47:28 +0800 Subject: [PATCH] Add HRANDFIELD and ZRANDMEMBER. improvements to SRANDMEMBER (#8297) New commands: `HRANDFIELD [ [WITHVALUES]]` `ZRANDMEMBER [ [WITHSCORES]]` Algorithms are similar to the one in SRANDMEMBER. Both return a simple bulk response when no arguments are given, and an array otherwise. In case values/scores are requested, RESP2 returns a long array, and RESP3 a nested array. note: in all 3 commands, the only option that also provides random order is the one with negative count. Changes to SRANDMEMBER * Optimization when count is 1, we can use the more efficient algorithm of non-unique random * optimization: work with sds strings rather than robj Other changes: * zzlGetScore: when zset needs to convert string to double, we use safer memcpy (in case the buffer is too small) * Solve a "bug" in SRANDMEMBER test: it intended to test a positive count (case 3 or case 4) and by accident used a negative count Co-authored-by: xinluton Co-authored-by: Oran Agra --- src/server.c | 19 +++ src/server.h | 3 + src/t_hash.c | 253 ++++++++++++++++++++++++++++++++++++ src/t_set.c | 36 +++--- src/t_zset.c | 269 ++++++++++++++++++++++++++++++++++++++- src/ziplist.c | 77 +++++++++++ src/ziplist.h | 11 ++ tests/support/util.tcl | 4 +- tests/unit/type/hash.tcl | 175 +++++++++++++++++++++++++ tests/unit/type/set.tcl | 2 +- tests/unit/type/zset.tcl | 176 +++++++++++++++++++++++++ 11 files changed, 1002 insertions(+), 23 deletions(-) diff --git a/src/server.c b/src/server.c index 0a12e2e31..79d6d813d 100644 --- a/src/server.c +++ b/src/server.c @@ -555,6 +555,10 @@ struct redisCommand redisCommandTable[] = { "write no-script fast @sortedset @blocking", 0,NULL,1,-2,1,0,0,0}, + {"zrandmember",zrandmemberCommand,-2, + "read-only random @sortedset", + 0,NULL,1,1,1,0,0,0}, + {"hset",hsetCommand,-4, "write use-memory fast @hash", 0,NULL,1,1,1,0,0,0}, @@ -611,6 +615,10 @@ struct redisCommand redisCommandTable[] = { "read-only fast @hash", 0,NULL,1,1,1,0,0,0}, + {"hrandfield",hrandfieldCommand,-2, + "read-only random @hash", + 0,NULL,1,1,1,0,0,0}, + {"hscan",hscanCommand,-3, "read-only random @hash", 0,NULL,1,1,1,0,0,0}, @@ -1456,6 +1464,17 @@ dictType hashDictType = { NULL /* allow to expand */ }; +/* Dict type without destructor */ +dictType sdsReplyDictType = { + dictSdsHash, /* hash function */ + NULL, /* key dup */ + NULL, /* val dup */ + dictSdsKeyCompare, /* key compare */ + NULL, /* key destructor */ + NULL, /* val destructor */ + NULL /* allow to expand */ +}; + /* Keylist hash table type has unencoded redis objects as keys and * lists as values. It's used for blocking operations (BLPOP) and to * map swapped keys to a list of clients waiting for this keys to be loaded. */ diff --git a/src/server.h b/src/server.h index 02ba8f7b9..1506c12a8 100644 --- a/src/server.h +++ b/src/server.h @@ -1718,6 +1718,7 @@ extern dictType hashDictType; extern dictType replScriptCacheDictType; extern dictType dbExpiresDictType; extern dictType modulesDictType; +extern dictType sdsReplyDictType; /*----------------------------------------------------------------------------- * Functions prototypes @@ -2555,6 +2556,7 @@ void zpopminCommand(client *c); void zpopmaxCommand(client *c); void bzpopminCommand(client *c); void bzpopmaxCommand(client *c); +void zrandmemberCommand(client *c); void multiCommand(client *c); void execCommand(client *c); void discardCommand(client *c); @@ -2588,6 +2590,7 @@ void hvalsCommand(client *c); void hgetallCommand(client *c); void hexistsCommand(client *c); void hscanCommand(client *c); +void hrandfieldCommand(client *c); void configCommand(client *c); void hincrbyCommand(client *c); void hincrbyfloatCommand(client *c); diff --git a/src/t_hash.c b/src/t_hash.c index 51c7d6758..9f7540a72 100644 --- a/src/t_hash.c +++ b/src/t_hash.c @@ -598,6 +598,42 @@ int hashZiplistValidateIntegrity(unsigned char *zl, size_t size, int deep) { return ret; } +/* Create a new sds string from the ziplist entry. */ +sds hashSdsFromZiplistEntry(ziplistEntry *e) { + return e->sval ? sdsnewlen(e->sval, e->slen) : sdsfromlonglong(e->lval); +} + +/* Reply with bulk string from the ziplist entry. */ +void hashReplyFromZiplistEntry(client *c, ziplistEntry *e) { + if (e->sval) + addReplyBulkCBuffer(c, e->sval, e->slen); + else + addReplyBulkLongLong(c, e->lval); +} + +/* Return random element from a non empty hash. + * 'key' and 'val' will be set to hold the element. + * The memory in them is not to be freed or modified by the caller. + * 'val' can be NULL in which case it's not extracted. */ +void hashTypeRandomElement(robj *hashobj, unsigned long hashsize, ziplistEntry *key, ziplistEntry *val) { + if (hashobj->encoding == OBJ_ENCODING_HT) { + dictEntry *de = dictGetFairRandomKey(hashobj->ptr); + sds s = dictGetKey(de); + key->sval = (unsigned char*)s; + key->slen = sdslen(s); + if (val) { + sds s = dictGetVal(de); + val->sval = (unsigned char*)s; + val->slen = sdslen(s); + } + } else if (hashobj->encoding == OBJ_ENCODING_ZIPLIST) { + ziplistRandomPair(hashobj->ptr, hashsize, key, val); + } else { + serverPanic("Unknown hash encoding"); + } +} + + /*----------------------------------------------------------------------------- * Hash type commands *----------------------------------------------------------------------------*/ @@ -922,3 +958,220 @@ void hscanCommand(client *c) { checkType(c,o,OBJ_HASH)) return; scanGenericCommand(c,o,cursor); } + +/* How many times bigger should be the hash compared to the requested size + * for us to not use the "remove elements" strategy? Read later in the + * implementation for more info. */ +#define HRANDFIELD_SUB_STRATEGY_MUL 3 + +void hrandfieldWithCountCommand(client *c, long l, int withvalues) { + unsigned long count, size; + int uniq = 1; + robj *hash; + + if ((hash = lookupKeyReadOrReply(c,c->argv[1],shared.null[c->resp])) + == NULL || checkType(c,hash,OBJ_HASH)) return; + size = hashTypeLength(hash); + + if(l >= 0) { + count = (unsigned long) l; + } else { + count = -l; + uniq = 0; + } + + /* If count is zero, serve it ASAP to avoid special cases later. */ + if (count == 0) { + addReply(c,shared.emptyarray); + return; + } + + /* CASE 1: The count was negative, so the extraction method is just: + * "return N random elements" sampling the whole set every time. + * This case is trivial and can be served without auxiliary data + * structures. This case is the only one that also needs to return the + * elements in random order. */ + if (!uniq || count == 1) { + if (withvalues && c->resp == 2) + addReplyArrayLen(c, count*2); + else + addReplyArrayLen(c, count); + if (hash->encoding == OBJ_ENCODING_HT) { + sds key, value; + while (count--) { + dictEntry *de = dictGetRandomKey(hash->ptr); + key = dictGetKey(de); + value = dictGetVal(de); + if (withvalues && c->resp > 2) + addReplyArrayLen(c,2); + addReplyBulkCBuffer(c, key, sdslen(key)); + if (withvalues) + addReplyBulkCBuffer(c, value, sdslen(value)); + } + } else if (hash->encoding == OBJ_ENCODING_ZIPLIST) { + ziplistEntry *keys, *vals = NULL; + keys = zmalloc(sizeof(ziplistEntry)*count); + if (withvalues) + vals = zmalloc(sizeof(ziplistEntry)*count); + ziplistRandomPairs(hash->ptr, count, keys, vals); + for (unsigned long i = 0; i < count; i++) { + if (withvalues && c->resp > 2) + addReplyArrayLen(c,2); + if (keys[i].sval) + addReplyBulkCBuffer(c, keys[i].sval, keys[i].slen); + else + addReplyBulkLongLong(c, keys[i].lval); + if (withvalues) { + if (vals[i].sval) + addReplyBulkCBuffer(c, vals[i].sval, vals[i].slen); + else + addReplyBulkLongLong(c, vals[i].lval); + } + } + zfree(keys); + zfree(vals); + } + return; + } + + /* Initiate reply count, RESP3 responds with nested array, RESP2 with flat one. */ + long reply_size = count < size ? count : size; + if (withvalues && c->resp == 2) + addReplyArrayLen(c, reply_size*2); + else + addReplyArrayLen(c, reply_size); + + /* CASE 2: + * The number of requested elements is greater than the number of + * elements inside the hash: simply return the whole hash. */ + if(count >= size) { + hashTypeIterator *hi = hashTypeInitIterator(hash); + while (hashTypeNext(hi) != C_ERR) { + if (withvalues && c->resp > 2) + addReplyArrayLen(c,2); + addHashIteratorCursorToReply(c, hi, OBJ_HASH_KEY); + if (withvalues) + addHashIteratorCursorToReply(c, hi, OBJ_HASH_VALUE); + } + hashTypeReleaseIterator(hi); + return; + } + + /* CASE 3: + * The number of elements inside the hash is not greater than + * HRANDFIELD_SUB_STRATEGY_MUL times the number of requested elements. + * In this case we create a hash from scratch with all the elements, and + * subtract random elements to reach the requested number of elements. + * + * This is done because if the number of requested elements is just + * a bit less than the number of elements in the hash, the natural approach + * used into CASE 4 is highly inefficient. */ + if (count*HRANDFIELD_SUB_STRATEGY_MUL > size) { + dict *d = dictCreate(&sdsReplyDictType, NULL); + hashTypeIterator *hi = hashTypeInitIterator(hash); + + /* Add all the elements into the temporary dictionary. */ + while ((hashTypeNext(hi)) != C_ERR) { + int ret = DICT_ERR; + sds key, value = NULL; + + key = hashTypeCurrentObjectNewSds(hi,OBJ_HASH_KEY); + if (withvalues) + value = hashTypeCurrentObjectNewSds(hi,OBJ_HASH_VALUE); + ret = dictAdd(d, key, value); + + serverAssert(ret == DICT_OK); + } + serverAssert(dictSize(d) == size); + hashTypeReleaseIterator(hi); + + /* Remove random elements to reach the right count. */ + while (size > count) { + dictEntry *de; + de = dictGetRandomKey(d); + dictUnlink(d,dictGetKey(de)); + sdsfree(dictGetKey(de)); + sdsfree(dictGetVal(de)); + dictFreeUnlinkedEntry(d,de); + size--; + } + + /* Reply with what's in the dict and release memory */ + dictIterator *di; + dictEntry *de; + di = dictGetIterator(d); + while ((de = dictNext(di)) != NULL) { + sds key = dictGetKey(de); + sds value = dictGetVal(de); + if (withvalues && c->resp > 2) + addReplyArrayLen(c,2); + addReplyBulkSds(c, key); + if (withvalues) + addReplyBulkSds(c, value); + } + + dictReleaseIterator(di); + dictRelease(d); + } + + /* CASE 4: We have a big hash compared to the requested number of elements. + * In this case we can simply get random elements from the hash and add + * to the temporary hash, trying to eventually get enough unique elements + * to reach the specified count. */ + else { + unsigned long added = 0; + ziplistEntry key, value; + dict *d = dictCreate(&hashDictType, NULL); + while(added < count) { + hashTypeRandomElement(hash, size, &key, withvalues? &value : NULL); + + /* Try to add the object to the dictionary. If it already exists + * free it, otherwise increment the number of objects we have + * in the result dictionary. */ + sds skey = hashSdsFromZiplistEntry(&key); + if (dictAdd(d,skey,NULL) != DICT_OK) { + sdsfree(skey); + continue; + } + added++; + + /* We can reply right away, so that we don't need to store the value in the dict. */ + if (withvalues && c->resp > 2) + addReplyArrayLen(c,2); + hashReplyFromZiplistEntry(c, &key); + if (withvalues) + hashReplyFromZiplistEntry(c, &value); + } + + /* Release memory */ + dictRelease(d); + } +} + +/* HRANDFIELD [ WITHVALUES] */ +void hrandfieldCommand(client *c) { + long l; + int withvalues = 0; + robj *hash; + ziplistEntry ele; + + if (c->argc >= 3) { + if (getLongFromObjectOrReply(c,c->argv[2],&l,NULL) != C_OK) return; + if (c->argc > 4 || (c->argc == 4 && strcasecmp(c->argv[3]->ptr,"withvalues"))) { + addReplyErrorObject(c,shared.syntaxerr); + return; + } else if (c->argc == 4) + withvalues = 1; + hrandfieldWithCountCommand(c, l, withvalues); + return; + } + + /* Handle variant without argument. Reply with simple bulk string */ + if ((hash = lookupKeyReadOrReply(c,c->argv[1],shared.null[c->resp]))== NULL || + checkType(c,hash,OBJ_HASH)) { + return; + } + + hashTypeRandomElement(hash,hashTypeLength(hash),&ele,NULL); + hashReplyFromZiplistEntry(c, &ele); +} diff --git a/src/t_set.c b/src/t_set.c index 64bbbd3a0..de0a9f954 100644 --- a/src/t_set.c +++ b/src/t_set.c @@ -690,8 +690,9 @@ void srandmemberWithCountCommand(client *c) { /* CASE 1: The count was negative, so the extraction method is just: * "return N random elements" sampling the whole set every time. * This case is trivial and can be served without auxiliary data - * structures. */ - if (!uniq) { + * structures. This case is the only one that also needs to return the + * elements in random order. */ + if (!uniq || count == 1) { addReplySetLen(c,count); while(count--) { encoding = setTypeRandomElement(set,&ele,&llele); @@ -713,7 +714,7 @@ void srandmemberWithCountCommand(client *c) { } /* For CASE 3 and CASE 4 we need an auxiliary dictionary. */ - d = dictCreate(&objectKeyPointerValueDictType,NULL); + d = dictCreate(&sdsReplyDictType,NULL); /* CASE 3: * The number of elements inside the set is not greater than @@ -729,13 +730,13 @@ void srandmemberWithCountCommand(client *c) { /* Add all the elements into the temporary dictionary. */ si = setTypeInitIterator(set); - while((encoding = setTypeNext(si,&ele,&llele)) != -1) { + while ((encoding = setTypeNext(si,&ele,&llele)) != -1) { int retval = DICT_ERR; if (encoding == OBJ_ENCODING_INTSET) { - retval = dictAdd(d,createStringObjectFromLongLong(llele),NULL); + retval = dictAdd(d,sdsfromlonglong(llele),NULL); } else { - retval = dictAdd(d,createStringObject(ele,sdslen(ele)),NULL); + retval = dictAdd(d,sdsdup(ele),NULL); } serverAssert(retval == DICT_OK); } @@ -743,11 +744,12 @@ void srandmemberWithCountCommand(client *c) { serverAssert(dictSize(d) == size); /* Remove random elements to reach the right count. */ - while(size > count) { + while (size > count) { dictEntry *de; - de = dictGetRandomKey(d); - dictDelete(d,dictGetKey(de)); + dictUnlink(d,dictGetKey(de)); + sdsfree(dictGetKey(de)); + dictFreeUnlinkedEntry(d,de); size--; } } @@ -758,22 +760,22 @@ void srandmemberWithCountCommand(client *c) { * to reach the specified count. */ else { unsigned long added = 0; - robj *objele; + sds sdsele; - while(added < count) { + while (added < count) { encoding = setTypeRandomElement(set,&ele,&llele); if (encoding == OBJ_ENCODING_INTSET) { - objele = createStringObjectFromLongLong(llele); + sdsele = sdsfromlonglong(llele); } else { - objele = createStringObject(ele,sdslen(ele)); + sdsele = sdsdup(ele); } /* Try to add the object to the dictionary. If it already exists * free it, otherwise increment the number of objects we have * in the result dictionary. */ - if (dictAdd(d,objele,NULL) == DICT_OK) + if (dictAdd(d,sdsele,NULL) == DICT_OK) added++; else - decrRefCount(objele); + sdsfree(sdsele); } } @@ -785,12 +787,13 @@ void srandmemberWithCountCommand(client *c) { addReplySetLen(c,count); di = dictGetIterator(d); while((de = dictNext(di)) != NULL) - addReplyBulk(c,dictGetKey(de)); + addReplyBulkSds(c,dictGetKey(de)); dictReleaseIterator(di); dictRelease(d); } } +/* SRANDMEMBER [] */ void srandmemberCommand(client *c) { robj *set; sds ele; @@ -805,6 +808,7 @@ void srandmemberCommand(client *c) { return; } + /* Handle variant without argument. Reply with simple bulk string */ if ((set = lookupKeyReadOrReply(c,c->argv[1],shared.null[c->resp])) == NULL || checkType(c,set,OBJ_SET)) return; diff --git a/src/t_zset.c b/src/t_zset.c index 6851ac86c..b55fc169e 100644 --- a/src/t_zset.c +++ b/src/t_zset.c @@ -721,20 +721,26 @@ zskiplistNode *zslLastInLexRange(zskiplist *zsl, zlexrangespec *range) { * Ziplist-backed sorted set API *----------------------------------------------------------------------------*/ +double zzlStrtod(unsigned char *vstr, unsigned int vlen) { + char buf[128]; + if (vlen > sizeof(buf)) + vlen = sizeof(buf); + memcpy(buf,vstr,vlen); + buf[vlen] = '\0'; + return strtod(buf,NULL); + } + double zzlGetScore(unsigned char *sptr) { unsigned char *vstr; unsigned int vlen; long long vlong; - char buf[128]; double score; serverAssert(sptr != NULL); serverAssert(ziplistGet(sptr,&vstr,&vlen,&vlong)); if (vstr) { - memcpy(buf,vstr,vlen); - buf[vlen] = '\0'; - score = strtod(buf,NULL); + score = zzlStrtod(vstr,vlen); } else { score = vlong; } @@ -1653,6 +1659,48 @@ int zsetZiplistValidateIntegrity(unsigned char *zl, size_t size, int deep) { return ret; } +/* Create a new sds string from the ziplist entry. */ +sds zsetSdsFromZiplistEntry(ziplistEntry *e) { + return e->sval ? sdsnewlen(e->sval, e->slen) : sdsfromlonglong(e->lval); +} + +/* Reply with bulk string from the ziplist entry. */ +void zsetReplyFromZiplistEntry(client *c, ziplistEntry *e) { + if (e->sval) + addReplyBulkCBuffer(c, e->sval, e->slen); + else + addReplyBulkLongLong(c, e->lval); +} + + +/* Return random element from a non empty zset. + * 'key' and 'val' will be set to hold the element. + * The memory in `key` is not to be freed or modified by the caller. + * 'score' can be NULL in which case it's not extracted. */ +void zsetTypeRandomElement(robj *zsetobj, unsigned long zsetsize, ziplistEntry *key, double *score) { + if (zsetobj->encoding == OBJ_ENCODING_SKIPLIST) { + zset *zs = zsetobj->ptr; + dictEntry *de = dictGetFairRandomKey(zs->dict); + sds s = dictGetKey(de); + key->sval = (unsigned char*)s; + key->slen = sdslen(s); + if (score) + *score = *(double*)dictGetVal(de); + } else if (zsetobj->encoding == OBJ_ENCODING_ZIPLIST) { + ziplistEntry val; + ziplistRandomPair(zsetobj->ptr, zsetsize, key, &val); + if (score) { + if (val.sval) { + *score = zzlStrtod(val.sval,val.slen); + } else { + *score = (double)val.lval; + } + } + } else { + serverPanic("Unknown zset encoding"); + } +} + /*----------------------------------------------------------------------------- * Sorted set commands *----------------------------------------------------------------------------*/ @@ -3907,3 +3955,216 @@ void bzpopminCommand(client *c) { void bzpopmaxCommand(client *c) { blockingGenericZpopCommand(c,ZSET_MAX); } + +/* How many times bigger should be the zset compared to the requested size + * for us to not use the "remove elements" strategy? Read later in the + * implementation for more info. */ +#define ZRANDMEMBER_SUB_STRATEGY_MUL 3 + +void zrandmemberWithCountCommand(client *c, long l, int withscores) { + unsigned long count, size; + int uniq = 1; + robj *zsetobj; + + if ((zsetobj = lookupKeyReadOrReply(c, c->argv[1], shared.null[c->resp])) + == NULL || checkType(c, zsetobj, OBJ_ZSET)) return; + size = zsetLength(zsetobj); + + if(l >= 0) { + count = (unsigned long) l; + } else { + count = -l; + uniq = 0; + } + + /* If count is zero, serve it ASAP to avoid special cases later. */ + if (count == 0) { + addReply(c,shared.emptyarray); + return; + } + + /* CASE 1: The count was negative, so the extraction method is just: + * "return N random elements" sampling the whole set every time. + * This case is trivial and can be served without auxiliary data + * structures. This case is the only one that also needs to return the + * elements in random order. */ + if (!uniq || count == 1) { + if (withscores && c->resp == 2) + addReplyArrayLen(c, count*2); + else + addReplyArrayLen(c, count); + if (zsetobj->encoding == OBJ_ENCODING_SKIPLIST) { + zset *zs = zsetobj->ptr; + while (count--) { + dictEntry *de = dictGetFairRandomKey(zs->dict); + sds key = dictGetKey(de); + if (withscores && c->resp > 2) + addReplyArrayLen(c,2); + addReplyBulkCBuffer(c, key, sdslen(key)); + if (withscores) + addReplyDouble(c, dictGetDoubleVal(de)); + } + } else if (zsetobj->encoding == OBJ_ENCODING_ZIPLIST) { + ziplistEntry *keys, *vals = NULL; + keys = zmalloc(sizeof(ziplistEntry)*count); + if (withscores) + vals = zmalloc(sizeof(ziplistEntry)*count); + ziplistRandomPairs(zsetobj->ptr, count, keys, vals); + for (unsigned long i = 0; i < count; i++) { + if (withscores && c->resp > 2) + addReplyArrayLen(c,2); + if (keys[i].sval) + addReplyBulkCBuffer(c, keys[i].sval, keys[i].slen); + else + addReplyBulkLongLong(c, keys[i].lval); + if (withscores) { + if (vals[i].sval) { + addReplyDouble(c, zzlStrtod(vals[i].sval,vals[i].slen)); + } else + addReplyDouble(c, vals[i].lval); + } + } + zfree(keys); + zfree(vals); + } + return; + } + + zsetopsrc src; + zsetopval zval; + src.subject = zsetobj; + src.type = zsetobj->type; + src.encoding = zsetobj->encoding; + zuiInitIterator(&src); + memset(&zval, 0, sizeof(zval)); + + /* Initiate reply count, RESP3 responds with nested array, RESP2 with flat one. */ + long reply_size = count < size ? count : size; + if (withscores && c->resp == 2) + addReplyArrayLen(c, reply_size*2); + else + addReplyArrayLen(c, reply_size); + + /* CASE 2: + * The number of requested elements is greater than the number of + * elements inside the zset: simply return the whole zset. */ + if (count >= size) { + while (zuiNext(&src, &zval)) { + if (withscores && c->resp > 2) + addReplyArrayLen(c,2); + addReplyBulkSds(c, zuiNewSdsFromValue(&zval)); + if (withscores) + addReplyDouble(c, zval.score); + } + return; + } + + /* CASE 3: + * The number of elements inside the zset is not greater than + * ZRANDMEMBER_SUB_STRATEGY_MUL times the number of requested elements. + * In this case we create a dict from scratch with all the elements, and + * subtract random elements to reach the requested number of elements. + * + * This is done because if the number of requested elements is just + * a bit less than the number of elements in the set, the natural approach + * used into CASE 4 is highly inefficient. */ + if (count*ZRANDMEMBER_SUB_STRATEGY_MUL > size) { + dict *d = dictCreate(&sdsReplyDictType, NULL); + /* Add all the elements into the temporary dictionary. */ + while (zuiNext(&src, &zval)) { + sds key = zuiNewSdsFromValue(&zval); + dictEntry *de = dictAddRaw(d, key, NULL); + serverAssert(de); + if (withscores) + dictSetDoubleVal(de, zval.score); + } + serverAssert(dictSize(d) == size); + + /* Remove random elements to reach the right count. */ + while (size > count) { + dictEntry *de; + de = dictGetRandomKey(d); + dictUnlink(d,dictGetKey(de)); + sdsfree(dictGetKey(de)); + dictFreeUnlinkedEntry(d,de); + size--; + } + + /* Reply with what's in the dict and release memory */ + dictIterator *di; + dictEntry *de; + di = dictGetIterator(d); + while ((de = dictNext(di)) != NULL) { + if (withscores && c->resp > 2) + addReplyArrayLen(c,2); + addReplyBulkSds(c, dictGetKey(de)); + if (withscores) + addReplyDouble(c, dictGetDoubleVal(de)); + } + + dictReleaseIterator(di); + dictRelease(d); + } + + /* CASE 4: We have a big zset compared to the requested number of elements. + * In this case we can simply get random elements from the zset and add + * to the temporary set, trying to eventually get enough unique elements + * to reach the specified count. */ + else { + unsigned long added = 0; + dict *d = dictCreate(&hashDictType, NULL); + + while (added < count) { + ziplistEntry key; + double score; + zsetTypeRandomElement(zsetobj, size, &key, withscores ? &score: NULL); + + /* Try to add the object to the dictionary. If it already exists + * free it, otherwise increment the number of objects we have + * in the result dictionary. */ + sds skey = zsetSdsFromZiplistEntry(&key); + if (dictAdd(d,skey,NULL) != DICT_OK) { + sdsfree(skey); + continue; + } + added++; + + if (withscores && c->resp > 2) + addReplyArrayLen(c,2); + zsetReplyFromZiplistEntry(c, &key); + if (withscores) + addReplyDouble(c, score); + } + + /* Release memory */ + dictRelease(d); + } +} + +/* ZRANDMEMBER [ WITHSCORES] */ +void zrandmemberCommand(client *c) { + long l; + int withscores = 0; + robj *zset; + ziplistEntry ele; + + if (c->argc >= 3) { + if (getLongFromObjectOrReply(c,c->argv[2],&l,NULL) != C_OK) return; + if (c->argc > 4 || (c->argc == 4 && strcasecmp(c->argv[3]->ptr,"withscores"))) { + addReplyErrorObject(c,shared.syntaxerr); + return; + } else if (c->argc == 4) + withscores = 1; + zrandmemberWithCountCommand(c, l, withscores); + return; + } + + /* Handle variant without argument. Reply with simple bulk string */ + if ((zset = lookupKeyReadOrReply(c,c->argv[1],shared.null[c->resp]))== NULL || + checkType(c,zset,OBJ_ZSET)) { + return; + } + + zsetTypeRandomElement(zset, zsetLength(zset), &ele,NULL); + zsetReplyFromZiplistEntry(c,&ele); +} diff --git a/src/ziplist.c b/src/ziplist.c index a4f38c5e8..62ffcb93d 100644 --- a/src/ziplist.c +++ b/src/ziplist.c @@ -1498,6 +1498,83 @@ int ziplistValidateIntegrity(unsigned char *zl, size_t size, int deep, return 1; } +/* Randomly select a pair of key and value. + * total_count is a pre-computed length/2 of the ziplist (to avoid calls to ziplistLen) + * 'key' and 'val' are used to store the result key value pair. + * 'val' can be NULL if the value is not needed. */ +void ziplistRandomPair(unsigned char *zl, unsigned long total_count, ziplistEntry *key, ziplistEntry *val) { + int ret; + unsigned char *p; + + /* Generate even numbers, because ziplist saved K-V pair */ + int r = (rand() % total_count) * 2; + p = ziplistIndex(zl, r); + ret = ziplistGet(p, &key->sval, &key->slen, &key->lval); + assert(ret != 0); + + if (!val) + return; + p = ziplistNext(zl, p); + ret = ziplistGet(p, &val->sval, &val->slen, &val->lval); + assert(ret != 0); +} + +/* int compare for qsort */ +int intCompare(const void *a, const void *b) { + return (*(int *) a - *(int *) b); +} + +/* Helper method to store a string into from val or lval into dest */ +static inline void ziplistSaveValue(unsigned char *val, unsigned int len, long long lval, ziplistEntry *dest) { + dest->sval = val; + dest->slen = len; + dest->lval = lval; +} + +/* Randomly select unique count of key value pairs and store into 'keys' and + * 'vals' args. The order of the picked entries is random. + * The 'vals' arg can be NULL in which case we skip these. */ +void ziplistRandomPairs(unsigned char *zl, int count, ziplistEntry *keys, ziplistEntry *vals) { + unsigned char *p, *key, *value; + unsigned int klen, vlen; + long long klval, vlval; + typedef struct { + int index; + int order; + } rand_pick; + rand_pick *picks = zmalloc(sizeof(rand_pick)*count); + unsigned long total_size = ziplistLen(zl)/2; + + /* create a pool of random indexes (some may be duplicate). */ + for (int i = 0; i < count; i++) { + picks[i].index = (rand() % total_size) * 2; /* Generate even indexes */ + /* keep track of the order we picked them */ + picks[i].order = i; + } + + /* sort by indexes. */ + qsort(picks, count, sizeof(rand_pick), intCompare); + + /* fetch the elements form the ziplist into a output array respecting the original order. */ + int zipindex = 0, pickindex = 0; + p = ziplistIndex(zl, 0); + while (ziplistGet(p, &key, &klen, &klval) && pickindex < count) { + p = ziplistNext(zl, p); + ziplistGet(p, &value, &vlen, &vlval); + while (pickindex < count && zipindex == picks[pickindex].index) { + int storeorder = picks[pickindex].order; + ziplistSaveValue(key, klen, klval, &keys[storeorder]); + if (vals) + ziplistSaveValue(value, vlen, vlval, &vals[storeorder]); + pickindex++; + } + zipindex += 2; + p = ziplistNext(zl, p); + } + + zfree(picks); +} + #ifdef REDIS_TEST #include #include "adlist.h" diff --git a/src/ziplist.h b/src/ziplist.h index 5153951dc..5fb8fd46a 100644 --- a/src/ziplist.h +++ b/src/ziplist.h @@ -34,6 +34,15 @@ #define ZIPLIST_HEAD 0 #define ZIPLIST_TAIL 1 +/* Each entry in the ziplist is either a string or an integer. */ +typedef struct { + /* When string is used, it is provided with the length (slen). */ + unsigned char *sval; + unsigned int slen; + /* When integer is used, 'sval' is NULL, and lval holds the value. */ + long long lval; +} ziplistEntry; + unsigned char *ziplistNew(void); unsigned char *ziplistMerge(unsigned char **first, unsigned char **second); unsigned char *ziplistPush(unsigned char *zl, unsigned char *s, unsigned int slen, int where); @@ -52,6 +61,8 @@ void ziplistRepr(unsigned char *zl); typedef int (*ziplistValidateEntryCB)(unsigned char* p, void* userdata); int ziplistValidateIntegrity(unsigned char *zl, size_t size, int deep, ziplistValidateEntryCB entry_cb, void *cb_userdata); +void ziplistRandomPair(unsigned char *zl, unsigned long total_count, ziplistEntry *key, ziplistEntry *val); +void ziplistRandomPairs(unsigned char *zl, int count, ziplistEntry *keys, ziplistEntry *vals); #ifdef REDIS_TEST int ziplistTest(int argc, char *argv[]); diff --git a/tests/support/util.tcl b/tests/support/util.tcl index e69395994..943f45a2b 100644 --- a/tests/support/util.tcl +++ b/tests/support/util.tcl @@ -569,8 +569,8 @@ proc generate_fuzzy_traffic_on_key {key duration} { # Commands per type, blocking commands removed # TODO: extract these from help.h or elsewhere, and improve to include other types set string_commands {APPEND BITCOUNT BITFIELD BITOP BITPOS DECR DECRBY GET GETBIT GETRANGE GETSET INCR INCRBY INCRBYFLOAT MGET MSET MSETNX PSETEX SET SETBIT SETEX SETNX SETRANGE STRALGO STRLEN} - set hash_commands {HDEL HEXISTS HGET HGETALL HINCRBY HINCRBYFLOAT HKEYS HLEN HMGET HMSET HSCAN HSET HSETNX HSTRLEN HVALS} - set zset_commands {ZADD ZCARD ZCOUNT ZINCRBY ZINTERSTORE ZLEXCOUNT ZPOPMAX ZPOPMIN ZRANGE ZRANGEBYLEX ZRANGEBYSCORE ZRANK ZREM ZREMRANGEBYLEX ZREMRANGEBYRANK ZREMRANGEBYSCORE ZREVRANGE ZREVRANGEBYLEX ZREVRANGEBYSCORE ZREVRANK ZSCAN ZSCORE ZUNIONSTORE} + set hash_commands {HDEL HEXISTS HGET HGETALL HINCRBY HINCRBYFLOAT HKEYS HLEN HMGET HMSET HSCAN HSET HSETNX HSTRLEN HVALS HRANDFIELD} + set zset_commands {ZADD ZCARD ZCOUNT ZINCRBY ZINTERSTORE ZLEXCOUNT ZPOPMAX ZPOPMIN ZRANGE ZRANGEBYLEX ZRANGEBYSCORE ZRANK ZREM ZREMRANGEBYLEX ZREMRANGEBYRANK ZREMRANGEBYSCORE ZREVRANGE ZREVRANGEBYLEX ZREVRANGEBYSCORE ZREVRANK ZSCAN ZSCORE ZUNIONSTORE ZRANDMEMBER} set list_commands {LINDEX LINSERT LLEN LPOP LPOS LPUSH LPUSHX LRANGE LREM LSET LTRIM RPOP RPOPLPUSH RPUSH RPUSHX} set set_commands {SADD SCARD SDIFF SDIFFSTORE SINTER SINTERSTORE SISMEMBER SMEMBERS SMOVE SPOP SRANDMEMBER SREM SSCAN SUNION SUNIONSTORE} set stream_commands {XACK XADD XCLAIM XDEL XGROUP XINFO XLEN XPENDING XRANGE XREAD XREADGROUP XREVRANGE XTRIM} diff --git a/tests/unit/type/hash.tcl b/tests/unit/type/hash.tcl index 79e58301a..2f3ea37c2 100644 --- a/tests/unit/type/hash.tcl +++ b/tests/unit/type/hash.tcl @@ -18,6 +18,181 @@ start_server {tags {"hash"}} { assert_encoding ziplist smallhash } + proc create_hash {key entries} { + r del $key + foreach entry $entries { + r hset $key [lindex $entry 0] [lindex $entry 1] + } + } + + proc get_keys {l} { + set res {} + foreach entry $l { + set key [lindex $entry 0] + lappend res $key + } + return $res + } + + foreach {type contents} "ziplist {{a 1} {b 2} {c 3}} hashtable {{a 1} {b 2} {[randstring 70 90 alpha] 3}}" { + set original_max_value [lindex [r config get hash-max-ziplist-value] 1] + r config set hash-max-ziplist-value 10 + create_hash myhash $contents + assert_encoding $type myhash + + test "HRANDFIELD - $type" { + unset -nocomplain myhash + array set myhash {} + for {set i 0} {$i < 100} {incr i} { + set key [r hrandfield myhash] + set myhash($key) 1 + } + assert_equal [lsort [get_keys $contents]] [lsort [array names myhash]] + } + r config set hash-max-ziplist-value $original_max_value + } + + test "HRANDFIELD with RESP3" { + r hello 3 + set res [r hrandfield myhash 3 withvalues] + assert_equal [llength $res] 3 + assert_equal [llength [lindex $res 1]] 2 + + set res [r hrandfield myhash 3] + assert_equal [llength $res] 3 + assert_equal [llength [lindex $res 1]] 1 + } + r hello 2 + + test "HRANDFIELD count of 0 is handled correctly" { + r hrandfield myhash 0 + } {} + + test "HRANDFIELD with against non existing key" { + r hrandfield nonexisting_key 100 + } {} + + foreach {type contents} " + hashtable {{a 1} {b 2} {c 3} {d 4} {e 5} {6 f} {7 g} {8 h} {9 i} {[randstring 70 90 alpha] 10}} + ziplist {{a 1} {b 2} {c 3} {d 4} {e 5} {6 f} {7 g} {8 h} {9 i} {10 j}} " { + test "HRANDFIELD with - $type" { + set original_max_value [lindex [r config get hash-max-ziplist-value] 1] + r config set hash-max-ziplist-value 10 + create_hash myhash $contents + assert_encoding $type myhash + + # create a dict for easy lookup + unset -nocomplain mydict + foreach {k v} [r hgetall myhash] { + dict append mydict $k $v + } + + # We'll stress different parts of the code, see the implementation + # of HRANDFIELD for more information, but basically there are + # four different code paths. + + # PATH 1: Use negative count. + + # 1) Check that it returns repeated elements with and without values. + set res [r hrandfield myhash -20] + assert_equal [llength $res] 20 + # again with WITHVALUES + set res [r hrandfield myhash -20 withvalues] + assert_equal [llength $res] 40 + + # 2) Check that all the elements actually belong to the original hash. + foreach {key val} $res { + assert {[dict exists $mydict $key]} + } + + # 3) Check that eventually all the elements are returned. + # Use both WITHVALUES and without + unset -nocomplain auxset + set iterations 1000 + while {$iterations != 0} { + incr iterations -1 + if {[expr {$iterations % 2}] == 0} { + set res [r hrandfield myhash -3 withvalues] + foreach {key val} $res { + dict append auxset $key $val + } + } else { + set res [r hrandfield myhash -3] + foreach key $res { + dict append auxset $key $val + } + } + if {[lsort [dict keys $mydict]] eq + [lsort [dict keys $auxset]]} { + break; + } + } + assert {$iterations != 0} + + # PATH 2: positive count (unique behavior) with requested size + # equal or greater than set size. + foreach size {10 20} { + set res [r hrandfield myhash $size] + assert_equal [llength $res] 10 + assert_equal [lsort $res] [lsort [dict keys $mydict]] + + # again with WITHVALUES + set res [r hrandfield myhash $size withvalues] + assert_equal [llength $res] 20 + assert_equal [lsort $res] [lsort $mydict] + } + + # PATH 3: Ask almost as elements as there are in the set. + # In this case the implementation will duplicate the original + # set and will remove random elements up to the requested size. + # + # PATH 4: Ask a number of elements definitely smaller than + # the set size. + # + # We can test both the code paths just changing the size but + # using the same code. + foreach size {8 2} { + set res [r hrandfield myhash $size] + assert_equal [llength $res] $size + # again with WITHVALUES + set res [r hrandfield myhash $size withvalues] + assert_equal [llength $res] [expr {$size * 2}] + + # 1) Check that all the elements actually belong to the + # original set. + foreach ele [dict keys $res] { + assert {[dict exists $mydict $ele]} + } + + # 2) Check that eventually all the elements are returned. + # Use both WITHVALUES and without + unset -nocomplain auxset + set iterations 1000 + while {$iterations != 0} { + incr iterations -1 + if {[expr {$iterations % 2}] == 0} { + set res [r hrandfield myhash $size withvalues] + foreach {key value} $res { + dict append auxset $key $value + } + } else { + set res [r hrandfield myhash $size] + foreach key $res { + dict append auxset $key + } + } + if {[lsort [dict keys $mydict]] eq + [lsort [dict keys $auxset]]} { + break; + } + } + assert {$iterations != 0} + } + } + r config set hash-max-ziplist-value $original_max_value + } + + test {HSET/HLEN - Big hash creation} { array set bighash {} for {set i 0} {$i < 1024} {incr i} { diff --git a/tests/unit/type/set.tcl b/tests/unit/type/set.tcl index 84c31c4e4..091ef7f0f 100644 --- a/tests/unit/type/set.tcl +++ b/tests/unit/type/set.tcl @@ -501,7 +501,7 @@ start_server { set iterations 1000 while {$iterations != 0} { incr iterations -1 - set res [r srandmember myset -10] + set res [r srandmember myset $size] foreach ele $res { set auxset($ele) 1 } diff --git a/tests/unit/type/zset.tcl b/tests/unit/type/zset.tcl index f3c952f0e..c657c1e4e 100644 --- a/tests/unit/type/zset.tcl +++ b/tests/unit/type/zset.tcl @@ -7,6 +7,8 @@ start_server {tags {"zset"}} { } proc basics {encoding} { + set original_max_entries [lindex [r config get zset-max-ziplist-entries] 1] + set original_max_value [lindex [r config get zset-max-ziplist-value] 1] if {$encoding == "ziplist"} { r config set zset-max-ziplist-entries 128 r config set zset-max-ziplist-value 64 @@ -925,6 +927,9 @@ start_server {tags {"zset"}} { assert_equal 0 [r zcard z1] assert_equal 1 [r zcard z2] } + + r config set zset-max-ziplist-entries $original_max_entries + r config set zset-max-ziplist-value $original_max_value } basics ziplist @@ -1022,6 +1027,8 @@ start_server {tags {"zset"}} { } proc stressers {encoding} { + set original_max_entries [lindex [r config get zset-max-ziplist-entries] 1] + set original_max_value [lindex [r config get zset-max-ziplist-value] 1] if {$encoding == "ziplist"} { # Little extra to allow proper fuzzing in the sorting stresser r config set zset-max-ziplist-entries 256 @@ -1446,6 +1453,8 @@ start_server {tags {"zset"}} { r zadd zset 0 foo assert_equal {zset foo 0} [$rd read] } + r config set zset-max-ziplist-entries $original_max_entries + r config set zset-max-ziplist-value $original_max_value } tags {"slow"} { @@ -1566,4 +1575,171 @@ start_server {tags {"zset"}} { catch {r zrangebyscore z1 0 -1 REV} err assert_match "*syntax*" $err } + + proc get_keys {l} { + set res {} + foreach {score key} $l { + lappend res $key + } + return $res + } + + foreach {type contents} "ziplist {1 a 2 b 3 c} skiplist {1 a 2 b 3 [randstring 70 90 alpha]}" { + set original_max_value [lindex [r config get zset-max-ziplist-value] 1] + r config set zset-max-ziplist-value 10 + create_zset myzset $contents + assert_encoding $type myzset + + test "ZRANDMEMBER - $type" { + unset -nocomplain myzset + array set myzset {} + for {set i 0} {$i < 100} {incr i} { + set key [r zrandmember myzset] + set myzset($key) 1 + } + assert_equal [lsort [get_keys $contents]] [lsort [array names myzset]] + } + r config set zset-max-ziplist-value $original_max_value + } + + test "ZRANDMEMBER with RESP3" { + r hello 3 + set res [r zrandmember myzset 3 withscores] + assert_equal [llength $res] 3 + assert_equal [llength [lindex $res 1]] 2 + + set res [r zrandmember myzset 3] + assert_equal [llength $res] 3 + assert_equal [llength [lindex $res 1]] 1 + } + r hello 2 + + test "ZRANDMEMBER count of 0 is handled correctly" { + r zrandmember myzset 0 + } {} + + test "ZRANDMEMBER with against non existing key" { + r zrandmember nonexisting_key 100 + } {} + + foreach {type contents} " + skiplist {1 a 2 b 3 c 4 d 5 e 6 f 7 g 7 h 9 i 10 [randstring 70 90 alpha]} + ziplist {1 a 2 b 3 c 4 d 5 e 6 f 7 g 7 h 9 i 10 j} " { + test "ZRANDMEMBER with - $type" { + set original_max_value [lindex [r config get zset-max-ziplist-value] 1] + r config set zset-max-ziplist-value 10 + create_zset myzset $contents + assert_encoding $type myzset + + # create a dict for easy lookup + unset -nocomplain mydict + foreach {k v} [r zrange myzset 0 -1 withscores] { + dict append mydict $k $v + } + + # We'll stress different parts of the code, see the implementation + # of ZRANDMEMBER for more information, but basically there are + # four different code paths. + + # PATH 1: Use negative count. + + # 1) Check that it returns repeated elements with and without values. + set res [r zrandmember myzset -20] + assert_equal [llength $res] 20 + # again with WITHSCORES + set res [r zrandmember myzset -20 withscores] + assert_equal [llength $res] 40 + + # 2) Check that all the elements actually belong to the original zset. + foreach {key val} $res { + assert {[dict exists $mydict $key]} + } + + # 3) Check that eventually all the elements are returned. + # Use both WITHSCORES and without + unset -nocomplain auxset + set iterations 1000 + while {$iterations != 0} { + incr iterations -1 + if {[expr {$iterations % 2}] == 0} { + set res [r zrandmember myzset -3 withscores] + foreach {key val} $res { + dict append auxset $key $val + } + } else { + set res [r zrandmember myzset -3] + foreach key $res { + dict append auxset $key $val + } + } + if {[lsort [dict keys $mydict]] eq + [lsort [dict keys $auxset]]} { + break; + } + } + assert {$iterations != 0} + + # PATH 2: positive count (unique behavior) with requested size + # equal or greater than set size. + foreach size {10 20} { + set res [r zrandmember myzset $size] + assert_equal [llength $res] 10 + assert_equal [lsort $res] [lsort [dict keys $mydict]] + + # again with WITHSCORES + set res [r zrandmember myzset $size withscores] + assert_equal [llength $res] 20 + assert_equal [lsort $res] [lsort $mydict] + } + + # PATH 3: Ask almost as elements as there are in the set. + # In this case the implementation will duplicate the original + # set and will remove random elements up to the requested size. + # + # PATH 4: Ask a number of elements definitely smaller than + # the set size. + # + # We can test both the code paths just changing the size but + # using the same code. + foreach size {8 2} { + set res [r zrandmember myzset $size] + assert_equal [llength $res] $size + # again with WITHSCORES + set res [r zrandmember myzset $size withscores] + assert_equal [llength $res] [expr {$size * 2}] + + # 1) Check that all the elements actually belong to the + # original set. + foreach ele [dict keys $res] { + assert {[dict exists $mydict $ele]} + } + + # 2) Check that eventually all the elements are returned. + # Use both WITHSCORES and without + unset -nocomplain auxset + set iterations 1000 + while {$iterations != 0} { + incr iterations -1 + if {[expr {$iterations % 2}] == 0} { + set res [r zrandmember myzset $size withscores] + foreach {key value} $res { + dict append auxset $key $value + } + } else { + set res [r zrandmember myzset $size] + foreach key $res { + dict append auxset $key + } + } + if {[lsort [dict keys $mydict]] eq + [lsort [dict keys $auxset]]} { + break; + } + } + assert {$iterations != 0} + } + } + r config set zset-max-ziplist-value $original_max_value + } + }