X-Git-Url: https://git.saurik.com/redis.git/blobdiff_plain/c772d9c6e7a4b65075c6efd15a53e84bb3c7ba3f..6c9897f6cfdd3d9dbf73d09a9f8a9c5706020717:/src/t_zset.c diff --git a/src/t_zset.c b/src/t_zset.c index e9da9fdb..f5cdec3e 100644 --- a/src/t_zset.c +++ b/src/t_zset.c @@ -1,6 +1,32 @@ -#include "redis.h" - -#include +/* + * Copyright (c) 2009-2012, Salvatore Sanfilippo + * Copyright (c) 2009-2012, Pieter Noordhuis + * All rights reserved. + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * * Redistributions of source code must retain the above copyright notice, + * this list of conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and/or other materials provided with the distribution. + * * Neither the name of Redis nor the names of its contributors may be used + * to endorse or promote products derived from this software without + * specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE + * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF + * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE + * POSSIBILITY OF SUCH DAMAGE. + */ /*----------------------------------------------------------------------------- * Sorted set API @@ -17,12 +43,15 @@ /* This skiplist implementation is almost a C translation of the original * algorithm described by William Pugh in "Skip Lists: A Probabilistic * Alternative to Balanced Trees", modified in three ways: - * a) this implementation allows for repeated values. + * a) this implementation allows for repeated scores. * b) the comparison is not just by key (our 'score') but by satellite data. * c) there is a back pointer, so it's a doubly linked list with the back * pointers being only at "level 1". This allows to traverse the list * from tail to head, useful for ZREVRANGE. */ +#include "redis.h" +#include + zskiplistNode *zslCreateNode(int level, double score, robj *obj) { zskiplistNode *zn = zmalloc(sizeof(*zn)+level*sizeof(struct zskiplistLevel)); zn->score = score; @@ -64,6 +93,10 @@ void zslFree(zskiplist *zsl) { zfree(zsl); } +/* Returns a random level for the new skiplist node we are going to create. + * The return value of this function is between 1 and ZSKIPLIST_MAXLEVEL + * (both inclusive), with a powerlaw-alike distribution where higher + * levels are less likely to be returned. */ int zslRandomLevel(void) { int level = 1; while ((random()&0xFFFF) < (ZSKIPLIST_P * 0xFFFF)) @@ -76,6 +109,7 @@ zskiplistNode *zslInsert(zskiplist *zsl, double score, robj *obj) { unsigned int rank[ZSKIPLIST_MAXLEVEL]; int i, level; + redisAssert(!isnan(score)); x = zsl->header; for (i = zsl->level-1; i >= 0; i--) { /* store rank that is crossed to reach the insert position */ @@ -497,7 +531,7 @@ int zzlIsInRange(unsigned char *zl, zrangespec *range) { return 0; p = ziplistIndex(zl,-1); /* Last score. */ - redisAssert(p != NULL); + if (p == NULL) return 0; /* Empty sorted set */ score = zzlGetScore(p); if (!zslValueGteMin(score,range)) return 0; @@ -578,7 +612,7 @@ unsigned char *zzlFind(unsigned char *zl, robj *ele, double *score) { ele = getDecodedObject(ele); while (eptr != NULL) { sptr = ziplistNext(zl,eptr); - redisAssert(sptr != NULL); + redisAssertWithInfo(NULL,ele,sptr != NULL); if (ziplistCompare(eptr,ele->ptr,sdslen(ele->ptr))) { /* Matching element, pull out score. */ @@ -612,7 +646,7 @@ unsigned char *zzlInsertAt(unsigned char *zl, unsigned char *eptr, robj *ele, do int scorelen; size_t offset; - redisAssert(ele->encoding == REDIS_ENCODING_RAW); + redisAssertWithInfo(NULL,ele,ele->encoding == REDIS_ENCODING_RAW); scorelen = d2string(scorebuf,sizeof(scorebuf),score); if (eptr == NULL) { zl = ziplistPush(zl,ele->ptr,sdslen(ele->ptr),ZIPLIST_TAIL); @@ -624,7 +658,7 @@ unsigned char *zzlInsertAt(unsigned char *zl, unsigned char *eptr, robj *ele, do eptr = zl+offset; /* Insert score after the element. */ - redisAssert((sptr = ziplistNext(zl,eptr)) != NULL); + redisAssertWithInfo(NULL,ele,(sptr = ziplistNext(zl,eptr)) != NULL); zl = ziplistInsert(zl,sptr,(unsigned char*)scorebuf,scorelen); } @@ -640,7 +674,7 @@ unsigned char *zzlInsert(unsigned char *zl, robj *ele, double score) { ele = getDecodedObject(ele); while (eptr != NULL) { sptr = ziplistNext(zl,eptr); - redisAssert(sptr != NULL); + redisAssertWithInfo(NULL,ele,sptr != NULL); s = zzlGetScore(sptr); if (s > score) { @@ -745,13 +779,13 @@ void zsetConvert(robj *zobj, int encoding) { zs->zsl = zslCreate(); eptr = ziplistIndex(zl,0); - redisAssert(eptr != NULL); + redisAssertWithInfo(NULL,zobj,eptr != NULL); sptr = ziplistNext(zl,eptr); - redisAssert(sptr != NULL); + redisAssertWithInfo(NULL,zobj,sptr != NULL); while (eptr != NULL) { score = zzlGetScore(sptr); - redisAssert(ziplistGet(eptr,&vstr,&vlen,&vlong)); + redisAssertWithInfo(NULL,zobj,ziplistGet(eptr,&vstr,&vlen,&vlong)); if (vstr == NULL) ele = createStringObjectFromLongLong(vlong); else @@ -759,7 +793,7 @@ void zsetConvert(robj *zobj, int encoding) { /* Has incremented refcount since it was just created. */ node = zslInsert(zs->zsl,score,ele); - redisAssert(dictAdd(zs->dict,ele,&node->score) == DICT_OK); + redisAssertWithInfo(NULL,zobj,dictAdd(zs->dict,ele,&node->score) == DICT_OK); incrRefCount(ele); /* Added to dictionary. */ zzlNext(zl,&eptr,&sptr); } @@ -810,11 +844,29 @@ void zaddGenericCommand(redisClient *c, int incr) { robj *ele; robj *zobj; robj *curobj; - double score, curscore = 0.0; + double score = 0, *scores, curscore = 0.0; + int j, elements = (c->argc-2)/2; + int added = 0; - if (getDoubleFromObjectOrReply(c,c->argv[2],&score,NULL) != REDIS_OK) + if (c->argc % 2) { + addReply(c,shared.syntaxerr); return; + } + + /* Start parsing all the scores, we need to emit any syntax error + * before executing additions to the sorted set, as the command should + * either execute fully or nothing at all. */ + scores = zmalloc(sizeof(double)*elements); + for (j = 0; j < elements; j++) { + if (getDoubleFromObjectOrReply(c,c->argv[2+j*2],&scores[j],NULL) + != REDIS_OK) + { + zfree(scores); + return; + } + } + /* Lookup the key and create the sorted set if does not exist. */ zobj = lookupKeyWrite(c->db,key); if (zobj == NULL) { if (server.zset_max_ziplist_entries == 0 || @@ -828,111 +880,105 @@ void zaddGenericCommand(redisClient *c, int incr) { } else { if (zobj->type != REDIS_ZSET) { addReply(c,shared.wrongtypeerr); + zfree(scores); return; } } - if (zobj->encoding == REDIS_ENCODING_ZIPLIST) { - unsigned char *eptr; + for (j = 0; j < elements; j++) { + score = scores[j]; - /* Prefer non-encoded element when dealing with ziplists. */ - ele = c->argv[3]; - if ((eptr = zzlFind(zobj->ptr,ele,&curscore)) != NULL) { - if (incr) { - score += curscore; - if (isnan(score)) { - addReplyError(c,nanerr); - /* Don't need to check if the sorted set is empty, because - * we know it has at least one element. */ - return; + if (zobj->encoding == REDIS_ENCODING_ZIPLIST) { + unsigned char *eptr; + + /* Prefer non-encoded element when dealing with ziplists. */ + ele = c->argv[3+j*2]; + if ((eptr = zzlFind(zobj->ptr,ele,&curscore)) != NULL) { + if (incr) { + score += curscore; + if (isnan(score)) { + addReplyError(c,nanerr); + /* Don't need to check if the sorted set is empty + * because we know it has at least one element. */ + zfree(scores); + return; + } } - } - /* Remove and re-insert when score changed. */ - if (score != curscore) { - zobj->ptr = zzlDelete(zobj->ptr,eptr); + /* Remove and re-insert when score changed. */ + if (score != curscore) { + zobj->ptr = zzlDelete(zobj->ptr,eptr); + zobj->ptr = zzlInsert(zobj->ptr,ele,score); + + signalModifiedKey(c->db,key); + server.dirty++; + } + } else { + /* Optimize: check if the element is too large or the list + * becomes too long *before* executing zzlInsert. */ zobj->ptr = zzlInsert(zobj->ptr,ele,score); + if (zzlLength(zobj->ptr) > server.zset_max_ziplist_entries) + zsetConvert(zobj,REDIS_ENCODING_SKIPLIST); + if (sdslen(ele->ptr) > server.zset_max_ziplist_value) + zsetConvert(zobj,REDIS_ENCODING_SKIPLIST); signalModifiedKey(c->db,key); server.dirty++; + if (!incr) added++; } + } else if (zobj->encoding == REDIS_ENCODING_SKIPLIST) { + zset *zs = zobj->ptr; + zskiplistNode *znode; + dictEntry *de; - if (incr) /* ZINCRBY */ - addReplyDouble(c,score); - else /* ZADD */ - addReply(c,shared.czero); - } else { - /* Optimize: check if the element is too large or the list becomes - * too long *before* executing zzlInsert. */ - zobj->ptr = zzlInsert(zobj->ptr,ele,score); - if (zzlLength(zobj->ptr) > server.zset_max_ziplist_entries) - zsetConvert(zobj,REDIS_ENCODING_SKIPLIST); - if (sdslen(ele->ptr) > server.zset_max_ziplist_value) - zsetConvert(zobj,REDIS_ENCODING_SKIPLIST); - - signalModifiedKey(c->db,key); - server.dirty++; - - if (incr) /* ZINCRBY */ - addReplyDouble(c,score); - else /* ZADD */ - addReply(c,shared.cone); - } - } else if (zobj->encoding == REDIS_ENCODING_SKIPLIST) { - zset *zs = zobj->ptr; - zskiplistNode *znode; - dictEntry *de; - - ele = c->argv[3] = tryObjectEncoding(c->argv[3]); - de = dictFind(zs->dict,ele); - if (de != NULL) { - curobj = dictGetEntryKey(de); - curscore = *(double*)dictGetEntryVal(de); - - if (incr) { - score += curscore; - if (isnan(score)) { - addReplyError(c,nanerr); - /* Don't need to check if the sorted set is empty, because - * we know it has at least one element. */ - return; + ele = c->argv[3+j*2] = tryObjectEncoding(c->argv[3+j*2]); + de = dictFind(zs->dict,ele); + if (de != NULL) { + curobj = dictGetKey(de); + curscore = *(double*)dictGetVal(de); + + if (incr) { + score += curscore; + if (isnan(score)) { + addReplyError(c,nanerr); + /* Don't need to check if the sorted set is empty + * because we know it has at least one element. */ + zfree(scores); + return; + } } - } - /* Remove and re-insert when score changed. We can safely delete - * the key object from the skiplist, since the dictionary still has - * a reference to it. */ - if (score != curscore) { - redisAssert(zslDelete(zs->zsl,curscore,curobj)); - znode = zslInsert(zs->zsl,score,curobj); - incrRefCount(curobj); /* Re-inserted in skiplist. */ - dictGetEntryVal(de) = &znode->score; /* Update score ptr. */ + /* Remove and re-insert when score changed. We can safely + * delete the key object from the skiplist, since the + * dictionary still has a reference to it. */ + if (score != curscore) { + redisAssertWithInfo(c,curobj,zslDelete(zs->zsl,curscore,curobj)); + znode = zslInsert(zs->zsl,score,curobj); + incrRefCount(curobj); /* Re-inserted in skiplist. */ + dictGetVal(de) = &znode->score; /* Update score ptr. */ + + signalModifiedKey(c->db,key); + server.dirty++; + } + } else { + znode = zslInsert(zs->zsl,score,ele); + incrRefCount(ele); /* Inserted in skiplist. */ + redisAssertWithInfo(c,NULL,dictAdd(zs->dict,ele,&znode->score) == DICT_OK); + incrRefCount(ele); /* Added to dictionary. */ signalModifiedKey(c->db,key); server.dirty++; + if (!incr) added++; } - - if (incr) /* ZINCRBY */ - addReplyDouble(c,score); - else /* ZADD */ - addReply(c,shared.czero); } else { - znode = zslInsert(zs->zsl,score,ele); - incrRefCount(ele); /* Inserted in skiplist. */ - redisAssert(dictAdd(zs->dict,ele,&znode->score) == DICT_OK); - incrRefCount(ele); /* Added to dictionary. */ - - signalModifiedKey(c->db,key); - server.dirty++; - - if (incr) /* ZINCRBY */ - addReplyDouble(c,score); - else /* ZADD */ - addReply(c,shared.cone); + redisPanic("Unknown sorted set encoding"); } - } else { - redisPanic("Unknown sorted set encoding"); } + zfree(scores); + if (incr) /* ZINCRBY */ + addReplyDouble(c,score); + else /* ZADD */ + addReplyLongLong(c,added); } void zaddCommand(redisClient *c) { @@ -945,8 +991,8 @@ void zincrbyCommand(redisClient *c) { void zremCommand(redisClient *c) { robj *key = c->argv[1]; - robj *ele = c->argv[2]; robj *zobj; + int deleted = 0, j; if ((zobj = lookupKeyWriteOrReply(c,key,shared.czero)) == NULL || checkType(c,zobj,REDIS_ZSET)) return; @@ -954,39 +1000,48 @@ void zremCommand(redisClient *c) { if (zobj->encoding == REDIS_ENCODING_ZIPLIST) { unsigned char *eptr; - if ((eptr = zzlFind(zobj->ptr,ele,NULL)) != NULL) { - zobj->ptr = zzlDelete(zobj->ptr,eptr); - if (zzlLength(zobj->ptr) == 0) dbDelete(c->db,key); - } else { - addReply(c,shared.czero); - return; + for (j = 2; j < c->argc; j++) { + if ((eptr = zzlFind(zobj->ptr,c->argv[j],NULL)) != NULL) { + deleted++; + zobj->ptr = zzlDelete(zobj->ptr,eptr); + if (zzlLength(zobj->ptr) == 0) { + dbDelete(c->db,key); + break; + } + } } } else if (zobj->encoding == REDIS_ENCODING_SKIPLIST) { zset *zs = zobj->ptr; dictEntry *de; double score; - de = dictFind(zs->dict,ele); - if (de != NULL) { - /* Delete from the skiplist */ - score = *(double*)dictGetEntryVal(de); - redisAssert(zslDelete(zs->zsl,score,ele)); - - /* Delete from the hash table */ - dictDelete(zs->dict,ele); - if (htNeedsResize(zs->dict)) dictResize(zs->dict); - if (dictSize(zs->dict) == 0) dbDelete(c->db,key); - } else { - addReply(c,shared.czero); - return; + for (j = 2; j < c->argc; j++) { + de = dictFind(zs->dict,c->argv[j]); + if (de != NULL) { + deleted++; + + /* Delete from the skiplist */ + score = *(double*)dictGetVal(de); + redisAssertWithInfo(c,c->argv[j],zslDelete(zs->zsl,score,c->argv[j])); + + /* Delete from the hash table */ + dictDelete(zs->dict,c->argv[j]); + if (htNeedsResize(zs->dict)) dictResize(zs->dict); + if (dictSize(zs->dict) == 0) { + dbDelete(c->db,key); + break; + } + } } } else { redisPanic("Unknown sorted set encoding"); } - signalModifiedKey(c->db,key); - server.dirty++; - addReply(c,shared.cone); + if (deleted) { + signalModifiedKey(c->db,key); + server.dirty += deleted; + } + addReplyLongLong(c,deleted); } void zremrangebyscoreCommand(redisClient *c) { @@ -997,7 +1052,7 @@ void zremrangebyscoreCommand(redisClient *c) { /* Parse the range arguments. */ if (zslParseRange(c->argv[2],c->argv[3],&range) != REDIS_OK) { - addReplyError(c,"min or max is not a double"); + addReplyError(c,"min or max is not a float"); return; } @@ -1228,13 +1283,16 @@ int zuiNext(zsetopsrc *op, zsetopval *val) { if (val->flags & OPVAL_DIRTY_ROBJ) decrRefCount(val->ele); - bzero(val,sizeof(zsetopval)); + memset(val,0,sizeof(zsetopval)); if (op->type == REDIS_SET) { iterset *it = &op->iter.set; if (op->encoding == REDIS_ENCODING_INTSET) { - if (!intsetGet(it->is.is,it->is.ii,&val->ell)) + int64_t ell; + + if (!intsetGet(it->is.is,it->is.ii,&ell)) return 0; + val->ell = ell; val->score = 1.0; /* Move to next element. */ @@ -1242,7 +1300,7 @@ int zuiNext(zsetopsrc *op, zsetopval *val) { } else if (op->encoding == REDIS_ENCODING_HT) { if (it->ht.de == NULL) return 0; - val->ele = dictGetEntryKey(it->ht.de); + val->ele = dictGetKey(it->ht.de); val->score = 1.0; /* Move to next element. */ @@ -1376,7 +1434,7 @@ int zuiFind(zsetopsrc *op, zsetopval *val, double *score) { } else if (op->encoding == REDIS_ENCODING_SKIPLIST) { dictEntry *de; if ((de = dictFind(it->sl.zs->dict,val->ele)) != NULL) { - *score = *(double*)dictGetEntryVal(de); + *score = *(double*)dictGetVal(de); return 1; } else { return 0; @@ -1396,7 +1454,7 @@ int zuiCompareByCardinality(const void *s1, const void *s2) { #define REDIS_AGGR_SUM 1 #define REDIS_AGGR_MIN 2 #define REDIS_AGGR_MAX 3 -#define zunionInterDictValue(_e) (dictGetEntryVal(_e) == NULL ? 1.0 : *(double*)dictGetEntryVal(_e)) +#define zunionInterDictValue(_e) (dictGetVal(_e) == NULL ? 1.0 : *(double*)dictGetVal(_e)) inline static void zunionInterAggregate(double *target, double val, int aggregate) { if (aggregate == REDIS_AGGR_SUM) { @@ -1416,7 +1474,8 @@ inline static void zunionInterAggregate(double *target, double val, int aggregat } void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) { - int i, j, setnum; + int i, j; + long setnum; int aggregate = REDIS_AGGR_SUM; zsetopsrc *src; zsetopval zval; @@ -1428,7 +1487,9 @@ void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) { int touched = 0; /* expect setnum input keys to be given */ - setnum = atoi(c->argv[2]->ptr); + if ((getLongFromObjectOrReply(c, c->argv[2], &setnum, NULL) != REDIS_OK)) + return; + if (setnum < 1) { addReplyError(c, "at least 1 input key is needed for ZUNIONSTORE/ZINTERSTORE"); @@ -1436,7 +1497,7 @@ void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) { } /* test if the expected number of keys would overflow */ - if (3+setnum > c->argc) { + if (setnum > c->argc-3) { addReply(c,shared.syntaxerr); return; } @@ -1472,7 +1533,7 @@ void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) { j++; remaining--; for (i = 0; i < setnum; i++, j++, remaining--) { if (getDoubleFromObjectOrReply(c,c->argv[j],&src[i].weight, - "weight value is not a double") != REDIS_OK) + "weight value is not a float") != REDIS_OK) { zfree(src); return; @@ -1520,8 +1581,15 @@ void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) { double score, value; score = src[0].weight * zval.score; + if (isnan(score)) score = 0; + for (j = 1; j < setnum; j++) { - if (zuiFind(&src[j],&zval,&value)) { + /* It is not safe to access the zset we are + * iterating, so explicitly check for equal object. */ + if (src[j].subject == src[0].subject) { + value = zval.score*src[j].weight; + zunionInterAggregate(&score,value,aggregate); + } else if (zuiFind(&src[j],&zval,&value)) { value *= src[j].weight; zunionInterAggregate(&score,value,aggregate); } else { @@ -1545,7 +1613,7 @@ void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) { } } else if (op == REDIS_OP_UNION) { for (i = 0; i < setnum; i++) { - if (zuiLength(&src[0]) == 0) + if (zuiLength(&src[i]) == 0) continue; while (zuiNext(&src[i],&zval)) { @@ -1557,11 +1625,17 @@ void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) { /* Initialize score */ score = src[i].weight * zval.score; + if (isnan(score)) score = 0; /* Because the inputs are sorted by size, it's only possible * for sets at larger indices to hold this element. */ for (j = (i+1); j < setnum; j++) { - if (zuiFind(&src[j],&zval,&value)) { + /* It is not safe to access the zset we are + * iterating, so explicitly check for equal object. */ + if(src[j].subject == src[i].subject) { + value = zval.score*src[j].weight; + zunionInterAggregate(&score,value,aggregate); + } else if (zuiFind(&src[j],&zval,&value)) { value *= src[j].weight; zunionInterAggregate(&score,value,aggregate); } @@ -1667,12 +1741,12 @@ void zrangeGenericCommand(redisClient *c, int reverse) { else eptr = ziplistIndex(zl,2*start); - redisAssert(eptr != NULL); + redisAssertWithInfo(c,zobj,eptr != NULL); sptr = ziplistNext(zl,eptr); while (rangelen--) { - redisAssert(eptr != NULL && sptr != NULL); - redisAssert(ziplistGet(eptr,&vstr,&vlen,&vlong)); + redisAssertWithInfo(c,zobj,eptr != NULL && sptr != NULL); + redisAssertWithInfo(c,zobj,ziplistGet(eptr,&vstr,&vlen,&vlong)); if (vstr == NULL) addReplyBulkLongLong(c,vlong); else @@ -1705,7 +1779,7 @@ void zrangeGenericCommand(redisClient *c, int reverse) { } while(rangelen--) { - redisAssert(ln != NULL); + redisAssertWithInfo(c,zobj,ln != NULL); ele = ln->obj; addReplyBulk(c,ele); if (withscores) @@ -1725,13 +1799,12 @@ void zrevrangeCommand(redisClient *c) { zrangeGenericCommand(c,1); } -/* This command implements ZRANGEBYSCORE, ZREVRANGEBYSCORE and ZCOUNT. - * If "justcount", only the number of elements in the range is returned. */ -void genericZrangebyscoreCommand(redisClient *c, int reverse, int justcount) { +/* This command implements ZRANGEBYSCORE, ZREVRANGEBYSCORE. */ +void genericZrangebyscoreCommand(redisClient *c, int reverse) { zrangespec range; robj *key = c->argv[1]; - robj *emptyreply, *zobj; - int offset = 0, limit = -1; + robj *zobj; + long offset = 0, limit = -1; int withscores = 0; unsigned long rangelen = 0; void *replylen = NULL; @@ -1747,7 +1820,7 @@ void genericZrangebyscoreCommand(redisClient *c, int reverse, int justcount) { } if (zslParseRange(c->argv[minidx],c->argv[maxidx],&range) != REDIS_OK) { - addReplyError(c,"min or max is not a double"); + addReplyError(c,"min or max is not a float"); return; } @@ -1762,8 +1835,8 @@ void genericZrangebyscoreCommand(redisClient *c, int reverse, int justcount) { pos++; remaining--; withscores = 1; } else if (remaining >= 3 && !strcasecmp(c->argv[pos]->ptr,"limit")) { - offset = atoi(c->argv[pos+1]->ptr); - limit = atoi(c->argv[pos+2]->ptr); + if ((getLongFromObjectOrReply(c, c->argv[pos+1], &offset, NULL) != REDIS_OK) || + (getLongFromObjectOrReply(c, c->argv[pos+2], &limit, NULL) != REDIS_OK)) return; pos += 3; remaining -= 3; } else { addReply(c,shared.syntaxerr); @@ -1773,8 +1846,7 @@ void genericZrangebyscoreCommand(redisClient *c, int reverse, int justcount) { } /* Ok, lookup the key and get the range */ - emptyreply = justcount ? shared.czero : shared.emptymultibulk; - if ((zobj = lookupKeyReadOrReply(c,key,emptyreply)) == NULL || + if ((zobj = lookupKeyReadOrReply(c,key,shared.emptymultibulk)) == NULL || checkType(c,zobj,REDIS_ZSET)) return; if (zobj->encoding == REDIS_ENCODING_ZIPLIST) { @@ -1786,34 +1858,36 @@ void genericZrangebyscoreCommand(redisClient *c, int reverse, int justcount) { double score; /* If reversed, get the last node in range as starting point. */ - if (reverse) + if (reverse) { eptr = zzlLastInRange(zl,range); - else + } else { eptr = zzlFirstInRange(zl,range); + } /* No "first" element in the specified interval. */ if (eptr == NULL) { - addReply(c,emptyreply); + addReply(c, shared.emptymultibulk); return; } /* Get score pointer for the first element. */ - redisAssert(eptr != NULL); + redisAssertWithInfo(c,zobj,eptr != NULL); sptr = ziplistNext(zl,eptr); /* We don't know in advance how many matching elements there are in the * list, so we push this object that will represent the multi-bulk * length in the output buffer, and will "fix" it later */ - if (!justcount) - replylen = addDeferredMultiBulkLength(c); + replylen = addDeferredMultiBulkLength(c); /* If there is an offset, just traverse the number of elements without * checking the score because that is done in the next loop. */ - while (eptr && offset--) - if (reverse) + while (eptr && offset--) { + if (reverse) { zzlPrev(zl,&eptr,&sptr); - else + } else { zzlNext(zl,&eptr,&sptr); + } + } while (eptr && limit--) { score = zzlGetScore(sptr); @@ -1825,24 +1899,26 @@ void genericZrangebyscoreCommand(redisClient *c, int reverse, int justcount) { if (!zslValueLteMax(score,&range)) break; } - /* Do our magic */ + /* We know the element exists, so ziplistGet should always succeed */ + redisAssertWithInfo(c,zobj,ziplistGet(eptr,&vstr,&vlen,&vlong)); + rangelen++; - if (!justcount) { - redisAssert(ziplistGet(eptr,&vstr,&vlen,&vlong)); - if (vstr == NULL) - addReplyBulkLongLong(c,vlong); - else - addReplyBulkCBuffer(c,vstr,vlen); - - if (withscores) - addReplyDouble(c,score); + if (vstr == NULL) { + addReplyBulkLongLong(c,vlong); + } else { + addReplyBulkCBuffer(c,vstr,vlen); + } + + if (withscores) { + addReplyDouble(c,score); } /* Move to next node */ - if (reverse) + if (reverse) { zzlPrev(zl,&eptr,&sptr); - else + } else { zzlNext(zl,&eptr,&sptr); + } } } else if (zobj->encoding == REDIS_ENCODING_SKIPLIST) { zset *zs = zobj->ptr; @@ -1850,27 +1926,32 @@ void genericZrangebyscoreCommand(redisClient *c, int reverse, int justcount) { zskiplistNode *ln; /* If reversed, get the last node in range as starting point. */ - if (reverse) + if (reverse) { ln = zslLastInRange(zsl,range); - else + } else { ln = zslFirstInRange(zsl,range); + } /* No "first" element in the specified interval. */ if (ln == NULL) { - addReply(c,emptyreply); + addReply(c, shared.emptymultibulk); return; } /* We don't know in advance how many matching elements there are in the * list, so we push this object that will represent the multi-bulk * length in the output buffer, and will "fix" it later */ - if (!justcount) - replylen = addDeferredMultiBulkLength(c); + replylen = addDeferredMultiBulkLength(c); /* If there is an offset, just traverse the number of elements without * checking the score because that is done in the next loop. */ - while (ln && offset--) - ln = reverse ? ln->backward : ln->level[0].forward; + while (ln && offset--) { + if (reverse) { + ln = ln->backward; + } else { + ln = ln->level[0].forward; + } + } while (ln && limit--) { /* Abort when the node is no longer in range. */ @@ -1880,39 +1961,114 @@ void genericZrangebyscoreCommand(redisClient *c, int reverse, int justcount) { if (!zslValueLteMax(ln->score,&range)) break; } - /* Do our magic */ rangelen++; - if (!justcount) { - addReplyBulk(c,ln->obj); - if (withscores) - addReplyDouble(c,ln->score); + addReplyBulk(c,ln->obj); + + if (withscores) { + addReplyDouble(c,ln->score); } /* Move to next node */ - ln = reverse ? ln->backward : ln->level[0].forward; + if (reverse) { + ln = ln->backward; + } else { + ln = ln->level[0].forward; + } } } else { redisPanic("Unknown sorted set encoding"); } - if (justcount) { - addReplyLongLong(c,(long)rangelen); - } else { - if (withscores) rangelen *= 2; - setDeferredMultiBulkLength(c,replylen,rangelen); + if (withscores) { + rangelen *= 2; } + + setDeferredMultiBulkLength(c, replylen, rangelen); } void zrangebyscoreCommand(redisClient *c) { - genericZrangebyscoreCommand(c,0,0); + genericZrangebyscoreCommand(c,0); } void zrevrangebyscoreCommand(redisClient *c) { - genericZrangebyscoreCommand(c,1,0); + genericZrangebyscoreCommand(c,1); } void zcountCommand(redisClient *c) { - genericZrangebyscoreCommand(c,0,1); + robj *key = c->argv[1]; + robj *zobj; + zrangespec range; + int count = 0; + + /* Parse the range arguments */ + if (zslParseRange(c->argv[2],c->argv[3],&range) != REDIS_OK) { + addReplyError(c,"min or max is not a float"); + return; + } + + /* Lookup the sorted set */ + if ((zobj = lookupKeyReadOrReply(c, key, shared.czero)) == NULL || + checkType(c, zobj, REDIS_ZSET)) return; + + if (zobj->encoding == REDIS_ENCODING_ZIPLIST) { + unsigned char *zl = zobj->ptr; + unsigned char *eptr, *sptr; + double score; + + /* Use the first element in range as the starting point */ + eptr = zzlFirstInRange(zl,range); + + /* No "first" element */ + if (eptr == NULL) { + addReply(c, shared.czero); + return; + } + + /* First element is in range */ + sptr = ziplistNext(zl,eptr); + score = zzlGetScore(sptr); + redisAssertWithInfo(c,zobj,zslValueLteMax(score,&range)); + + /* Iterate over elements in range */ + while (eptr) { + score = zzlGetScore(sptr); + + /* Abort when the node is no longer in range. */ + if (!zslValueLteMax(score,&range)) { + break; + } else { + count++; + zzlNext(zl,&eptr,&sptr); + } + } + } else if (zobj->encoding == REDIS_ENCODING_SKIPLIST) { + zset *zs = zobj->ptr; + zskiplist *zsl = zs->zsl; + zskiplistNode *zn; + unsigned long rank; + + /* Find first element in range */ + zn = zslFirstInRange(zsl, range); + + /* Use rank of first element, if any, to determine preliminary count */ + if (zn != NULL) { + rank = zslGetRank(zsl, zn->score, zn->obj); + count = (zsl->length - (rank - 1)); + + /* Find last element in range */ + zn = zslLastInRange(zsl, range); + + /* Use rank of last element, if any, to determine the actual count */ + if (zn != NULL) { + rank = zslGetRank(zsl, zn->score, zn->obj); + count -= (zsl->length - rank); + } + } + } else { + redisPanic("Unknown sorted set encoding"); + } + + addReplyLongLong(c, count); } void zcardCommand(redisClient *c) { @@ -1945,7 +2101,7 @@ void zscoreCommand(redisClient *c) { c->argv[2] = tryObjectEncoding(c->argv[2]); de = dictFind(zs->dict,c->argv[2]); if (de != NULL) { - score = *(double*)dictGetEntryVal(de); + score = *(double*)dictGetVal(de); addReplyDouble(c,score); } else { addReply(c,shared.nullbulk); @@ -1966,15 +2122,15 @@ void zrankGenericCommand(redisClient *c, int reverse) { checkType(c,zobj,REDIS_ZSET)) return; llen = zsetLength(zobj); - redisAssert(ele->encoding == REDIS_ENCODING_RAW); + redisAssertWithInfo(c,ele,ele->encoding == REDIS_ENCODING_RAW); if (zobj->encoding == REDIS_ENCODING_ZIPLIST) { unsigned char *zl = zobj->ptr; unsigned char *eptr, *sptr; eptr = ziplistIndex(zl,0); - redisAssert(eptr != NULL); + redisAssertWithInfo(c,zobj,eptr != NULL); sptr = ziplistNext(zl,eptr); - redisAssert(sptr != NULL); + redisAssertWithInfo(c,zobj,sptr != NULL); rank = 1; while(eptr != NULL) { @@ -2001,9 +2157,9 @@ void zrankGenericCommand(redisClient *c, int reverse) { ele = c->argv[2] = tryObjectEncoding(c->argv[2]); de = dictFind(zs->dict,ele); if (de != NULL) { - score = *(double*)dictGetEntryVal(de); + score = *(double*)dictGetVal(de); rank = zslGetRank(zsl,score,ele); - redisAssert(rank); /* Existing elements always have a rank. */ + redisAssertWithInfo(c,ele,rank); /* Existing elements always have a rank. */ if (reverse) addReplyLongLong(c,llen-rank); else