diff --git a/src/Makefile b/src/Makefile index 0138be0137..852b536c42 100644 --- a/src/Makefile +++ b/src/Makefile @@ -401,7 +401,7 @@ endif REDIS_SERVER_NAME=redis-server$(PROG_SUFFIX) REDIS_SENTINEL_NAME=redis-sentinel$(PROG_SUFFIX) -REDIS_SERVER_OBJ=threads_mngr.o memory_prefetch.o adlist.o quicklist.o ae.o anet.o dict.o ebuckets.o eventnotifier.o iothread.o mstr.o entry.o kvstore.o fwtree.o estore.o server.o sds.o zmalloc.o lzf_c.o lzf_d.o lzf_d_2.o pqsort.o zipmap.o sha1.o ziplist.o release.o networking.o util.o object.o db.o replication.o rdb.o t_string.o t_list.o t_set.o t_set_hashtable.o t_set_listpack.o t_set_intset.o t_zset.o t_hash.o bless.o t_array.o sparsearray.o config.o aof.o pubsub.o multi.o debug.o sort.o intset.o syncio.o cluster.o cluster_asm.o cluster_legacy.o cluster_slot_stats.o crc16.o endianconv.o slowlog.o eval.o bio.o rio.o rand.o memtest.o syscheck.o crcspeed.o crccombine.o crc64.o bitops.o bitroar.o sentinel.o notify.o setproctitle.o blocked.o hyperloglog.o latency.o sparkline.o redis-check-rdb.o redis-check-aof.o geo.o lazyfree.o module.o evict.o expire.o geohash.o geohash_helper.o childinfo.o defrag.o siphash.o rax.o t_stream.o listpack.o localtime.o lolwut.o lolwut5.o lolwut6.o lolwut8.o acl.o tracking.o socket.o tls.o sha256.o timeout.o setcpuaffinity.o monotonic.o mt19937-64.o resp_parser.o call_reply.o script_lua.o script.o functions.o function_lua.o commands.o strl.o connection.o unix.o logreqres.o keymeta.o chk.o hotkeys.o gcra.o vector.o fast_float_strtod.o client_comp.o +REDIS_SERVER_OBJ=threads_mngr.o memory_prefetch.o adlist.o quicklist.o ae.o anet.o dict.o ebuckets.o eventnotifier.o iothread.o mstr.o entry.o kvstore.o fwtree.o estore.o server.o sds.o zmalloc.o lzf_c.o lzf_d.o lzf_d_2.o pqsort.o zipmap.o sha1.o ziplist.o release.o networking.o util.o object.o db.o replication.o rdb.o t_string.o t_list.o t_set.o t_set_hashtable.o t_set_listpack.o t_set_intset.o t_zset.o zset_btree.o t_hash.o bless.o t_array.o sparsearray.o config.o aof.o pubsub.o multi.o debug.o sort.o intset.o syncio.o cluster.o cluster_asm.o cluster_legacy.o cluster_slot_stats.o crc16.o endianconv.o slowlog.o eval.o bio.o rio.o rand.o memtest.o syscheck.o crcspeed.o crccombine.o crc64.o bitops.o bitroar.o sentinel.o notify.o setproctitle.o blocked.o hyperloglog.o latency.o sparkline.o redis-check-rdb.o redis-check-aof.o geo.o lazyfree.o module.o evict.o expire.o geohash.o geohash_helper.o childinfo.o defrag.o siphash.o rax.o t_stream.o listpack.o localtime.o lolwut.o lolwut5.o lolwut6.o lolwut8.o acl.o tracking.o socket.o tls.o sha256.o timeout.o setcpuaffinity.o monotonic.o mt19937-64.o resp_parser.o call_reply.o script_lua.o script.o functions.o function_lua.o commands.o strl.o connection.o unix.o logreqres.o keymeta.o chk.o hotkeys.o gcra.o vector.o fast_float_strtod.o client_comp.o REDIS_CLI_NAME=redis-cli$(PROG_SUFFIX) REDIS_CLI_OBJ=anet.o adlist.o dict.o redis-cli.o zmalloc.o release.o ae.o redisassert.o crcspeed.o crccombine.o crc64.o siphash.o crc16.o monotonic.o cli_common.o mt19937-64.o strl.o cli_commands.o REDIS_BENCHMARK_NAME=redis-benchmark$(PROG_SUFFIX) diff --git a/src/aof.c b/src/aof.c index 92625a0308..d7dae7fc98 100644 --- a/src/aof.c +++ b/src/aof.c @@ -12,6 +12,7 @@ #include "rio.h" #include "functions.h" #include "cluster_asm.h" +#include "zset_btree.h" #include #include @@ -2409,6 +2410,27 @@ int rewriteSortedSetObject(rio *r, robj *key, robj *o) { items--; } dictResetIterator(&di); + } else if (o->encoding == OBJ_ENCODING_BTREE) { + zbtreeIterator iter; + const unsigned char *ele; + size_t len; + double score; + zbtreeIteratorStart(o->ptr, 0, &iter); + while (zbtreeIteratorNext(&iter, 0, &ele, &len, &score)) { + if (count == 0) { + int cmd_items = (items > AOF_REWRITE_ITEMS_PER_CMD) ? + AOF_REWRITE_ITEMS_PER_CMD : items; + if (!rioWriteBulkCount(r,'*',2+cmd_items*2) || + !rioWriteBulkString(r,"ZADD",4) || + !rioWriteBulkObject(r,key)) + return 0; + } + if (!rioWriteBulkDouble(r,score) || + !rioWriteBulkString(r,(char *)ele,len)) + return 0; + if (++count == AOF_REWRITE_ITEMS_PER_CMD) count = 0; + items--; + } } else { serverPanic("Unknown sorted zset encoding"); } diff --git a/src/db.c b/src/db.c index 8794f3ee25..3d5280626c 100644 --- a/src/db.c +++ b/src/db.c @@ -27,6 +27,7 @@ #include #include "bio.h" #include "keymeta.h" +#include "zset_btree.h" /*----------------------------------------------------------------------------- * C-level DB API @@ -1670,6 +1671,59 @@ typedef struct { redisDb *db; /* database reference for expiration checks */ } scanData; +typedef struct { + client *c; + sds pattern; + int pattern_len; + int use_pattern; + unsigned long emitted; + int streaming; + void *cursor_reply; + void *replylen; + struct { + const unsigned char *ele; + size_t len; + double score; + } buffered[32]; +} zsetBtreeScanData; + +static void zsetBtreeScanAddReply(client *c, const unsigned char *ele, size_t len, double score) { + addReplyBulkCBuffer(c, ele, len); + char buf[MAX_D2STRING_CHARS]; + int scorelen = d2string(buf, sizeof(buf), score); + addReplyBulkCBuffer(c, buf, scorelen); +} + +/* Buffer small replies so their headers need no deferred reply nodes. The + * borrowed members remain valid throughout this read-only command. Larger + * replies switch to streaming, keeping temporary storage bounded. */ +static void zsetBtreeScanReply(void *privdata, const unsigned char *ele, + size_t len, double score) +{ + zsetBtreeScanData *data = privdata; + if (data->use_pattern && + !stringmatchlen(data->pattern, data->pattern_len, (char *)ele, len, 0)) + return; + size_t capacity = sizeof(data->buffered) / sizeof(data->buffered[0]); + if (!data->streaming && data->emitted / 2 < capacity) { + size_t i = data->emitted / 2; + data->buffered[i].ele = ele; + data->buffered[i].len = len; + data->buffered[i].score = score; + } else { + if (!data->streaming) { + data->streaming = 1; + data->cursor_reply = addReplyDeferredLen(data->c); + data->replylen = addReplyDeferredLen(data->c); + for (size_t i = 0; i < capacity; i++) + zsetBtreeScanAddReply(data->c, data->buffered[i].ele, + data->buffered[i].len, data->buffered[i].score); + } + zsetBtreeScanAddReply(data->c, ele, len, score); + } + data->emitted += 2; +} + /* Helper function to compare key type in scan commands */ int objectTypeCompare(robj *o, long long target) { if (o->type != OBJ_MODULE) { @@ -2015,6 +2069,26 @@ void scanGenericCommand(client *c, robj *o, unsigned long long cursor) { else serverAssert(cur_length == array_reply_len); /* fail on corrupt data */ return; + } else if (o->type == OBJ_ZSET && o->encoding == OBJ_ENCODING_BTREE) { + vecRelease(&keys); + addReplyArrayLen(c, 2); + zsetBtreeScanData data = { + .c = c, .pattern = pat, .pattern_len = patlen, .use_pattern = use_pattern + }; + uint64_t next = zbtreeScan(o->ptr, cursor, count, zsetBtreeScanReply, &data); + char cursor_buf[LONG_STR_SIZE]; + int cursor_len = ull2string(cursor_buf, sizeof(cursor_buf), next); + if (data.streaming) { + setDeferredReplyBulkSds(c, data.cursor_reply, sdsnewlen(cursor_buf, cursor_len)); + setDeferredArrayLen(c, data.replylen, data.emitted); + } else { + addReplyBulkCBuffer(c, cursor_buf, cursor_len); + addReplyArrayLen(c, data.emitted); + for (size_t i = 0; i < data.emitted / 2; i++) + zsetBtreeScanAddReply(c, data.buffered[i].ele, + data.buffered[i].len, data.buffered[i].score); + } + return; } else if ((o->type == OBJ_HASH || o->type == OBJ_ZSET) && o->encoding == OBJ_ENCODING_LISTPACK) { diff --git a/src/debug.c b/src/debug.c index f408cc1363..40f634333a 100644 --- a/src/debug.c +++ b/src/debug.c @@ -22,6 +22,7 @@ #include "fpconv_dtoa.h" #include "fast_float_strtod.h" #include "cluster.h" +#include "zset_btree.h" #include "threads_mngr.h" #include "script.h" #include "cluster_asm.h" @@ -341,6 +342,20 @@ void xorObjectDigest(redisDb *db, robj *keyobj, unsigned char *digest, robj *o) xorDigest(digest,eledigest,20); } dictResetIterator(&di); + } else if (o->encoding == OBJ_ENCODING_BTREE) { + zbtreeIterator iter; + const unsigned char *ele; + size_t elelen; + double score; + zbtreeIteratorStart(o->ptr, 0, &iter); + while (zbtreeIteratorNext(&iter, 0, &ele, &elelen, &score)) { + const int len = fpconv_dtoa(score, buf); + buf[len] = '\0'; + memset(eledigest,0,20); + mixDigest(eledigest,ele,elelen); + mixDigest(eledigest,buf,strlen(buf)); + xorDigest(digest,eledigest,20); + } } else { serverPanic("Unknown sorted set encoding"); } diff --git a/src/defrag.c b/src/defrag.c index 6de324316d..2d1cf6f459 100644 --- a/src/defrag.c +++ b/src/defrag.c @@ -19,6 +19,7 @@ */ #include "server.h" +#include "zset_btree.h" #include "bitroar.h" #include #include @@ -116,6 +117,7 @@ typedef struct { * before "later" and we will search by key name to find the entry when we defrag the item later. */ list *defrag_later; unsigned long defrag_later_cursor; + zbtreeDefragState zset_btree_state; } defragKeysCtx; static_assert(offsetof(defragKeysCtx, kvstate) == 0, "defragStageKvstoreHelper requires this"); @@ -729,6 +731,26 @@ void defragZsetSkiplist(defragKeysCtx *ctx, kvobj *ob) { zs->dict = newdict; } +/* Defragment a B+ tree now, or schedule its pages for incremental work. */ +void defragZsetBtree(defragKeysCtx *ctx, kvobj *ob) { + zbtreeSet *zs = ob->ptr; + serverAssert(ob->type == OBJ_ZSET && ob->encoding == OBJ_ENCODING_BTREE); + + ob->ptr = zs = zbtreeDefragStart(zs, activeDefragAlloc); + if (zbtreeLength(zs) > server.active_defrag_max_scan_fields) { + defragLater(ctx, ob); + } else { + zbtreeDefragState state = {0}; + unsigned long scanned; + while (!zbtreeDefragStep(zs, &state, + activeDefragAllocWithoutFree, + activeDefragFree, &scanned)) + { + server.stat_active_defrag_scanned += scanned; + } + } +} + void defragHash(defragKeysCtx *ctx, kvobj *ob) { dict *d, *newd; serverAssert(ob->type == OBJ_HASH && ob->encoding == OBJ_ENCODING_HT); @@ -1218,6 +1240,8 @@ void defragKey(defragKeysCtx *ctx, dictEntry *de, dictEntryLink link) { ob->ptr = newzl; } else if (ob->encoding == OBJ_ENCODING_SKIPLIST) { defragZsetSkiplist(ctx, ob); + } else if (ob->encoding == OBJ_ENCODING_BTREE) { + defragZsetBtree(ctx, ob); } else { serverPanic("Unknown sorted set encoding"); } @@ -1358,7 +1382,15 @@ void defragPubsubScanCallback(void *privdata, const dictEntry *de, dictEntryLink /* returns 0 more work may or may not be needed (see non-zero cursor), * and 1 if time is up and more work is needed. */ -int defragLaterItem(kvobj *ob, unsigned long *cursor, monotime endtime, int dbid) { +int defragLaterItem(defragKeysCtx *ctx, kvobj *ob, monotime endtime) { + unsigned long *cursor = &ctx->defrag_later_cursor; + zbtreeDefragState *zset_btree_state = &ctx->zset_btree_state; + int is_btree = ob && ob->type == OBJ_ZSET && ob->encoding == OBJ_ENCODING_BTREE; + if (!is_btree && zset_btree_state->tree) { + zbtreeDefragStateReset(zset_btree_state); + *cursor = 0; + } + if (ob) { if (ob->type == OBJ_LIST && ob->encoding == OBJ_ENCODING_QUICKLIST) { return scanLaterList(ob, cursor, endtime); @@ -1366,6 +1398,15 @@ int defragLaterItem(kvobj *ob, unsigned long *cursor, monotime endtime, int dbid scanLaterSet(ob, cursor); } else if (ob->type == OBJ_ZSET && ob->encoding == OBJ_ENCODING_SKIPLIST) { scanLaterZset(ob, cursor); + } else if (ob->type == OBJ_ZSET && ob->encoding == OBJ_ENCODING_BTREE) { + if (*cursor == 0) + zbtreeDefragStateReset(zset_btree_state); + unsigned long scanned; + int done = zbtreeDefragStep(ob->ptr, zset_btree_state, + activeDefragAllocWithoutFree, + activeDefragFree, &scanned); + server.stat_active_defrag_scanned += scanned; + *cursor = done ? 0 : 1; } else if (ob->type == OBJ_HASH && ob->encoding == OBJ_ENCODING_HT) { scanLaterHash(ob, cursor); } else if (ob->type == OBJ_HASH && ob->encoding == OBJ_ENCODING_TMPL_ARRAY) { @@ -1375,7 +1416,7 @@ int defragLaterItem(kvobj *ob, unsigned long *cursor, monotime endtime, int dbid } else if (ob->type == OBJ_MODULE) { robj keyobj; initStaticStringObject(keyobj, kvobjGetKey(ob)); - return moduleLateDefrag(&keyobj, ob, cursor, endtime, dbid); + return moduleLateDefrag(&keyobj, ob, cursor, endtime, ctx->dbid); } else if (ob->type == OBJ_ARRAY) { redisArray *ar = ob->ptr; *cursor = arDefragIncremental(&ar, *cursor, activeDefragAlloc); @@ -1384,9 +1425,11 @@ int defragLaterItem(kvobj *ob, unsigned long *cursor, monotime endtime, int dbid *cursor = bitroarDefragIncremental(ob, *cursor); } else { *cursor = 0; /* object type/encoding may have changed since we schedule it for later */ + zbtreeDefragStateReset(zset_btree_state); } } else { *cursor = 0; /* object may have been deleted already */ + zbtreeDefragStateReset(zset_btree_state); } return 0; } @@ -1415,7 +1458,7 @@ static doneStatus defragLaterStep(void *ctx, monotime endtime) { long long key_defragged = server.stat_active_defrag_hits; if (server.memory_tracking_enabled && kv) oldsize = kvobjAllocSize(kv); - int timeout = (defragLaterItem(kv, &defrag_keys_ctx->defrag_later_cursor, endtime, defrag_keys_ctx->dbid) == 1); + int timeout = (defragLaterItem(defrag_keys_ctx, kv, endtime) == 1); if (server.memory_tracking_enabled && kv) updateSlotAllocSize(db, slot, kv, oldsize, kvobjAllocSize(kv)); if (key_defragged != server.stat_active_defrag_hits) { diff --git a/src/dict.c b/src/dict.c index 67a7c80b20..be0cef664f 100644 --- a/src/dict.c +++ b/src/dict.c @@ -75,7 +75,7 @@ static int _dictInit(dict *d, dictType *type); static dictEntryLink dictGetNextLink(dictEntry *de); static void dictSetNext(dictEntry *de, dictEntry *next); static int dictDefaultCompare(dictCmpCache *cache, const void *key1, const void *key2); -static dictEntryLink dictFindLinkInternal(dict *d, const void *key, dictEntryLink *bucket); +static dictEntryLink dictFindLinkInternal(dict *d, const void *key, uint64_t hash, dictEntryLink *bucket); dictEntryLink dictFindLinkForInsert(dict *d, const void *key, dictEntry **existing); static dictEntry *dictInsertKeyAtLink(dict *d, void *key __stored_key, dictEntryLink link); @@ -763,20 +763,19 @@ void dictRelease(dict *d) * * bucket - return pointer to bucket that the key was mapped. unless dict is empty. */ -static dictEntryLink dictFindLinkInternal(dict *d, const void *key, dictEntryLink *bucket) { +static dictEntryLink dictFindLinkInternal(dict *d, const void *key, uint64_t hash, dictEntryLink *bucket) { dictCmpCache cmpCache = {0}; dictEntryLink link; uint64_t idx; int table; - + if (bucket) { *bucket = NULL; } else { /* If dict is empty and no need to find bucket, return NULL */ - if (dictSize(d) == 0) return NULL; + if (dictSize(d) == 0) return NULL; } - const uint64_t hash = dictGetHash(d, key); idx = hash & DICTHT_SIZE_MASK(d->ht_size_exp[0]); keyCmpFunc cmpFunc = dictGetCmpFunc(d); @@ -865,10 +864,23 @@ dictEntryLink dictFindLink(dict *d, const void *key, dictEntryLink *bucket) { if (unlikely(dictSize(d) == 0)) return NULL; - return dictFindLinkInternal(d, key, bucket); + return dictFindLinkInternal(d, key, dictGetHash(d, key), bucket); } -/* Set the key with link +/* Same as dictFindLink(), but for a caller that already computed this exact + * key's hash (via dictGetHash(d, key)) for another purpose and wants to avoid + * paying for the dict type's hashFunction() a second time. The hash is not + * validated against key -- passing a mismatched hash silently searches the + * wrong bucket. */ +dictEntryLink dictFindLinkWithHash(dict *d, const void *key, uint64_t hash, dictEntryLink *bucket) { + if (bucket) *bucket = NULL; + if (unlikely(dictSize(d) == 0)) + return NULL; + + return dictFindLinkInternal(d, key, hash, bucket); +} + +/* Set the key with link * * link: - When `newItem` is set, `link` points to the bucket of the key. * - When `newItem` is not set, `link` points to the link of the key. @@ -893,7 +905,8 @@ void dictSetKeyAtLink(dict *d, void *key __stored_key, dictEntryLink *link, int if (snap[0] != d->ht_size_exp[0] || snap[1] != d->ht_size_exp[1] || *link == NULL) { dictEntryLink bucket; /* Bypass dictFindLink() to search bucket even if dict is empty!!! */ - *link = dictFindLinkInternal(d, dictStoredKey2Key(d, key), &bucket); + const void *lookup_key = dictStoredKey2Key(d, key); + *link = dictFindLinkInternal(d, lookup_key, dictGetHash(d, lookup_key), &bucket); assert(bucket != NULL); assert(*link == NULL); *link = bucket; /* On newItem the link should be the bucket */ diff --git a/src/dict.h b/src/dict.h index 08460727ad..2074e84409 100644 --- a/src/dict.h +++ b/src/dict.h @@ -304,6 +304,7 @@ void dictCombineStats(dictStats *from, dictStats *into); void dictFreeStats(dictStats *stats); dictEntryLink dictFindLink(dict *d, const void *key, dictEntryLink *bucket); +dictEntryLink dictFindLinkWithHash(dict *d, const void *key, uint64_t hash, dictEntryLink *bucket); void dictSetKeyAtLink(dict *d, void *key __stored_key, dictEntryLink *link, int newItem); /* API relevant only when dict is used as a hash-map (no_value=0) */ diff --git a/src/geo.c b/src/geo.c index 978270bb68..b85a99705f 100644 --- a/src/geo.c +++ b/src/geo.c @@ -30,6 +30,7 @@ #include "geo.h" #include "geohash_helper.h" +#include "zset_btree.h" #include "debugmacro.h" #include "pqsort.h" @@ -319,6 +320,21 @@ int geoGetPointsInRange(robj *zobj, double min, double max, GeoShape *shape, geo if (ga->used && limit && ga->used >= limit) break; ln = ln->level[0].forward; } + } else if (zobj->encoding == OBJ_ENCODING_BTREE) { + zbtreeIterator iter; + const unsigned char *ele; + size_t len; + double score; + if (!zbtreeIteratorSeekScore(zobj->ptr, range.min, range.minex, 0, &iter, NULL)) + return 0; + while (zbtreeIteratorNext(&iter, 0, &ele, &len, &score)) { + double xy[2]; + double distance = 0; + if (!zslValueLteMax(score, &range)) break; + if (geoWithinShape(shape, score, xy, &distance) == C_OK) + geoArrayAppend(ga, xy, distance, score, sdsnewlen(ele, len)); + if (ga->used && limit && ga->used >= limit) break; + } } return ga->used - origincount; } @@ -804,32 +820,23 @@ void georadiusGeneric(client *c, int srcKeyIndex, int flags) { } else { /* Target key, create a sorted set with the results. */ robj *zobj; - zset *zs; int i; - size_t maxelelen = 0, totelelen = 0; - if (returned_items) { - zobj = createZsetObject(); - zs = zobj->ptr; - } + if (returned_items) + zobj = zsetTypeCreate(returned_items, 0); for (i = 0; i < returned_items; i++) { - zskiplistNode *znode; geoPoint *gp = ga->array+i; gp->dist /= shape.conversion; /* Fix according to unit. */ double score = storedist ? gp->dist : gp->score; - size_t elelen = sdslen(gp->member); - - if (maxelelen < elelen) maxelelen = elelen; - totelelen += elelen; - znode = zslInsert(zs->zsl,score,gp->member); - serverAssert(dictAdd(zs->dict, znode, NULL) == DICT_OK); - sdsfree(gp->member); /* zslInsert copies the sds, so free the original */ + int out_flags; + serverAssert(zsetAdd(zobj, score, gp->member, ZADD_IN_NONE, &out_flags, NULL)); + serverAssert(out_flags & ZADD_OUT_ADDED); + sdsfree(gp->member); gp->member = NULL; } if (returned_items) { - zsetConvertToListpackIfNeeded(zobj,maxelelen,totelelen); setKey(c,c->db,storekey,&zobj,0); notifyKeyspaceEvent(NOTIFY_ZSET,flags & GEOSEARCH ? "geosearchstore" : "georadiusstore",storekey, c->db->id); diff --git a/src/lazyfree.c b/src/lazyfree.c index f4f5bd5a3f..4d707ebab7 100644 --- a/src/lazyfree.c +++ b/src/lazyfree.c @@ -155,6 +155,8 @@ size_t lazyfreeGetFreeEffort(robj *key, robj *obj, int dbid) { } else if (obj->type == OBJ_ZSET && obj->encoding == OBJ_ENCODING_SKIPLIST){ zset *zs = obj->ptr; return zs->zsl->length; + } else if (obj->type == OBJ_ZSET && obj->encoding == OBJ_ENCODING_BTREE) { + return zsetLength(obj); } else if (obj->type == OBJ_HASH && obj->encoding == OBJ_ENCODING_HT) { dict *ht = obj->ptr; return dictSize(ht); diff --git a/src/module.c b/src/module.c index 7b3cfc9ed3..7a9f493a7e 100644 --- a/src/module.c +++ b/src/module.c @@ -46,6 +46,7 @@ #include "call_reply.h" #include "hdr_histogram.h" #include "crc16_slottable.h" +#include "zset_btree.h" #include #include #include @@ -198,7 +199,7 @@ struct RedisModuleKey { zlexrangespec lrs; /* Lex range. */ uint32_t start; /* Start pos for positional ranges. */ uint32_t end; /* End pos for positional ranges. */ - void *current; /* Zset iterator current node. */ + void *current; /* Node, or B+ tree rank plus one. */ int er; /* Zset iterator end reached flag (true if end was reached). */ } zset; @@ -5455,6 +5456,18 @@ void zsetKeyReset(RedisModuleKey *key) { key->u.zset.er = 1; } +/* Read one B+ tree member by rank. Seeking for every call keeps no leaf + * pointer across module code, which may modify the sorted set. */ +static int moduleZsetBtreeElementAtRank(RedisModuleKey *key, + unsigned long rank, + const unsigned char **ele, + size_t *len, double *score) +{ + zbtreeIterator iter; + if (!zbtreeIteratorSeekRank(key->kv->ptr, rank, &iter)) return 0; + return zbtreeIteratorNext(&iter, 0, ele, len, score); +} + /* Stop a sorted set iteration. */ void RM_ZsetRangeStop(RedisModuleKey *key) { if (!key->kv || key->kv->type != OBJ_ZSET) return; @@ -5502,6 +5515,19 @@ int zsetInitScoreRange(RedisModuleKey *key, double min, double max, int minex, i zskiplist *zsl = zs->zsl; key->u.zset.current = first ? zslNthInRange(zsl, zrs, 0, NULL) : zslNthInRange(zsl, zrs, -1, NULL); + } else if (key->kv->encoding == OBJ_ENCODING_BTREE) { + zbtreeIterator iter; + unsigned long rank; + double score; + int found = first ? + zbtreeIteratorSeekScore(key->kv->ptr, min, minex, 0, &iter, &rank) : + zbtreeIteratorSeekScore(key->kv->ptr, max, maxex, 1, &iter, &rank); + if (found) { + serverAssert(moduleZsetBtreeElementAtRank(key, rank, NULL, NULL, &score)); + /* Module callers can pass NaN bounds, so check both sides. */ + if (zslValueGteMin(score, zrs) && zslValueLteMax(score, zrs)) + key->u.zset.current = (void *)(uintptr_t)(rank + 1); + } } else { serverPanic("Unsupported zset encoding"); } @@ -5566,6 +5592,23 @@ int zsetInitLexRange(RedisModuleKey *key, RedisModuleString *min, RedisModuleStr zskiplist *zsl = zs->zsl; key->u.zset.current = first ? zslNthInLexRange(zsl,zlrs,0,NULL) : zslNthInLexRange(zsl,zlrs,-1,NULL); + } else if (key->kv->encoding == OBJ_ENCODING_BTREE) { + zbtreeIterator iter; + unsigned long rank; + int found = first ? + zbtreeIteratorSeekLex(key->kv->ptr, zlrs->min, zlrs->minex, 0, &iter, &rank) : + zbtreeIteratorSeekLex(key->kv->ptr, zlrs->max, zlrs->maxex, 1, &iter, &rank); + if (found) { + const unsigned char *ele; + size_t len; + serverAssert(moduleZsetBtreeElementAtRank(key, rank, &ele, &len, NULL)); + sds value = sdsnewlen(ele, len); + int in_range = zslLexValueGteMin(value, zlrs) && + zslLexValueLteMax(value, zlrs); + sdsfree(value); + if (in_range) + key->u.zset.current = (void *)(uintptr_t)(rank + 1); + } } else { serverPanic("Unsupported zset encoding"); } @@ -5618,6 +5661,17 @@ RedisModuleString *RM_ZsetRangeCurrentElement(RedisModuleKey *key, double *score if (score) *score = ln->score; sds ele = zslGetNodeElement(ln); str = createStringObject(ele,sdslen(ele)); + } else if (key->kv->encoding == OBJ_ENCODING_BTREE) { + const unsigned char *ele; + size_t len; + double value; + unsigned long rank = (uintptr_t)key->u.zset.current - 1; + if (!moduleZsetBtreeElementAtRank(key, rank, &ele, &len, &value)) { + key->u.zset.er = 1; + return NULL; + } + if (score) *score = value; + str = createStringObject((char *)ele, len); } else { serverPanic("Unsupported zset encoding"); } @@ -5684,9 +5738,31 @@ int RM_ZsetRangeNext(RedisModuleKey *key) { key->u.zset.current = next; return 1; } + } else if (key->kv->encoding == OBJ_ENCODING_BTREE) { + unsigned long rank = (uintptr_t)key->u.zset.current - 1; + if (++rank >= zbtreeLength(key->kv->ptr)) goto end; + const unsigned char *ele; + size_t len; + double score; + if (!moduleZsetBtreeElementAtRank(key, rank, &ele, &len, &score)) + goto end; + if (key->u.zset.type == REDISMODULE_ZSET_RANGE_SCORE) { + if (!zslValueLteMax(score, &key->u.zset.rs)) goto end; + } else if (key->u.zset.type == REDISMODULE_ZSET_RANGE_LEX) { + sds value = sdsnewlen(ele, len); + int valid = zslLexValueLteMax(value, &key->u.zset.lrs); + sdsfree(value); + if (!valid) goto end; + } + key->u.zset.current = (void *)(uintptr_t)(rank + 1); + return 1; } else { serverPanic("Unsupported zset encoding"); } + +end: + key->u.zset.er = 1; + return 0; } /* Go to the previous element of the sorted set iterator. Returns 1 if there was @@ -5748,9 +5824,32 @@ int RM_ZsetRangePrev(RedisModuleKey *key) { key->u.zset.current = prev; return 1; } + } else if (key->kv->encoding == OBJ_ENCODING_BTREE) { + unsigned long rank = (uintptr_t)key->u.zset.current - 1; + if (rank == 0) goto end; + rank--; + const unsigned char *ele; + size_t len; + double score; + if (!moduleZsetBtreeElementAtRank(key, rank, &ele, &len, &score)) + goto end; + if (key->u.zset.type == REDISMODULE_ZSET_RANGE_SCORE) { + if (!zslValueGteMin(score, &key->u.zset.rs)) goto end; + } else if (key->u.zset.type == REDISMODULE_ZSET_RANGE_LEX) { + sds value = sdsnewlen(ele, len); + int valid = zslLexValueGteMin(value, &key->u.zset.lrs); + sdsfree(value); + if (!valid) goto end; + } + key->u.zset.current = (void *)(uintptr_t)(rank + 1); + return 1; } else { serverPanic("Unsupported zset encoding"); } + +end: + key->u.zset.er = 1; + return 0; } /* -------------------------------------------------------------------------- @@ -12486,6 +12585,34 @@ typedef struct { RedisModuleScanKeyCB fn; } ScanKeyCBData; +typedef struct { + robj *field; + double score; +} BtreeScanItem; + +typedef struct { + BtreeScanItem *items; + size_t count; + size_t capacity; +} BtreeScanData; + +/* B+ tree leaves may move if module code changes the sorted set. Collect one + * small batch before calling the module so no tree pointer survives a + * callback. */ +static void moduleBtreeScanCollect(void *privdata, const unsigned char *ele, + size_t len, double score) +{ + BtreeScanData *data = privdata; + if (data->count == data->capacity) { + data->capacity = data->capacity ? data->capacity * 2 : 64; + data->items = zrealloc(data->items, data->capacity * sizeof(*data->items)); + } + + data->items[data->count].field = createStringObject((char *)ele, len); + data->items[data->count].score = score; + data->count++; +} + static void moduleScanKeyCallback(void *privdata, const dictEntry *de, dictEntryLink plink) { UNUSED(plink); ScanKeyCBData *data = privdata; @@ -12636,6 +12763,34 @@ int RM_ScanKey(RedisModuleKey *key, RedisModuleScanCursor *cursor, RedisModuleSc cursor->cursor = 1; cursor->done = 1; ret = 0; + } else if (kv->type == OBJ_ZSET && kv->encoding == OBJ_ENCODING_BTREE) { + BtreeScanData data = {0}; + cursor->cursor = zbtreeScan(kv->ptr, cursor->cursor, 64, moduleBtreeScanCollect, &data); + + size_t i = 0; + for (; i < data.count; i++) { + robj *value = createStringObjectFromLongDouble(data.items[i].score, 0); + fn(key, data.items[i].field, value, privdata); + decrRefCount(value); + decrRefCount(data.items[i].field); + + /* Removing the current member is safe. If the callback removes the + * key or changes its encoding, this B+ tree cursor has no meaning. */ + if (key->kv != kv || kv->encoding != OBJ_ENCODING_BTREE) { + cursor->cursor = 0; + cursor->done = 1; + ret = 0; + i++; + break; + } + } + while (i < data.count) decrRefCount(data.items[i++].field); + zfree(data.items); + + if (cursor->cursor == 0) { + cursor->done = 1; + ret = 0; + } } else if (kv->type == OBJ_ZSET || kv->type == OBJ_HASH) { unsigned char *lp, *p; /* is hash with expiry on fields, then lp tuples are [field][value][expire] */ diff --git a/src/object.c b/src/object.c index e89823868e..7e1b02aa24 100644 --- a/src/object.c +++ b/src/object.c @@ -16,6 +16,7 @@ #include "functions.h" #include "intset.h" /* Compact integer set structure */ #include "cluster_asm.h" +#include "zset_btree.h" #include #include @@ -510,6 +511,13 @@ robj *createZsetObject(void) { return o; } +/* Create an empty sorted set using the packed B+ tree encoding. */ +robj *createZsetBtreeObject(void) { + robj *o = createObject(OBJ_ZSET, zbtreeCreate()); + o->encoding = OBJ_ENCODING_BTREE; + return o; +} + robj *createZsetListpackObject(void) { unsigned char *lp = lpNew(0); robj *o = createObject(OBJ_ZSET,lp); @@ -586,6 +594,9 @@ void freeZsetObject(robj *o) { zslFree(zs->zsl); zfree(zs); break; + case OBJ_ENCODING_BTREE: + zbtreeFree(o->ptr); + break; case OBJ_ENCODING_LISTPACK: zfree(o->ptr); break; @@ -769,6 +780,8 @@ void dismissZsetObject(robj *o, size_t size_hint) { /* Dismiss hash table memory. */ dismissDictBucketsMemory(zs->dict); + } else if (o->encoding == OBJ_ENCODING_BTREE) { + zbtreeDismissMemory(o->ptr); } else if (o->encoding == OBJ_ENCODING_LISTPACK) { dismissMemory(o->ptr, lpBytes((unsigned char*)o->ptr)); } else { @@ -1322,6 +1335,7 @@ char *strEncoding(int encoding) { case OBJ_ENCODING_TMPL_LP: return "template-listpack"; case OBJ_ENCODING_TMPL_ARRAY: return "template-array"; case OBJ_ENCODING_BITMAP_ROARING: return "bitmap-roaring"; + case OBJ_ENCODING_BTREE: return "btree"; default: return "unknown"; } } diff --git a/src/object.h b/src/object.h index 36896f125d..e04b8261a7 100644 --- a/src/object.h +++ b/src/object.h @@ -78,7 +78,7 @@ struct RedisModuleType; #define OBJ_ENCODING_HT 2 /* Encoded as hash table */ #define OBJ_ENCODING_ZIPMAP 3 /* No longer used: old hash encoding. */ #define OBJ_ENCODING_LINKEDLIST 4 /* No longer used: old list encoding. */ -#define OBJ_ENCODING_ZIPLIST 5 /* No longer used: old list/hash/zset encoding. */ +#define OBJ_ENCODING_BTREE 5 /* Sorted set as a packed B+ tree */ #define OBJ_ENCODING_INTSET 6 /* Encoded as intset */ #define OBJ_ENCODING_SKIPLIST 7 /* Encoded as skiplist */ #define OBJ_ENCODING_EMBSTR 8 /* Embedded sds string encoding */ @@ -180,6 +180,7 @@ robj *createIntsetObject(void); robj *createSetListpackObject(void); robj *createHashObject(void); robj *createZsetObject(void); +robj *createZsetBtreeObject(void); robj *createZsetListpackObject(void); robj *createStreamObject(void); robj *createGCRAObject(long long value); diff --git a/src/rdb.c b/src/rdb.c index 7625327fb7..b41a9430d2 100644 --- a/src/rdb.c +++ b/src/rdb.c @@ -24,6 +24,7 @@ #include "bio.h" #include "cluster_asm.h" #include "keymeta.h" +#include "zset_btree.h" #include #include @@ -713,7 +714,8 @@ int rdbSaveObjectType(rio *rdb, robj *o) { case OBJ_ZSET: if (o->encoding == OBJ_ENCODING_LISTPACK) return rdbSaveType(rdb,RDB_TYPE_ZSET_LISTPACK); - else if (o->encoding == OBJ_ENCODING_SKIPLIST) + else if (o->encoding == OBJ_ENCODING_SKIPLIST || + o->encoding == OBJ_ENCODING_BTREE) return rdbSaveType(rdb,RDB_TYPE_ZSET_2); else serverPanic("Unknown sorted set encoding"); @@ -1283,7 +1285,9 @@ ssize_t rdbSaveObject(rio *rdb, robj *o, robj *key, int dbid) { serverPanic("Unknown set encoding"); } } else if (o->type == OBJ_ZSET) { - /* Save a sorted set value */ + /* Save a sorted set value. SKIPLIST is only a temporary encoding, + * keyspace zsets are LISTPACK or BTREE. */ + debugServerAssert(o->encoding != OBJ_ENCODING_SKIPLIST); if (o->encoding == OBJ_ENCODING_LISTPACK) { size_t l = lpBytes((unsigned char*)o->ptr); @@ -1316,6 +1320,24 @@ ssize_t rdbSaveObject(rio *rdb, robj *o, robj *key, int dbid) { nwritten += n; zn = zn->backward; } + } else if (o->encoding == OBJ_ENCODING_BTREE) { + zbtreeIterator iter; + const unsigned char *ele; + size_t len; + double score; + + if ((n = rdbSaveLen(rdb, zbtreeLength(o->ptr))) == -1) + return -1; + nwritten += n; + zbtreeIteratorStart(o->ptr, 1, &iter); + while (zbtreeIteratorNext(&iter, 1, &ele, &len, &score)) { + if ((n = rdbSaveRawString(rdb, (unsigned char *)ele, len)) == -1) + return -1; + nwritten += n; + if ((n = rdbSaveBinaryDoubleValue(rdb, score)) == -1) + return -1; + nwritten += n; + } } else { serverPanic("Unknown sorted set encoding"); } @@ -2543,23 +2565,27 @@ static int _listZiplistEntryConvertAndValidate(unsigned char *p, unsigned int he return 1; } -/* callback for to check the listpack doesn't have duplicate records */ -static int _lpEntryValidation(unsigned char *p, unsigned int head_count, void *userdata) { - struct { - int tuple_len; - long count; - dict *fields; - long long last_expireat; - } *data = userdata; +typedef struct { + int tuple_len; + long count; + dict *fields; + long long last_expireat; + int check_duplicates; + int validate_scores; +} lpValidationData; - if (data->fields == NULL) { +/* Validate tuple fields after lpValidateIntegrity checked the entry. */ +static int _lpEntryValidation(unsigned char *p, unsigned int head_count, void *userdata) { + lpValidationData *data = userdata; + + if (data->check_duplicates && data->fields == NULL) { data->fields = dictCreate(&hashDictType); dictExpand(data->fields, head_count/data->tuple_len); } /* If we're checking pairs, then even records are field names. Otherwise * we're checking all elements. Add to dict and check that's not a dup */ - if (data->count % data->tuple_len == 0) { + if (data->check_duplicates && data->count % data->tuple_len == 0) { unsigned char *str; int64_t slen; unsigned char buf[LP_INTBUF_SIZE]; @@ -2573,6 +2599,17 @@ static int _lpEntryValidation(unsigned char *p, unsigned int head_count, void *u } } + /* A zset's second entry is a score, not an arbitrary hash value. Parse + * it before any conversion can store an invalid or NaN score. Integer + * entries are always valid, and infinity remains a supported score. */ + if (data->validate_scores && data->count % data->tuple_len == 1) { + unsigned int len; + long long integer; + unsigned char *str = lpGetValue(p, &len, &integer); + double score; + if (str && !string2d((char *)str, len, &score)) return 0; + } + /* Validate TTL field, only for listpackex. */ if (data->count % data->tuple_len == 2) { long long expire_at; @@ -2596,17 +2633,15 @@ static int _lpEntryValidation(unsigned char *p, unsigned int head_count, void *u * tuple_len indicates what is a logical entry tuple size. * Whether tuple is of size 1 (set), 2 (field-value) or 3 (field-value[-ttl]), * first element in the tuple must be unique */ -int lpValidateIntegrityAndDups(unsigned char *lp, size_t size, int deep, int tuple_len) { - if (!deep) +static int lpValidateIntegrityAndDupsEx(unsigned char *lp, size_t size, + int deep, int tuple_len, int validate_scores) +{ + if (!deep && !validate_scores) return lpValidateIntegrity(lp, size, 0, NULL, NULL); - /* Keep track of the field names to locate duplicate ones */ - struct { - int tuple_len; - long count; - dict *fields; /* Initialisation at the first callback. */ - long long last_expireat; /* Last field's expiry time to ensure order in TTL fields. */ - } data = {tuple_len, 0, NULL, -1}; + /* Score validation must visit every entry even when duplicate checking + * is disabled. The integrity walker bounds each entry before decoding. */ + lpValidationData data = {tuple_len, 0, NULL, -1, deep, validate_scores}; int ret = lpValidateIntegrity(lp, size, 1, _lpEntryValidation, &data); @@ -2618,6 +2653,10 @@ int lpValidateIntegrityAndDups(unsigned char *lp, size_t size, int deep, int tup return ret; } +int lpValidateIntegrityAndDups(unsigned char *lp, size_t size, int deep, int tuple_len) { + return lpValidateIntegrityAndDupsEx(lp, size, deep, tuple_len, 0); +} + /* ---- shared SDS-array helpers ---- */ /* Free 'count' SDS strings and the array holding them. Safe on NULL. */ @@ -3116,16 +3155,16 @@ robj *rdbLoadObject(int rdbtype, rio *rdb, sds key, int dbid, int *error) /* Read sorted set value. */ uint64_t zsetlen; size_t maxelelen = 0, totelelen = 0; - zset *zs; if ((zsetlen = rdbLoadLen(rdb,NULL)) == RDB_LENERR) return NULL; if (zsetlen == 0) goto emptykey; - - o = createZsetObject(); - zs = o->ptr; - - if (zsetlen > DICT_HT_INITIAL_SIZE && dictTryExpand(zs->dict,zsetlen) != DICT_OK) { - rdbReportCorruptRDB("OOM in dictTryExpand %llu", (unsigned long long)zsetlen); + /* Decide listpack eligibility only after seeing every member. A + * temporary listpack would normalize -0 before a later long member + * promotes the object, changing the final B+tree's score bits. */ + o = createZsetBtreeObject(); + if (!zbtreeTryReserve(o->ptr, zsetlen)) { + rdbReportCorruptRDB("Unable to reserve zset index for %llu elements", + (unsigned long long)zsetlen); decrRefCount(o); return NULL; } @@ -3134,7 +3173,6 @@ robj *rdbLoadObject(int rdbtype, rio *rdb, sds key, int dbid, int *error) while(zsetlen--) { sds sdsele; double score; - zskiplistNode *znode; if ((sdsele = rdbGenericLoadStringObject(rdb,RDB_LOAD_SDS,NULL)) == NULL) { decrRefCount(o); @@ -3162,27 +3200,23 @@ robj *rdbLoadObject(int rdbtype, rio *rdb, sds key, int dbid, int *error) return NULL; } - /* Don't care about integer-encoded strings. */ - if (sdslen(sdsele) > maxelelen) maxelelen = sdslen(sdsele); - totelelen += sdslen(sdsele); - - znode = zslInsert(zs->zsl,score,sdsele); - if (dictAdd(zs->dict, znode, NULL) != DICT_OK) { + int out_flags; + serverAssert(zsetAdd(o, score, sdsele, ZADD_IN_NX, &out_flags, NULL)); + if (!(out_flags & ZADD_OUT_ADDED)) { rdbReportCorruptRDB("Duplicate zset fields detected"); decrRefCount(o); - sdsfree(sdsele); /* zslInsert copies the sds, so we need to free the original */ + sdsfree(sdsele); return NULL; } - sdsfree(sdsele); /* zslInsert copies the sds into the node, so free the original */ - } - - /* Convert *after* loading, since sorted sets are not stored ordered. */ - if (zsetLength(o) <= server.zset_max_listpack_entries && - maxelelen <= server.zset_max_listpack_value && - lpSafeToAdd(NULL, totelelen)) - { - zsetConvert(o, OBJ_ENCODING_LISTPACK); + size_t elelen = sdslen(sdsele); + if (elelen > maxelelen) maxelelen = elelen; + if (totelelen <= SIZE_MAX - elelen) + totelelen += elelen; + else + totelelen = SIZE_MAX; + sdsfree(sdsele); } + zsetConvertAfterBulkInsert(o, maxelelen, totelelen); } else if (rdbtype == RDB_TYPE_HASH) { uint64_t len, original_len; int ret; @@ -3820,7 +3854,12 @@ robj *rdbLoadObject(int rdbtype, rio *rdb, sds key, int dbid, int *error) case RDB_TYPE_ZSET_ZIPLIST: { unsigned char *lp = lpNew(encoded_len); - if (!ziplistPairsConvertAndValidateIntegrity(encoded, encoded_len, &lp)) { + int valid = ziplistPairsConvertAndValidateIntegrity(encoded, encoded_len, &lp); + /* See RDB_TYPE_ZSET_LISTPACK below for when scores are parsed. */ + if (valid && (deep_integrity_validation || + lpLength(lp) / 2 > server.zset_max_listpack_entries)) + valid = lpValidateIntegrityAndDupsEx(lp, lpBytes(lp), 0, 2, 1); + if (!valid) { rdbReportCorruptRDB("Zset ziplist integrity check failed."); zfree(lp); zfree(encoded); @@ -3839,14 +3878,16 @@ robj *rdbLoadObject(int rdbtype, rio *rdb, sds key, int dbid, int *error) } if (zsetLength(o) > server.zset_max_listpack_entries) - zsetConvert(o, OBJ_ENCODING_SKIPLIST); + zsetConvert(o, OBJ_ENCODING_BTREE); else o->ptr = lpShrinkToFit(o->ptr); break; } case RDB_TYPE_ZSET_LISTPACK: if (deep_integrity_validation) server.stat_dump_payload_sanitizations++; - if (!lpValidateIntegrityAndDups(encoded, encoded_len, deep_integrity_validation, 2)) { + if (!lpValidateIntegrityAndDupsEx(encoded, encoded_len, + deep_integrity_validation, 2, + deep_integrity_validation)) { rdbReportCorruptRDB("Zset listpack integrity check failed."); zfree(encoded); o->ptr = NULL; @@ -3860,8 +3901,21 @@ robj *rdbLoadObject(int rdbtype, rio *rdb, sds key, int dbid, int *error) goto emptykey; } - if (zsetLength(o) > server.zset_max_listpack_entries) - zsetConvert(o, OBJ_ENCODING_SKIPLIST); + if (zsetLength(o) > server.zset_max_listpack_entries) { + /* Parsing every score costs a full walk, so a listpack + * that stays one is checked only by deep sanitization. + * A conversion parses every score anyway and must not + * store an invalid one, so check them first. */ + if (!deep_integrity_validation && + !lpValidateIntegrityAndDupsEx(encoded, encoded_len, 0, 2, 1)) { + rdbReportCorruptRDB("Zset listpack integrity check failed."); + zfree(encoded); + o->ptr = NULL; + decrRefCount(o); + return NULL; + } + zsetConvert(o, OBJ_ENCODING_BTREE); + } break; case RDB_TYPE_HASH_ZIPLIST: { diff --git a/src/server.c b/src/server.c index 7134a2f889..3dd2052cd4 100644 --- a/src/server.c +++ b/src/server.c @@ -8338,6 +8338,7 @@ int __test_num = 0; typedef int redisTestProc(int argc, char **argv, int flags); int bitopsTest(int argc, char **argv, int flags); int zsetTest(int argc, char **argv, int flags); +int zsetBtreeTest(int argc, char **argv, int flags); int vectorTest(int argc, char **argv, int flags); int lzfTest(int argc, char **argv, int flags); struct redisTest { @@ -8367,6 +8368,7 @@ struct redisTest { {"bitmap", bitopsTest}, {"rax", raxTest}, {"zset", zsetTest}, + {"zsetbtree", zsetBtreeTest}, {"topk", chkTopKTest}, {"fastfloat", fastFloatTest}, {"lzf", lzfTest}, diff --git a/src/server.h b/src/server.h index 3199032674..bd26ee4551 100644 --- a/src/server.h +++ b/src/server.h @@ -3895,8 +3895,10 @@ unsigned char *zzlFirstInRange(unsigned char *zl, zrangespec *range); unsigned char *zzlLastInRange(unsigned char *zl, zrangespec *range); unsigned long zsetLength(const robj *zobj); size_t zsetAllocSize(const robj *o); +robj *zsetTypeCreate(size_t size_hint, size_t val_len_hint); void zsetConvert(robj *zobj, int encoding); -void zsetConvertToListpackIfNeeded(robj *zobj, size_t maxelelen, size_t totelelen); +void zsetConvertAndExpand(robj *zobj, int encoding, unsigned long cap); +void zsetConvertAfterBulkInsert(robj *zobj, size_t maxelelen, size_t totelelen); int zsetScore(robj *zobj, sds member, double *score); unsigned long zslGetRank(zskiplist *zsl, double score, sds o); int zsetAdd(robj *zobj, double score, sds ele, int in_flags, int *out_flags, double *newscore); diff --git a/src/sort.c b/src/sort.c index 0d8dcdd9b4..abbef680d0 100644 --- a/src/sort.c +++ b/src/sort.c @@ -10,6 +10,7 @@ #include "fast_float_strtod.h" #include "server.h" +#include "zset_btree.h" #include "pqsort.h" /* Partial qsort for SORT+LIMIT */ #include /* isnan() */ #include "cluster.h" @@ -338,11 +339,11 @@ void sortCommandGeneric(client *c, int readonly) { sortby = NULL; } - /* Destructively convert encoded sorted sets for SORT. */ - if (sortval->type == OBJ_ZSET) { + /* Destructively convert listpack encoded sorted sets for SORT. */ + if (sortval->type == OBJ_ZSET && sortval->encoding == OBJ_ENCODING_LISTPACK) { if (server.memory_tracking_enabled) oldsize = kvobjAllocSize(sortval); - zsetConvert(sortval, OBJ_ENCODING_SKIPLIST); + zsetConvert(sortval, OBJ_ENCODING_BTREE); if (server.memory_tracking_enabled) updateSlotAllocSize(c->db, getKeySlot(c->argv[1]->ptr), sortval, oldsize, kvobjAllocSize(sortval)); } @@ -351,7 +352,7 @@ void sortCommandGeneric(client *c, int readonly) { switch(sortval->type) { case OBJ_LIST: vectorlen = listTypeLength(sortval); break; case OBJ_SET: vectorlen = setTypeSize(sortval); break; - case OBJ_ZSET: vectorlen = dictSize(((zset*)sortval->ptr)->dict); break; + case OBJ_ZSET: vectorlen = zsetLength(sortval); break; default: vectorlen = 0; serverPanic("Bad SORT type"); /* Avoid GCC warning */ } @@ -446,54 +447,79 @@ void sortCommandGeneric(client *c, int readonly) { * Note that in this case we also handle LIMIT here in a direct * way, just getting the required range, as an optimization. */ - zset *zs = sortval->ptr; - zskiplist *zsl = zs->zsl; - zskiplistNode *ln; - sds sdsele; int rangelen = vectorlen; - - /* Check if starting point is trivial, before doing log(N) lookup. */ - if (desc) { - long zsetlen = dictSize(((zset*)sortval->ptr)->dict); - - ln = zsl->tail; - if (start > 0) - ln = zslGetElementByRank(zsl,zsetlen-start); + if (sortval->encoding == OBJ_ENCODING_BTREE) { + serverAssert(zsetLength(sortval) > 0); + zbtreeIterator iter; + const unsigned char *ele; + size_t len; + unsigned long rank = + desc ? zsetLength(sortval) - (unsigned long)start - 1 : + (unsigned long)start; + serverAssert(zbtreeIteratorSeekRank(sortval->ptr, rank, &iter)); + while(rangelen--) { + serverAssert(zbtreeIteratorNext(&iter, desc, &ele, &len, NULL)); + vector[j].obj = createStringObject((char *)ele,len); + vector[j].u.score = 0; + vector[j].u.cmpobj = NULL; + j++; + } } else { - ln = zsl->header->level[0].forward; - if (start > 0) - ln = zslGetElementByRank(zsl,start+1); - } + zset *zs = sortval->ptr; + zskiplist *zsl = zs->zsl; + zskiplistNode *ln; - while(rangelen--) { - serverAssertWithInfo(c,sortval,ln != NULL); - sdsele = zslGetNodeElement(ln); - vector[j].obj = createStringObject(sdsele,sdslen(sdsele)); - vector[j].u.score = 0; - vector[j].u.cmpobj = NULL; - j++; - ln = desc ? ln->backward : ln->level[0].forward; + if (desc) { + ln = zsl->tail; + if (start > 0) + ln = zslGetElementByRank(zsl,zsetLength(sortval)-start); + } else { + ln = zsl->header->level[0].forward; + if (start > 0) + ln = zslGetElementByRank(zsl,start+1); + } + + while(rangelen--) { + serverAssertWithInfo(c,sortval,ln != NULL); + sds ele = zslGetNodeElement(ln); + vector[j].obj = createStringObject(ele,sdslen(ele)); + vector[j].u.score = 0; + vector[j].u.cmpobj = NULL; + j++; + ln = desc ? ln->backward : ln->level[0].forward; + } } /* Fix start/end: output code is not aware of this optimization. */ end -= start; start = 0; } else if (sortval->type == OBJ_ZSET) { - dict *set = ((zset*)sortval->ptr)->dict; - dictIterator di; - dictEntry *setele; - sds sdsele; - if (server.memory_tracking_enabled) oldsize = kvobjAllocSize(sortval); - dictInitIterator(&di, set); - while((setele = dictNext(&di)) != NULL) { - sdsele = zslGetNodeElement(dictGetKey(setele)); - vector[j].obj = createStringObject(sdsele,sdslen(sdsele)); - vector[j].u.score = 0; - vector[j].u.cmpobj = NULL; - j++; + if (sortval->encoding == OBJ_ENCODING_BTREE) { + zbtreeIterator iter; + const unsigned char *ele; + size_t len; + zbtreeIteratorStart(sortval->ptr, 0, &iter); + while(zbtreeIteratorNext(&iter, 0, &ele, &len, NULL)) { + vector[j].obj = createStringObject((char *)ele,len); + vector[j].u.score = 0; + vector[j].u.cmpobj = NULL; + j++; + } + } else { + dict *set = ((zset*)sortval->ptr)->dict; + dictIterator di; + dictEntry *setele; + dictInitIterator(&di, set); + while((setele = dictNext(&di)) != NULL) { + sds ele = zslGetNodeElement(dictGetKey(setele)); + vector[j].obj = createStringObject(ele,sdslen(ele)); + vector[j].u.score = 0; + vector[j].u.cmpobj = NULL; + j++; + } + dictResetIterator(&di); } - dictResetIterator(&di); if (server.memory_tracking_enabled) updateSlotAllocSize(c->db, getKeySlot(c->argv[1]->ptr), sortval, oldsize, kvobjAllocSize(sortval)); } else { diff --git a/src/t_zset.c b/src/t_zset.c index d12c6c8ff1..7a0b77287a 100644 --- a/src/t_zset.c +++ b/src/t_zset.c @@ -17,13 +17,13 @@ * Sorted set API *----------------------------------------------------------------------------*/ -/* ZSETs are ordered sets using two data structures to hold the same elements - * in order to get O(log(N)) INSERT and REMOVE operations into a sorted - * data structure. +/* Large ZSETs use the packed B+ tree and compact member hash index + * implemented in zset_btree.c. The skiplist below is only used as a + * temporary structure by ZUNION/ZINTER/ZDIFF and their STORE variants: + * a SKIPLIST-encoded zset is never stored in the keyspace. * - * The elements are added to a hash table mapping Redis objects to scores. - * At the same time the elements are added to a skip list mapping scores - * to Redis objects (so objects are sorted by scores in this "view"). + * The skiplist holds every element in both a hash table and a skiplist. The + * hash table finds a member directly, while the skiplist keeps score order. * * Note that the SDS string representing the element is the same in both * the hash table and skiplist in order to save memory. What we do in order @@ -42,6 +42,7 @@ * from tail to head, useful for ZREVRANGE. */ #include "fast_float_strtod.h" #include "server.h" +#include "zset_btree.h" #include "intset.h" /* Compact integer set structure */ #include @@ -69,7 +70,6 @@ dictType zsetDictType = { int zslLexValueGteMin(sds value, zlexrangespec *spec); int zslLexValueLteMax(sds value, zlexrangespec *spec); -void zsetConvertAndExpand(robj *zobj, int encoding, unsigned long cap); static zskiplistNode *zslGetElementByRankFromNode(zskiplistNode *start_node, int start_level, unsigned long rank); static inline unsigned long zslGetNodeSpanAtLevel(zskiplistNode *x, int level) { @@ -842,6 +842,31 @@ int zslLexValueLteMax(sds value, zlexrangespec *spec) { (sdscmplex(value,spec->max) <= 0); } +/* Compare a packed B+ tree member with a lexical range bound. */ +static int zbtLexCompare(const unsigned char *value, size_t len, sds bound) { + if (bound == shared.minstring) return 1; + if (bound == shared.maxstring) return -1; + size_t boundlen = sdslen(bound); + size_t minlen = len < boundlen ? len : boundlen; + int cmp = memcmp(value, bound, minlen); + if (cmp != 0) return cmp; + if (len < boundlen) return -1; + if (len > boundlen) return 1; + return 0; +} + +/* Test a packed B+ tree member against the lower lexical bound. */ +static int zbtLexValueGteMin(const unsigned char *value, size_t len, zlexrangespec *spec) { + int cmp = zbtLexCompare(value, len, spec->min); + return spec->minex ? cmp > 0 : cmp >= 0; +} + +/* Test a packed B+ tree member against the upper lexical bound. */ +static int zbtLexValueLteMax(const unsigned char *value, size_t len, zlexrangespec *spec) { + int cmp = zbtLexCompare(value, len, spec->max); + return spec->maxex ? cmp < 0 : cmp <= 0; +} + /* Returns if there is a part of the zset is in the lex range. */ static int zslIsInLexRange(zskiplist *zsl, zlexrangespec *range) { zskiplistNode *x; @@ -1380,6 +1405,8 @@ unsigned long zsetLength(const robj *zobj) { length = zzlLength(zobj->ptr); } else if (zobj->encoding == OBJ_ENCODING_SKIPLIST) { length = ((const zset*)zobj->ptr)->zsl->length; + } else if (zobj->encoding == OBJ_ENCODING_BTREE) { + length = zbtreeLength(zobj->ptr); } else { serverPanic("Unknown sorted set encoding"); } @@ -1396,6 +1423,8 @@ size_t zsetAllocSize(const robj *o) { zskiplist *zsl = ((zset*)o->ptr)->zsl; size = sizeof(zset) + zslAllocSize(zsl) + sizeof(dict) + dictMemUsage(d); + } else if (o->encoding == OBJ_ENCODING_BTREE) { + size = zbtreeAllocSize(o->ptr); } else { serverPanic("Unknown sorted set encoding"); } @@ -1417,10 +1446,7 @@ robj *zsetTypeCreate(size_t size_hint, size_t val_len_hint) { return createZsetListpackObject(); } - robj *zobj = createZsetObject(); - zset *zs = zobj->ptr; - dictExpand(zs->dict, size_hint); - return zobj; + return createZsetBtreeObject(); } /* Check if the existing zset should be converted to another encoding based off the @@ -1429,7 +1455,7 @@ void zsetTypeMaybeConvert(robj *zobj, size_t size_hint) { if (zobj->encoding == OBJ_ENCODING_LISTPACK && size_hint > server.zset_max_listpack_entries) { - zsetConvertAndExpand(zobj, OBJ_ENCODING_SKIPLIST, size_hint); + zsetConvertAndExpand(zobj, OBJ_ENCODING_BTREE, size_hint); } } @@ -1440,6 +1466,69 @@ void zsetConvert(robj *zobj, int encoding) { zsetConvertAndExpand(zobj, encoding, zsetLength(zobj)); } +/* Decode one listpack zset member into a flat (raw bytes, length) view. + * Integer members are formatted into the caller-supplied buffer (must be at + * least LONG_STR_SIZE bytes); string members are returned as a direct + * pointer into the listpack itself, valid as long as the listpack isn't + * modified. */ +static const unsigned char *zzlDecodeElement(unsigned char *eptr, char *buf, + size_t bufsize, size_t *rawlen) +{ + unsigned int vlen; + long long vlong; + unsigned char *vstr = lpGetValue(eptr, &vlen, &vlong); + if (vstr == NULL) { + *rawlen = ll2string(buf, bufsize, vlong); + return (unsigned char *)buf; + } + *rawlen = vlen; + return vstr; +} + +/* Verify a listpack-encoded zset is genuinely sorted (score, then member) + * and duplicate-free -- its invariant when uncorrupted. A single linear + * pass here lets the LISTPACK->BTREE conversion below use the descent-free + * append path for every member instead of paying a full B+tree search (the + * duplicate check) plus a second full descent (the insert) per element. + * Only a corrupt listpack should make this return false, in which case the + * caller falls back to the original per-element-checked, slower path. */ +static int zzlIsSortedNoDup(unsigned char *zl) { + unsigned char *preveptr = NULL; + unsigned char *eptr, *sptr; + double prevscore = 0; + + eptr = lpSeek(zl, 0); + if (eptr == NULL) return 1; + sptr = lpNext(zl, eptr); + if (sptr == NULL) return 0; + + while (eptr != NULL) { + double score = zzlGetScore(sptr); + if (isnan(score)) return 0; + if (preveptr != NULL) { + int cmp; + if (prevscore != score) { + cmp = prevscore < score ? -1 : 1; + } else { + char prevbuf[LONG_STR_SIZE], buf[LONG_STR_SIZE]; + size_t prevrawlen, rawlen; + const unsigned char *prevraw = zzlDecodeElement( + preveptr, prevbuf, sizeof(prevbuf), &prevrawlen); + const unsigned char *raw = zzlDecodeElement(eptr, buf, sizeof(buf), &rawlen); + size_t minlen = prevrawlen < rawlen ? prevrawlen : rawlen; + cmp = memcmp(prevraw, raw, minlen); + if (cmp == 0) + cmp = prevrawlen < rawlen ? -1 : (prevrawlen > rawlen ? 1 : 0); + } + if (cmp >= 0) return 0; + } + preveptr = eptr; + prevscore = score; + zzlNext(zl, &eptr, &sptr); + } + return 1; +} + /* Converts a zset to the specified encoding, pre-sizing it for 'cap' elements. */ void zsetConvertAndExpand(robj *zobj, int encoding, unsigned long cap) { zset *zs; @@ -1450,11 +1539,91 @@ void zsetConvertAndExpand(robj *zobj, int encoding, unsigned long cap) { if (zobj->encoding == encoding) return; if (zobj->encoding == OBJ_ENCODING_LISTPACK) { unsigned char *zl = zobj->ptr; - unsigned char *eptr, *sptr; + unsigned char *eptr, *sptr = NULL; unsigned char *vstr; unsigned int vlen; long long vlong; + if (encoding == OBJ_ENCODING_BTREE) { + zbtreeSet *bt = zbtreeCreate(); + /* ZADD's hint counts its arguments, which can be fewer than the + * existing members after a listpack limit is lowered. */ + unsigned long length = zzlLength(zl); + if (cap < length) cap = length; + /* Presize the member index to avoid rehashing, mirroring the + * dictExpand(zs->dict, cap) a few lines below for the + * LISTPACK->SKIPLIST branch. cap here is always a real element + * count, never an unvalidated RDB/client length: zsetConvert()'s + * callers (SORT/SORT_RO, RDB listpack-promotion loads) pass + * zsetLength(zobj), the listpack's own already-decoded length; + * zsetTypeMaybeConvert()'s ZADD caller passes elements, this + * command's own argc-derived count of score-member pairs; + * and ZADD's listpack-overflow path below passes + * zsetLength(zobj) + 1. + * zsetConvertAfterBulkInsert()'s zsetConvert() calls never reach + * here: it returns immediately on OBJ_ENCODING_LISTPACK, so its + * BTREE conversion always lands in the SKIPLIST->BTREE branch + * below instead. */ + zbtreeReserve(bt, cap); + + eptr = lpSeek(zl,0); + if (eptr != NULL) { + sptr = lpNext(zl,eptr); + serverAssertWithInfo(NULL,zobj,sptr != NULL); + } + + if (zzlIsSortedNoDup(zl)) { + /* Uncorrupted listpack (the common case, verified above): + * known sorted and duplicate-free, so every insert can use + * the descent-free append path -- no per-element duplicate + * search, and no separate full-tree descent to insert. + * Batched: buf is a per-iteration stack scratch reused every + * loop pass, so raw's bytes are only valid until the next + * zzlDecodeElement() call -- zbtreeAppendBatchAdd() copies + * them out synchronously before returning, same contract the + * unbatched zbtreeInsertNewAppend() call it replaces had. */ + zbtreeAppendBatch *batch = zbtreeAppendBatchCreate(bt); + while (eptr != NULL) { + score = zzlGetScore(sptr); + serverAssert(!isnan(score)); + char buf[LONG_STR_SIZE]; + size_t rawlen; + const unsigned char *raw = zzlDecodeElement(eptr, buf, sizeof(buf), &rawlen); + zbtreeAppendBatchAdd(batch, score, raw, rawlen); + zzlNext(zl,&eptr,&sptr); + } + zbtreeAppendBatchFinish(batch); + } else { + while (eptr != NULL) { + score = zzlGetScore(sptr); + serverAssert(!isnan(score)); + vstr = lpGetValue(eptr,&vlen,&vlong); + char buf[LONG_STR_SIZE]; + const unsigned char *raw; + size_t rawlen; + if (vstr == NULL) { + rawlen = ll2string(buf,sizeof(buf),vlong); + raw = (unsigned char *)buf; + } else { + raw = vstr; + rawlen = vlen; + } + /* A corrupt listpack may contain a duplicate. dictAdd() + * used to catch this when the target was always a + * skiplist. */ + int duplicate = zbtreeScoreRaw(bt,raw,rawlen,NULL); + serverAssert(!duplicate); + zbtreeInsertNewRaw(bt,score,raw,rawlen,NULL); + zzlNext(zl,&eptr,&sptr); + } + } + + zfree(zobj->ptr); + zobj->ptr = bt; + zobj->encoding = OBJ_ENCODING_BTREE; + return; + } + if (encoding != OBJ_ENCODING_SKIPLIST) serverPanic("Unknown target encoding"); @@ -1489,13 +1658,40 @@ void zsetConvertAndExpand(robj *zobj, int encoding, unsigned long cap) { zobj->ptr = zs; zobj->encoding = OBJ_ENCODING_SKIPLIST; } else if (zobj->encoding == OBJ_ENCODING_SKIPLIST) { - unsigned char *zl = lpNew(0); - + if (encoding == OBJ_ENCODING_BTREE) { + zs = zobj->ptr; + zbtreeSet *bt = zbtreeCreate(); + /* Same reserve mechanism as the LISTPACK->BTREE branch above. + * This branch is reached only via zsetConvert(), so cap is + * always zsetLength(zobj) here -- the skiplist's own real + * length, not a size_hint/overflow-derived count. */ + zbtreeReserve(bt, cap); + node = zs->zsl->header->level[0].forward; + zbtreeAppendBatch *batch = zbtreeAppendBatchCreate(bt); + while (node) { + /* The skiplist is Redis's own, always sorted, never + * corrupted -- unlike the listpack branch above, safe to + * use the descent-free append path (zbtScoreInsertAppend's + * precondition holds by construction: zslInsert never + * produces (score, member) out of order or duplicated). */ + sds ele = zslGetNodeElement(node); + zbtreeAppendBatchAdd(batch, node->score, (unsigned char *)ele, sdslen(ele)); + node = node->level[0].forward; + } + zbtreeAppendBatchFinish(batch); + dictRelease(zs->dict); + zslFree(zs->zsl); + zfree(zs); + zobj->ptr = bt; + zobj->encoding = OBJ_ENCODING_BTREE; + return; + } if (encoding != OBJ_ENCODING_LISTPACK) serverPanic("Unknown target encoding"); /* Approach similar to zslFree(), since we want to free the skiplist at * the same time as creating the listpack. */ + unsigned char *zl = lpNew(0); zs = zobj->ptr; dictRelease(zs->dict); node = zs->zsl->header->level[0].forward; @@ -1512,23 +1708,62 @@ void zsetConvertAndExpand(robj *zobj, int encoding, unsigned long cap) { zfree(zs); zobj->ptr = zl; zobj->encoding = OBJ_ENCODING_LISTPACK; + } else if (zobj->encoding == OBJ_ENCODING_BTREE) { + if (encoding == OBJ_ENCODING_SKIPLIST) { + zs = zmalloc(sizeof(*zs)); + zs->dict = dictCreate(&zsetDictType); + zs->zsl = zslCreate(); + dictExpand(zs->dict, cap); + + zbtreeIterator iter; + const unsigned char *raw; + size_t len; + zbtreeIteratorStart(zobj->ptr, 0, &iter); + while (zbtreeIteratorNext(&iter, 0, &raw, &len, &score)) { + ele = sdsnewlen(raw, len); + node = zslInsert(zs->zsl, score, ele); + serverAssert(dictAdd(zs->dict, node, NULL) == DICT_OK); + sdsfree(ele); + } + zbtreeFree(zobj->ptr); + zobj->ptr = zs; + zobj->encoding = OBJ_ENCODING_SKIPLIST; + return; + } + + if (encoding != OBJ_ENCODING_LISTPACK) + serverPanic("Unknown target encoding"); + + unsigned char *zl = lpNew(0); + zbtreeIterator iter; + const unsigned char *raw; + size_t len; + zbtreeIteratorStart(zobj->ptr, 0, &iter); + while (zbtreeIteratorNext(&iter, 0, &raw, &len, &score)) { + ele = sdsnewlen(raw, len); + zl = zzlInsertAt(zl, NULL, ele, score); + sdsfree(ele); + } + zbtreeFree(zobj->ptr); + zobj->ptr = zl; + zobj->encoding = OBJ_ENCODING_LISTPACK; } else { serverPanic("Unknown sorted set encoding"); } } -/* Convert the sorted set object into a listpack if it is not already a listpack - * and if the number of elements and the maximum element size and total elements size - * are within the expected ranges. */ -void zsetConvertToListpackIfNeeded(robj *zobj, size_t maxelelen, size_t totelelen) { +/* Choose the best encoding after a sorted set was built in bulk. Small sets + * become listpacks and larger skiplist sets become B+ trees. */ +void zsetConvertAfterBulkInsert(robj *zobj, size_t maxelelen, size_t totelelen) { if (zobj->encoding == OBJ_ENCODING_LISTPACK) return; - zset *zset = zobj->ptr; - - if (zset->zsl->length <= server.zset_max_listpack_entries && + unsigned long length = zsetLength(zobj); + if (length <= server.zset_max_listpack_entries && maxelelen <= server.zset_max_listpack_value && lpSafeToAdd(NULL, totelelen)) { zsetConvert(zobj,OBJ_ENCODING_LISTPACK); + } else if (zobj->encoding == OBJ_ENCODING_SKIPLIST) { + zsetConvert(zobj,OBJ_ENCODING_BTREE); } } @@ -1547,6 +1782,8 @@ int zsetScore(robj *zobj, sds member, double *score) { if (de == NULL) return C_ERR; zskiplistNode *znode = dictGetKey(de); *score = znode->score; + } else if (zobj->encoding == OBJ_ENCODING_BTREE) { + if (!zbtreeScore(zobj->ptr, member, score)) return C_ERR; } else { serverPanic("Unknown sorted set encoding"); } @@ -1656,7 +1893,7 @@ int zsetAdd(robj *zobj, double score, sds ele, int in_flags, int *out_flags, dou sdslen(ele) > server.zset_max_listpack_value || !lpSafeToAdd(zobj->ptr, sdslen(ele))) { - zsetConvertAndExpand(zobj, OBJ_ENCODING_SKIPLIST, zsetLength(zobj) + 1); + zsetConvertAndExpand(zobj, OBJ_ENCODING_BTREE, zsetLength(zobj) + 1); } else { zobj->ptr = zzlInsert(zobj->ptr,ele,score); if (newscore) *newscore = score; @@ -1669,8 +1906,7 @@ int zsetAdd(robj *zobj, double score, sds ele, int in_flags, int *out_flags, dou } } - /* Note that the above block handling listpack would have either returned or - * converted the key to skiplist. */ + /* The listpack block either returned or converted the key. */ if (zobj->encoding == OBJ_ENCODING_SKIPLIST) { zset *zs = zobj->ptr; zskiplistNode *znode; @@ -1728,6 +1964,42 @@ int zsetAdd(robj *zobj, double score, sds ele, int in_flags, int *out_flags, dou /* Add node pointer to dict using the bucket we already found */ dictSetKeyAtLink(zs->dict, znode, &bucket, 1); + *out_flags |= ZADD_OUT_ADDED; + if (newscore) *newscore = score; + return 1; + } else { + *out_flags |= ZADD_OUT_NOP; + return 1; + } + } else if (zobj->encoding == OBJ_ENCODING_BTREE) { + zbtreeSet *zs = zobj->ptr; + zbtreeInsertPosition position; + int exists = zbtreeFindForAdd(zs, ele, &curscore, &position); + + if (exists) { + if (nx) { + *out_flags |= ZADD_OUT_NOP; + return 1; + } + if (incr) { + score += curscore; + if (isnan(score)) { + *out_flags |= ZADD_OUT_NAN; + return 0; + } + } + if ((lt && score >= curscore) || (gt && score <= curscore)) { + *out_flags |= ZADD_OUT_NOP; + return 1; + } + if (newscore) *newscore = score; + if (score != curscore) { + zbtreeUpdateScore(zs, ele, score, &position); + *out_flags |= ZADD_OUT_UPDATED; + } + return 1; + } else if (!xx) { + zbtreeInsertNew(zs, score, ele, &position); *out_flags |= ZADD_OUT_ADDED; if (newscore) *newscore = score; return 1; @@ -1784,6 +2056,8 @@ int zsetDel(robj *zobj, sds ele) { if (zsetRemoveFromSkiplist(zs, ele)) { return 1; } + } else if (zobj->encoding == OBJ_ENCODING_BTREE) { + return zbtreeDelete(zobj->ptr, ele); } else { serverPanic("Unknown sorted set encoding"); } @@ -1856,6 +2130,8 @@ long zsetRank(robj *zobj, sds ele, int reverse, double *output_score) { } else { return -1; } + } else if (zobj->encoding == OBJ_ENCODING_BTREE) { + return zbtreeRank(zobj->ptr, ele, reverse, output_score); } else { serverPanic("Unknown sorted set encoding"); } @@ -1904,6 +2180,9 @@ robj *zsetDup(robj *o) { dictAdd(new_zs->dict, znode, NULL); ln = ln->backward; } + } else if (o->encoding == OBJ_ENCODING_BTREE) { + zobj = createObject(OBJ_ZSET, zbtreeDup(o->ptr)); + zobj->encoding = OBJ_ENCODING_BTREE; } else { serverPanic("Unknown sorted set encoding"); } @@ -1949,6 +2228,15 @@ void zsetTypeRandomElement(robj *zsetobj, unsigned long zsetsize, listpackEntry *score = (double)val.lval; } } + } else if (zsetobj->encoding == OBJ_ENCODING_BTREE) { + zbtreeIterator iter; + const unsigned char *ele; + size_t len; + unsigned long rank = randomULong() % zsetsize; + serverAssert(zbtreeIteratorSeekRank(zsetobj->ptr, rank, &iter)); + serverAssert(zbtreeIteratorNext(&iter, 0, &ele, &len, score)); + key->sval = (unsigned char *)ele; + key->slen = len; } else { serverPanic("Unknown zset encoding"); } @@ -2256,6 +2544,34 @@ void zremrangeGenericCommand(client *c, zrange_type rangetype) { } else { dictShrinkIfNeeded(zs->dict); } + } else if (zobj->encoding == OBJ_ENCODING_BTREE) { + zbtreeSet *zs = zobj->ptr; + switch(rangetype) { + case ZRANGE_AUTO: + case ZRANGE_RANK: + deleted = zbtreeDeleteRangeByRank(zs, start, end); + break; + case ZRANGE_SCORE: + deleted = zbtreeDeleteRangeByScore(zs, range.min, range.minex, range.max, range.maxex); + break; + case ZRANGE_LEX: { + zbtreeIterator first_iter, last_iter; + unsigned long first, last; + if (zbtreeIteratorSeekLex(zs, lexrange.min, lexrange.minex, 0, &first_iter, &first) && + zbtreeIteratorSeekLex(zs, lexrange.max, lexrange.maxex, 1, &last_iter, &last) && + last >= first) + { + deleted = zbtreeDeleteRangeByRank(zs, first, last); + } + break; + } + } + if (zbtreeLength(zs) == 0) { + if (server.memory_tracking_enabled) + updateSlotAllocSize(c->db, getKeySlot(key->ptr), zobj, oldsize, kvobjAllocSize(zobj)); + dbDeleteSkipKeysizesUpdate(c->db, key); + keyremoved = 1; + } } else { serverPanic("Unknown sorted set encoding"); } @@ -2333,6 +2649,9 @@ typedef struct { zset *zs; zskiplistNode *node; } sl; + struct { + zbtreeIterator iter; + } bt; } zset; } iter; } zsetopsrc; @@ -2354,7 +2673,7 @@ typedef struct { unsigned char _buf[32]; /* Private buffer. */ sds ele; unsigned char *estr; - unsigned int elen; + size_t elen; long long ell; double score; } zsetopval; @@ -2396,6 +2715,8 @@ void zuiInitIterator(zsetopsrc *op) { } else if (op->encoding == OBJ_ENCODING_SKIPLIST) { it->sl.zs = op->subject->ptr; it->sl.node = it->sl.zs->zsl->tail; + } else if (op->encoding == OBJ_ENCODING_BTREE) { + zbtreeIteratorStart(op->subject->ptr, 1, &it->bt.iter); } else { serverPanic("Unknown sorted set encoding"); } @@ -2425,6 +2746,8 @@ void zuiClearIterator(zsetopsrc *op) { UNUSED(it); /* skip */ } else if (op->encoding == OBJ_ENCODING_SKIPLIST) { UNUSED(it); /* skip */ + } else if (op->encoding == OBJ_ENCODING_BTREE) { + UNUSED(it); /* skip */ } else { serverPanic("Unknown sorted set encoding"); } @@ -2453,6 +2776,8 @@ unsigned long zuiLength(zsetopsrc *op) { } else if (op->encoding == OBJ_ENCODING_SKIPLIST) { zset *zs = op->subject->ptr; return zs->zsl->length; + } else if (op->encoding == OBJ_ENCODING_BTREE) { + return zbtreeLength(op->subject->ptr); } else { serverPanic("Unknown sorted set encoding"); } @@ -2495,7 +2820,9 @@ int zuiNext(zsetopsrc *op, zsetopval *val) { } else if (op->encoding == OBJ_ENCODING_LISTPACK) { if (it->lp.p == NULL) return 0; - val->estr = lpGetValue(it->lp.p, &val->elen, &val->ell); + unsigned int len = 0; + val->estr = lpGetValue(it->lp.p, &len, &val->ell); + val->elen = len; val->score = 1.0; /* Move to next element. */ @@ -2509,7 +2836,9 @@ int zuiNext(zsetopsrc *op, zsetopval *val) { /* No need to check both, but better be explicit. */ if (it->zl.eptr == NULL || it->zl.sptr == NULL) return 0; - val->estr = lpGetValue(it->zl.eptr,&val->elen,&val->ell); + unsigned int len = 0; + val->estr = lpGetValue(it->zl.eptr,&len,&val->ell); + val->elen = len; val->score = zzlGetScore(it->zl.sptr); /* Move to next element (going backwards, see zuiInitIterator). */ @@ -2522,6 +2851,13 @@ int zuiNext(zsetopsrc *op, zsetopval *val) { /* Move to next element. (going backwards, see zuiInitIterator) */ it->sl.node = it->sl.node->backward; + } else if (op->encoding == OBJ_ENCODING_BTREE) { + const unsigned char *ele; + size_t len; + if (!zbtreeIteratorNext(&it->bt.iter, 1, &ele, &len, &val->score)) + return 0; + val->estr = (unsigned char *)ele; + val->elen = len; } else { serverPanic("Unknown sorted set encoding"); } @@ -2608,6 +2944,11 @@ int zuiFind(zsetopsrc *op, zsetopval *val, double *score) { return 0; } } else if (op->type == OBJ_ZSET) { + if (op->encoding == OBJ_ENCODING_BTREE) { + zuiBufferFromValue(val); + return zbtreeScoreRaw(op->subject->ptr, val->estr, val->elen, score); + } + zuiSdsFromValue(val); if (op->encoding == OBJ_ENCODING_LISTPACK) { @@ -2643,6 +2984,27 @@ int zuiCompareByCardinality(const void *s1, const void *s2) { return 0; } +/* An unlinked zskiplistNode plus the member hash ZUNIONSTORE's UNION path + * already computed for its dict lookup -- carried through to a direct-to- + * B+tree bulk build so zbtreeInsertNewAppendWithHash() doesn't need to + * recompute it. */ +typedef struct { + zskiplistNode *node; + uint32_t hash; +} zsetUnionMember; + +/* Orders zsetUnionMembers the same way zslInsert() and the B+tree both order + * entries: by score, then by member (sdscmp, matching zbtCompareElements() + * in zset_btree.c byte-for-byte). Used to bulk-sort ZUNIONSTORE's result + * dict before a direct-to-B+tree build. */ +static int zsetUnionMemberScoreCompare(const void *a, const void *b) { + const zskiplistNode *na = ((const zsetUnionMember *)a)->node; + const zskiplistNode *nb = ((const zsetUnionMember *)b)->node; + if (na->score < nb->score) return -1; + if (na->score > nb->score) return 1; + return sdscmp(zslGetNodeElement(na), zslGetNodeElement(nb)); +} + static int zuiCompareByRevCardinality(const void *s1, const void *s2) { return zuiCompareByCardinality(s1, s2) * -1; } @@ -2897,6 +3259,9 @@ void zunionInterDiffGenericCommand(client *c, robj *dstkey, int numkeysIndex, in int withscores = 0; unsigned long cardinality = 0; long limit = 0; /* Stop searching after reaching the limit. 0 means unlimited. */ + int union_built_as_btree = 0; /* SET_OP_UNION only: result already converted + * to OBJ_ENCODING_BTREE in-line, skip the + * generic post-build conversion below. */ /* expect setnum input keys to be given */ if ((getLongFromObjectOrReply(c, c->argv[numkeysIndex], &setnum, NULL) != C_OK)) @@ -3080,6 +3445,12 @@ void zunionInterDiffGenericCommand(client *c, robj *dstkey, int numkeysIndex, in dictExpand(dstzset->dict,zuiLength(&src[setnum-1])); } + /* Only STORE needs the node/hash array for the B+tree bulk build. + * Grow it with distinct members: repeated or overlapping inputs must + * not allocate space proportional to the sum of their lengths. */ + zsetUnionMember *members = NULL; + size_t nmembers = 0, memberscap = 0; + /* Step 1: Iterate all sorted sets and aggregate scores. * For each element, either insert into skiplist (new) or update score (existing). */ for (i = 0; i < setnum; i++) { @@ -3092,9 +3463,11 @@ void zunionInterDiffGenericCommand(client *c, robj *dstkey, int numkeysIndex, in if (isnan(score)) score = 0; /* Search for this element in the dict (which stores node pointers). */ + sds elesds = zuiSdsFromValue(&zval); + uint64_t hash = dictGetHash(dstzset->dict, elesds); dictEntryLink bucket, link; - link = dictFindLink(dstzset->dict, zuiSdsFromValue(&zval), &bucket); - + link = dictFindLinkWithHash(dstzset->dict, elesds, hash, &bucket); + if (link == NULL) { /* if not exists */ /* New element: create node and insert into dict */ tmp = zuiNewSdsFromValue(&zval); @@ -3108,6 +3481,16 @@ void zunionInterDiffGenericCommand(client *c, robj *dstkey, int numkeysIndex, in znode = zslCreateNode(dstzset->zsl, zslRandomLevel(), score, tmp); /* Add node pointer to dict using the bucket we already found */ dictSetKeyAtLink(dstzset->dict, znode, &bucket, 1); + if (dstkey) { + if (nmembers == memberscap) { + serverAssert(memberscap <= SIZE_MAX / sizeof(*members) / 2); + memberscap = memberscap ? memberscap * 2 : 16; + members = zrealloc(members, memberscap * sizeof(*members)); + } + members[nmembers].node = znode; + members[nmembers].hash = (uint32_t)hash; + nmembers++; + } sdsfree(tmp); /* zslCreateNode copied it, we can free our copy */ } else { /* Existing element: aggregate score */ @@ -3121,14 +3504,65 @@ void zunionInterDiffGenericCommand(client *c, robj *dstkey, int numkeysIndex, in zuiClearIterator(&src[i]); } - /* Step 2: Done filling dict with nodes and updating scores. Now insert skiplist */ - dictInitIterator(&di, dstzset->dict); + /* Step 2: dict is fully filled and scores are final, so length, + * maxelelen and totelelen already reflect the final result -- decide + * the encoding now, before paying for a skiplist link pass. + * + * If the result won't fit a listpack, it is going to end up + * BTREE-encoded either way (zsetConvertAfterBulkInsert() always + * converts a non-listpack-eligible SKIPLIST result to BTREE). The + * normal route gets there via zslInsertNode() here (a real + * search-and-link skiplist build, one O(log n) descent per node in + * dict hash order) followed by zsetConvertAndExpand()'s SKIPLIST-> + * BTREE branch re-walking that skiplist to bulk-append it into a + * B+tree, then tearing the skiplist down again. Skip both passes: + * sort the members captured above once (matches zslInsert()'s own + * order -- score, then sdscmp(member)), and bulk-append directly + * into a fresh B+tree, reusing each member's already-known hash. + * + * Only safe when dstkey is set (ZUNIONSTORE): the STORE-less ZUNION + * variant reuses this exact accumulation code but replies straight + * off dstzset->zsl's linked list a few lines below, so it must keep + * getting a real, linked skiplist. */ + unsigned long unioncount = dictSize(dstzset->dict); + if (dstkey && !(unioncount <= server.zset_max_listpack_entries && + maxelelen <= server.zset_max_listpack_value && + lpSafeToAdd(NULL, totelelen))) + { + qsort(members, unioncount, sizeof(zsetUnionMember), zsetUnionMemberScoreCompare); - while((de = dictNext(&di)) != NULL) { - zskiplistNode *znode = dictGetKey(de); - zslInsertNode(dstzset->zsl, znode); + zbtreeSet *bt = zbtreeCreate(); + zbtreeReserve(bt, unioncount); + zbtreeAppendBatch *batch = zbtreeAppendBatchCreate(bt); + for (unsigned long n = 0; n < unioncount; n++) { + zskiplistNode *node = members[n].node; + sds ele = zslGetNodeElement(node); + zbtreeAppendBatchAddWithHash(batch, node->score, + (unsigned char *)ele, sdslen(ele), + members[n].hash); + zslFreeNode(dstzset->zsl, node); + } + zbtreeAppendBatchFinish(batch); + dictRelease(dstzset->dict); + zfree(dstzset->zsl->header); + zfree(dstzset->zsl); + zfree(dstzset); + dstzset = NULL; + dstobj->ptr = bt; + dstobj->encoding = OBJ_ENCODING_BTREE; + union_built_as_btree = 1; + } else { + /* Small result: still listpack-eligible, keep the normal + * skiplist-link path -- zsetConvertAfterBulkInsert() below will + * convert it to a listpack, never touching the B+tree side. */ + dictInitIterator(&di, dstzset->dict); + while((de = dictNext(&di)) != NULL) { + zskiplistNode *znode = dictGetKey(de); + zslInsertNode(dstzset->zsl, znode); + } + dictResetIterator(&di); } - dictResetIterator(&di); + zfree(members); } else if (op == SET_OP_DIFF) { zdiff(src, setnum, dstzset, &maxelelen, &totelelen); } else { @@ -3144,8 +3578,12 @@ void zunionInterDiffGenericCommand(client *c, robj *dstkey, int numkeysIndex, in } if (dstkey) { - if (dstzset->zsl->length) { - zsetConvertToListpackIfNeeded(dstobj, maxelelen, totelelen); + if (zsetLength(dstobj)) { + /* Already BTREE-encoded via the ZUNIONSTORE fast path above -- + * dstzset (and its zsl) were freed there, nothing left to + * convert. */ + if (!union_built_as_btree) + zsetConvertAfterBulkInsert(dstobj, maxelelen, totelelen); setKey(c, c->db, dstkey, &dstobj, 0); addReplyLongLong(c, zsetLength(dstobj)); notifyKeyspaceEvent(NOTIFY_ZSET, @@ -3257,7 +3695,9 @@ struct zrange_result_handler { robj *dstobj; void *userdata; int withscores; + int reverse; int should_emit_array_length; + zbtreeAppendBatch *append_batch; zrangeResultBeginFunction beginResultEmission; zrangeResultFinalizeFunction finalizeResultEmission; zrangeResultEmitCBufferFunction emitResultFromCBuffer; @@ -3331,11 +3771,37 @@ static void zrangeResultFinalizeClient(zrange_result_handler *handler, static void zrangeResultBeginStore(zrange_result_handler *handler, long length) { handler->dstobj = zsetTypeCreate(length >= 0 ? length : 0, 0); + /* The plain-rank ZRANGESTORE path (genericZrangebyrankCommand) computes + * an exact result count here; BYSCORE/BYLEX pass length<0 and this is a + * no-op, same as before this reserve existed. This length is + * server-computed, satisfying the trust requirement documented on + * zbtreeReserve() itself. */ + if (length > 0 && handler->dstobj->encoding == OBJ_ENCODING_BTREE) + zbtreeReserve(handler->dstobj->ptr, (unsigned long)length); } static void zrangeResultEmitCBufferForStore(zrange_result_handler *handler, const void *value, size_t value_length_in_bytes, double score) { + /* ZRANGESTORE (forward direction only, see zrangeResultHandlerDirectionSet()) + * emits a contiguous slice of the SOURCE zset's own sorted order into a + * destination that starts empty -- every element is provably greater + * than the last one stored and provably distinct (a zset can't hold a + * duplicate member). Once the destination is BTREE-encoded, that's + * exactly zbtreeAppendBatch's precondition (same as the + * zbtreeInsertNewAppend() it batches), so skip zsetAdd()'s generic + * zbtreeFindForAdd() existence/position search (which can only ever + * report "not found") entirely. Small zset (still LISTPACK) and + * reverse-order emission both keep the original, always-correct path. + * The batch is created lazily on the first fast-path emit (the + * destination may start as LISTPACK and convert mid-stream) and flushed + * once in zrangeResultFinalizeStore(). */ + if (!handler->reverse && handler->dstobj->encoding == OBJ_ENCODING_BTREE) { + if (!handler->append_batch) + handler->append_batch = zbtreeAppendBatchCreate(handler->dstobj->ptr); + zbtreeAppendBatchAdd(handler->append_batch, score, value, value_length_in_bytes); + return; + } double newscore; int retflags = 0; sds ele = sdsnewlen(value, value_length_in_bytes); @@ -3347,6 +3813,18 @@ static void zrangeResultEmitCBufferForStore(zrange_result_handler *handler, static void zrangeResultEmitLongLongForStore(zrange_result_handler *handler, long long value, double score) { + /* Same fast path as zrangeResultEmitCBufferForStore() above, formatting + * the integer member directly into a stack buffer instead of an sds -- + * zbtreeAppendBatchAdd() copies the bytes out synchronously, so the + * stack buffer doesn't need to outlive this call. */ + if (!handler->reverse && handler->dstobj->encoding == OBJ_ENCODING_BTREE) { + char buf[LONG_STR_SIZE]; + size_t len = ll2string(buf, sizeof(buf), value); + if (!handler->append_batch) + handler->append_batch = zbtreeAppendBatchCreate(handler->dstobj->ptr); + zbtreeAppendBatchAdd(handler->append_batch, score, (unsigned char *)buf, len); + return; + } double newscore; int retflags = 0; sds ele = sdsfromlonglong(value); @@ -3357,6 +3835,10 @@ static void zrangeResultEmitLongLongForStore(zrange_result_handler *handler, static void zrangeResultFinalizeStore(zrange_result_handler *handler, size_t result_count) { + if (handler->append_batch) { + zbtreeAppendBatchFinish(handler->append_batch); + handler->append_batch = NULL; + } if (result_count) { setKey(handler->client, handler->client->db, handler->dstkey, &handler->dstobj, 0); addReplyLongLong(handler->client, result_count); @@ -3409,6 +3891,14 @@ static void zrangeResultHandlerDestinationKeySet (zrange_result_handler *handler handler->dstkey = dstkey; } +/* Records the emission direction so the STORE handler methods know whether + * results arrive in ascending (score, member) order -- the only direction + * the descent-free B+tree append fast path is valid for. Harmless (and + * unused) for the CLIENT-reply handler. */ +static void zrangeResultHandlerDirectionSet(zrange_result_handler *handler, int reverse) { + handler->reverse = reverse; +} + /* This command implements ZRANGE, ZREVRANGE. */ void genericZrangebyrankCommand(zrange_result_handler *handler, robj *zobj, long start, long end, int withscores, int reverse) { @@ -3494,6 +3984,17 @@ void genericZrangebyrankCommand(zrange_result_handler *handler, handler->emitResultFromCBuffer(handler, ele, sdslen(ele), ln->score); ln = reverse ? ln->backward : ln->level[0].forward; } + } else if (zobj->encoding == OBJ_ENCODING_BTREE) { + zbtreeIterator iter; + const unsigned char *ele; + size_t len; + double score; + unsigned long rank = reverse ? llen - start - 1 : start; + serverAssert(zbtreeIteratorSeekRank(zobj->ptr, rank, &iter)); + while (rangelen--) { + serverAssert(zbtreeIteratorNext(&iter, reverse, &ele, &len, &score)); + handler->emitResultFromCBuffer(handler, ele, len, score); + } } else { serverPanic("Unknown sorted set encoding"); } @@ -3622,6 +4123,38 @@ void genericZrangebyscoreCommand(zrange_result_handler *handler, ln = ln->level[0].forward; } } + } else if (zobj->encoding == OBJ_ENCODING_BTREE) { + zbtreeIterator iter; + const unsigned char *ele; + size_t len; + double score; + unsigned long rank; + int found = reverse ? + zbtreeIteratorSeekScore(zobj->ptr, range->max, range->maxex, 1, &iter, &rank) : + zbtreeIteratorSeekScore(zobj->ptr, range->min, range->minex, 0, &iter, &rank); + + if (found && offset) { + if ((!reverse && rank + offset >= zsetLength(zobj)) || + (reverse && (unsigned long)offset > rank)) + { + found = 0; + } else { + rank = reverse ? rank - offset : rank + offset; + found = zbtreeIteratorSeekRank(zobj->ptr, rank, &iter); + } + } + + while (found && limit-- && + zbtreeIteratorNext(&iter, reverse, &ele, &len, &score)) + { + if (reverse) { + if (!zslValueGteMin(score, range)) break; + } else { + if (!zslValueLteMax(score, range)) break; + } + rangelen++; + handler->emitResultFromCBuffer(handler, ele, len, score); + } } else { serverPanic("Unknown sorted set encoding"); } @@ -3711,6 +4244,8 @@ void zcountCommand(client *c) { count -= (zsl->length - rank); } } + } else if (zobj->encoding == OBJ_ENCODING_BTREE) { + count = zbtreeCountByScore(zobj->ptr, range.min, range.minex, range.max, range.maxex); } else { serverPanic("Unknown sorted set encoding"); } @@ -3787,6 +4322,17 @@ void zlexcountCommand(client *c) { count -= (zsl->length - rank); } } + } else if (zobj->encoding == OBJ_ENCODING_BTREE) { + zbtreeIterator first_iter, last_iter; + unsigned long first, last; + if (zbtreeIteratorSeekLex(zobj->ptr, range.min, range.minex, + 0, &first_iter, &first) && + zbtreeIteratorSeekLex(zobj->ptr, range.max, range.maxex, + 1, &last_iter, &last) && + last >= first) + { + count = last - first + 1; + } } else { serverPanic("Unknown sorted set encoding"); } @@ -3896,6 +4442,37 @@ void genericZrangebylexCommand(zrange_result_handler *handler, ln = ln->level[0].forward; } } + } else if (zobj->encoding == OBJ_ENCODING_BTREE) { + zbtreeIterator iter; + const unsigned char *ele; + size_t len; + double score; + unsigned long rank; + int found = reverse ? + zbtreeIteratorSeekLex(zobj->ptr, range->max, range->maxex, 1, &iter, &rank) : + zbtreeIteratorSeekLex(zobj->ptr, range->min, range->minex, 0, &iter, &rank); + + if (found && offset) { + if ((!reverse && rank + offset >= zsetLength(zobj)) || + (reverse && (unsigned long)offset > rank)) + { + found = 0; + } else { + rank = reverse ? rank - offset : rank + offset; + found = zbtreeIteratorSeekRank(zobj->ptr, rank, &iter); + } + } + + while (found && limit && zbtreeIteratorNext(&iter, reverse, &ele, &len, &score)) { + /* The seek already placed the iterator inside the near bound, so + * only the bound we are moving towards has to be rechecked. */ + if (reverse ? !zbtLexValueGteMin(ele, len, range) : + !zbtLexValueLteMax(ele, len, range)) + break; + limit--; + rangelen++; + handler->emitResultFromCBuffer(handler, ele, len, score); + } } else { serverPanic("Unknown sorted set encoding"); } @@ -4032,6 +4609,7 @@ void zrangeGenericCommand(zrange_result_handler *handler, int argc_start, int st if (opt_withscores || store) { zrangeResultHandlerScoreEmissionEnable(handler); } + zrangeResultHandlerDirectionSet(handler, direction == ZRANGE_DIRECTION_REVERSE); /* Step 3: Lookup the key and get the range. */ kvobj *zobj = lookupKeyRead(c->db, key); @@ -4119,12 +4697,28 @@ void zmscoreCommand(client *c) { if (server.memory_tracking_enabled && zobj != NULL) oldsize = kvobjAllocSize(zobj); addReplyArrayLen(c,c->argc - 2); - for (int j = 2; j < c->argc; j++) { - /* Treat a missing set the same way as an empty set */ - if (zobj == NULL || zsetScore(zobj,c->argv[j]->ptr,&score) == C_ERR) { - addReplyNull(c); - } else { - addReplyDouble(c,score); + if (zobj && zobj->encoding == OBJ_ENCODING_BTREE) { + for (int j = 2; j < c->argc;) { + sds members[ZBT_SCORE_BATCH_SIZE]; + double scores[ZBT_SCORE_BATCH_SIZE]; + unsigned int count = c->argc - j; + if (count > ZBT_SCORE_BATCH_SIZE) count = ZBT_SCORE_BATCH_SIZE; + for (unsigned int i = 0; i < count; i++) members[i] = c->argv[j+i]->ptr; + unsigned int found = zbtreeScores(zobj->ptr, members, count, scores); + for (unsigned int i = 0; i < count; i++) { + if (found & (1U << i)) addReplyDouble(c, scores[i]); + else addReplyNull(c); + } + j += count; + } + } else { + for (int j = 2; j < c->argc; j++) { + /* Treat a missing set the same way as an empty set */ + if (zobj == NULL || zsetScore(zobj,c->argv[j]->ptr,&score) == C_ERR) { + addReplyNull(c); + } else { + addReplyDouble(c,score); + } } } if (server.memory_tracking_enabled && zobj != NULL) @@ -4323,11 +4917,23 @@ void genericZpopCommand(client *c, robj **keyv, int keyc, int where, int emitkey serverAssertWithInfo(c,zobj,zln != NULL); ele = sdsdup(zslGetNodeElement(zln)); score = zln->score; + } else if (zobj->encoding == OBJ_ENCODING_BTREE) { + zbtreeIterator iter; + const unsigned char *raw; + size_t len; + serverAssert(zbtreeIteratorStart(zobj->ptr, where == ZSET_MAX, &iter)); + serverAssert(zbtreeIteratorNext(&iter, where == ZSET_MAX, &raw, &len, &score)); + /* Copy the reply before deletion invalidates the leaf bytes. */ + if (use_nested_array) addReplyArrayLen(c, 2); + addReplyBulkCBuffer(c, raw, len); + addReplyDouble(c, score); + zbtreeDeleteEdge(zobj->ptr, where == ZSET_MAX); + ele = NULL; } else { serverPanic("Unknown sorted set encoding"); } - serverAssertWithInfo(c,zobj,zsetDel(zobj,ele)); + if (ele) serverAssertWithInfo(c,zobj,zsetDel(zobj,ele)); server.dirty++; if (result_count == 0) { /* Do this only for the first iteration. */ @@ -4335,13 +4941,37 @@ void genericZpopCommand(client *c, robj **keyv, int keyc, int where, int emitkey notifyKeyspaceEvent(NOTIFY_ZSET,events[where],key,c->db->id); } - if (use_nested_array) { - addReplyArrayLen(c,2); + if (ele) { + if (use_nested_array) addReplyArrayLen(c,2); + addReplyBulkCBuffer(c,ele,sdslen(ele)); + addReplyDouble(c,score); + sdsfree(ele); } - addReplyBulkCBuffer(c,ele,sdslen(ele)); - addReplyDouble(c,score); - sdsfree(ele); ++result_count; + + /* The first pop above preserves the usual notification point. Read + * the rest before deleting their rank range, so each B+tree leaf is + * edited once instead of shifting its packed arrays for every item. */ + if (zobj->encoding == OBJ_ENCODING_BTREE && rangelen > 1) { + zbtreeSet *zs = zobj->ptr; + long remaining = rangelen - 1; + zbtreeIterator iter; + serverAssert(zbtreeIteratorStart(zs, where == ZSET_MAX, &iter)); + for (long i = 0; i < remaining; i++) { + const unsigned char *raw; + size_t len; + serverAssert(zbtreeIteratorNext(&iter, where == ZSET_MAX, &raw, &len, &score)); + if (use_nested_array) addReplyArrayLen(c, 2); + addReplyBulkCBuffer(c, raw, len); + addReplyDouble(c, score); + } + unsigned long start = where == ZSET_MAX ? zbtreeLength(zs) - remaining : 0; + serverAssert(zbtreeDeleteRangeByRank(zs, start, + start + remaining - 1) == (unsigned long)remaining); + result_count += remaining; + server.dirty += remaining; + break; + } } while(--rangelen); if (server.memory_tracking_enabled) @@ -4570,6 +5200,19 @@ void zrandmemberWithCountCommand(client *c, long l, int withscores) { } zfree(keys); zfree(vals); + } else if (zsetobj->encoding == OBJ_ENCODING_BTREE) { + while (count--) { + listpackEntry key; + double score; + zsetTypeRandomElement(zsetobj, size, &key, withscores ? &score : NULL); + if (withscores && c->resp > 2) + addReplyArrayLen(c,2); + addReplyBulkCBuffer(c, key.sval, key.slen); + if (withscores) + addReplyDouble(c, score); + if (c->flags & CLIENT_CLOSE_ASAP) + break; + } } goto out; } @@ -5027,7 +5670,108 @@ int zsetTest(int argc, char **argv, int flags) { zfree(elements); zslFree(zsl); - + + printf("Testing LISTPACK->BTREE conversion order-check fast path\n"); + { + /* Member strings are named by (49 - i), the OPPOSITE of insertion/ + * score order, so score order and member-lexicographic order are + * inversely correlated throughout -- a comparator that silently + * fell back to comparing members instead of scores (or vice versa) + * would disagree with the true order and get caught, instead of + * accidentally agreeing because the two happened to correlate. */ + unsigned char *zl = lpNew(0); + char namebuf[32]; + for (int i = 0; i < 50; i++) { + snprintf(namebuf, sizeof(namebuf), "member:%04d", 49 - i); + sds ele = sdsnew(namebuf); + zl = zzlInsert(zl, ele, (double)i); + sdsfree(ele); + } + test_cond("zzlIsSortedNoDup accepts a genuinely sorted, duplicate-free listpack", + zzlIsSortedNoDup(zl) == 1); + + /* Swap the last two entries' relative order by appending them in + * reverse (zzlInsertAt(..., NULL, ...) always appends, ignoring + * sort order) onto an otherwise-sorted prefix -- no duplicates, just + * misordered. */ + unsigned char *zl2 = lpNew(0); + for (int i = 0; i < 48; i++) { + snprintf(namebuf, sizeof(namebuf), "member:%04d", 49 - i); + sds ele = sdsnew(namebuf); + zl2 = zzlInsertAt(zl2, NULL, ele, (double)i); + sdsfree(ele); + } + sds elehigh = sdsnew("member:0000"); /* i=49 */ + sds elelow = sdsnew("member:0001"); /* i=48 */ + zl2 = zzlInsertAt(zl2, NULL, elehigh, 49.0); + zl2 = zzlInsertAt(zl2, NULL, elelow, 48.0); + sdsfree(elehigh); + sdsfree(elelow); + test_cond("zzlIsSortedNoDup rejects an out-of-order (but duplicate-free) listpack", + zzlIsSortedNoDup(zl2) == 0); + + /* Exact duplicate of the LAST element appended right after it, so + * the only thing wrong is the (score, member) tie -- unlike + * duplicating an earlier element (which would also read as + * out-of-order via the score comparison alone), this specifically + * exercises the equal-keys boundary of the comparison. */ + unsigned char *zl3 = lpNew(0); + for (int i = 0; i < 50; i++) { + snprintf(namebuf, sizeof(namebuf), "member:%04d", 49 - i); + sds ele = sdsnew(namebuf); + zl3 = zzlInsert(zl3, ele, (double)i); + sdsfree(ele); + } + sds dup = sdsnew("member:0000"); /* i=49's member */ + zl3 = zzlInsertAt(zl3, NULL, dup, 49.0); + sdsfree(dup); + test_cond("zzlIsSortedNoDup rejects a listpack containing a duplicate", + zzlIsSortedNoDup(zl3) == 0); + + /* End-to-end: a sorted source exercises the fast (append) path; an + * out-of-order-but-duplicate-free source exercises the fallback + * (zbtreeInsertNewRaw's own descent-based positioning handles + * arbitrary insertion order correctly, only true duplicates ever + * hit the assert this whole check exists to avoid paying for + * unnecessarily) -- both must still produce a correct, fully sorted + * B+tree either way. */ + robj sortedobj = {.encoding = OBJ_ENCODING_LISTPACK, .ptr = zl, .type = OBJ_ZSET}; + zsetConvertAndExpand(&sortedobj, OBJ_ENCODING_BTREE, 50); + robj unsortedobj = {.encoding = OBJ_ENCODING_LISTPACK, .ptr = zl2, .type = OBJ_ZSET}; + zsetConvertAndExpand(&unsortedobj, OBJ_ENCODING_BTREE, 50); + + int sorted_ok = 1, unsorted_ok = 1; + int sorted_len = 0, unsorted_len = 0; + double prevscore; + zbtreeIterator iter; + const unsigned char *ele; + size_t elelen; + double score; + + zbtreeIteratorStart(sortedobj.ptr, 0, &iter); + prevscore = -1; + while (zbtreeIteratorNext(&iter, 0, &ele, &elelen, &score)) { + if (score < prevscore) sorted_ok = 0; + prevscore = score; + sorted_len++; + } + zbtreeIteratorStart(unsortedobj.ptr, 0, &iter); + prevscore = -1; + while (zbtreeIteratorNext(&iter, 0, &ele, &elelen, &score)) { + if (score < prevscore) unsorted_ok = 0; + prevscore = score; + unsorted_len++; + } + test_cond("Fast-path (sorted source) conversion yields a correctly sorted, complete B+tree", + sorted_ok && sorted_len == 50); + test_cond("Fallback-path (unsorted source) conversion also yields a correctly sorted, complete B+tree", + unsorted_ok && unsorted_len == 50); + + zbtreeFree(sortedobj.ptr); + zbtreeFree(unsortedobj.ptr); + lpFree(zl3); + } + return 0; } #endif diff --git a/src/valgrind.sup b/src/valgrind.sup index 5d6367e3f3..d901820070 100644 --- a/src/valgrind.sup +++ b/src/valgrind.sup @@ -24,3 +24,15 @@ fun:ztrymalloc_usable fun:ztrymalloc } + +# The only pointer to a long zset member stored outside a B+ tree leaf lives +# unaligned inside a packed leaf record, where the leak checker does not look. +# Every such member is allocated by zbtExternalDup(), so match only that frame. +{ + + Memcheck:Leak + match-leak-kinds: definite + ... + fun:zbtAlloc + fun:zbtExternalDup +} diff --git a/src/zset_btree.c b/src/zset_btree.c new file mode 100644 index 0000000000..7d57bc265b --- /dev/null +++ b/src/zset_btree.c @@ -0,0 +1,6195 @@ +/* zset_btree.c -- memory efficient sorted set implementation. + * + * Copyright (c) 2026-Present, Redis Ltd. + * All rights reserved. + * + * Licensed under your choice of (a) the Redis Source Available License 2.0 + * (RSALv2); or (b) the Server Side Public License v1 (SSPLv1); or (c) the + * GNU Affero General Public License v3 (AGPLv3). + */ + +#include "server.h" +#include "zset_btree.h" + +#include +#include +#include +#include + +/*----------------------------------------------------------------------------- + * B+ tree implementation of the low level sorted set API + *----------------------------------------------------------------------------*/ + +/* A large sorted set needs to answer two different questions efficiently: + * + * 1. Given a member, what is its score? + * 2. Given a score and member, where is it in the sorted order? + * + * The traditional Redis encoding uses a hash table for the first question + * and a skiplist for the second. This file uses a small hash table and one B+ + * tree. Both are made for this exact pair of indexes, so they share + * information and omit everything the sorted set does not need. + * + * THE SCORE TREE + * ============== + * + * The score tree is ordered by (score, member), with the member breaking ties + * between equal scores. It owns all member data. The leaves are linked in + * both directions, so a range starts with one tree search and then reads + * consecutive memory pages. + * + * A score leaf is one allocation with the following layout: + * + * +--------+---------------+-------------+------+------------+---------+ + * | header | packed scores | u16 offsets | tags | free space | records | + * +--------+---------------+-------------+------+------------+---------+ + * + * A leaf holds at most ZBT_SCORE_LEAF_MAX elements and requests at most 4095 + * bytes from the allocator. Each offset locates a member record relative to + * the leaf. A parallel byte stores the high eight bits of the member hash. A + * lookup checks these small tags first and normally reads only the requested + * member. + * + * Records grow backwards from the end of the allocation. Scores and offsets + * grow forwards. Spare space between them lets common inserts modify the leaf + * without allocating another one. Short members follow a small length prefix + * inside the record. Members of 320 bytes or more use a separate allocation; + * their record holds its pointer, length, and hash. This keeps leaves dense + * and lets a member keep the same allocation when its score changes. + * + * Scores are stored without losing precision. An IEEE double is first changed + * into an unsigned integer having the same order. A leaf stores one base and + * bit-packed differences from that base. If every difference has the same + * low zero bits, the leaf records that shift once and omits those bits from + * every score. Consecutive integer scores benefit especially from this. + * + * Inner score pages contain child pointers, the maximum score below each + * child, and the number of elements below each child. The counts are what + * make rank searches possible in the same tree. When maximum scores are + * equal, the maximum member below the child is read from its last leaf only + * when needed. + * + * THE MEMBER INDEX + * ================ + * + * The by-name index is an open addressed hash table. Eight entries share one + * bucket. Every entry needs only a one byte piece of the member hash and a + * small number identifying its score leaf. The number normally uses two + * bytes and grows to four before its range is exhausted. An entry does not + * store the member pointer, score, complete hash, or a separate allocation. + * + * A lookup uses the complete hash to select the first bucket, but reads the + * eight one byte pieces together. Usually none match and the more expensive + * leaf access is avoided. On a match, the score leaf number leads to a small + * packed page where the real member is compared. The one byte value is only a + * quick rejection test, so hash collisions can cost work but cannot produce + * a false result. + * + * Empty slots have a zero tag. Deleted slots retain their tag and use a + * reserved leaf number, so searches that passed through them still reach + * later buckets. A small filter for each starting bucket records which tags + * may have started there. It lets most missing-member lookups stop before + * walking a nearly full table. The filter may say "perhaps", never "present". + * + * The table is resized incrementally. It moves one complete score leaf at a + * time. Inline members provide their full hashes, while external records keep + * the hash beside the pointer so a long member is not read again. The active + * resize has a non-zero number. A leaf records that number after its entries + * are copied, which tells lookups which of the two tables contains its current + * entries. + * + * SCORE LEAF IDs + * ============== + * + * Hash entries refer to score leaves by a small number, not by a pointer. + * The score_leaf_by_id array translates the number to the current address. + * This matters because adding, deleting, compacting, or defragmenting a score + * leaf can move its allocation. A replacement keeps the old ID, a split keeps + * it on the left page and gives the right page a new one, and a merge keeps + * the ID of the page with more members. Only members that move to another + * leaf need their member record changed. + * + * Released IDs form a free list inside score_leaf_by_id itself. A marked + * pointer stores the next free ID. Real allocations are aligned, so their low + * bit is zero and can never be confused with a free-list link. + * + * PAGE CHANGES AND TEMPORARY POINTERS + * =================================== + * + * A failed ZADD lookup records its empty hash slot, but score-tree work before + * insertion may advance a resize and move the table. Index changes update + * member_revision whenever a saved slot address can become stale; the + * position is used only if its revision still matches. + * + * The same rule applies to callers: iterators and member pointers returned by + * this file belong to the score tree. They are valid only until the next + * sorted set modification and must not be retained across one. + * + * Score leaves split according to both element count and byte size. Small + * leaves are joined with a neighbor when the result fits. Inner pages are + * removed when empty and a one-child root is collapsed; otherwise sparse + * inner pages are left alone. Their high fanout keeps their number small. + * + * Long members have one owner: the leaf record that points to them. When a + * leaf is rebuilt, the new record takes the pointer and clears it in the old + * record. Freeing the old leaf then releases only members that were deleted. + * COPY makes independent allocations, and active defragmentation moves each + * allocation and repairs its one pointer. + */ + +/* Page limits. Score leaves have both a count limit and a byte limit because + * they contain variable length records. Keeping a leaf below 4 KB bounds the + * bytes moved by an edit and keeps its offsets small. A newly built non-full + * leaf asks for a little spare room, bounded by the same limit. */ +#define ZBT_SCORE_LEAF_MAX 96 +#define ZBT_SCORE_LEAF_BYTES 4095 +#define ZBT_SCORE_LEAF_RESERVE 128 +#define ZBT_SCORE_INNER_MAX 32 +#define ZBT_SCORE_LEAF_MERGE (ZBT_SCORE_LEAF_MAX / 2) +/* Score changes count what they move out of a leaf in 1/96 of a page, by + * count or bytes, whichever is larger. The leaf is checked after a quarter. */ +#define ZBT_SCORE_MOVED_CHECK (ZBT_SCORE_LEAF_MAX / 4) +/* Member lookup uses an open addressed table with eight slots per bucket. */ +#define ZBT_INDEX_BUCKET_ITEMS 8 +#define ZBT_INDEX_INITIAL_BUCKETS 4 +/* Leave headroom for bounded probe runs, including the overflow walk used + * by ZSCAN. Tombstones need a separate limit because they also extend runs. */ +#define ZBT_INDEX_MAX_LOAD_NUM 7 +#define ZBT_INDEX_MAX_LOAD_DEN 8 +#define ZBT_INDEX_MIN_LOAD_NUM 1 +#define ZBT_INDEX_MIN_LOAD_DEN 8 +#define ZBT_INDEX_MAX_FILLED_NUM 15 +#define ZBT_INDEX_MAX_FILLED_DEN 16 +#define ZBT_INDEX_DELETED_ID UINT32_MAX +#define ZBT_INDEX_WIDE_ID_AT (UINT16_MAX / 2) +/* Narrow buckets store leaf IDs below UINT16_MAX, which marks a deleted slot. + * One modification creates at most one score leaf; keep some margin. */ +#define ZBT_INDEX_NARROW_ID_LIMIT (UINT16_MAX - 16) +/* Copy steps per leaf-creating modification while a narrow resize delays the + * widening. It then ends long before ZBT_INDEX_NARROW_ID_LIMIT. */ +#define ZBT_INDEX_WIDEN_STEPS 4 + +/* UINT32_MAX asks for a new leaf ID or marks the absence of a leaf. Use the + * name that describes the meaning at each call site. */ +#define ZBT_NEW_LEAF_ID UINT32_MAX +#define ZBT_NO_LEAF_ID UINT32_MAX +/* Longer members get their own allocation. This keeps score leaves dense and + * avoids copying large strings whenever a leaf changes shape. */ +#define ZBT_EXTERNAL_MEMBER_AT 320 + +#define ZBT_PACKED_SHORT 254 +#define ZBT_PACKED_EXTERNAL 255 +#define ZBT_EXTERNAL_HASH_OFFSET 1 +#define ZBT_EXTERNAL_LENGTH_OFFSET \ + (ZBT_EXTERNAL_HASH_OFFSET + sizeof(uint32_t)) +#define ZBT_EXTERNAL_POINTER_OFFSET \ + (ZBT_EXTERNAL_LENGTH_OFFSET + sizeof(uint64_t)) +#define ZBT_EXTERNAL_RECORD_BYTES \ + (ZBT_EXTERNAL_POINTER_OFFSET + sizeof(void *)) + +/* ZSCAN visits groups of this many adjacent buckets, which share one walk of + * the probe run that follows them. At the usual table fill a group already + * holds more members than the default COUNT. It must be a power of two no + * larger than ZBT_INDEX_INITIAL_BUCKETS. Member hashes have 32 bits, so no + * home bucket lies beyond ZBT_SCAN_MAX_GROUPS groups. */ +#define ZBT_SCAN_BUCKETS_PER_STEP 4 +#define ZBT_SCAN_MAX_GROUPS \ + ((UINT64_C(1) << 32) / ZBT_SCAN_BUCKETS_PER_STEP) + +/* Released score leaf IDs form a list in the leaf table. Allocations are + * aligned, so the low bit distinguishes a list link from a live pointer. */ +#define ZBT_FREE_LEAF_ID(next) \ + ((zbtScoreLeaf *)((((uintptr_t)(next)) << 1) | 1)) +#define ZBT_IS_FREE_LEAF_ID(ptr) (((uintptr_t)(ptr) & 1) != 0) +#define ZBT_NEXT_FREE_LEAF_ID(ptr) ((uint32_t)((uintptr_t)(ptr) >> 1)) + +/* ----------------------------- Page headers ----------------------------- */ + +typedef struct zbtScoreNode zbtScoreNode; +typedef struct zbtScoreLeaf zbtScoreLeaf; +typedef struct zbtScoreInner zbtScoreInner; + +struct zbtScoreNode { + zbtScoreInner *parent; /* NULL only for the root. */ + uint64_t subtree; /* Elements represented by this node. */ + uint8_t parent_index; /* Our position in parent->child[]. */ + uint8_t count; /* Elements in a leaf, children in an inner node. */ + uint8_t isleaf; /* Non-zero for zbtScoreLeaf. */ + /* Only score leaves use these. They live in this header's padding so that + * a score leaf still costs 56 bytes before its packed arrays. */ + uint8_t score_shift; + uint8_t has_external; + uint8_t moved_out; /* Page share moved out by score changes. */ + uint16_t index_resize; /* Resize that copied this leaf, zero if none. */ +}; + +/* The data[] array starts with packed scores and offsets. Member records + * occupy the end and record_start marks their first used byte. */ +struct zbtScoreLeaf { + zbtScoreNode n; + zbtScoreLeaf *prev; + zbtScoreLeaf *next; + uint64_t score_base; + uint32_t id; + uint8_t score_bits; + uint8_t reversed; + uint16_t record_start; + unsigned char data[]; +}; + +/* max_score[] and child_count[] describe the child at the same index. */ +struct zbtScoreInner { + zbtScoreNode n; + zbtScoreNode *child[ZBT_SCORE_INNER_MAX]; + uint64_t child_count[ZBT_SCORE_INNER_MAX]; + double max_score[ZBT_SCORE_INNER_MAX]; +}; + +/* Most sets use 16 bit leaf numbers, making one bucket 32 bytes. The small + * home_tags summary rejects absent members before a table search. Keeping it + * beside the bucket also puts both checks in the same cache line. Very large + * sets switch incrementally to the 48 byte form before 16 bit IDs fill up. */ +typedef struct zbtIndexBucket16 { + uint64_t tags; + uint64_t home_tags; + uint16_t id[ZBT_INDEX_BUCKET_ITEMS]; +} zbtIndexBucket16; + +typedef struct zbtIndexBucket32 { + uint64_t tags; + uint64_t home_tags; + uint32_t id[ZBT_INDEX_BUCKET_ITEMS]; +} zbtIndexBucket32; + +typedef union zbtIndexBucket { + zbtIndexBucket16 narrow; + zbtIndexBucket32 wide; +} zbtIndexBucket; + +typedef struct zbtIndexTable { + zbtIndexBucket *buckets; + unsigned long size; /* Number of buckets, always a power of two. */ + unsigned long used; /* Live entries. */ + unsigned long filled; /* Live entries plus tombstones. */ + int wide_ids; /* IDs are 32 rather than 16 bits. */ +} zbtIndexTable; + +typedef struct zbtIndexRehash { + zbtIndexTable table; /* Table being built. */ + uint32_t next_leaf_id; /* Where the next copy step resumes. */ + uint16_t resize_id; /* Non-zero number of this resize. */ +} zbtIndexRehash; + +struct zbtreeSet { + zbtScoreNode *score_root; + zbtScoreLeaf *score_first; + zbtScoreLeaf *score_last; + zbtIndexTable member_index; + zbtIndexRehash *member_rehash; + + /* + * Hash entries store these small numbers instead of pointers. A score + * leaf can then move during compaction without changing its member + * entries. Released numbers form a free list inside this same array. + */ + zbtScoreLeaf **score_leaf_by_id; + uint32_t score_leaf_cap; + uint32_t next_score_leaf_id; + uint32_t free_score_leaf_id; + uint32_t member_revision; /* Invalidates saved hash slots. */ + unsigned long length; /* Number of sorted set members. */ + size_t alloc_size; /* Usable bytes of all allocations, including us. */ +}; + +/* Temporary view used while a score leaf is being built. An external member + * may already belong to an old record. The builder moves that ownership to + * the new record instead of copying the member bytes. */ +typedef struct { + const unsigned char *ptr; + size_t len; + uint32_t hash; + unsigned char *external; + unsigned char *owner_record; +} zbtBuildElement; + +/* ------------------------- Packed score leaves -------------------------- */ + +/* Turn an IEEE double into an unsigned integer with the same ordering. This + * lets us subtract nearby scores and still reconstruct every bit exactly. + * NaN never reaches this encoding: the sorted set command layer rejects it. */ +static uint64_t zbtScoreToOrdered(double score) { + uint64_t bits; + memcpy(&bits, &score, sizeof(bits)); + if (bits & (UINT64_C(1) << 63)) return ~bits; + return bits ^ (UINT64_C(1) << 63); +} + +/* Inverse of zbtScoreToOrdered(). */ +static double zbtOrderedToScore(uint64_t ordered) { + uint64_t bits; + double score; + if (ordered & (UINT64_C(1) << 63)) + bits = ordered ^ (UINT64_C(1) << 63); + else + bits = ~ordered; + memcpy(&score, &bits, sizeof(score)); + return score; +} + +/* Return the number of bits needed to store 'value'. */ +static unsigned int zbtBitsNeeded(uint64_t value) { + if (value == 0) return 0; + return 64 - __builtin_clzll(value); +} + +/* Eight spare bytes let the bit readers use one unaligned 64 bit load even + * for the last score. The final rounding keeps the offset array aligned. */ +static size_t zbtScoreLeafScoreBytes(unsigned int count, + unsigned int score_bits) +{ + if (score_bits == 0) return 0; + size_t bytes = ((size_t)count * score_bits + 7) / 8; + return (bytes + 8 + 1) & ~(size_t)1; +} + +/* Return the offset array that follows the packed scores. */ +static inline uint16_t *zbtScoreLeafOffsets(zbtScoreLeaf *leaf) { + return (uint16_t *)(leaf->data + + zbtScoreLeafScoreBytes(leaf->n.count, leaf->score_bits)); +} + +/* Return the hash-tag array that follows the offsets. */ +static inline uint8_t *zbtScoreLeafHashTags(zbtScoreLeaf *leaf) { + return (uint8_t *)(zbtScoreLeafOffsets(leaf) + leaf->n.count); +} + +/* Left-edge leaves keep their arrays in reverse physical order. Logical + * positions remain sorted, but prepending then inserts at the physical end + * and does not move the existing packed scores or offsets. */ +static inline unsigned int zbtScoreLeafPhysicalPos(zbtScoreLeaf *leaf, + unsigned int pos) +{ + return leaf->reversed ? leaf->n.count - pos - 1 : pos; +} + +/* Return the member record offset stored at logical position 'pos'. */ +static inline uint16_t zbtScoreLeafOffset(zbtScoreLeaf *leaf, + unsigned int pos) +{ + return zbtScoreLeafOffsets(leaf)[zbtScoreLeafPhysicalPos(leaf, pos)]; +} + +/* The packed arrays form one bit stream: bit p lives in byte p/8 at bit p%8. + * A word is therefore loaded and stored in little endian order whatever the + * host is. Without that, a value written through one byte offset lands on + * top of a value written through another, because a big endian load places + * the first byte in the most significant position. These conversions compile + * away on a little endian build. + */ +static uint64_t zbtReadBits(const unsigned char *data, size_t bitpos, + unsigned int width) +{ + size_t byte = bitpos / 8; + unsigned int shift = bitpos & 7; + unsigned int lowbits = width < 64 - shift ? width : 64 - shift; + uint64_t word, value; + + memcpy(&word, data + byte, sizeof(word)); + value = intrev64ifbe(word) >> shift; + if (width > lowbits) + value |= (uint64_t)data[byte + 8] << lowbits; + if (width < 64) + value &= (UINT64_C(1) << width) - 1; + return value; +} + +/* Replace one field in the packed bit stream. */ +static void zbtWriteBits(unsigned char *data, size_t bitpos, + unsigned int width, uint64_t value) +{ + size_t byte = bitpos / 8; + unsigned int shift = bitpos & 7; + unsigned int lowbits = width < 64 - shift ? width : 64 - shift; + uint64_t mask = lowbits == 64 ? UINT64_MAX : + (UINT64_C(1) << lowbits) - 1; + uint64_t word; + + memcpy(&word, data + byte, sizeof(word)); + word = intrev64ifbe(word); + word &= ~(mask << shift); + word |= (value & mask) << shift; + word = intrev64ifbe(word); + memcpy(data + byte, &word, sizeof(word)); + + if (width > lowbits) { + unsigned int highbits = width - lowbits; + unsigned char highmask = (1U << highbits) - 1; + data[byte + 8] = (data[byte + 8] & ~highmask) | + ((value >> lowbits) & highmask); + } +} + +/* Insert one field in the packed score stream. Copy backwards because source + * and destination overlap. leaf->n.count is still the old count. */ +static void zbtInsertScoreBits(zbtScoreLeaf *leaf, unsigned int pos, + uint64_t value) +{ + unsigned int bits = leaf->score_bits; + if (bits == 0) return; + + size_t insert_at = (size_t)pos * bits; + size_t remaining = (size_t)(leaf->n.count - pos) * bits; + while (remaining) { + unsigned int chunk = remaining > 64 ? 64 : remaining; + size_t source = insert_at + remaining - chunk; + uint64_t part = zbtReadBits(leaf->data, source, chunk); + zbtWriteBits(leaf->data, source + bits, chunk, part); + remaining -= chunk; + } + zbtWriteBits(leaf->data, insert_at, bits, value); +} + +/* Delete one field from the packed score stream. Copy forwards because source + * and destination overlap. The final field is cleared because later writes + * add bits to a field that is expected to start at zero. */ +static void zbtDeleteScoreBits(zbtScoreLeaf *leaf, unsigned int pos) { + unsigned int bits = leaf->score_bits; + if (bits == 0) return; + + size_t destination = (size_t)pos * bits; + size_t remaining = (size_t)(leaf->n.count - pos - 1) * bits; + size_t moved = 0; + while (moved < remaining) { + unsigned int chunk = remaining - moved > 64 ? + 64 : remaining - moved; + uint64_t part = zbtReadBits(leaf->data, + destination + bits + moved, chunk); + zbtWriteBits(leaf->data, destination + moved, chunk, part); + moved += chunk; + } + zbtWriteBits(leaf->data, (size_t)(leaf->n.count - 1) * bits, + bits, 0); +} + +/* Write one score into a zeroed packed field. The base, shift and width must + * already have been selected by zbtScoreEncoding(). */ +static void zbtScoreLeafWriteScore(zbtScoreLeaf *leaf, unsigned int pos, + double score) +{ + unsigned int bits = leaf->score_bits; + if (bits == 0) return; + + uint64_t ordered = zbtScoreToOrdered(score); + uint64_t delta = leaf->reversed ? leaf->score_base - ordered : + ordered - leaf->score_base; + delta >>= leaf->n.score_shift; + size_t bitpos = (size_t)zbtScoreLeafPhysicalPos(leaf, pos) * bits; + size_t byte = bitpos / 8; + unsigned int shift = bitpos & 7; + uint64_t word; + + memcpy(&word, leaf->data + byte, sizeof(word)); + word = intrev64ifbe(intrev64ifbe(word) | (delta << shift)); + memcpy(leaf->data + byte, &word, sizeof(word)); + if (shift && bits > 64 - shift) + leaf->data[byte + 8] |= delta >> (64 - shift); +} + +/* Decode the score at logical position 'pos'. */ +static double zbtScoreLeafScore(zbtScoreLeaf *leaf, unsigned int pos) { + unsigned int bits = leaf->score_bits; + uint64_t delta = 0; + + if (bits != 0) { + size_t bitpos = (size_t)zbtScoreLeafPhysicalPos(leaf, pos) * bits; + size_t byte = bitpos / 8; + unsigned int shift = bitpos & 7; + memcpy(&delta, leaf->data + byte, sizeof(delta)); + delta = intrev64ifbe(delta) >> shift; + if (shift && bits > 64 - shift) + delta |= (uint64_t)leaf->data[byte + 8] << (64 - shift); + if (bits < 64) + delta &= (UINT64_C(1) << bits) - 1; + delta <<= leaf->n.score_shift; + } + uint64_t ordered = leaf->reversed ? leaf->score_base - delta : + leaf->score_base + delta; + return zbtOrderedToScore(ordered); +} + +/* Return the raw hash byte stored with a member. */ +static inline uint8_t zbtScoreLeafTag(zbtScoreLeaf *leaf, unsigned int pos) { + return zbtScoreLeafHashTags(leaf)[zbtScoreLeafPhysicalPos(leaf, pos)]; +} + +/* ---------------------------- Packed members ---------------------------- */ + +/* Most members live directly in the leaf and have only a one byte length for + * the common case. Long members live in their own allocation. Their leaf + * record keeps the length, hash, and pointer needed to find and own them. + * + * short inline member: length byte, member bytes + * longer inline member: byte 254, uint16 length, member bytes + * external member: byte 255, uint32 hash, uint64 length, pointer + */ +static size_t zbtPackedLengthBytes(size_t len) { + return len < ZBT_PACKED_SHORT ? 1 : 3; +} + +/* Write an inline member into a score leaf. */ +static void zbtPackedWriteInline(unsigned char *dst, + const unsigned char *src, size_t len) +{ + serverAssert(len < ZBT_EXTERNAL_MEMBER_AT && len <= UINT16_MAX); + if (len < ZBT_PACKED_SHORT) { + dst[0] = len; + dst++; + } else { + uint16_t shortlen = len; + dst[0] = ZBT_PACKED_SHORT; + memcpy(dst + 1, &shortlen, sizeof(shortlen)); + dst += 3; + } + memcpy(dst, src, len); +} + +/* These fields are deliberately unaligned. Access them through memcpy so the + * records remain tightly packed on every supported CPU. */ +static uint32_t zbtExternalHash(const unsigned char *record) { + uint32_t hash; + serverAssert(record[0] == ZBT_PACKED_EXTERNAL); + memcpy(&hash, record + ZBT_EXTERNAL_HASH_OFFSET, sizeof(hash)); + return hash; +} + +static size_t zbtExternalLength(const unsigned char *record) { + uint64_t len; + serverAssert(record[0] == ZBT_PACKED_EXTERNAL); + memcpy(&len, record + ZBT_EXTERNAL_LENGTH_OFFSET, sizeof(len)); + serverAssert(len <= SIZE_MAX); + return (size_t)len; +} + +static unsigned char *zbtExternalPointer(const unsigned char *record) { + unsigned char *ptr; + serverAssert(record[0] == ZBT_PACKED_EXTERNAL); + memcpy(&ptr, record + ZBT_EXTERNAL_POINTER_OFFSET, sizeof(ptr)); + return ptr; +} + +static void zbtExternalSetPointer(unsigned char *record, unsigned char *ptr) { + serverAssert(record[0] == ZBT_PACKED_EXTERNAL); + memcpy(record + ZBT_EXTERNAL_POINTER_OFFSET, &ptr, sizeof(ptr)); +} + +/* Return a member and store its length in 'len'. */ +static const unsigned char *zbtPackedData(const unsigned char *record, + size_t *len) +{ + if (record[0] < ZBT_PACKED_SHORT) { + *len = record[0]; + return record + 1; + } else if (record[0] == ZBT_PACKED_SHORT) { + uint16_t shortlen; + memcpy(&shortlen, record + 1, sizeof(shortlen)); + *len = shortlen; + return record + 3; + } else { + unsigned char *ptr = zbtExternalPointer(record); + serverAssert(ptr != NULL); + *len = zbtExternalLength(record); + return ptr; + } +} + +/* Return the record at logical position 'pos'. */ +static inline unsigned char *zbtScoreLeafRecord(zbtScoreLeaf *leaf, + unsigned int pos) +{ + return (unsigned char *)leaf + zbtScoreLeafOffset(leaf, pos); +} + +/* Return the member at logical position 'pos'. */ +static inline const unsigned char *zbtScoreLeafElement(zbtScoreLeaf *leaf, + unsigned int pos, + size_t *len) +{ + return zbtPackedData(zbtScoreLeafRecord(leaf, pos), len); +} + +/* Return the complete member hash. External records keep it because hashing a + * long member again during a split or table resize would make that small + * maintenance step depend on the member size. */ +static uint32_t zbtScoreLeafHash(zbtScoreLeaf *leaf, unsigned int pos) { + unsigned char *record = zbtScoreLeafRecord(leaf, pos); + if (record[0] == ZBT_PACKED_EXTERNAL) + return zbtExternalHash(record); + size_t len; + const unsigned char *ele = zbtPackedData(record, &len); + return (uint32_t)dictGenHashFunction(ele, len); +} + +/* Fill a build entry from a live leaf. The old record remains the owner until + * a new record adopts its external pointer. */ +static void zbtBuildElementFromLeaf(zbtBuildElement *ele, + zbtScoreLeaf *leaf, unsigned int pos) +{ + unsigned char *record = zbtScoreLeafRecord(leaf, pos); + ele->ptr = zbtPackedData(record, &ele->len); + if (record[0] == ZBT_PACKED_EXTERNAL) { + ele->hash = zbtExternalHash(record); + ele->external = (unsigned char *)ele->ptr; + ele->owner_record = record; + } else { + debugServerAssert(ele->len < ZBT_EXTERNAL_MEMBER_AT); + ele->hash = 0; + ele->external = NULL; + ele->owner_record = NULL; + } +} + +/* Fill a build entry from bytes supplied by the caller. */ +static void zbtBuildElementFromBytes(zbtBuildElement *ele, + const unsigned char *ptr, size_t len, + uint32_t hash) +{ + ele->ptr = ptr; + ele->len = len; + ele->hash = hash; + ele->external = NULL; + ele->owner_record = NULL; +} + +/* Compare two members without constructing temporary SDS strings. */ +static int zbtCompareElements(const unsigned char *a, size_t alen, + const unsigned char *b, size_t blen) +{ + size_t minlen = alen < blen ? alen : blen; + int cmp = memcmp(a, b, minlen); + if (cmp != 0) return cmp; + if (alen < blen) return -1; + if (alen > blen) return 1; + return 0; +} + +/* Compare two sorted-set entries by score and then by member. */ +static inline int zbtScoreCompare(double score, + const unsigned char *ele, size_t elelen, + double other_score, + const unsigned char *other_ele, + size_t other_elelen) +{ + if (score < other_score) return -1; + if (score > other_score) return 1; + return zbtCompareElements(ele, elelen, other_ele, other_elelen); +} + +/* ---------------------- Allocation and score leaf IDs ------------------- */ + +/* Allocate memory and add the allocator's usable size to the exact object + * accounting returned by MEMORY USAGE. */ +static void *zbtAlloc(zbtreeSet *zs, size_t bytes) { + size_t usable; + void *ptr = zmalloc_usable(bytes, &usable); + zs->alloc_size += usable; + return ptr; +} + +/* Cached insertion positions use these numbers to notice that a member table + * changed. A process-wide counter also distinguishes a deleted set from a new + * one that happens to reuse the same address. */ +static uint32_t zbtIndexNextRevision(void) { + static uint32_t revision = 0; + if (++revision == 0) revision++; + return revision; +} + +/* Free one accounted allocation. */ +static void zbtFreeAllocation(zbtreeSet *zs, void *ptr) { + size_t usable; + zfree_usable(ptr, &usable); + zs->alloc_size -= usable; +} + +/* The only pointer to an external member lives unaligned inside a packed + * leaf record, where LeakSanitizer does not look, so it would report every + * external member still in a set as leaked. Valgrind is handled by + * src/valgrind.sup. */ +#if defined(__SANITIZE_ADDRESS__) +#define ZBT_LSAN_IGNORE 1 +#elif defined(__has_feature) +#if __has_feature(address_sanitizer) +#define ZBT_LSAN_IGNORE 1 +#endif +#endif + +#ifdef ZBT_LSAN_IGNORE +#include +#define zbtExternalIgnoreLeak(ptr) __lsan_ignore_object(ptr) +#else +#define zbtExternalIgnoreLeak(ptr) ((void)(ptr)) +#endif + +/* Copy a long member into its own external allocation. Every new external + * member is allocated here, and src/valgrind.sup suppresses only this frame, + * so no other allocation is hidden from Valgrind. */ +__attribute__((noinline)) +static unsigned char *zbtExternalDup(zbtreeSet *zs, const unsigned char *src, size_t len) { + unsigned char *external = zbtAlloc(zs, len); + zbtExternalIgnoreLeak(external); + memcpy(external, src, len); + return external; +} + +/* Return true when a build entry already refers to an external allocation or + * is long enough that a new one should be made. Existing external records + * retain their form when a leaf is rebuilt. */ +static int zbtBuildElementIsExternal(const zbtBuildElement *ele) { + return ele->external != NULL || ele->len >= ZBT_EXTERNAL_MEMBER_AT; +} + +/* Return the leaf bytes needed by one member record. */ +static size_t zbtElementStorageBytes(const zbtBuildElement *ele) { + if (zbtBuildElementIsExternal(ele)) return ZBT_EXTERNAL_RECORD_BYTES; + return zbtPackedLengthBytes(ele->len) + ele->len; +} + +/* Return the bytes occupied by an existing record without following an + * external pointer. This is used while deleting the record itself. */ +static size_t zbtRecordStorageBytes(const unsigned char *record) { + if (record[0] == ZBT_PACKED_EXTERNAL) + return ZBT_EXTERNAL_RECORD_BYTES; + size_t len; + zbtPackedData(record, &len); + return zbtPackedLengthBytes(len) + len; +} + +/* Write one build entry into a leaf. Existing external allocations move from + * the old record to the new one: clearing the old pointer makes ownership + * visible and lets the normal leaf-free path distinguish moved members from + * deleted ones. */ +static int zbtBuildElementWrite(zbtreeSet *zs, unsigned char *record, + zbtBuildElement *ele) +{ + if (!zbtBuildElementIsExternal(ele)) { + zbtPackedWriteInline(record, ele->ptr, ele->len); + return 0; + } + + unsigned char *external = ele->external; + if (external == NULL) external = zbtExternalDup(zs, ele->ptr, ele->len); + + uint64_t len = ele->len; + record[0] = ZBT_PACKED_EXTERNAL; + memcpy(record + ZBT_EXTERNAL_HASH_OFFSET, &ele->hash, sizeof(ele->hash)); + memcpy(record + ZBT_EXTERNAL_LENGTH_OFFSET, &len, sizeof(len)); + memcpy(record + ZBT_EXTERNAL_POINTER_OFFSET, &external, sizeof(external)); + + if (ele->owner_record) { + debugServerAssert(zbtExternalPointer(ele->owner_record) == external); + zbtExternalSetPointer(ele->owner_record, NULL); + } + ele->external = NULL; + ele->owner_record = NULL; + return 1; +} + +/* Detach an external member before removing its old record. The caller can + * then reinsert the same allocation after a score change without copying the + * member bytes. */ +static void zbtBuildElementDetach(zbtBuildElement *ele) { + if (ele->owner_record == NULL) return; + debugServerAssert(zbtExternalPointer(ele->owner_record) == ele->external); + zbtExternalSetPointer(ele->owner_record, NULL); + ele->owner_record = NULL; +} + +/* Free a leaf and every external member it still owns. A record copied into a + * replacement leaf has a NULL pointer here because the builder cleared it. + * Records for deleted members remain non-NULL and are released now. */ +static void zbtScoreLeafFreeOwned(zbtreeSet *zs, zbtScoreLeaf *leaf) { + if (leaf->n.has_external) { + for (unsigned int i = 0; i < leaf->n.count; i++) { + unsigned char *record = zbtScoreLeafRecord(leaf, i); + if (record[0] != ZBT_PACKED_EXTERNAL) continue; + unsigned char *external = zbtExternalPointer(record); + if (external) zbtFreeAllocation(zs, external); + } + } + zbtFreeAllocation(zs, leaf); +} + +/* Give 'leaf' an ID and publish its address in score_leaf_by_id. Passing + * ZBT_NEW_LEAF_ID first reuses a released ID, then grows the table if all + * existing IDs are busy. Passing an existing ID replaces its address. */ +static void zbtScoreRegisterLeaf(zbtreeSet *zs, zbtScoreLeaf *leaf, uint32_t id) { + if (id == ZBT_NEW_LEAF_ID) { + /* Old hash entries remain in place during a resize. Do not reuse a + * released ID until that old table is gone, or a stale entry could + * appear to refer to an unrelated new leaf. */ + if (zs->free_score_leaf_id && zs->member_rehash == NULL) { + id = zs->free_score_leaf_id - 1; + serverAssert(ZBT_IS_FREE_LEAF_ID(zs->score_leaf_by_id[id])); + zs->free_score_leaf_id = + ZBT_NEXT_FREE_LEAF_ID(zs->score_leaf_by_id[id]); + } else { + id = zs->next_score_leaf_id++; + if (id == zs->score_leaf_cap) { + uint32_t newcap = zs->score_leaf_cap ? + zs->score_leaf_cap * 2 : 16; + size_t usable, old_usable = 0; + zs->score_leaf_by_id = zrealloc_usable( + zs->score_leaf_by_id, + newcap * sizeof(*zs->score_leaf_by_id), + &usable, &old_usable); + memset(zs->score_leaf_by_id + zs->score_leaf_cap, 0, + (newcap - zs->score_leaf_cap) * + sizeof(*zs->score_leaf_by_id)); + zs->score_leaf_cap = newcap; + zs->alloc_size += usable - old_usable; + } + } + } + leaf->id = id; + zs->score_leaf_by_id[id] = leaf; +} + +/* Return an ID to the in-table free list. free_score_leaf_id stores ID+1 so + * zero can mean that the list is empty. */ +static void zbtScoreReleaseLeafId(zbtreeSet *zs, uint32_t id) { + serverAssert(id < zs->next_score_leaf_id); + zs->score_leaf_by_id[id] = + ZBT_FREE_LEAF_ID(zs->free_score_leaf_id); + zs->free_score_leaf_id = id + 1; +} + +/* --------------------------- Score tree pages --------------------------- */ + +/* Return the final member below a score node. Inner nodes deliberately store + * only maximum scores; this walk supplies the tie breaker when equal scores + * make it necessary. */ +static const unsigned char *zbtScoreNodeMaxElement(zbtScoreNode *node, + size_t *len) +{ + while (!node->isleaf) { + zbtScoreInner *inner = (zbtScoreInner *)node; + node = inner->child[inner->n.count - 1]; + } + zbtScoreLeaf *leaf = (zbtScoreLeaf *)node; + return zbtScoreLeafElement(leaf, leaf->n.count - 1, len); +} + +/* Return the greatest score below one score-tree node. */ +static double zbtScoreNodeMaxScore(zbtScoreNode *node) { + if (node->isleaf) { + zbtScoreLeaf *leaf = (zbtScoreLeaf *)node; + return zbtScoreLeafScore(leaf, leaf->n.count - 1); + } + zbtScoreInner *inner = (zbtScoreInner *)node; + return inner->max_score[inner->n.count - 1]; +} + +/* Recompute all derived fields of one score inner page. Besides maximums + * and counts, this repairs each child's parent pointer and parent position. + * Call it after the child array has been rearranged. */ +static void zbtScoreRefreshNode(zbtScoreInner *inner) { + uint64_t total = 0; + for (unsigned int i = 0; i < inner->n.count; i++) { + zbtScoreNode *child = inner->child[i]; + child->parent = inner; + child->parent_index = i; + inner->child_count[i] = child->subtree; + inner->max_score[i] = zbtScoreNodeMaxScore(child); + total += child->subtree; + } + inner->n.subtree = total; +} + +/* Recompute complete inner pages from 'node' up to the root. This is used + * after a reshape, where several child positions may have changed. */ +static void zbtScoreRefreshParents(zbtScoreNode *node) { + while (node->parent) { + zbtScoreInner *parent = node->parent; + zbtScoreRefreshNode(parent); + node = &parent->n; + } +} + +/* Update only the path fields affected by a local leaf edit or by replacing a + * leaf at the same position. length_change is the change in the leaf's + * element count: +1 for insertion, -1 for deletion, zero for a score change. */ +static void zbtScoreUpdatePath(zbtScoreNode *node, long long length_change) { + while (node->parent) { + zbtScoreInner *parent = node->parent; + unsigned int pos = node->parent_index; + double max = zbtScoreNodeMaxScore(node); + /* An in place score change that leaves the node maximum alone leaves + * every page above it alone too. */ + if (length_change == 0 && parent->max_score[pos] == max) return; + parent->child_count[pos] = node->subtree; + parent->max_score[pos] = max; + parent->n.subtree += length_change; + node = &parent->n; + } +} + +/* Scores that share low bits waste those bits in every packed difference. + * Integer scores are the common example: consecutive integers of the same + * magnitude are many ulps apart, so their ordered forms all end in the same + * long run of zero bits. The leaf records that run once as a shift and stores + * only the bits above it. Nothing is lost, because a difference is a multiple + * of 2^shift exactly when every score agrees on those low bits. + */ +static void zbtScoreEncoding(unsigned int count, double *scores, int desc, + uint64_t *base, uint8_t *bits, uint8_t *shift) +{ + uint64_t first = zbtScoreToOrdered(scores[0]); + uint64_t min = first, max = first, differing = 0; + for (unsigned int i = 1; i < count; i++) { + uint64_t ordered = zbtScoreToOrdered(scores[i]); + differing |= ordered ^ first; + if (ordered < min) min = ordered; + if (ordered > max) max = ordered; + } + *base = desc ? max : min; + if (differing == 0) { + *bits = 0; + *shift = 0; + return; + } + *shift = __builtin_ctzll(differing); + *bits = zbtBitsNeeded((max - min) >> *shift); +} + +/* Return the smallest leaf allocation that holds the supplied arrays. */ +static size_t zbtScoreLeafRequestBytes(unsigned int count, double *scores, + zbtBuildElement *elems) +{ + uint64_t base; + uint8_t bits, shift; + /* The shift and the number of packed bits do not depend on the base, so + * an ascending encoding gives the size of a reversed leaf as well. */ + zbtScoreEncoding(count, scores, 0, &base, &bits, &shift); + size_t bytes = offsetof(zbtScoreLeaf, data) + + zbtScoreLeafScoreBytes(count, bits) + + count * (sizeof(uint16_t) + sizeof(uint8_t)); + for (unsigned int i = 0; i < count; i++) + bytes += zbtElementStorageBytes(&elems[i]); + return bytes; +} + +/* Build a complete score leaf from sorted arrays. Inline members are copied. + * External members are either allocated here or adopted from their old leaf. + * + * A normal non-full leaf asks the allocator for some reserve, allowing later + * inserts to use the same allocation. An edge page asks for the full 4095 + * bytes immediately: monotonic insertion will fill that page, so repeated + * growth would only copy it needlessly. + * + * 'id' may be an existing ID when rebuilding a leaf, or ZBT_NEW_LEAF_ID when + * creating another one. The function publishes the new address in the ID + * table but does not link the leaf into the tree or leaf list. + */ +static zbtScoreLeaf *zbtScoreLeafBuild(zbtreeSet *zs, unsigned int count, + double *scores, zbtBuildElement *elems, + uint8_t *tags, uint32_t id, + int edge_page) +{ + uint64_t score_base; + uint8_t score_bits, score_shift; + zbtScoreEncoding(count, scores, 0, &score_base, &score_bits, &score_shift); + size_t request = zbtScoreLeafRequestBytes(count, scores, elems); + size_t alloc_request = request; + if (edge_page && request <= ZBT_SCORE_LEAF_BYTES) { + alloc_request = ZBT_SCORE_LEAF_BYTES; + } else if (request <= ZBT_SCORE_LEAF_BYTES && + count < ZBT_SCORE_LEAF_MAX) + { + alloc_request += ZBT_SCORE_LEAF_RESERVE; + if (alloc_request > ZBT_SCORE_LEAF_BYTES) + alloc_request = ZBT_SCORE_LEAF_BYTES; + } + zbtScoreLeaf *leaf = zbtAlloc(zs, alloc_request); + uint16_t *offsets; + uint8_t *hash_tags; + unsigned char *records; + + serverAssert(count > 0 && count <= ZBT_SCORE_LEAF_MAX); + serverAssert(request <= ZBT_SCORE_LEAF_BYTES); + + memset(leaf, 0, offsetof(zbtScoreLeaf, data)); + leaf->n.isleaf = 1; + leaf->n.count = count; + leaf->n.subtree = count; + leaf->n.score_shift = score_shift; + leaf->score_base = score_base; + leaf->score_bits = score_bits; + memset(leaf->data, 0, zbtScoreLeafScoreBytes(count, score_bits)); + offsets = zbtScoreLeafOffsets(leaf); + hash_tags = zbtScoreLeafHashTags(leaf); + size_t capacity = zmalloc_usable_size(leaf); + if (capacity > ZBT_SCORE_LEAF_BYTES) + capacity = ZBT_SCORE_LEAF_BYTES; + records = (unsigned char *)leaf + capacity; + + for (unsigned int i = 0; i < count; i++) { + size_t item_bytes = zbtElementStorageBytes(&elems[i]); + + zbtScoreLeafWriteScore(leaf, i, scores[i]); + debugServerAssert(zbtScoreLeafScore(leaf, i) == scores[i]); + records -= item_bytes; + serverAssert((char *)records - (char *)leaf <= + ZBT_SCORE_LEAF_BYTES); + offsets[i] = (uint16_t)((char *)records - (char *)leaf); + hash_tags[i] = tags[i]; + if (zbtBuildElementWrite(zs, records, &elems[i])) + leaf->n.has_external = 1; + } + leaf->record_start = records - (unsigned char *)leaf; + zbtScoreRegisterLeaf(zs, leaf, id); + return leaf; +} + +/* Allocate an empty inner score-tree page. */ +static zbtScoreInner *zbtScoreInnerCreate(zbtreeSet *zs) { + zbtScoreInner *inner = zbtAlloc(zs, sizeof(*inner)); + memset(inner, 0, sizeof(*inner)); + return inner; +} + +/* Put a rebuilt leaf in every place that named oldleaf: parent, neighboring + * leaves, first/last pointers, and possibly the root. Its ID table entry was + * already changed by zbtScoreLeafBuild(). */ +static void zbtScoreReplaceLeaf(zbtreeSet *zs, zbtScoreLeaf *oldleaf, zbtScoreLeaf *newleaf) { + newleaf->n.parent = oldleaf->n.parent; + newleaf->n.parent_index = oldleaf->n.parent_index; + newleaf->n.index_resize = oldleaf->n.index_resize; + newleaf->prev = oldleaf->prev; + newleaf->next = oldleaf->next; + + if (newleaf->prev) newleaf->prev->next = newleaf; + else zs->score_first = newleaf; + if (newleaf->next) newleaf->next->prev = newleaf; + else zs->score_last = newleaf; + + if (oldleaf->n.parent) + oldleaf->n.parent->child[oldleaf->n.parent_index] = &newleaf->n; + else + zs->score_root = &newleaf->n; + /* The leaf keeps its position, so only its path can change. Refreshing + * every child of each ancestor would read up to 32 cold sibling leaves. */ + zbtScoreUpdatePath(&newleaf->n, (long long)newleaf->n.subtree - + (long long)oldleaf->n.subtree); + zbtScoreLeafFreeOwned(zs, oldleaf); +} + +/* Rebuild a leaf without spare or abandoned record bytes, retaining its ID. */ +static zbtScoreLeaf *zbtScoreCompactLeaf(zbtreeSet *zs, + zbtScoreLeaf *oldleaf) +{ + double scores[ZBT_SCORE_LEAF_MAX]; + zbtBuildElement elems[ZBT_SCORE_LEAF_MAX]; + uint8_t tags[ZBT_SCORE_LEAF_MAX]; + + for (unsigned int i = 0; i < oldleaf->n.count; i++) { + scores[i] = zbtScoreLeafScore(oldleaf, i); + zbtBuildElementFromLeaf(&elems[i], oldleaf, i); + tags[i] = zbtScoreLeafTag(oldleaf, i); + } + zbtScoreLeaf *newleaf = zbtScoreLeafBuild(zs, oldleaf->n.count, + scores, elems, tags, oldleaf->id, 0); + zbtScoreReplaceLeaf(zs, oldleaf, newleaf); + return newleaf; +} + +/* Detach a score node from the tree. Empty ancestors are removed recursively, + * and a root with one child is replaced by that child. The node itself is + * still owned and freed by the caller. */ +static void zbtScoreDetachNode(zbtreeSet *zs, zbtScoreNode *node) { + zbtScoreInner *parent = node->parent; + + if (parent == NULL) { + serverAssert(zs->score_root == node); + zs->score_root = NULL; + return; + } + + unsigned int pos = node->parent_index; + memmove(&parent->child[pos], &parent->child[pos + 1], + (parent->n.count - pos - 1) * sizeof(parent->child[0])); + parent->n.count--; + + if (parent->n.count == 0) { + zbtScoreDetachNode(zs, &parent->n); + zbtFreeAllocation(zs, parent); + } else if (parent->n.parent == NULL && parent->n.count == 1) { + zbtScoreNode *child = parent->child[0]; + child->parent = NULL; + child->parent_index = 0; + zs->score_root = child; + zbtFreeAllocation(zs, parent); + } else { + zbtScoreRefreshNode(parent); + zbtScoreRefreshParents(&parent->n); + } +} + +/* Unlink and free a score leaf, including its ID. Member-index records using + * that ID must be removed or redirected before another lookup can run. */ +static void zbtScoreRemoveLeaf(zbtreeSet *zs, zbtScoreLeaf *leaf) { + if (zs->member_rehash && + zs->member_rehash->next_leaf_id == leaf->id) + { + zs->member_rehash->next_leaf_id = + leaf->next ? leaf->next->id : ZBT_NO_LEAF_ID; + } + if (leaf->prev) leaf->prev->next = leaf->next; + else zs->score_first = leaf->next; + if (leaf->next) leaf->next->prev = leaf->prev; + else zs->score_last = leaf->prev; + zbtScoreDetachNode(zs, &leaf->n); + zbtScoreReleaseLeafId(zs, leaf->id); + zbtScoreLeafFreeOwned(zs, leaf); +} + +static void zbtScoreInsertSibling(zbtreeSet *zs, zbtScoreNode *left, zbtScoreNode *right); + +/* Insert one child into a full inner page, split the resulting children in + * half, then insert the new right page into the level above. */ +static void zbtScoreSplitInner(zbtreeSet *zs, zbtScoreInner *inner, + unsigned int insert_at, zbtScoreNode *right) +{ + zbtScoreNode *children[ZBT_SCORE_INNER_MAX + 1]; + unsigned int src = 0; + + for (unsigned int i = 0; i < ZBT_SCORE_INNER_MAX + 1; i++) { + if (i == insert_at) children[i] = right; + else children[i] = inner->child[src++]; + } + + unsigned int left_count = (ZBT_SCORE_INNER_MAX + 1) / 2; + unsigned int right_count = ZBT_SCORE_INNER_MAX + 1 - left_count; + zbtScoreInner *newright = zbtScoreInnerCreate(zs); + inner->n.count = left_count; + newright->n.count = right_count; + + for (unsigned int i = 0; i < left_count; i++) + inner->child[i] = children[i]; + for (unsigned int i = 0; i < right_count; i++) + newright->child[i] = children[left_count + i]; + + zbtScoreRefreshNode(inner); + zbtScoreRefreshNode(newright); + zbtScoreInsertSibling(zs, &inner->n, &newright->n); +} + +/* Link 'right' immediately after 'left' at the same tree level. A new root is + * created when left was the old root; a full parent is split recursively. */ +static void zbtScoreInsertSibling(zbtreeSet *zs, zbtScoreNode *left, zbtScoreNode *right) { + zbtScoreInner *parent = left->parent; + + if (parent == NULL) { + zbtScoreInner *root = zbtScoreInnerCreate(zs); + root->n.count = 2; + root->child[0] = left; + root->child[1] = right; + zbtScoreRefreshNode(root); + zs->score_root = &root->n; + return; + } + + unsigned int pos = left->parent_index + 1; + if (parent->n.count == ZBT_SCORE_INNER_MAX) { + zbtScoreSplitInner(zs, parent, pos, right); + return; + } + + memmove(&parent->child[pos + 1], &parent->child[pos], + (parent->n.count - pos) * sizeof(parent->child[0])); + parent->child[pos] = right; + parent->n.count++; + zbtScoreRefreshNode(parent); + zbtScoreRefreshParents(&parent->n); +} + +/* Same operation as zbtScoreInsertSibling(), but put 'left' immediately + * before the existing node 'right'. This is used by descending edge builds. */ +static void zbtScoreInsertBefore(zbtreeSet *zs, zbtScoreNode *right, + zbtScoreNode *left) +{ + zbtScoreInner *parent = right->parent; + + if (parent == NULL) { + zbtScoreInner *root = zbtScoreInnerCreate(zs); + root->n.count = 2; + root->child[0] = left; + root->child[1] = right; + zbtScoreRefreshNode(root); + zs->score_root = &root->n; + return; + } + + unsigned int pos = right->parent_index; + if (parent->n.count == ZBT_SCORE_INNER_MAX) { + zbtScoreSplitInner(zs, parent, pos, left); + return; + } + + memmove(&parent->child[pos + 1], &parent->child[pos], + (parent->n.count - pos) * sizeof(parent->child[0])); + parent->child[pos] = left; + parent->n.count++; + zbtScoreRefreshNode(parent); + zbtScoreRefreshParents(&parent->n); +} + +/* ---------------------- Score tree search and insert -------------------- */ + +/* Select the first child whose greatest (score, member) is not below the + * requested pair. max_score[] answers the common case. Only equal scores need + * the member comparison that walks to the child's final leaf. */ +static unsigned int zbtScoreChild(zbtScoreInner *inner, double score, + const unsigned char *ele, size_t elelen) +{ + for (unsigned int i = 0; i < inner->n.count; i++) { + if (score < inner->max_score[i]) return i; + if (score == inner->max_score[i]) { + size_t maxlen; + const unsigned char *maxele = + zbtScoreNodeMaxElement(inner->child[i], &maxlen); + if (zbtCompareElements(ele, elelen, maxele, maxlen) <= 0) + return i; + } + } + return inner->n.count - 1; +} + +/* Find the score leaf where a (score, member) pair belongs. */ +static zbtScoreLeaf *zbtScoreFindLeaf(const zbtreeSet *zs, double score, + const unsigned char *ele, size_t elelen) +{ + zbtScoreNode *node = zs->score_root; + if (node == NULL) return NULL; + while (!node->isleaf) { + zbtScoreInner *inner = (zbtScoreInner *)node; + node = inner->child[zbtScoreChild(inner, score, ele, elelen)]; + } + return (zbtScoreLeaf *)node; +} + +/* Return the first logical position not below (score, member). */ +static unsigned int zbtScoreLeafLowerBound(zbtScoreLeaf *leaf, double score, + const unsigned char *ele, + size_t elelen) +{ + unsigned int lo = 0, hi = leaf->n.count; + while (lo < hi) { + unsigned int mid = lo + (hi - lo) / 2; + size_t midlen; + const unsigned char *midele = + zbtScoreLeafElement(leaf, mid, &midlen); + int cmp = zbtScoreCompare(zbtScoreLeafScore(leaf, mid), + midele, midlen, score, ele, elelen); + if (cmp < 0) lo = mid + 1; + else hi = mid; + } + return lo; +} + +/* Look for an exact member in one score leaf. The byte tags are contiguous, + * so memchr rejects almost every unrelated member before its packed bytes + * are touched. */ +static int zbtScoreLeafFindMember(zbtScoreLeaf *leaf, uint32_t hash, + const unsigned char *ele, size_t elelen, + double *score, + unsigned int *position) +{ + uint8_t tag = hash >> 24; + uint16_t *offsets = zbtScoreLeafOffsets(leaf); + uint8_t *tags = zbtScoreLeafHashTags(leaf); + uint8_t *next = tags; + uint8_t *end = tags + leaf->n.count; + + while (next < end) { + uint8_t *found = memchr(next, tag, end - next); + if (found == NULL) break; + unsigned int physical = found - tags; + uint16_t offset = offsets[physical]; + unsigned char *record = (unsigned char *)leaf + offset; + + if (record[0] == ZBT_PACKED_EXTERNAL && + zbtExternalHash(record) != hash) + { + next = found + 1; + continue; + } + + size_t len; + const unsigned char *stored = zbtPackedData(record, &len); + if (len == elelen && memcmp(stored, ele, len) == 0) + { + unsigned int logical = + zbtScoreLeafPhysicalPos(leaf, physical); + if (score) *score = zbtScoreLeafScore(leaf, logical); + if (position) *position = logical; + return 1; + } + next = found + 1; + } + return 0; +} + +/* Choose a score-leaf split that balances bytes, not just element count. + * Both sides must fit in the 4 KB leaf limit. */ +static unsigned int zbtScoreChooseSplit(unsigned int count, double *scores, + zbtBuildElement *elems) +{ + size_t total = 0, prefix = 0; + unsigned int candidate = count / 2; + + /* Eight bytes is an upper bound for one packed score. Use that bound to + * find a balanced split without repeatedly decoding every score. + */ + for (unsigned int i = 0; i < count; i++) + total += zbtElementStorageBytes(&elems[i]) + + sizeof(double) + sizeof(uint16_t) + sizeof(uint8_t); + for (unsigned int i = 1; i < count; i++) { + prefix += zbtElementStorageBytes(&elems[i - 1]) + + sizeof(double) + sizeof(uint16_t) + sizeof(uint8_t); + if (prefix * 2 >= total) { + candidate = i; + break; + } + } + + if (zbtScoreLeafRequestBytes(candidate, scores, elems) <= + ZBT_SCORE_LEAF_BYTES && + zbtScoreLeafRequestBytes(count - candidate, scores + candidate, + elems + candidate) <= ZBT_SCORE_LEAF_BYTES) + return candidate; + + /* The byte estimate is conservative, but verify every boundary if an + * unusual mix of member lengths did not fit at the estimated midpoint. */ + unsigned int best = 0; + size_t best_largest = SIZE_MAX; + + for (unsigned int i = 1; i < count; i++) { + size_t left = zbtScoreLeafRequestBytes(i, scores, elems); + size_t right = zbtScoreLeafRequestBytes(count - i, scores + i, + elems + i); + if (left > ZBT_SCORE_LEAF_BYTES || + right > ZBT_SCORE_LEAF_BYTES) + continue; + size_t largest = left > right ? left : right; + if (largest < best_largest) { + best = i; + best_largest = largest; + } + } + serverAssert(best != 0); + return best; +} + +/* Insert after the new score encoding has been computed, while retaining the + * same leaf allocation. Existing records remain where they are; only the new + * record is added and the packed scores and offsets are rewritten. + * + * 'locations' contains the old record offsets in the new logical order. + * Return zero when the changed arrays and new record no longer fit in the + * free space, so the caller can rebuild or split the leaf. + */ +static int zbtScoreInsertInPlace(zbtreeSet *zs, zbtScoreLeaf *leaf, + unsigned int count, + double *scores, zbtBuildElement *elems, + uint8_t *tags, uint16_t *locations, + unsigned int inserted) +{ + uint64_t score_base; + uint8_t score_bits, score_shift; + zbtScoreEncoding(count, scores, leaf->reversed, + &score_base, &score_bits, &score_shift); + size_t score_bytes = zbtScoreLeafScoreBytes(count, score_bits); + size_t index_end = offsetof(zbtScoreLeaf, data) + score_bytes + + count * (sizeof(uint16_t) + sizeof(uint8_t)); + size_t item_bytes = zbtElementStorageBytes(&elems[inserted]); + if (index_end + item_bytes > leaf->record_start) + return 0; + + uint16_t new_offset = leaf->record_start - item_bytes; + if (zbtBuildElementWrite(zs, (unsigned char *)leaf + new_offset, + &elems[inserted])) + leaf->n.has_external = 1; + + leaf->n.count = count; + leaf->n.subtree = count; + leaf->n.score_shift = score_shift; + leaf->score_base = score_base; + leaf->score_bits = score_bits; + leaf->record_start = new_offset; + memset(leaf->data, 0, score_bytes); + + uint16_t *offsets = zbtScoreLeafOffsets(leaf); + uint8_t *hash_tags = zbtScoreLeafHashTags(leaf); + for (unsigned int i = 0; i < count; i++) { + zbtScoreLeafWriteScore(leaf, i, scores[i]); + uint16_t offset = i == inserted ? new_offset : locations[i]; + unsigned int physical = zbtScoreLeafPhysicalPos(leaf, i); + offsets[physical] = offset; + hash_tags[physical] = tags[i]; + debugServerAssert(zbtScoreLeafScore(leaf, i) == scores[i]); + } + zbtScoreUpdatePath(&leaf->n, 1); + return 1; +} + +/* Return the packed representation of a score when the leaf's current base, + * shift and width can hold it exactly. */ +static int zbtScoreLeafPackScore(zbtScoreLeaf *leaf, double score, + uint64_t *packed) +{ + uint64_t ordered = zbtScoreToOrdered(score); + if ((!leaf->reversed && ordered < leaf->score_base) || + (leaf->reversed && ordered > leaf->score_base)) + return 0; + uint64_t delta = leaf->reversed ? leaf->score_base - ordered : + ordered - leaf->score_base; + unsigned int shift = leaf->n.score_shift; + if (delta & ((UINT64_C(1) << shift) - 1)) return 0; + delta >>= shift; + if (leaf->score_bits == 0) { + if (delta != 0) return 0; + } else if (leaf->score_bits < 64 && + delta >= (UINT64_C(1) << leaf->score_bits)) + { + return 0; + } + *packed = delta; + return 1; +} + +/* Fast insertion when the current base, shift and width can encode the new + * score. Shift the packed score fields and offsets in place instead of + * decoding and rewriting the whole leaf. Return zero if the score needs a + * different encoding or the allocation has insufficient free space. + */ +static int zbtScoreInsertFast(zbtreeSet *zs, zbtScoreLeaf *leaf, double score, + zbtBuildElement *ele, + uint8_t tag, unsigned int pos) +{ + uint64_t delta; + if (!zbtScoreLeafPackScore(leaf, score, &delta)) return 0; + + unsigned int oldcount = leaf->n.count; + unsigned int count = oldcount + 1; + size_t score_bytes = + zbtScoreLeafScoreBytes(count, leaf->score_bits); + size_t item_bytes = zbtElementStorageBytes(ele); + size_t index_end = offsetof(zbtScoreLeaf, data) + score_bytes + + count * (sizeof(uint16_t) + sizeof(uint8_t)); + if (index_end + item_bytes > leaf->record_start) + return 0; + + unsigned int physical = leaf->reversed ? oldcount - pos : pos; + + /* Adding a score can move both arrays that follow the packed bits. Save + * the tags, shift the bits while the old offsets are still available, + * then rebuild the offsets and tags at their new addresses. */ + uint16_t *old_offsets = zbtScoreLeafOffsets(leaf); + uint8_t old_tags[ZBT_SCORE_LEAF_MAX]; + memcpy(old_tags, zbtScoreLeafHashTags(leaf), oldcount); + zbtInsertScoreBits(leaf, physical, delta); + uint16_t *offsets = (uint16_t *)(leaf->data + score_bytes); + memmove(offsets, old_offsets, oldcount * sizeof(uint16_t)); + memmove(offsets + physical + 1, offsets + physical, + (oldcount - physical) * sizeof(uint16_t)); + uint8_t *hash_tags = (uint8_t *)(offsets + count); + memcpy(hash_tags, old_tags, physical); + hash_tags[physical] = tag; + memcpy(hash_tags + physical + 1, old_tags + physical, + oldcount - physical); + uint16_t new_offset = leaf->record_start - item_bytes; + if (zbtBuildElementWrite(zs, (unsigned char *)leaf + new_offset, ele)) + leaf->n.has_external = 1; + + leaf->n.count = count; + leaf->n.subtree = count; + leaf->record_start = new_offset; + offsets[physical] = new_offset; + zbtScoreUpdatePath(&leaf->n, 1); + return 1; +} + +static int zbtIndexMove(zbtreeSet *zs, uint32_t hash, + uint32_t old_leaf_id, uint32_t new_leaf_id, + int target_is_incomplete, int *target_was_copied); + +/* Start a new leaf at one end of the tree. The old edge leaf was allocated + * at full page size while growing; compact it now that it will no longer + * receive monotonic inserts. The new left edge uses reversed arrays. */ +static zbtScoreLeaf *zbtScoreInsertEdgeLeaf(zbtreeSet *zs, + zbtScoreLeaf *oldleaf, + double score, + zbtBuildElement *ele, + uint8_t tag, int prepend) +{ + if (zmalloc_usable_size(oldleaf) >= ZBT_SCORE_LEAF_BYTES) + oldleaf = zbtScoreCompactLeaf(zs, oldleaf); + + double newscore[1] = {score}; + uint8_t newtag[1] = {tag}; + zbtScoreLeaf *newleaf = zbtScoreLeafBuild(zs, 1, newscore, ele, + newtag, ZBT_NEW_LEAF_ID, 1); + newleaf->n.index_resize = oldleaf->n.index_resize; + if (prepend) { + newleaf->reversed = 1; + newleaf->prev = oldleaf->prev; + newleaf->next = oldleaf; + oldleaf->prev = newleaf; + zs->score_first = newleaf; + zbtScoreInsertBefore(zs, &oldleaf->n, &newleaf->n); + } else { + newleaf->prev = oldleaf; + newleaf->next = oldleaf->next; + oldleaf->next = newleaf; + zs->score_last = newleaf; + zbtScoreInsertSibling(zs, &oldleaf->n, &newleaf->n); + } + return newleaf; +} + +/* Insert a new (score, member) pair into the score tree and return the leaf + * that received it. The member index is updated separately by the caller. + * 'oldleaf'/'pos' are the target leaf and the position the pair belongs at + * within it; 'at_left_edge'/'at_right_edge' say whether that position is the + * very start/end of the whole tree. Callers that already know all four + * (bulk-append, see zbtScoreInsertAppend below) skip the descent that would + * otherwise recompute them. + * + * The paths, in increasing order of work, are: + * + * 1. Shift packed fields in the current leaf without changing its encoding. + * 2. Re-encode the arrays while keeping the same allocation. + * 3. Rebuild the leaf, preserving its stable ID. + * 4. Split it and repair the member records of elements moved to the right. + * + * Monotonic insertion has another path that fills an edge page and then + * starts a new one, avoiding the half-full leaves produced by normal splits. + */ +static zbtScoreLeaf *zbtScoreInsertAt(zbtreeSet *zs, zbtScoreLeaf *oldleaf, + unsigned int pos, int at_left_edge, + int at_right_edge, double score, + zbtBuildElement *ele) +{ + uint32_t hash = ele->hash; + uint8_t tag = hash >> 24; + unsigned int oldcount = oldleaf->n.count; + + /* Monotonic inserts are common when a sorted set is loaded. Once an edge + * leaf reaches its element limit, start another full-size leaf instead of + * splitting it. At the left edge the arrays are kept in reverse physical + * order, so prepending has the same cheap path. + */ + if (oldcount == ZBT_SCORE_LEAF_MAX && + (at_left_edge || at_right_edge)) + { + return zbtScoreInsertEdgeLeaf(zs, oldleaf, score, ele, tag, + at_left_edge); + } + + if (oldcount < ZBT_SCORE_LEAF_MAX && + zbtScoreInsertFast(zs, oldleaf, score, ele, tag, pos)) + return oldleaf; + + double scores[ZBT_SCORE_LEAF_MAX + 1]; + zbtBuildElement elems[ZBT_SCORE_LEAF_MAX + 1]; + uint8_t tags[ZBT_SCORE_LEAF_MAX + 1]; + uint16_t locations[ZBT_SCORE_LEAF_MAX + 1]; + + for (unsigned int i = 0, src = 0; i < oldcount + 1; i++) { + if (i == pos) { + scores[i] = score; + elems[i] = *ele; + tags[i] = tag; + locations[i] = 0; + } else { + scores[i] = zbtScoreLeafScore(oldleaf, src); + zbtBuildElementFromLeaf(&elems[i], oldleaf, src); + tags[i] = zbtScoreLeafTag(oldleaf, src); + locations[i] = zbtScoreLeafOffset(oldleaf, src); + src++; + } + } + + size_t request = zbtScoreLeafRequestBytes(oldcount + 1, scores, elems); + + /* Long members can fill the page before its element limit is reached. + * Preserve a full edge page in that case for the same reason as above. */ + if (request > ZBT_SCORE_LEAF_BYTES && + (at_left_edge || at_right_edge)) + { + return zbtScoreInsertEdgeLeaf(zs, oldleaf, score, ele, tag, + at_left_edge); + } + + if (oldcount < ZBT_SCORE_LEAF_MAX && + request <= ZBT_SCORE_LEAF_BYTES) + { + if (zbtScoreInsertInPlace(zs, oldleaf, oldcount + 1, scores, elems, + tags, locations, pos)) + return oldleaf; + zbtScoreLeaf *newleaf = zbtScoreLeafBuild(zs, oldcount + 1, + scores, elems, tags, + oldleaf->id, 0); + zbtScoreReplaceLeaf(zs, oldleaf, newleaf); + return newleaf; + } + + unsigned int left_count = + zbtScoreChooseSplit(oldcount + 1, scores, elems); + unsigned int right_count = oldcount + 1 - left_count; + zbtScoreLeaf *left = zbtScoreLeafBuild(zs, left_count, scores, elems, + tags, oldleaf->id, 0); + zbtScoreLeaf *right = zbtScoreLeafBuild(zs, right_count, + scores + left_count, + elems + left_count, + tags + left_count, + ZBT_NEW_LEAF_ID, 0); + left->n.index_resize = right->n.index_resize = oldleaf->n.index_resize; + + left->n.parent = oldleaf->n.parent; + left->n.parent_index = oldleaf->n.parent_index; + left->prev = oldleaf->prev; + left->next = right; + right->prev = left; + right->next = oldleaf->next; + if (left->prev) left->prev->next = left; + else zs->score_first = left; + if (right->next) right->next->prev = right; + else zs->score_last = right; + + if (oldleaf->n.parent) + oldleaf->n.parent->child[oldleaf->n.parent_index] = &left->n; + else + zs->score_root = &left->n; + + /* Existing member-index entries still name the old leaf. The left leaf + * keeps that ID. Entries for members copied to the right leaf only need + * their small ID field changed. The hash-table position does not depend + * on that ID. + * The newly inserted member has no index entry yet, so skip it here. + */ + int right_was_copied = 0; + for (unsigned int i = left_count; i < oldcount + 1; i++) { + if (i == pos) continue; + uint32_t moved_hash = + zbtScoreLeafHash(right, i - left_count); + serverAssert(zbtIndexMove(zs, moved_hash, oldleaf->id, right->id, + pos >= left_count, &right_was_copied)); + } + + zbtScoreLeafFreeOwned(zs, oldleaf); + zbtScoreInsertSibling(zs, &left->n, &right->n); + zbtScoreRefreshParents(&left->n); + return pos < left_count ? left : right; +} + +/* Insert a new (score, member) pair anywhere in the tree: find its target + * leaf and position, then hand off to zbtScoreInsertAt(). */ +static zbtScoreLeaf *zbtScoreInsert(zbtreeSet *zs, double score, + zbtBuildElement *ele) +{ + if (zs->score_root == NULL) { + double scores[1] = {score}; + uint8_t tags[1] = {(uint8_t)(ele->hash >> 24)}; + zbtScoreLeaf *leaf = zbtScoreLeafBuild(zs, 1, scores, ele, tags, + ZBT_NEW_LEAF_ID, 0); + zs->score_root = &leaf->n; + zs->score_first = zs->score_last = leaf; + return leaf; + } + + zbtScoreLeaf *oldleaf = zbtScoreFindLeaf(zs, score, ele->ptr, ele->len); + unsigned int pos = zbtScoreLeafLowerBound(oldleaf, score, + ele->ptr, ele->len); + int at_left_edge = pos == 0 && oldleaf == zs->score_first; + int at_right_edge = pos == oldleaf->n.count && oldleaf == zs->score_last; + return zbtScoreInsertAt(zs, oldleaf, pos, at_left_edge, at_right_edge, + score, ele); +} + +/* Insert a new (score, member) pair that the caller guarantees sorts after + * every pair already in the tree -- the shape zsetConvertAndExpand()'s + * BTREE-target loops produce, walking an already-sorted source (a listpack + * or skiplist) and inserting members known to be absent. Skips + * zbtScoreFindLeaf()'s top-down descent entirely: the target is always the + * current rightmost leaf, appended at its end, which zbtScoreInsertAt() + * already has a cheap edge-leaf path for. Precondition checked in debug + * builds, not derived -- callers outside a from-scratch bulk build must not + * use this. */ +static zbtScoreLeaf *zbtScoreInsertAppend(zbtreeSet *zs, double score, + zbtBuildElement *ele) +{ + if (zs->score_root == NULL) { + double scores[1] = {score}; + uint8_t tags[1] = {(uint8_t)(ele->hash >> 24)}; + zbtScoreLeaf *leaf = zbtScoreLeafBuild(zs, 1, scores, ele, tags, + ZBT_NEW_LEAF_ID, 0); + zs->score_root = &leaf->n; + zs->score_first = zs->score_last = leaf; + return leaf; + } + +#ifdef DEBUG_ASSERTIONS + { + zbtScoreLeaf *last = zs->score_last; + unsigned int lastpos = last->n.count - 1; + double lastscore = zbtScoreLeafScore(last, lastpos); + int after_last = score > lastscore; + if (!after_last && score == lastscore) { + size_t len; + const unsigned char *maxele = + zbtScoreLeafElement(last, lastpos, &len); + after_last = zbtCompareElements(ele->ptr, ele->len, + maxele, len) > 0; + } + debugServerAssert(after_last); + } +#endif + + zbtScoreLeaf *oldleaf = zs->score_last; + unsigned int pos = oldleaf->n.count; + return zbtScoreInsertAt(zs, oldleaf, pos, 0, 1, score, ele); +} + +/* --------------------------- Member index ------------------------------- */ + +/* Return the number of slots in a member-index table. */ +static inline unsigned long zbtIndexSlots(const zbtIndexTable *table) { + return table->size * ZBT_INDEX_BUCKET_ITEMS; +} + +/* Return the byte size of one narrow or wide bucket. */ +static inline size_t zbtIndexBucketBytes(const zbtIndexTable *table) { + return table->wide_ids ? sizeof(zbtIndexBucket32) : + sizeof(zbtIndexBucket16); +} + +/* Return the bytes occupied by all buckets in a table. */ +static inline size_t zbtIndexTableBytes(const zbtIndexTable *table) { + return table->size * zbtIndexBucketBytes(table); +} + +/* Return one bucket from a table whose bucket width is selected at runtime. */ +static inline zbtIndexBucket *zbtIndexBucketAt(const zbtIndexTable *table, + unsigned long index) +{ + return (zbtIndexBucket *)((unsigned char *)table->buckets + + index * zbtIndexBucketBytes(table)); +} + +/* Return the missing-member filter for one home bucket. */ +static inline uint64_t *zbtIndexHomeTagsAt(const zbtIndexTable *table, + unsigned long index) +{ + return &zbtIndexBucketAt(table, index)->narrow.home_tags; +} + +/* Read a leaf ID from either bucket form. */ +static inline uint32_t zbtIndexGetId(const zbtIndexTable *table, + const zbtIndexBucket *bucket, + unsigned int pos) +{ + if (table->wide_ids) return bucket->wide.id[pos]; + uint16_t id = bucket->narrow.id[pos]; + return id == UINT16_MAX ? ZBT_INDEX_DELETED_ID : id; +} + +/* Both bucket forms have tags as their common first field. */ +static inline uint64_t zbtIndexTags(const zbtIndexBucket *bucket) { + return bucket->narrow.tags; +} + +/* Store a leaf ID in either bucket form. */ +static inline void zbtIndexSetId(const zbtIndexTable *table, + zbtIndexBucket *bucket, unsigned int pos, + uint32_t id) +{ + if (table->wide_ids) { + bucket->wide.id[pos] = id; + } else { + serverAssert(id == ZBT_INDEX_DELETED_ID || id < UINT16_MAX); + bucket->narrow.id[pos] = id == ZBT_INDEX_DELETED_ID ? + UINT16_MAX : id; + } +} + +/* Fold tag zero into one because zero marks an unused table slot. */ +static inline uint8_t zbtIndexTag(uint32_t hash) { + uint8_t tag = hash >> 24; + return tag ? tag : 1; +} + +/* Three bits summarize one tag. Collisions only cause a normal table search. */ +static inline uint64_t zbtIndexTagBits(uint8_t tag) { + uint64_t mixed = (uint64_t)tag * UINT64_C(0x9e3779b97f4a7c15); + return (UINT64_C(1) << (tag & 63)) | + (UINT64_C(1) << (mixed >> 58)) | + (UINT64_C(1) << ((mixed >> 36) & 63)); +} + +/* Record one member in the missing-member filter of its home bucket. */ +static inline void zbtIndexRecordHomeTag(zbtIndexTable *table, uint32_t hash) { + unsigned long home = hash & (table->size - 1); + *zbtIndexHomeTagsAt(table, home) |= zbtIndexTagBits(zbtIndexTag(hash)); +} + +/* Return zero only when the home bucket proves that a hash is absent. */ +static inline int zbtIndexHomeMayContain(const zbtIndexTable *table, + uint32_t hash) +{ + unsigned long home = hash & (table->size - 1); + uint64_t bits = zbtIndexTagBits(zbtIndexTag(hash)); + return (*zbtIndexHomeTagsAt(table, home) & bits) == bits; +} + +/* Return a mask with the high bit set in every byte equal to tag. As with + * the usual zero-byte test, a bit above a real match may also be set. Every + * candidate is therefore checked against the stored byte before use. */ +static inline uint64_t zbtIndexTagMask(uint64_t tags, uint8_t tag) { + uint64_t x = tags ^ (UINT64_C(0x0101010101010101) * tag); + return (x - UINT64_C(0x0101010101010101)) & ~x & + UINT64_C(0x8080808080808080); +} + +/* Return the first matching byte position in a non-zero tag mask. */ +static inline unsigned int zbtIndexFirstTag(uint64_t mask) { + return __builtin_ctzll(mask) >> 3; +} + +/* Replace one tag byte without disturbing the other seven. */ +static inline void zbtIndexSetTag(zbtIndexBucket *bucket, unsigned int pos, + uint8_t tag) +{ + uint64_t shift = pos * 8; + bucket->narrow.tags = + (zbtIndexTags(bucket) & ~(UINT64_C(0xff) << shift)) | + ((uint64_t)tag << shift); +} + +/* Round a requested bucket count up to the table's minimum power of two. */ +static unsigned long zbtIndexNextPower(unsigned long size) { + unsigned long result = ZBT_INDEX_INITIAL_BUCKETS; + while (result < size) result <<= 1; + return result; +} + +/* Choose enough buckets to keep the requested elements below maximum load. */ +static unsigned long zbtIndexBucketsForElements(unsigned long elements) { + if (elements == 0) return ZBT_INDEX_INITIAL_BUCKETS; + unsigned long slots = + (elements * ZBT_INDEX_MAX_LOAD_DEN + ZBT_INDEX_MAX_LOAD_NUM - 1) / + ZBT_INDEX_MAX_LOAD_NUM; + return zbtIndexNextPower( + (slots + ZBT_INDEX_BUCKET_ITEMS - 1) / ZBT_INDEX_BUCKET_ITEMS); +} + +/* Allocate and clear a narrow or wide member-index table. A fallible call + * returns zero instead of aborting when the table cannot be allocated, and + * leaves 'table' unchanged. */ +static int zbtIndexTableAlloc(zbtreeSet *zs, zbtIndexTable *table, + unsigned long buckets, int wide_ids, + int fallible) +{ + zbtIndexTable t = {.size = zbtIndexNextPower(buckets), + .wide_ids = wide_ids}; + if (t.size > SIZE_MAX / zbtIndexBucketBytes(&t)) return 0; + size_t bytes = zbtIndexTableBytes(&t), usable; + t.buckets = fallible ? ztrycalloc_usable(bytes, &usable) : + zcalloc_usable(bytes, &usable); + if (t.buckets == NULL) return 0; + zs->alloc_size += usable; + *table = t; + return 1; +} + +/* Allocate a member-index table, aborting on out of memory. */ +static void zbtIndexTableInit(zbtreeSet *zs, zbtIndexTable *table, + unsigned long buckets, int wide_ids) +{ + serverAssert(zbtIndexTableAlloc(zs, table, buckets, wide_ids, 0)); +} + +/* Release a member-index table and clear its descriptor. */ +static void zbtIndexTableRelease(zbtreeSet *zs, zbtIndexTable *table) { + if (table->buckets) zbtFreeAllocation(zs, table->buckets); + memset(table, 0, sizeof(*table)); +} + +/* Insert an entry without checking for duplicates or changing table size. + * Reusing a tombstone keeps the slot occupied, so every existing probe path + * remains valid and filled does not increase. */ +static void zbtIndexTableInsertRaw(zbtIndexTable *table, uint32_t hash, + uint32_t id) +{ + uint8_t tag = zbtIndexTag(hash); + unsigned long mask = table->size - 1; + unsigned long index = hash & mask; + + for (unsigned long probes = 0; probes < table->size; probes++) { + zbtIndexBucket *bucket = zbtIndexBucketAt(table, index); + for (unsigned int pos = 0; pos < ZBT_INDEX_BUCKET_ITEMS; pos++) { + uint8_t oldtag = zbtIndexTags(bucket) >> (pos * 8); + uint32_t oldid = zbtIndexGetId(table, bucket, pos); + if (oldtag == 0 || oldid == ZBT_INDEX_DELETED_ID) { + if (oldtag == 0) table->filled++; + zbtIndexSetTag(bucket, pos, tag); + zbtIndexSetId(table, bucket, pos, id); + table->used++; + zbtIndexRecordHomeTag(table, hash); + return; + } + } + index = (index + 1) & mask; + } + serverPanic("B+ tree member index has no free slot"); +} + +/* Search one table. A matching hash still has to be checked against the + * member in its score leaf because the stored hash is not a unique key. + * Inline so read-only callers can omit the insertion/cursor bookkeeping. */ +static inline __attribute__((always_inline)) +int zbtIndexTableFind(zbtreeSet *zs, zbtIndexTable *table, + uint32_t hash, const unsigned char *ele, + size_t elelen, double *score, + zbtScoreLeaf **score_leaf, + unsigned int *score_pos, + zbtIndexBucket **found_bucket, + unsigned int *found_pos, + zbtIndexBucket **insert_bucket, + unsigned int *insert_pos) +{ + if (table->size == 0) return 0; + uint8_t tag = zbtIndexTag(hash); + unsigned long mask = table->size - 1; + unsigned long index = hash & mask; + zbtIndexBucket *first_deleted = NULL; + unsigned int first_deleted_pos = 0; + + for (unsigned long probes = 0; probes < table->size; probes++) { + zbtIndexBucket *bucket = zbtIndexBucketAt(table, index); + uint64_t matches = zbtIndexTagMask(zbtIndexTags(bucket), tag); + while (matches) { + unsigned int pos = zbtIndexFirstTag(matches); + if ((uint8_t)(zbtIndexTags(bucket) >> (pos * 8)) != tag) { + matches &= matches - 1; + continue; + } + uint32_t id = zbtIndexGetId(table, bucket, pos); + if (id != ZBT_INDEX_DELETED_ID) + { + zbtScoreLeaf *leaf = id < zs->next_score_leaf_id ? + zs->score_leaf_by_id[id] : NULL; + if (leaf == NULL || ZBT_IS_FREE_LEAF_ID(leaf)) { + serverAssert(zs->member_rehash && + table == &zs->member_index); + matches &= matches - 1; + continue; + } + serverAssert(leaf->id == id); + if (zs->member_rehash) { + int in_new = table == &zs->member_rehash->table; + int migrated = leaf->n.index_resize == + zs->member_rehash->resize_id; + if (in_new != migrated) { + matches &= matches - 1; + continue; + } + } + if (zbtScoreLeafFindMember(leaf, hash, ele, elelen, score, + score_pos)) + { + if (score_leaf) *score_leaf = leaf; + if (found_bucket) *found_bucket = bucket; + if (found_pos) *found_pos = pos; + return 1; + } + } + matches &= matches - 1; + } + + /* Hits in the home bucket need no missing-member filter. Check it + * only after those candidates, before following the probe run. */ + if (probes == 0 && !zbtIndexHomeMayContain(table, hash)) return 0; + + /* Save the first tombstone for a possible insertion. An unused slot, + * unlike a tombstone, proves that the member is not farther ahead. */ + for (unsigned int pos = 0; pos < ZBT_INDEX_BUCKET_ITEMS; pos++) { + uint8_t oldtag = zbtIndexTags(bucket) >> (pos * 8); + if (oldtag == 0) { + if (insert_bucket) { + *insert_bucket = first_deleted ? first_deleted : bucket; + *insert_pos = first_deleted ? first_deleted_pos : pos; + } + return 0; + } + if (first_deleted == NULL && + zbtIndexGetId(table, bucket, pos) == ZBT_INDEX_DELETED_ID) + { + first_deleted = bucket; + first_deleted_pos = pos; + } + } + index = (index + 1) & mask; + } + + if (insert_bucket && first_deleted) { + *insert_bucket = first_deleted; + *insert_pos = first_deleted_pos; + } + return 0; +} + +/* Find an entry whose leaf number has to change. Entries with the same tag + * and leaf number look alike, but any reachable one is sufficient. With + * linear probing, two search paths stay together after they meet, so + * exchanging such entries cannot make either member unreachable. */ +static int zbtIndexTableFindReference(zbtIndexTable *table, uint32_t hash, + uint32_t id, + zbtIndexBucket **found_bucket, + unsigned int *found_pos) +{ + if (table->size == 0) return 0; + if (!zbtIndexHomeMayContain(table, hash)) return 0; + uint8_t tag = zbtIndexTag(hash); + unsigned long mask = table->size - 1; + unsigned long index = hash & mask; + + for (unsigned long probes = 0; probes < table->size; probes++) { + zbtIndexBucket *bucket = zbtIndexBucketAt(table, index); + uint64_t matches = zbtIndexTagMask(zbtIndexTags(bucket), tag); + while (matches) { + unsigned int pos = zbtIndexFirstTag(matches); + if ((uint8_t)(zbtIndexTags(bucket) >> (pos * 8)) == tag && + zbtIndexGetId(table, bucket, pos) == id) + { + *found_bucket = bucket; + *found_pos = pos; + return 1; + } + matches &= matches - 1; + } + if (zbtIndexTagMask(zbtIndexTags(bucket), 0)) return 0; + index = (index + 1) & mask; + } + return 0; +} + +/* Start copying the member index into a table sized for 'elements'. + * + * Every later insertion copies at least one more score leaf, so the table + * also leaves room for one insertion per leaf. The copy then finishes before + * the destination fills, and zbtIndexExpandIfNeeded() does not have to + * complete it at once. Without this, a tombstone clean-up of a set close to + * the maximum load gets a same-sized table with almost no room. Leaf IDs + * bound the number of leaves, and so does the member count. */ +static int zbtIndexStartResize(zbtreeSet *zs, unsigned long elements, + int allow_same_size) +{ + if (zs->member_rehash) return 0; + unsigned long room = 0; + if (zs->length) { + room = zs->next_score_leaf_id; + if (room > elements) room = elements; + } + unsigned long buckets = zbtIndexBucketsForElements(elements + room); + if (buckets == zs->member_index.size && !allow_same_size) return 0; + + /* A size hint may underestimate the first batch. There are no leaves + * to migrate yet, so replace the empty table directly. */ + if (zs->length == 0) { + int wide_ids = zs->member_index.wide_ids || + zs->next_score_leaf_id >= ZBT_INDEX_WIDE_ID_AT; + zbtIndexTableRelease(zs, &zs->member_index); + zbtIndexTableInit(zs, &zs->member_index, buckets, wide_ids); + zs->member_revision = zbtIndexNextRevision(); + return 1; + } + + zs->member_rehash = zbtAlloc(zs, sizeof(*zs->member_rehash)); + memset(zs->member_rehash, 0, sizeof(*zs->member_rehash)); + int wide_ids = zs->member_index.wide_ids || + zs->next_score_leaf_id >= ZBT_INDEX_WIDE_ID_AT; + zbtIndexTableInit(zs, &zs->member_rehash->table, buckets, wide_ids); + serverAssert(zs->score_first != NULL); + zs->member_rehash->next_leaf_id = zs->score_first->id; + uint16_t resize_id = zs->score_first->n.index_resize + 1; + /* Zero is what a newly allocated leaf contains. Never use it for an + * active resize, including when this counter wraps. */ + if (resize_id == 0) resize_id = 1; + zs->member_rehash->resize_id = resize_id; + zs->member_revision = zbtIndexNextRevision(); + return 1; +} + +/* Return whether a leaf has already been copied by the active resize. */ +static int zbtIndexLeafMigrated(const zbtreeSet *zs, + const zbtScoreLeaf *leaf) +{ + return zs->member_rehash && + leaf->n.index_resize == zs->member_rehash->resize_id; +} + +/* Copy one score leaf into the new table. Its old entries are left untouched + * and ignored from now on. This is much cheaper than locating and removing + * every old entry, and the complete old table can be freed at the end. */ +static void zbtIndexCopyLeaf(zbtreeSet *zs, zbtScoreLeaf *leaf) { + zbtIndexRehash *rehash = zs->member_rehash; + serverAssert(rehash != NULL); + if (zbtIndexLeafMigrated(zs, leaf)) return; + + for (unsigned int i = 0; i < leaf->n.count; i++) { + uint32_t hash = zbtScoreLeafHash(leaf, i); + zbtIndexTableInsertRaw(&rehash->table, hash, leaf->id); + } + leaf->n.index_resize = rehash->resize_id; + zs->member_revision = zbtIndexNextRevision(); +} + +/* Copy complete score leaves to the new table. The leaf bytes supply the full + * hashes which the compact member index deliberately does not store. */ +static int zbtIndexRehashStep(zbtreeSet *zs, int steps) { + zbtIndexRehash *rehash = zs->member_rehash; + if (rehash == NULL) return 0; + + while (steps--) { + zbtScoreLeaf *leaf = NULL; + if (rehash->next_leaf_id != ZBT_NO_LEAF_ID && + rehash->next_leaf_id < zs->next_score_leaf_id) + { + leaf = zs->score_leaf_by_id[rehash->next_leaf_id]; + if (leaf && ZBT_IS_FREE_LEAF_ID(leaf)) leaf = NULL; + } + /* A merge can release the saved ID between steps. In that case resume + * from the first leaf and skip all leaves already copied. */ + if (leaf == NULL) leaf = zs->score_first; + while (leaf && zbtIndexLeafMigrated(zs, leaf)) leaf = leaf->next; + if (leaf == NULL) { + /* Every live leaf is in the destination. Make it current and + * release the old table in one step. */ + serverAssert(rehash->table.used == zs->length); + zbtIndexTable old = zs->member_index; + zs->member_index = rehash->table; + zbtFreeAllocation(zs, rehash); + zs->member_rehash = NULL; + zbtIndexTableRelease(zs, &old); + zs->member_revision = zbtIndexNextRevision(); + return 0; + } + + rehash->next_leaf_id = + leaf->next ? leaf->next->id : ZBT_NO_LEAF_ID; + zbtIndexCopyLeaf(zs, leaf); + } + return 1; +} + +/* Call before a modification that can create a score leaf. New leaf IDs are + * written to the resize destination while one is active, otherwise to the + * current table; during a resize the old table only receives deletions. If + * that receiving table is narrow, next_score_leaf_id must stay below + * UINT16_MAX. Released IDs are reused only below next_score_leaf_id, and one + * modification creates at most one leaf, so a check here is sufficient. + * + * Widening starts incrementally at ZBT_INDEX_WIDE_ID_AT, as for insertions. + * A narrow resize that is already running is advanced faster so that the + * widening can follow it. Finishing it at once is only a last resort. */ +static void zbtIndexReserveLeafId(zbtreeSet *zs, unsigned long add) { + if (zs->next_score_leaf_id < ZBT_INDEX_WIDE_ID_AT) return; + if (zs->member_rehash && !zs->member_rehash->table.wide_ids) { + if (zs->next_score_leaf_id < ZBT_INDEX_NARROW_ID_LIMIT) + zbtIndexRehashStep(zs, ZBT_INDEX_WIDEN_STEPS); + else + while (zs->member_rehash) zbtIndexRehashStep(zs, 64); + } + if (zs->member_rehash == NULL && zs->member_index.size && + !zs->member_index.wide_ids) + { + zbtIndexStartResize(zs, zs->member_index.used + add, 1); + } +} + +/* Grow, widen, or clean the member index before adding more entries. */ +static void zbtIndexExpandIfNeeded(zbtreeSet *zs, unsigned long add) { + if (zs->member_index.size == 0) { + /* A set emptied by deletions keeps its leaf ID counter. */ + zbtIndexTableInit(zs, &zs->member_index, + zbtIndexBucketsForElements(add), + zs->next_score_leaf_id >= ZBT_INDEX_WIDE_ID_AT); + zs->member_revision = zbtIndexNextRevision(); + return; + } + zbtIndexReserveLeafId(zs, add); + if (zs->member_rehash) { + zbtIndexTable *target = &zs->member_rehash->table; + unsigned long slots = zbtIndexSlots(target); + + /* A tombstone cleanup can use a same-sized destination. If many new + * members arrive before that copy finishes, the destination itself + * may need to grow. Finish the current copy before it runs out of + * room; the normal checks below will then start the larger table. */ + if ((zs->length + add) * ZBT_INDEX_MAX_LOAD_DEN <= + slots * ZBT_INDEX_MAX_LOAD_NUM && + (target->filled + add) * ZBT_INDEX_MAX_FILLED_DEN < + slots * ZBT_INDEX_MAX_FILLED_NUM) + return; + while (zs->member_rehash) + zbtIndexRehashStep(zs, 64); + } + + /* A narrow table can represent a little more than twice this many leaf + * IDs. Starting here leaves one operation per old leaf to finish the + * upgrade even if every operation also creates a new leaf. */ + if (!zs->member_index.wide_ids && + zs->next_score_leaf_id >= ZBT_INDEX_WIDE_ID_AT) + { + zbtIndexStartResize(zs, zs->member_index.used + add, 1); + return; + } + + unsigned long slots = zbtIndexSlots(&zs->member_index); + unsigned long live = zs->member_index.used + add; + unsigned long filled = zs->member_index.filled + add; + if (live * ZBT_INDEX_MAX_LOAD_DEN > + slots * ZBT_INDEX_MAX_LOAD_NUM || + filled * ZBT_INDEX_MAX_FILLED_DEN >= + slots * ZBT_INDEX_MAX_FILLED_NUM) + { + zbtIndexStartResize(zs, live, 1); + } +} + +/* Start a smaller table when deletions leave the current one mostly empty. */ +static void zbtIndexShrinkIfNeeded(zbtreeSet *zs) { + /* Redis will normally free the whole object after its final member is + * removed. Release the index here as well because range deletion finishes + * its internal work before the command layer can do that. There is no + * score leaf from which an incremental resize could start. */ + if (zs->length == 0) { + zbtIndexTableRelease(zs, &zs->member_index); + if (zs->member_rehash) { + zbtIndexTableRelease(zs, &zs->member_rehash->table); + zbtFreeAllocation(zs, zs->member_rehash); + zs->member_rehash = NULL; + } + zs->member_revision = zbtIndexNextRevision(); + return; + } + if (zs->member_rehash || + zs->member_index.size <= ZBT_INDEX_INITIAL_BUCKETS) + return; + unsigned long used = zs->member_index.used; + if (used * ZBT_INDEX_MIN_LOAD_DEN <= + zbtIndexSlots(&zs->member_index) * ZBT_INDEX_MIN_LOAD_NUM) + { + zbtIndexStartResize(zs, used ? used : 1, 0); + } +} + +/* Return whether a saved bucket address belongs to this table. */ +static int zbtIndexOwnsBucket(zbtIndexTable *table, + zbtIndexBucket *bucket) +{ + if (table->buckets == NULL) return 0; + uintptr_t address = (uintptr_t)bucket; + uintptr_t first = (uintptr_t)table->buckets; + uintptr_t last = first + table->size * zbtIndexBucketBytes(table); + return address >= first && address < last; +} + +/* Look in both tables and optionally return score and insertion positions. */ +static inline __attribute__((always_inline)) +int zbtIndexFind(zbtreeSet *zs, uint32_t hash, + const unsigned char *ele, size_t elelen, + double *score, + zbtScoreLeaf **found_score_leaf, + unsigned int *found_score_position, + zbtIndexBucket **found_bucket, + unsigned int *found_bucket_position) +{ + zbtIndexBucket *bucket = NULL, *hint_bucket = NULL; + unsigned int pos = 0, hint_pos = 0; + if (zs->member_rehash) zbtIndexRehashStep(zs, 1); + + if (zs->member_rehash && + zbtIndexTableFind(zs, &zs->member_rehash->table, hash, ele, elelen, + score, + found_score_leaf, found_score_position, + found_bucket ? &bucket : NULL, + found_bucket_position ? &pos : NULL, + found_bucket ? &hint_bucket : NULL, + found_bucket ? &hint_pos : NULL)) + goto found; + + /* During a resize new entries go to the new table, so an insertion slot + * found in the old table would not be useful to the caller. */ + if (zbtIndexTableFind(zs, &zs->member_index, hash, ele, elelen, score, + found_score_leaf, found_score_position, + found_bucket ? &bucket : NULL, + found_bucket_position ? &pos : NULL, + zs->member_rehash || !found_bucket ? NULL : &hint_bucket, + zs->member_rehash || !found_bucket ? NULL : &hint_pos)) + goto found; + + if (found_bucket) *found_bucket = hint_bucket; + if (found_bucket_position) *found_bucket_position = hint_pos; + return 0; + +found: + if (found_bucket) *found_bucket = bucket; + if (found_bucket_position) *found_bucket_position = pos; + return 1; +} + +/* Find a table entry by hash tag and leaf ID without reading its member. */ +static int zbtIndexFindReference(zbtreeSet *zs, uint32_t hash, + uint32_t score_leaf_id, + zbtIndexBucket **found_bucket, + unsigned int *found_pos) +{ + zbtIndexBucket *bucket; + if (zs->member_rehash && + zbtIndexTableFindReference(&zs->member_rehash->table, hash, + score_leaf_id, &bucket, found_pos)) + { + *found_bucket = bucket; + return 1; + } + if (zbtIndexTableFindReference(&zs->member_index, hash, score_leaf_id, + &bucket, found_pos)) + { + *found_bucket = bucket; + return 1; + } + return 0; +} + +/* Change the leaf ID named by one member entry, also during a resize. + * + * A resize may place the source and destination leaves in different tables. + * If the destination leaf is incomplete, copying it would miss the member + * currently being inserted, so the caller copies it after finishing the leaf. + * Until then only the old entry is removed: the copy adds every member of the + * destination. Relabeling the old entry instead could put a wide ID in a + * narrow old table, and during a resize that table only receives deletions. + * When several members move to one leaf, target_was_copied records that its + * complete contents were already copied; later old entries are then removed + * instead of creating duplicates. */ +static int zbtIndexMove(zbtreeSet *zs, uint32_t hash, + uint32_t old_leaf_id, uint32_t new_leaf_id, + int target_is_incomplete, int *target_was_copied) +{ + zbtIndexBucket *bucket; + unsigned int pos; + if (zs->member_rehash == NULL) { + if (!zbtIndexTableFindReference(&zs->member_index, hash, + old_leaf_id, &bucket, &pos)) + return 0; + zbtIndexSetId(&zs->member_index, bucket, pos, new_leaf_id); + return 1; + } + + serverAssert(new_leaf_id < zs->next_score_leaf_id); + zbtScoreLeaf *newleaf = zs->score_leaf_by_id[new_leaf_id]; + serverAssert(newleaf && !ZBT_IS_FREE_LEAF_ID(newleaf)); + if (!zbtIndexLeafMigrated(zs, newleaf) && target_is_incomplete) { + /* A split may put the newly inserted, not-yet-indexed member in the + * right leaf. zbtIndexInsert(), or the zbtIndexMove() that follows a + * score change, copies that complete leaf after the split. A split + * leaf passes its copy state to both halves, and the new table holds + * no entries for leaves not yet copied, so the entry is old. */ + if (!zbtIndexFindReference(zs, hash, old_leaf_id, &bucket, &pos)) + return 0; + serverAssert(zbtIndexOwnsBucket(&zs->member_index, bucket)); + zbtIndexSetId(&zs->member_index, bucket, pos, ZBT_INDEX_DELETED_ID); + zs->member_index.used--; + return 1; + } + int copied_target = target_was_copied && *target_was_copied; + if (!zbtIndexLeafMigrated(zs, newleaf)) { + zbtIndexCopyLeaf(zs, newleaf); + copied_target = 1; + if (target_was_copied) *target_was_copied = 1; + } + + if (!zbtIndexFindReference(zs, hash, old_leaf_id, &bucket, &pos)) + return 0; + + zbtIndexTable *source = &zs->member_index; + if (!zbtIndexOwnsBucket(source, bucket)) { + serverAssert(zbtIndexOwnsBucket(&zs->member_rehash->table, bucket)); + source = &zs->member_rehash->table; + } + if (copied_target) { + zbtIndexSetId(source, bucket, pos, ZBT_INDEX_DELETED_ID); + source->used--; + } else if (source == &zs->member_rehash->table) { + zbtIndexSetId(source, bucket, pos, new_leaf_id); + } else { + zbtIndexSetId(source, bucket, pos, ZBT_INDEX_DELETED_ID); + source->used--; + zbtIndexTableInsertRaw(&zs->member_rehash->table, + hash, new_leaf_id); + } + return 1; +} + +/* Turn one occupied member-index slot into a tombstone. */ +static void zbtIndexDeleteAt(zbtreeSet *zs, zbtIndexBucket *bucket, + unsigned int pos) +{ + zbtIndexTable *table = &zs->member_index; + if (!zbtIndexOwnsBucket(table, bucket)) { + serverAssert(zs->member_rehash && + zbtIndexOwnsBucket(&zs->member_rehash->table, bucket)); + table = &zs->member_rehash->table; + } + serverAssert(zbtIndexGetId(table, bucket, pos) != ZBT_INDEX_DELETED_ID); + zbtIndexSetId(table, bucket, pos, ZBT_INDEX_DELETED_ID); + table->used--; +} + +/* Add one member, reusing a valid lookup position when possible. */ +static void zbtIndexInsert(zbtreeSet *zs, uint32_t hash, uint32_t leaf_id, + const zbtreeInsertPosition *position) +{ + if (zs->member_rehash) { + serverAssert(leaf_id < zs->next_score_leaf_id); + zbtScoreLeaf *leaf = zs->score_leaf_by_id[leaf_id]; + serverAssert(leaf && !ZBT_IS_FREE_LEAF_ID(leaf)); + if (!zbtIndexLeafMigrated(zs, leaf)) { + /* The new member is already in the score leaf, so copying the + * complete leaf inserts its index entry as well. */ + zbtIndexCopyLeaf(zs, leaf); + zbtIndexRehashStep(zs, 1); + return; + } + } + + zbtIndexTable *table = zs->member_rehash ? + &zs->member_rehash->table : &zs->member_index; + zbtIndexBucket *bucket = position ? + (zbtIndexBucket *)position->bucket : NULL; + unsigned int pos = position ? position->pos : 0; + if (position && position->hash == hash && + position->revision == zs->member_revision && bucket && + zbtIndexOwnsBucket(table, bucket) && + pos < ZBT_INDEX_BUCKET_ITEMS) + { + uint8_t oldtag = zbtIndexTags(bucket) >> (pos * 8); + uint32_t oldid = zbtIndexGetId(table, bucket, pos); + if (oldtag == 0 || oldid == ZBT_INDEX_DELETED_ID) { + if (oldtag == 0) table->filled++; + zbtIndexSetTag(bucket, pos, zbtIndexTag(hash)); + zbtIndexSetId(table, bucket, pos, leaf_id); + table->used++; + zbtIndexRecordHomeTag(table, hash); + zbtIndexRehashStep(zs, 1); + return; + } + } + + table = zs->member_rehash ? + &zs->member_rehash->table : &zs->member_index; + zbtIndexTableInsertRaw(table, hash, leaf_id); + zbtIndexRehashStep(zs, 1); +} + +/* ------------------------------ Deletion -------------------------------- */ + +/* Return the bytes used by a leaf's member records. Gaps left by score + * updates are not counted. */ +static size_t zbtScoreLeafRecordBytes(zbtScoreLeaf *leaf) { + uint16_t *offsets = zbtScoreLeafOffsets(leaf); + size_t bytes = 0; + for (unsigned int i = 0; i < leaf->n.count; i++) + bytes += zbtRecordStorageBytes((unsigned char *)leaf + offsets[i]); + return bytes; +} + +/* Return the live bytes excluding growth reserve and allocator rounding. + * The current score encoding can represent every remaining score, so a + * fresh encoding never needs more bits. */ +static size_t zbtScoreLeafLiveBytes(zbtScoreLeaf *leaf) { + return offsetof(zbtScoreLeaf, data) + + zbtScoreLeafScoreBytes(leaf->n.count, leaf->score_bits) + + leaf->n.count * (sizeof(uint16_t) + sizeof(uint8_t)) + + zbtScoreLeafRecordBytes(leaf); +} + +/* Join two adjacent leaves if the result stays within max_count elements and + * max_bytes, returning the merged leaf or NULL. The ID of the larger leaf + * survives. Before releasing the other ID, every member reference to it is + * redirected to the merged leaf. */ +static zbtScoreLeaf *zbtScoreMergePair(zbtreeSet *zs, zbtScoreLeaf *left, + zbtScoreLeaf *right, + unsigned int max_count, + size_t max_bytes) +{ + double scores[ZBT_SCORE_LEAF_MAX]; + zbtBuildElement elems[ZBT_SCORE_LEAF_MAX]; + uint8_t tags[ZBT_SCORE_LEAF_MAX]; + + unsigned int count = left->n.count + right->n.count; + if (count > max_count) return NULL; + /* Records, offsets and tags alone must fit before scores are decoded. */ + if (offsetof(zbtScoreLeaf, data) + + count * (sizeof(uint16_t) + sizeof(uint8_t)) + + zbtScoreLeafRecordBytes(left) + zbtScoreLeafRecordBytes(right) > + max_bytes) + return NULL; + + for (unsigned int i = 0; i < left->n.count; i++) { + scores[i] = zbtScoreLeafScore(left, i); + zbtBuildElementFromLeaf(&elems[i], left, i); + tags[i] = zbtScoreLeafTag(left, i); + } + int merged_was_copied = 0; + for (unsigned int i = 0; i < right->n.count; i++) { + unsigned int dst = left->n.count + i; + scores[dst] = zbtScoreLeafScore(right, i); + zbtBuildElementFromLeaf(&elems[dst], right, i); + tags[dst] = zbtScoreLeafTag(right, i); + } + if (zbtScoreLeafRequestBytes(count, scores, elems) > max_bytes) + return NULL; + + /* Keep the ID of the leaf with more members: each member of the other + * one needs its hash, and inline members have to be hashed again. */ + zbtScoreLeaf *kept = right->n.count > left->n.count ? right : left; + zbtScoreLeaf *gone = kept == left ? right : left; + uint32_t kept_id = kept->id; + uint32_t gone_id = gone->id; + zbtScoreLeaf *merged = zbtScoreLeafBuild(zs, count, scores, elems, + tags, kept_id, 0); + zbtScoreReplaceLeaf(zs, kept, merged); + + for (unsigned int i = 0; i < gone->n.count; i++) { + uint32_t hash = zbtScoreLeafHash(gone, i); + serverAssert(zbtIndexMove(zs, hash, gone_id, kept_id, 0, + &merged_was_copied)); + } + zbtScoreRemoveLeaf(zs, gone); + return merged; +} + +/* Join a small score leaf with its previous or next neighbor when the result + * stays within the supplied limits. Return the merged leaf, or 'leaf' if + * nothing changed. */ +static zbtScoreLeaf *zbtScoreMerge(zbtreeSet *zs, zbtScoreLeaf *leaf, + unsigned int max_count, size_t max_bytes) +{ + zbtScoreLeaf *merged = NULL; + if (leaf->prev) + merged = zbtScoreMergePair(zs, leaf->prev, leaf, max_count, max_bytes); + if (merged == NULL && leaf->next) + merged = zbtScoreMergePair(zs, leaf, leaf->next, max_count, max_bytes); + return merged ? merged : leaf; +} + +/* Remove one packed score and offset. The existing base, shift and width can + * still represent every remaining score, so no full re-encoding is needed. + * compact_records decides whether the member record gap is closed too. */ +static void zbtScoreDeleteInPlace(zbtreeSet *zs, zbtScoreLeaf *leaf, + unsigned int pos, int compact_records) +{ + unsigned int oldcount = leaf->n.count; + unsigned int count = oldcount - 1; + + serverAssert(count != 0 && pos < oldcount); + unsigned int physical = zbtScoreLeafPhysicalPos(leaf, pos); + uint16_t *old_offsets = zbtScoreLeafOffsets(leaf); + uint8_t old_tags[ZBT_SCORE_LEAF_MAX]; + memcpy(old_tags, zbtScoreLeafHashTags(leaf), oldcount); + uint16_t deleted_offset = old_offsets[physical]; + unsigned char *deleted_record = (unsigned char *)leaf + deleted_offset; + size_t deleted_bytes = zbtRecordStorageBytes(deleted_record); + + if (deleted_record[0] == ZBT_PACKED_EXTERNAL) { + unsigned char *external = zbtExternalPointer(deleted_record); + if (compact_records) { + serverAssert(external != NULL); + zbtFreeAllocation(zs, external); + } else { + debugServerAssert(external == NULL); + } + } + + /* Normal deletes close the record gap immediately. A score update is a + * delete followed by an insert, so it leaves the old bytes behind and + * avoids moving every record in the leaf. The next rebuild or compaction + * removes such gaps. + */ + /* The record of the first member of a rebuilt leaf is the last one in + * memory, so closing its gap would move every other record. When no + * record lies above the deleted one, leave the gap at the top instead; + * the next rebuild reclaims it. */ + if (compact_records && deleted_offset != leaf->record_start) { + int above = 0; + for (unsigned int i = 0; i < oldcount && !above; i++) + above = old_offsets[i] > deleted_offset; + if (!above) compact_records = 0; + } + if (compact_records) { + unsigned char *records = (unsigned char *)leaf + leaf->record_start; + memmove(records + deleted_bytes, records, + deleted_offset - leaf->record_start); + leaf->record_start += deleted_bytes; + } else if (deleted_offset == leaf->record_start) { + leaf->record_start += deleted_bytes; + } + + zbtDeleteScoreBits(leaf, physical); + memmove(old_offsets + physical, old_offsets + physical + 1, + (count - physical) * sizeof(uint16_t)); + size_t score_bytes = zbtScoreLeafScoreBytes(count, leaf->score_bits); + uint16_t *offsets = (uint16_t *)(leaf->data + score_bytes); + memmove(offsets, old_offsets, count * sizeof(uint16_t)); + uint8_t *hash_tags = (uint8_t *)(offsets + count); + memcpy(hash_tags, old_tags, physical); + memcpy(hash_tags + physical, old_tags + physical + 1, + count - physical); + + leaf->n.count = count; + leaf->n.subtree = count; + if (compact_records) { + for (unsigned int i = 0; i < count; i++) { + if (offsets[i] < deleted_offset) offsets[i] += deleted_bytes; + } + } + zbtScoreUpdatePath(&leaf->n, -1); +} + +/* Put the smallest score at the physical end before draining the first + * leaf. Reusing the existing reversed encoding avoids shifting all packed + * scores on every pop. Member records and index references stay in place. */ +static void zbtScoreLeafReverseForPop(zbtScoreLeaf *leaf) { + unsigned int count = leaf->n.count; + double scores[ZBT_SCORE_LEAF_MAX]; + for (unsigned int i = 0; i < count; i++) + scores[i] = zbtScoreLeafScore(leaf, i); + uint16_t *offsets = zbtScoreLeafOffsets(leaf); + uint8_t *tags = zbtScoreLeafHashTags(leaf); + for (unsigned int i = 0; i < count / 2; i++) { + unsigned int j = count - i - 1; + uint16_t offset = offsets[i]; offsets[i] = offsets[j]; offsets[j] = offset; + uint8_t tag = tags[i]; tags[i] = tags[j]; tags[j] = tag; + } + /* +0.0 and -0.0 compare equal and are ordered by member, so the last + * logical score need not have the largest ordered value. */ + uint64_t base = zbtScoreToOrdered(scores[0]); + for (unsigned int i = 1; i < count; i++) { + uint64_t ordered = zbtScoreToOrdered(scores[i]); + if (ordered > base) base = ordered; + } + leaf->score_base = base; + leaf->reversed = 1; + memset(leaf->data, 0, zbtScoreLeafScoreBytes(count, leaf->score_bits)); + for (unsigned int i = 0; i < count; i++) { + zbtScoreLeafWriteScore(leaf, i, scores[i]); + debugServerAssert(zbtScoreToOrdered(zbtScoreLeafScore(leaf, i)) == + zbtScoreToOrdered(scores[i])); + } +} + +/* Delete one score-tree element. Score updates disable merging because they + * are about to reinsert the same member; normal deletions also compact the + * records and try to join a small leaf. */ +static void zbtScoreDeleteAt(zbtreeSet *zs, zbtScoreLeaf *leaf, + unsigned int pos, int allow_merge) +{ + if (leaf->n.count == 1) { + zbtScoreRemoveLeaf(zs, leaf); + return; + } + + /* ZPOPMIN, ZPOPMAX and queues keep deleting from one end of the set. + * Merging that edge leaf would copy its neighbor into it only for those + * members to be deleted next, so let it empty and disappear instead. */ + int edge = (pos == 0 && leaf == zs->score_first) || + (pos + 1 == leaf->n.count && leaf == zs->score_last); + if (allow_merge && pos == 0 && leaf == zs->score_first && !leaf->reversed) + zbtScoreLeafReverseForPop(leaf); + zbtScoreDeleteInPlace(zs, leaf, pos, allow_merge); + if (!allow_merge || leaf->n.count > ZBT_SCORE_LEAF_MERGE) return; + if (edge) { + /* Check occasionally while draining an edge, including its final + * member, so stopping a pop sequence cannot leave a large page for + * one member. Powers of two keep the checks amortized. */ + if (leaf->n.count & (leaf->n.count - 1)) return; + } else if (zbtScoreMerge(zs, leaf, ZBT_SCORE_LEAF_MAX, + ZBT_SCORE_LEAF_BYTES) != leaf) + { + return; + } + + /* Deleting never shrinks the allocation. When a small leaf cannot join + * a neighbor, rebuild it at the real size once at least half of the + * allocation is unused, so sparse leaves return their memory. */ + if (zmalloc_usable_size(leaf) >= + 2 * zbtScoreLeafLiveBytes(leaf) + ZBT_SCORE_LEAF_RESERVE) + zbtScoreCompactLeaf(zs, leaf); +} + +/* Remove a run of consecutive elements from one score leaf. Their member + * entries are found by hash and leaf number, so no member string has to be + * copied and no second full lookup is needed. + * + * A run longer than one element rebuilds the leaf once instead of shifting + * its packed scores, offsets and records for every removed element. That is + * what makes a rank range deletion cost one leaf pass rather than one pass + * per element. + */ +static void zbtScoreDeleteRun(zbtreeSet *zs, zbtScoreLeaf *score_leaf, + unsigned int score_pos, unsigned int n, + int update_member_index, int allow_merge) +{ + unsigned int count = score_leaf->n.count; + serverAssert(n >= 1 && score_pos + n <= count); + int edge = (score_pos == 0 && score_leaf == zs->score_first) || + (score_pos + n == count && score_leaf == zs->score_last); + + if (update_member_index) { + for (unsigned int i = 0; i < n; i++) { + uint32_t hash = zbtScoreLeafHash(score_leaf, score_pos + i); + zbtIndexBucket *bucket; + unsigned int bucket_pos; + + serverAssert(zbtIndexFindReference(zs, hash, score_leaf->id, + &bucket, &bucket_pos)); + zbtIndexDeleteAt(zs, bucket, bucket_pos); + } + } + zs->length -= n; + + if (n == 1) { + if (allow_merge) { + zbtScoreDeleteAt(zs, score_leaf, score_pos, 1); + } else if (count == 1) { + zbtScoreRemoveLeaf(zs, score_leaf); + } else { + zbtScoreDeleteInPlace(zs, score_leaf, score_pos, 1); + } + return; + } + if (n == count) { + zbtScoreRemoveLeaf(zs, score_leaf); + return; + } + + double scores[ZBT_SCORE_LEAF_MAX]; + zbtBuildElement elems[ZBT_SCORE_LEAF_MAX]; + uint8_t tags[ZBT_SCORE_LEAF_MAX]; + unsigned int kept = 0; + + for (unsigned int i = 0; i < count; i++) { + if (i == score_pos) { + i += n - 1; + continue; + } + scores[kept] = zbtScoreLeafScore(score_leaf, i); + zbtBuildElementFromLeaf(&elems[kept], score_leaf, i); + tags[kept] = zbtScoreLeafTag(score_leaf, i); + kept++; + } + serverAssert(kept == count - n); + + zbtScoreLeaf *newleaf = zbtScoreLeafBuild(zs, kept, scores, elems, tags, + score_leaf->id, 0); + zbtScoreReplaceLeaf(zs, score_leaf, newleaf); + /* Like single-member pops, a range draining an end need not pull + * neighboring members into the leaf that will be emptied next. */ + if (allow_merge && !edge && kept <= ZBT_SCORE_LEAF_MERGE) + zbtScoreMerge(zs, newleaf, ZBT_SCORE_LEAF_MAX, ZBT_SCORE_LEAF_BYTES); +} + +/* ------------------------- Creation and release ------------------------- */ + +/* Free a tree recursively. Tree height is small because inner pages have high + * fanout, so recursion depth is bounded by the practical set size. */ +static void zbtScoreFreeTree(zbtreeSet *zs, zbtScoreNode *node) { + if (node->isleaf) { + zbtScoreLeafFreeOwned(zs, (zbtScoreLeaf *)node); + return; + } + + zbtScoreInner *inner = (zbtScoreInner *)node; + for (unsigned int i = 0; i < inner->n.count; i++) + zbtScoreFreeTree(zs, inner->child[i]); + zbtFreeAllocation(zs, node); +} + +/* Create an empty B+ tree sorted set. The first score leaf and member-index + * table are allocated lazily on insertion. */ +zbtreeSet *zbtreeCreate(void) { + size_t usable; + zbtreeSet *zs = zmalloc_usable(sizeof(*zs), &usable); + memset(zs, 0, sizeof(*zs)); + zs->member_revision = zbtIndexNextRevision(); + zs->alloc_size = usable; + return zs; +} + +/* Size the member index of a fresh set for 'count' insertions at once, so + * that building it starts no incremental resize. Only a set that never held + * a member is presized: a set emptied by deletion keeps its leaf ID counter, + * which decides the width of its next table. Other sets succeed unchanged. + * + * A leaf holds at most ZBT_SCORE_LEAF_MAX members, so a count that needs at + * least ZBT_INDEX_WIDE_ID_AT leaves will widen the table while the set is + * built. Such a table starts wide instead, which avoids copying it again. + * + * The count check keeps zbtIndexBucketsForElements() from overflowing before + * the count is narrowed to unsigned long; zbtIndexTableAlloc() checks the + * bucket bytes. A fallible call returns zero when the count is too large or + * the table cannot be allocated, and then leaves the set unchanged. */ +static int zbtIndexReserve(zbtreeSet *zs, uint64_t count, int fallible) { + if (zs->length != 0 || zs->member_index.size != 0 || + zs->next_score_leaf_id != 0 || count == 0) + return 1; + if (count > (ULONG_MAX - (ZBT_INDEX_MAX_LOAD_NUM - 1)) / + ZBT_INDEX_MAX_LOAD_DEN) + return 0; + int wide_ids = count / ZBT_SCORE_LEAF_MAX >= ZBT_INDEX_WIDE_ID_AT; + if (!zbtIndexTableAlloc(zs, &zs->member_index, + zbtIndexBucketsForElements((unsigned long)count), + wide_ids, fallible)) + return 0; + zs->member_revision = zbtIndexNextRevision(); + return 1; +} + +/* Presize a fresh set for a count the server computed itself. Out of memory + * aborts, so never pass a length read from a client or an RDB stream: use + * zbtreeTryReserve() for those. */ +void zbtreeReserve(zbtreeSet *zs, unsigned long count) { + zbtIndexReserve(zs, count, 0); +} + +/* Presize a fresh set for an untrusted count, such as an RDB length. Return + * zero, leaving the set unchanged, when the table cannot be allocated. */ +int zbtreeTryReserve(zbtreeSet *zs, uint64_t count) { + return zbtIndexReserve(zs, count, 1); +} + +/* Copy a score tree while repairing every pointer between its pages. Leaves + * are visited in score order, which also rebuilds their linked list and the + * leaf ID table. Inline records can be copied unchanged. External members + * need their own allocation so the two sorted sets remain independent. */ +static zbtScoreNode *zbtScoreDupTree(zbtreeSet *dst, + const zbtScoreNode *source, + zbtScoreInner *parent, + unsigned int parent_index) +{ + size_t bytes = source->isleaf ? zmalloc_usable_size((void *)source) : + sizeof(zbtScoreInner); + zbtScoreNode *copy = zbtAlloc(dst, bytes); + memcpy(copy, source, bytes); + copy->parent = parent; + copy->parent_index = parent_index; + + if (copy->isleaf) { + zbtScoreLeaf *leaf = (zbtScoreLeaf *)copy; + if (leaf->n.has_external) { + for (unsigned int i = 0; i < leaf->n.count; i++) { + unsigned char *record = zbtScoreLeafRecord(leaf, i); + if (record[0] != ZBT_PACKED_EXTERNAL) continue; + size_t len = zbtExternalLength(record); + unsigned char *oldptr = zbtExternalPointer(record); + serverAssert(oldptr != NULL); + zbtExternalSetPointer(record, zbtExternalDup(dst, oldptr, len)); + } + } + leaf->prev = dst->score_last; + leaf->next = NULL; + if (dst->score_last) dst->score_last->next = leaf; + else dst->score_first = leaf; + dst->score_last = leaf; + dst->score_leaf_by_id[leaf->id] = leaf; + } else { + zbtScoreInner *inner = (zbtScoreInner *)copy; + const zbtScoreInner *source_inner = + (const zbtScoreInner *)source; + for (unsigned int i = 0; i < inner->n.count; i++) + inner->child[i] = zbtScoreDupTree(dst, source_inner->child[i], + inner, i); + } + return copy; +} + +/* The member tables contain leaf IDs rather than page pointers. Once the + * score tree has been copied, duplicating a table is just a byte copy. */ +static void zbtIndexDupTable(zbtreeSet *dst, zbtIndexTable *copy, + const zbtIndexTable *source) +{ + *copy = *source; + if (source->buckets) { + size_t bytes = zbtIndexTableBytes(source); + copy->buckets = zbtAlloc(dst, bytes); + memcpy(copy->buckets, source->buckets, bytes); + } +} + +/* Duplicate an entire set without decoding and reinserting its members. */ +zbtreeSet *zbtreeDup(const zbtreeSet *source) { + zbtreeSet *copy = zbtreeCreate(); + + copy->score_leaf_cap = source->score_leaf_cap; + copy->next_score_leaf_id = source->next_score_leaf_id; + copy->free_score_leaf_id = source->free_score_leaf_id; + copy->length = source->length; + + if (source->score_leaf_by_id) { + size_t bytes = source->score_leaf_cap * + sizeof(*source->score_leaf_by_id); + copy->score_leaf_by_id = zbtAlloc(copy, bytes); + /* Free-list links contain IDs, not addresses, so they can be copied. + * Live addresses are replaced while the score tree is duplicated. */ + memcpy(copy->score_leaf_by_id, source->score_leaf_by_id, bytes); + } + if (source->score_root) + copy->score_root = zbtScoreDupTree(copy, source->score_root, NULL, 0); + + zbtIndexDupTable(copy, ©->member_index, &source->member_index); + if (source->member_rehash) { + copy->member_rehash = zbtAlloc(copy, sizeof(*copy->member_rehash)); + *copy->member_rehash = *source->member_rehash; + zbtIndexDupTable(copy, ©->member_rehash->table, + &source->member_rehash->table); + } + copy->member_revision = zbtIndexNextRevision(); + return copy; +} + +/* Free the score tree, member index, leaf ID table, and set header. Score + * leaves also release the long members they own. */ +void zbtreeFree(zbtreeSet *zs) { + size_t self; + if (zs->score_root) zbtScoreFreeTree(zs, zs->score_root); + zbtIndexTableRelease(zs, &zs->member_index); + if (zs->member_rehash) { + zbtIndexTableRelease(zs, &zs->member_rehash->table); + zbtFreeAllocation(zs, zs->member_rehash); + } + if (zs->score_leaf_by_id) zbtFreeAllocation(zs, zs->score_leaf_by_id); + /* Every allocation but the header has been returned. Leak checkers are + * told to ignore external members (see zbtExternalIgnoreLeak()), so this + * is what catches one that was dropped without being freed. */ + size_t leftover = zs->alloc_size; + zfree_usable(zs, &self); + debugServerAssert(leftover == self); + UNUSED(leftover); +} + +/* Return the number of members. */ +unsigned long zbtreeLength(const zbtreeSet *zs) { + return zs->length; +} + +/* Return allocator usable bytes owned by this sorted set. */ +size_t zbtreeAllocSize(const zbtreeSet *zs) { + return zs->alloc_size; +} + +/* The fork child no longer needs a sorted set after writing it. Release its + * page-sized allocations there so later writes in the parent do not copy + * those pages. Sub-page ranges are skipped by zmadvise_dontneed(). */ +void zbtreeDismissMemory(zbtreeSet *zs) { + if (zs->member_index.buckets) + dismissMemory(zs->member_index.buckets, + zbtIndexTableBytes(&zs->member_index)); + if (zs->member_rehash && zs->member_rehash->table.buckets) + dismissMemory(zs->member_rehash->table.buckets, + zbtIndexTableBytes(&zs->member_rehash->table)); + if (zs->score_leaf_by_id) + dismissMemory(zs->score_leaf_by_id, + zs->score_leaf_cap * sizeof(*zs->score_leaf_by_id)); + + zbtScoreLeaf *leaf = zs->score_first; + while (leaf) { + zbtScoreLeaf *next = leaf->next; + if (leaf->n.has_external) { + for (unsigned int i = 0; i < leaf->n.count; i++) { + unsigned char *record = zbtScoreLeafRecord(leaf, i); + if (record[0] == ZBT_PACKED_EXTERNAL) + dismissMemory(zbtExternalPointer(record), 0); + } + } + dismissMemory(leaf, 0); + leaf = next; + } +} + +/* ----------------------- Active defragmentation ------------------------- */ + +/* Active defragmentation may return a new address for any allocation. Moving + * a tree page is therefore more than copying bytes: every pointer to the old + * address must be changed before the old allocation is released. + * + * A score leaf is named by its parent, two leaf neighbors, the set's edge + * pointers, and its stable ID table entry. An inner page is named by its + * parent and by every child's parent pointer. zbtDefragScoreNode() repairs + * these references. + * + * The incremental scan walks linked leaves. External members are visited one + * allocation per call. After the last member of the last child below an inner + * page, the scan moves that page and continues upward while the same condition + * holds. This visits each page once during an uninterrupted pass without + * keeping a stack. + * + * Writes may happen between time-limited slices. The cursor is a stable leaf + * ID, never a possibly stale leaf pointer. This is a best-effort scan: a write + * can cause repeated work, but cannot make the next slice dereference a freed + * page. Hash-table allocations contain no pointers into themselves and are + * moved once before this incremental page walk starts. + */ + +/* Keep exact allocation accounting correct when the allocator moves an object + * to a size class with a different usable size. */ +static void zbtDefragAdjustSize(zbtreeSet *zs, size_t old_size, void *newptr) { + size_t new_size = zmalloc_usable_size(newptr); + if (new_size >= old_size) + zs->alloc_size += new_size - old_size; + else + zs->alloc_size -= old_size - new_size; +} + +/* Move an allocation which has no pointers back to its own address. */ +static void *zbtDefragPlainAllocation( + zbtreeSet *zs, void *ptr, zbtreeDefragAllocFunction *defrag_alloc) +{ + size_t old_size = zmalloc_usable_size(ptr); + void *newptr = defrag_alloc(ptr); + if (newptr == NULL) return ptr; + zbtDefragAdjustSize(zs, old_size, newptr); + return newptr; +} + +/* The header, leaf table, and hash-table allocations have no pointers back to + * their old addresses, so they can move before the page scan starts. */ +zbtreeSet *zbtreeDefragStart(zbtreeSet *zs, + zbtreeDefragAllocFunction *defrag_alloc) +{ + size_t old_size = zmalloc_usable_size(zs); + zbtreeSet *newzs = defrag_alloc(zs); + if (newzs) { + zs = newzs; + zbtDefragAdjustSize(zs, old_size, zs); + } + + if (zs->score_leaf_by_id) + zs->score_leaf_by_id = zbtDefragPlainAllocation( + zs, zs->score_leaf_by_id, defrag_alloc); + if (zs->member_index.buckets) + zs->member_index.buckets = zbtDefragPlainAllocation( + zs, zs->member_index.buckets, defrag_alloc); + if (zs->member_rehash) { + zs->member_rehash = zbtDefragPlainAllocation( + zs, zs->member_rehash, defrag_alloc); + if (zs->member_rehash->table.buckets) + zs->member_rehash->table.buckets = zbtDefragPlainAllocation( + zs, zs->member_rehash->table.buckets, defrag_alloc); + } + return zs; +} + +/* Forget an incremental scan. The next call starts at the first score leaf. */ +void zbtreeDefragStateReset(zbtreeDefragState *state) { + memset(state, 0, sizeof(*state)); +} + +/* Move one score page after saving all the places that refer to it. Leaves + * are referenced by their neighbors and leaf ID table. Inner pages are + * referenced by their parent and their children. */ +static zbtScoreNode *zbtDefragScoreNode( + zbtreeSet *zs, zbtScoreNode *node, + zbtreeDefragAllocFunction *defrag_alloc, + zbtreeDefragFreeFunction *defrag_free) +{ + size_t old_size = zmalloc_usable_size(node); + zbtScoreNode *newnode = defrag_alloc(node); + if (newnode == NULL) return node; + serverAssert(newnode != node); + zbtDefragAdjustSize(zs, old_size, newnode); + + if (newnode->parent) + newnode->parent->child[newnode->parent_index] = newnode; + else + zs->score_root = newnode; + + if (newnode->isleaf) { + zbtScoreLeaf *leaf = (zbtScoreLeaf *)newnode; + if (leaf->prev) leaf->prev->next = leaf; + else zs->score_first = leaf; + if (leaf->next) leaf->next->prev = leaf; + else zs->score_last = leaf; + zs->score_leaf_by_id[leaf->id] = leaf; + } else { + zbtScoreInner *inner = (zbtScoreInner *)newnode; + for (unsigned int i = 0; i < inner->n.count; i++) + inner->child[i]->parent = inner; + } + + defrag_free(node); + return newnode; +} + +/* Move one external member and repair its record. Keeping this separate from + * the leaf move lets the incremental scan stop after one possibly large copy. */ +static void zbtDefragExternalMember( + zbtreeSet *zs, unsigned char *record, + zbtreeDefragAllocFunction *defrag_alloc, + zbtreeDefragFreeFunction *defrag_free) +{ + unsigned char *oldptr = zbtExternalPointer(record); + serverAssert(oldptr != NULL); + size_t old_size = zmalloc_usable_size(oldptr); + unsigned char *newptr = defrag_alloc(oldptr); + if (newptr) { + zbtExternalIgnoreLeak(newptr); + zbtDefragAdjustSize(zs, old_size, newptr); + zbtExternalSetPointer(record, newptr); + defrag_free(oldptr); + } +} + +/* Move each inner page whose last leaf was reached. This visits every inner + * page once during an uninterrupted ordered scan. */ +static void zbtDefragFinishedLeaf( + zbtreeSet *zs, zbtScoreLeaf *leaf, + zbtreeDefragAllocFunction *defrag_alloc, + zbtreeDefragFreeFunction *defrag_free) +{ + zbtScoreNode *node = &leaf->n; + while (node->parent) { + zbtScoreInner *parent = node->parent; + if (node->parent_index + 1 != parent->n.count) break; + node = zbtDefragScoreNode(zs, &parent->n, + defrag_alloc, defrag_free); + } +} + +/* Resolve a saved leaf ID, such as a defrag cursor. Deleted IDs contain a + * marked free-list link rather than a live leaf and must be treated as + * missing. */ +static zbtScoreLeaf *zbtScoreLeafById(zbtreeSet *zs, uint32_t id) { + if (id >= zs->next_score_leaf_id) return NULL; + zbtScoreLeaf *leaf = zs->score_leaf_by_id[id]; + if (leaf == NULL || ZBT_IS_FREE_LEAF_ID(leaf)) return NULL; + serverAssert(leaf->id == id); + return leaf; +} + +/* Defrag one leaf incrementally. A call moves the leaf page if needed, scans + * its inline records, and moves at most one external member. After the final + * member it moves the inner pages which end at this leaf. The next call resumes + * through a stable leaf ID and a small position, never through a page pointer. */ +int zbtreeDefragStep(zbtreeSet *zs, zbtreeDefragState *state, + zbtreeDefragAllocFunction *defrag_alloc, + zbtreeDefragFreeFunction *defrag_free, + unsigned long *scanned) +{ + *scanned = 0; + if (state->tree != zs) { + zbtreeDefragStateReset(state); + state->tree = zs; + state->score_leaf_id = + zs->score_first ? zs->score_first->id : ZBT_NO_LEAF_ID; + } + + while (state->score_leaf_id != ZBT_NO_LEAF_ID) { + zbtScoreLeaf *leaf = + zbtScoreLeafById(zs, state->score_leaf_id); + if (leaf == NULL) { + /* A write removed the saved leaf between defrag slices. New + * allocations do not need defrag, so restarting is sufficient. */ + state->score_leaf_id = + zs->score_first ? zs->score_first->id : ZBT_NO_LEAF_ID; + state->score_pos = 0; + state->leaf_moved = 0; + continue; + } + + if (!state->leaf_moved) { + leaf = (zbtScoreLeaf *)zbtDefragScoreNode( + zs, &leaf->n, defrag_alloc, defrag_free); + state->leaf_moved = 1; + } + + while (state->score_pos < leaf->n.count) { + unsigned int pos = state->score_pos++; + unsigned char *record = zbtScoreLeafRecord(leaf, pos); + (*scanned)++; + if (record[0] != ZBT_PACKED_EXTERNAL) continue; + zbtDefragExternalMember(zs, record, defrag_alloc, defrag_free); + (*scanned)++; + return 0; + } + + uint32_t next_id = + leaf->next ? leaf->next->id : ZBT_NO_LEAF_ID; + zbtDefragFinishedLeaf(zs, leaf, defrag_alloc, defrag_free); + state->score_leaf_id = next_id; + state->score_pos = 0; + state->leaf_moved = 0; + return 0; + } + return 1; +} + +/* -------------------------- Lookup and update --------------------------- */ + +/* Find a member supplied as raw bytes, without making a temporary sds. */ +int zbtreeScoreRaw(zbtreeSet *zs, const unsigned char *ele, + size_t elelen, double *score) +{ + uint32_t hash = (uint32_t)dictGenHashFunction(ele, elelen); + return zbtIndexFind(zs, hash, ele, elelen, score, + NULL, NULL, NULL, NULL); +} + +/* Find an sds member and optionally return its score. */ +int zbtreeScore(zbtreeSet *zs, sds ele, double *score) { + return zbtreeScoreRaw(zs, (unsigned char *)ele, sdslen(ele), score); +} + +/* ZMSCORE can bring several independent index buckets into cache before it + * reads any of their members. Hash once, then reuse the normal lookup. */ +unsigned int zbtreeScores(zbtreeSet *zs, const sds *members, unsigned int count, + double *scores) +{ + serverAssert(count <= ZBT_SCORE_BATCH_SIZE); + uint32_t hashes[ZBT_SCORE_BATCH_SIZE]; + for (unsigned int i = 0; i < count; i++) { + hashes[i] = (uint32_t)dictGenHashFunction(members[i], sdslen(members[i])); + if (zs->member_index.size) + redis_prefetch_read(zbtIndexBucketAt(&zs->member_index, + hashes[i] & (zs->member_index.size - 1))); + if (zs->member_rehash) + redis_prefetch_read(zbtIndexBucketAt(&zs->member_rehash->table, + hashes[i] & (zs->member_rehash->table.size - 1))); + } + unsigned int found = 0; + for (unsigned int i = 0; i < count; i++) { + if (zbtIndexFind(zs, hashes[i], (unsigned char *)members[i], + sdslen(members[i]), &scores[i], NULL, NULL, NULL, NULL)) + found |= 1U << i; + } + return found; +} + +/* Lookup used by ZADD. In addition to the normal result, save the member hash + * and insertion position found on a miss. zbtreeInsertNew() can reuse it and + * avoid a second index lookup. The saved revision makes stale positions + * harmless if score insertion happened to advance a table resize first. + */ +int zbtreeFindForAdd(zbtreeSet *zs, sds ele, double *score, + zbtreeInsertPosition *position) +{ + uint32_t hash = (uint32_t)dictGenHashFunction(ele, sdslen(ele)); + zbtIndexBucket *bucket = NULL; + zbtScoreLeaf *score_leaf = NULL; + unsigned int pos = 0; + unsigned int score_pos = 0; + int found = zbtIndexFind(zs, hash, (unsigned char *)ele, + sdslen(ele), score, + &score_leaf, &score_pos, &bucket, &pos); + position->bucket = bucket; + position->pos = pos; + position->score_leaf = score_leaf; + position->score_pos = score_pos; + position->hash = hash; + position->revision = zs->member_revision; + return found; +} + +/* Compute a zero-based rank by adding complete left siblings while walking + * from a score leaf to the root. */ +static unsigned long zbtScoreRank(zbtScoreLeaf *leaf, unsigned int pos) { + unsigned long rank = pos; + zbtScoreNode *node = &leaf->n; + while (node->parent) { + zbtScoreInner *parent = node->parent; + for (unsigned int i = 0; i < node->parent_index; i++) + rank += parent->child_count[i]; + node = &parent->n; + } + return rank; +} + +/* Return the zero-based rank of a member, or -1 if it does not exist. + * 'reverse' counts from the greatest element instead. */ +long zbtreeRank(zbtreeSet *zs, sds ele, int reverse, double *score) { + uint32_t hash = (uint32_t)dictGenHashFunction(ele, sdslen(ele)); + zbtScoreLeaf *leaf; + unsigned int pos; + + if (!zbtIndexFind(zs, hash, (unsigned char *)ele, + sdslen(ele), score, &leaf, &pos, NULL, NULL)) + return -1; + + unsigned long rank = zbtScoreRank(leaf, pos); + return reverse ? zs->length - rank - 1 : rank; +} + +/* Insert a member that is known to be absent. Bulk builders pass the bytes + * straight out of another structure, so no temporary sds is created. */ +void zbtreeInsertNewRaw(zbtreeSet *zs, double score, const unsigned char *ele, + size_t elelen, const zbtreeInsertPosition *position) +{ + uint32_t hash = position ? position->hash : + (uint32_t)dictGenHashFunction(ele, elelen); + debugServerAssert(hash == (uint32_t)dictGenHashFunction(ele, elelen)); + zbtIndexExpandIfNeeded(zs, 1); + zbtBuildElement element; + zbtBuildElementFromBytes(&element, ele, elelen, hash); + zbtScoreLeaf *leaf = zbtScoreInsert(zs, score, &element); + zs->length++; + zbtIndexInsert(zs, hash, leaf->id, position); +} + +/* Insert an absent SDS member without taking ownership of its bytes. */ +void zbtreeInsertNew(zbtreeSet *zs, double score, sds ele, + const zbtreeInsertPosition *position) +{ + zbtreeInsertNewRaw(zs, score, (unsigned char *)ele, sdslen(ele), + position); +} + +/* Like zbtreeInsertNewRaw(), but for a caller building a tree from a source + * that is already fully sorted (a listpack or another zset's skiplist) and + * inserting members it knows are absent -- zsetConvertAndExpand()'s two + * BTREE-target loops, the only intended callers. Skips the top-down leaf + * search entirely via zbtScoreInsertAppend(); everything else (index + * expand/insert, length bump) matches zbtreeInsertNewRaw() exactly. */ +void zbtreeInsertNewAppend(zbtreeSet *zs, double score, + const unsigned char *ele, size_t elelen) +{ + zbtreeInsertNewAppendWithHash(zs, score, ele, elelen, + (uint32_t)dictGenHashFunction(ele, elelen)); +} + +/* Same as zbtreeInsertNewAppend(), but for a caller that already computed + * this exact member's hash for another purpose (e.g. a dict lookup using the + * same SipHash-based dictGenHashFunction()) and wants to avoid paying for it + * twice. The hash is not validated against ele/elelen -- passing a mismatched + * hash silently corrupts the member index. */ +void zbtreeInsertNewAppendWithHash(zbtreeSet *zs, double score, + const unsigned char *ele, size_t elelen, + uint32_t hash) +{ + zbtIndexExpandIfNeeded(zs, 1); + zbtBuildElement element; + zbtBuildElementFromBytes(&element, ele, elelen, hash); + zbtScoreLeaf *leaf = zbtScoreInsertAppend(zs, score, &element); + zs->length++; + zbtIndexInsert(zs, hash, leaf->id, NULL); +} + +/* zbtreeInsertNewAppend() calls zbtScoreInsertFast() once per element, which + * pays an O(leaf size) cost on every single insert: the offsets array grows + * by one entry each time, so the hash-tag array that follows it in the + * packed layout must shift forward and gets fully snapshotted and rebuilt on + * every call -- O(leaf size^2) total to fill one leaf. zbtreeAppendBatch + * buffers up to one leaf's worth of elements and flushes them with a single + * zbtScoreLeafBuild() call (the same single-pass builder zbtScoreCompactLeaf() + * and zbtScoreInsertEdgeLeaf() already trust) instead. + * + * Same precondition as zbtreeInsertNewAppend(): every added element must sort + * after everything already in the tree and be known absent. A batch never + * spans more than one leaf -- Add() flushes first whenever adding the next + * element could exceed either ZBT_SCORE_LEAF_MAX or a conservative (always + * >= the real zbtScoreLeafRequestBytes()) byte estimate, so every flush is a + * single build-and-splice, reusing the append-only edge-leaf machinery + * unmodified rather than needing a multi-leaf split path. */ +struct zbtreeAppendBatch { + zbtreeSet *zs; + unsigned int count; + size_t conservative_bytes; + double scores[ZBT_SCORE_LEAF_MAX]; + uint8_t tags[ZBT_SCORE_LEAF_MAX]; + zbtBuildElement elems[ZBT_SCORE_LEAF_MAX]; + unsigned char scratch[ZBT_SCORE_LEAF_BYTES]; + size_t scratch_used; +}; + +/* offsetof(zbtScoreLeaf, data) for the leaf header, plus slack for the + * from-scratch reserve zbtScoreLeafBuild() may add for a non-edge leaf. */ +#define ZBT_BATCH_HEADER_SLACK (offsetof(zbtScoreLeaf, data) + 16) +/* Per item: 8 bytes is an explicit upper bound for one packed score (see the + * comment at zbtScoreEncoding()'s bit-width selection), plus the offset + * (uint16_t) and hash tag (uint8_t) every element always carries. */ +#define ZBT_BATCH_ITEM_SLACK (sizeof(double) + sizeof(uint16_t) + sizeof(uint8_t)) +/* How many elements ahead to software-prefetch the member-index insert's + * target bucket. A few iterations' worth of other per-element work (the + * insert itself, its own probe-chain walk) needs to happen between issuing + * a prefetch and consuming it for the prefetch to have finished landing; + * too far ahead risks the cache line being evicted again before use on a + * table this size. Not tuned beyond this order-of-magnitude judgment call. */ +#define ZBT_INDEX_PREFETCH_AHEAD 4 + +zbtreeAppendBatch *zbtreeAppendBatchCreate(zbtreeSet *zs) { + zbtreeAppendBatch *b = zcalloc(sizeof(*b)); + b->zs = zs; + return b; +} + +/* Build and splice one leaf from everything buffered so far. A no-op on an + * empty batch so Finish() can call this unconditionally. */ +static void zbtreeAppendBatchFlush(zbtreeAppendBatch *b) { + if (b->count == 0) return; + zbtreeSet *zs = b->zs; + zbtIndexExpandIfNeeded(zs, b->count); + + /* edge_page=0: unlike zbtScoreInsertEdgeLeaf()'s single-element edge + * leaves (which ask for the full ZBT_SCORE_LEAF_BYTES up front because + * they expect many more zbtScoreInsertFast() calls to grow them + * in-place), a batch-built leaf is complete and final the moment + * zbtScoreLeafBuild() returns -- nothing ever appends into it + * incrementally afterward. Asking for the full page here would only + * waste memory and make zmalloc_usable_size(oldleaf) >= ZBT_SCORE_LEAF_BYTES + * true on essentially every subsequent flush, triggering a needless + * compact-and-rebuild of a leaf that was never actually over-allocated + * relative to its own real content. */ + zbtScoreLeaf *newleaf; + if (zs->score_root == NULL) { + newleaf = zbtScoreLeafBuild(zs, b->count, b->scores, b->elems, b->tags, + ZBT_NEW_LEAF_ID, 0); + zs->score_root = &newleaf->n; + zs->score_first = zs->score_last = newleaf; + } else { + /* Mirrors zbtScoreInsertEdgeLeaf()'s append (non-prepend) branch for + * the splice/wiring shape: compact an over-allocated old edge leaf + * first (using its own, possibly-reassigned return value from here + * on -- kept as a defensive check even though it should rarely fire + * for a tightly-allocated batch leaf), build the new leaf, then wire + * the doubly-linked leaf list by hand before handing off to + * zbtScoreInsertSibling() for the tree-side splice -- that function + * only ever touches zbtScoreNode fields (parent/child/rollup), never + * ->prev/->next/score_first/score_last. */ + zbtScoreLeaf *oldleaf = zs->score_last; + if (zmalloc_usable_size(oldleaf) >= ZBT_SCORE_LEAF_BYTES) + oldleaf = zbtScoreCompactLeaf(zs, oldleaf); + newleaf = zbtScoreLeafBuild(zs, b->count, b->scores, b->elems, b->tags, + ZBT_NEW_LEAF_ID, 0); + newleaf->prev = oldleaf; + newleaf->next = oldleaf->next; + oldleaf->next = newleaf; + zs->score_last = newleaf; + zbtScoreInsertSibling(zs, &oldleaf->n, &newleaf->n); + } + /* The leaf is complete, but none of its members have index entries yet. + * Keep its initial index_resize (zero, never an active resize ID) so a + * rehash copies the whole leaf once. Publish its length before indexing: + * zbtIndexInsert() can finish the rehash and check the total member count. */ + zs->length += b->count; + if (zs->member_rehash) { + zbtIndexInsert(zs, b->elems[0].hash, newleaf->id, NULL); + } else { + /* zbtIndexTableInsertRaw()'s bucket lookup is a near-guaranteed + * cache miss on a member_index too large to fit in cache. Every + * element's hash is already materialized in b->elems[] (a free read, + * unlike re-deriving a hash from a leaf's packed bytes elsewhere in + * this file), so prefetch a few elements ahead to hide that miss + * behind the current iteration's own work. Re-derives the live + * target table exactly as zbtIndexInsert() does; always safe to + * call on a stale guess (a rehash starting or advancing mid-batch) + * since a prefetch is a hint, never a memory access -- a wrong + * address is merely a wasted prefetch, not a correctness hazard. */ + for (unsigned int i = 0; i < b->count; i++) { + unsigned int ahead = i + ZBT_INDEX_PREFETCH_AHEAD; + if (ahead < b->count) { + zbtIndexTable *table = zs->member_rehash ? + &zs->member_rehash->table : &zs->member_index; + if (table->size > 0) { + unsigned long idx = b->elems[ahead].hash & (table->size - 1); + redis_prefetch_write(zbtIndexBucketAt(table, idx)); + } + } + zbtIndexInsert(zs, b->elems[i].hash, newleaf->id, NULL); + } + } + b->count = 0; + b->conservative_bytes = 0; + b->scratch_used = 0; +} + +void zbtreeAppendBatchAddWithHash(zbtreeAppendBatch *b, double score, + const unsigned char *ele, size_t elelen, + uint32_t hash) +{ + zbtBuildElement probe; + zbtBuildElementFromBytes(&probe, ele, elelen, hash); + size_t item_cost = zbtElementStorageBytes(&probe) + ZBT_BATCH_ITEM_SLACK; + + if (b->count == ZBT_SCORE_LEAF_MAX || + (b->count > 0 && ZBT_BATCH_HEADER_SLACK + b->conservative_bytes + + item_cost > ZBT_SCORE_LEAF_BYTES)) + zbtreeAppendBatchFlush(b); + + unsigned int i = b->count; + b->scores[i] = score; + b->tags[i] = (uint8_t)(hash >> 24); + + if (zbtBuildElementIsExternal(&probe)) { + /* A large member's bytes never go through the shared inline scratch + * arena -- that arena is only sized for the leaf's *inline* byte + * budget. Instead, allocate this member's own external storage right + * now, using the exact mechanism zbtBuildElementWrite() already uses + * for a fresh external member with no prior allocation (see its + * "external == NULL" branch). Pre-setting ->external here makes that + * later call adopt this allocation instead of copying again. */ + unsigned char *external = zbtExternalDup(b->zs, ele, elelen); + b->elems[i].ptr = external; + b->elems[i].len = elelen; + b->elems[i].hash = hash; + b->elems[i].external = external; + b->elems[i].owner_record = NULL; + } else { + memcpy(b->scratch + b->scratch_used, ele, elelen); + zbtBuildElementFromBytes(&b->elems[i], b->scratch + b->scratch_used, + elelen, hash); + b->scratch_used += elelen; + } + b->conservative_bytes += item_cost; + b->count = i + 1; +} + +void zbtreeAppendBatchAdd(zbtreeAppendBatch *b, double score, + const unsigned char *ele, size_t elelen) +{ + zbtreeAppendBatchAddWithHash(b, score, ele, elelen, + (uint32_t)dictGenHashFunction(ele, elelen)); +} + +/* Flush whatever remains and free the batch. The batch is unusable after + * this returns. */ +void zbtreeAppendBatchFinish(zbtreeAppendBatch *b) { + zbtreeAppendBatchFlush(b); + zfree(b); +} + +/* Delete a member, returning one if it existed and zero otherwise. */ +int zbtreeDelete(zbtreeSet *zs, sds ele) { + uint32_t hash = (uint32_t)dictGenHashFunction(ele, sdslen(ele)); + zbtScoreLeaf *score_leaf; + zbtIndexBucket *bucket; + unsigned int score_pos, bucket_pos; + + if (!zbtIndexFind(zs, hash, (unsigned char *)ele, sdslen(ele), NULL, + &score_leaf, &score_pos, &bucket, &bucket_pos)) + return 0; + zbtIndexDeleteAt(zs, bucket, bucket_pos); + zbtScoreDeleteAt(zs, score_leaf, score_pos, 1); + zs->length--; + zbtIndexShrinkIfNeeded(zs); + return 1; +} + +/* Delete an already selected end member without looking it up by name. */ +void zbtreeDeleteEdge(zbtreeSet *zs, int reverse) { + zbtIndexRehashStep(zs, 1); + zbtScoreLeaf *leaf = reverse ? zs->score_last : zs->score_first; + serverAssert(leaf != NULL); + zbtScoreDeleteRun(zs, leaf, reverse ? leaf->n.count - 1 : 0, 1, 1, 1); + zbtIndexShrinkIfNeeded(zs); +} + +/* Change a score without moving the element when two conditions hold: + * + * 1. The new (score, member) pair remains strictly between its neighbors. + * 2. The leaf's current base, shift and width represent the score exactly. + * + * In that common case only one packed field and possibly maximums above it + * change. Return zero when the caller must remove and reinsert the element. + */ +static int zbtScoreUpdateInPlace(zbtScoreLeaf *leaf, unsigned int pos, + double score, const unsigned char *ele, + size_t elelen) +{ + zbtScoreLeaf *neighbor = pos ? leaf : leaf->prev; + unsigned int neighbor_pos = pos ? pos - 1 : + (neighbor ? (unsigned int)neighbor->n.count - 1 : 0); + if (neighbor) { + double other_score = zbtScoreLeafScore(neighbor, neighbor_pos); + if (score < other_score) return 0; + if (score == other_score) { + size_t len; + const unsigned char *other = + zbtScoreLeafElement(neighbor, neighbor_pos, &len); + if (zbtCompareElements(ele, elelen, other, len) <= 0) return 0; + } + } + + neighbor = pos + 1 < leaf->n.count ? leaf : leaf->next; + neighbor_pos = pos + 1 < leaf->n.count ? pos + 1 : 0; + if (neighbor) { + double other_score = zbtScoreLeafScore(neighbor, neighbor_pos); + if (score > other_score) return 0; + if (score == other_score) { + size_t len; + const unsigned char *other = + zbtScoreLeafElement(neighbor, neighbor_pos, &len); + if (zbtCompareElements(ele, elelen, other, len) >= 0) return 0; + } + } + + uint64_t delta; + if (!zbtScoreLeafPackScore(leaf, score, &delta)) return 0; + + if (leaf->score_bits) { + unsigned int physical = zbtScoreLeafPhysicalPos(leaf, pos); + zbtWriteBits(leaf->data, (size_t)physical * leaf->score_bits, + leaf->score_bits, delta); + } + if (pos + 1 == leaf->n.count) zbtScoreUpdatePath(&leaf->n, 0); + debugServerAssert(zbtScoreLeafScore(leaf, pos) == score); + return 1; +} + +/* Record the page share that a score change moves out of a leaf, in 1/96 of + * a page by count or bytes, whichever is larger. A rebuilt leaf starts at + * zero. */ +static void zbtScoreLeafNoteMove(zbtScoreLeaf *leaf, unsigned int pos) { + size_t bytes = zbtRecordStorageBytes(zbtScoreLeafRecord(leaf, pos)) + + sizeof(uint16_t) + sizeof(uint8_t); + unsigned int share = bytes * ZBT_SCORE_LEAF_MAX / ZBT_SCORE_LEAF_BYTES; + if (share == 0) share = 1; + share += leaf->n.moved_out; + leaf->n.moved_out = share > UINT8_MAX ? UINT8_MAX : share; +} + +/* A score change deletes the member from its old leaf without merging: the + * member is detached from its record at that point and its index entry still + * names the old leaf. Once it has its new record and entry, the old leaf is + * an ordinary leaf again. After every quarter page moved out of it, join it + * with a neighbor or rebuild it at its real size, as deletion does. Moves + * also fill leaves, so a join must leave a quarter of a page free; otherwise + * the next few moves into the merged leaf could split it again. The leaf is + * found by ID because the insertion may have rebuilt or removed it. */ +static void zbtScoreMovedOut(zbtreeSet *zs, uint32_t id) { + zbtScoreLeaf *leaf = zbtScoreLeafById(zs, id); + if (leaf == NULL || leaf->n.moved_out < ZBT_SCORE_MOVED_CHECK) return; + leaf->n.moved_out = 0; + + /* A neighbor's records are rarely in cache. When this leaf is still more + * than half of a join's budget, estimate the neighbor from this leaf's + * average member and read it only if a join looks possible. A sparse + * leaf always checks the real sizes, so a neighbor with shorter members + * cannot be skipped forever. */ + const unsigned int max_count = ZBT_SCORE_LEAF_MAX * 3 / 4; + const size_t max_bytes = ZBT_SCORE_LEAF_BYTES * 3 / 4; + size_t live = zbtScoreLeafLiveBytes(leaf); + size_t member_bytes = (live - offsetof(zbtScoreLeaf, data)) / + leaf->n.count; + int estimate = live > max_bytes / 2; + for (int side = 0; side < 2; side++) { + zbtScoreLeaf *other = side == 0 ? leaf->prev : leaf->next; + if (other == NULL || leaf->n.count + other->n.count > max_count || + (estimate && live + other->n.count * member_bytes > max_bytes)) + continue; + if (zbtScoreMergePair(zs, side == 0 ? other : leaf, + side == 0 ? leaf : other, + max_count, max_bytes)) + return; + } + if (zmalloc_usable_size(leaf) >= 2 * live + ZBT_SCORE_LEAF_RESERVE) + zbtScoreCompactLeaf(zs, leaf); +} + +/* Update an existing member's score. Its hash entry remains in place unless + * the member lands in a different score leaf. */ +void zbtreeUpdateScore(zbtreeSet *zs, sds ele, double score, + const zbtreeInsertPosition *position) +{ + uint32_t hash = position->hash; + zbtScoreLeaf *oldleaf = position->score_leaf; + unsigned int oldpos = position->score_pos; + + serverAssert(position->revision == zs->member_revision); + serverAssert(oldleaf != NULL && oldpos < oldleaf->n.count); + + /* Most small score changes preserve the member's position. In that case + * only its packed score bits and the maxima on the path are changed. */ + if (zbtScoreUpdateInPlace(oldleaf, oldpos, score, + (unsigned char *)ele, sdslen(ele))) + return; + + /* A larger move removes and reinserts the member in the score tree. Its + * hash entry stays in place; only the score leaf number may need + * to change. Reinsertion can split a leaf, so first make sure that a new + * leaf ID fits the member index. That may advance a resize, which moves + * no score leaf. + */ + zbtIndexReserveLeafId(zs, 0); + zbtBuildElement moved; + zbtBuildElementFromLeaf(&moved, oldleaf, oldpos); + if (moved.external) { + serverAssert(moved.hash == hash); + zbtBuildElementDetach(&moved); + } else { + /* Inline bytes may move or disappear during deletion, so use the + * command argument until the member has its new record. */ + zbtBuildElementFromBytes(&moved, (unsigned char *)ele, + sdslen(ele), hash); + } + uint32_t oldid = oldleaf->id; + if (oldleaf->n.count > 1) zbtScoreLeafNoteMove(oldleaf, oldpos); + zbtScoreDeleteAt(zs, oldleaf, oldpos, 0); + zbtScoreLeaf *newleaf = zbtScoreInsert(zs, score, &moved); + if (oldid != newleaf->id) { + int target_was_copied = 0; + serverAssert(zbtIndexMove(zs, hash, oldid, newleaf->id, 0, + &target_was_copied)); + } + zbtScoreMovedOut(zs, oldid); +} + +/* Give surviving score leaves dense IDs and shrink their address table. The + * member index is rebuilt immediately afterwards, so no old leaf ID survives. */ +static void zbtScoreRenumberLeaves(zbtreeSet *zs) { + uint32_t count = 0; + for (zbtScoreLeaf *leaf = zs->score_first; leaf; leaf = leaf->next) + count++; + + zbtScoreLeaf **oldtable = zs->score_leaf_by_id; + if (count == 0) { + zs->score_leaf_by_id = NULL; + zs->score_leaf_cap = 0; + zs->next_score_leaf_id = 0; + zs->free_score_leaf_id = 0; + if (oldtable) zbtFreeAllocation(zs, oldtable); + return; + } + + uint32_t cap = 16; + while (cap < count) cap *= 2; + zbtScoreLeaf **table = zbtAlloc(zs, cap * sizeof(*table)); + memset(table, 0, cap * sizeof(*table)); + + uint32_t id = 0; + for (zbtScoreLeaf *leaf = zs->score_first; leaf; leaf = leaf->next) { + leaf->id = id; + leaf->n.index_resize = 0; + table[id++] = leaf; + } + zs->score_leaf_by_id = table; + zs->score_leaf_cap = cap; + zs->next_score_leaf_id = count; + zs->free_score_leaf_id = 0; + if (oldtable) zbtFreeAllocation(zs, oldtable); +} + +/* Rebuild the hash index from the surviving score leaves. This also removes + * tombstones and makes every entry refer to the newly dense leaf numbers. */ +static void zbtIndexRebuild(zbtreeSet *zs) { + zbtScoreRenumberLeaves(zs); + zbtIndexTable oldindex = zs->member_index; + zbtIndexRehash *oldrehash = zs->member_rehash; + memset(&zs->member_index, 0, sizeof(zs->member_index)); + zs->member_rehash = NULL; + + if (zs->length) + zbtIndexTableInit(zs, &zs->member_index, + zbtIndexBucketsForElements(zs->length), + zs->next_score_leaf_id >= ZBT_INDEX_WIDE_ID_AT); + unsigned long rebuilt = 0; + for (zbtScoreLeaf *leaf = zs->score_first; leaf; leaf = leaf->next) { + for (unsigned int i = 0; i < leaf->n.count; i++) { + uint32_t hash = zbtScoreLeafHash(leaf, i); + zbtIndexTableInsertRaw(&zs->member_index, hash, leaf->id); + rebuilt++; + } + } + serverAssert(rebuilt == zs->length); + zbtIndexTableRelease(zs, &oldindex); + if (oldrehash) { + zbtIndexTableRelease(zs, &oldrehash->table); + zbtFreeAllocation(zs, oldrehash); + } + zs->member_revision = zbtIndexNextRevision(); +} + +/* Delete the inclusive zero-based rank interval [start,end], returning the + * number removed. Large deletions rebuild the member index from the survivors; + * smaller deletions maintain it one score-leaf run at a time. */ +unsigned long zbtreeDeleteRangeByRank(zbtreeSet *zs, unsigned long start, + unsigned long end) +{ + if (start >= zs->length || start > end) return 0; + if (end >= zs->length) end = zs->length - 1; + unsigned long deleted = end - start + 1; + + if (deleted > 1024 && deleted >= zs->length / 10) { + unsigned long remaining = deleted; + while (remaining) { + zbtreeIterator iter; + serverAssert(zbtreeIteratorSeekRank(zs, start, &iter)); + zbtScoreLeaf *leaf = iter.leaf; + unsigned long run = leaf->n.count - iter.pos; + if (run > remaining) run = remaining; + /* Nothing reads the old member index before zbtIndexRebuild(), so + * maintaining entries for partial leaves would be wasted work. */ + zbtScoreDeleteRun(zs, leaf, iter.pos, run, 0, 0); + remaining -= run; + } + zbtIndexRebuild(zs); + return deleted; + } + + /* Take as long a run as one leaf holds, so a range that fits in a leaf + * is one rebuild rather than one shift per element. No member is looked + * up by name here, so advance a member index resize by one step for each + * deleted member, as ZREM does. Otherwise a shrink started by an earlier + * deletion may never finish. */ + unsigned long remaining = deleted; + while (remaining) { + zbtreeIterator iter; + serverAssert(zbtreeIteratorSeekRank(zs, start, &iter)); + zbtScoreLeaf *leaf = iter.leaf; + unsigned long run = leaf->n.count - iter.pos; + if (run > remaining) run = remaining; + zbtScoreDeleteRun(zs, leaf, iter.pos, run, 1, 1); + remaining -= run; + zbtIndexRehashStep(zs, (int)run); + } + zbtIndexShrinkIfNeeded(zs); + return deleted; +} + +/* Delete all elements inside a score interval. Convert its two bounds to + * ranks and use the rank deletion path above. */ +unsigned long zbtreeDeleteRangeByScore(zbtreeSet *zs, double min, int minex, + double max, int maxex) +{ + zbtreeIterator first, last; + unsigned long first_rank, last_rank; + if (!zbtreeIteratorSeekScore(zs, min, minex, 0, &first, &first_rank)) + return 0; + if (!zbtreeIteratorSeekScore(zs, max, maxex, 1, &last, &last_rank)) + return 0; + if (last_rank < first_rank) return 0; + return zbtreeDeleteRangeByRank(zs, first_rank, last_rank); +} + +/* ------------------------------ Iterators ------------------------------- */ + +/* Iterators are a score leaf plus a logical position. zbtreeIteratorNext() + * returns the current element and advances, which makes the usual loop: + * + * seek or start; + * while (zbtreeIteratorNext(...)) { ... } + * + * The returned member pointer belongs to the current score leaf. Both it and + * the iterator become invalid after any modification of this sorted set. + */ + +/* Seek the element having the specified zero-based rank. */ +int zbtreeIteratorSeekRank(const zbtreeSet *zs, unsigned long rank, + zbtreeIterator *iter) +{ + if (rank >= zs->length) { + iter->leaf = NULL; + iter->pos = 0; + return 0; + } + + zbtScoreNode *node = zs->score_root; + while (!node->isleaf) { + zbtScoreInner *inner = (zbtScoreInner *)node; + unsigned int pos = 0; + while (rank >= inner->child_count[pos]) { + rank -= inner->child_count[pos]; + pos++; + } + node = inner->child[pos]; + } + iter->leaf = node; + iter->pos = rank; + return 1; +} + +/* Find the first score not below the bound. If after_equal is true, equal + * scores are skipped too. A binary search avoids many packed score reads. */ +static unsigned int zbtScoreLeafScoreBound(zbtScoreLeaf *leaf, double score, + int after_equal) +{ + unsigned int lo = 0, hi = leaf->n.count; + while (lo < hi) { + unsigned int mid = lo + (hi - lo) / 2; + double current = zbtScoreLeafScore(leaf, mid); + if (current < score || (after_equal && current == score)) + lo = mid + 1; + else + hi = mid; + } + return lo; +} + +/* Seek a score bound. In forward order, find the first score greater than or + * equal to the bound (strictly greater when exclusive). In reverse order, + * find the last score less than or equal to it (strictly less when + * exclusive). Return zero if no such element exists. + * + * If 'rank' is not NULL, also return the zero-based forward rank. + */ +int zbtreeIteratorSeekScore(const zbtreeSet *zs, double score, int exclusive, + int reverse, zbtreeIterator *iter, + unsigned long *rank) +{ + zbtScoreNode *node = zs->score_root; + if (node == NULL) { + iter->leaf = NULL; + iter->pos = 0; + return 0; + } + + while (!node->isleaf) { + zbtScoreInner *inner = (zbtScoreInner *)node; + unsigned int pos = 0; + if (reverse) { + while (pos < inner->n.count && + (exclusive ? inner->max_score[pos] < score : + inner->max_score[pos] <= score)) + pos++; + if (pos == inner->n.count) pos--; + } else { + while (pos + 1 < inner->n.count && + (exclusive ? inner->max_score[pos] <= score : + inner->max_score[pos] < score)) + pos++; + } + node = inner->child[pos]; + } + + zbtScoreLeaf *leaf = (zbtScoreLeaf *)node; + if (reverse) { + unsigned int pos = zbtScoreLeafScoreBound(leaf, score, !exclusive); + if (pos == 0) { + leaf = leaf->prev; + if (leaf == NULL) goto notfound; + pos = leaf->n.count - 1; + } else { + pos--; + } + iter->leaf = leaf; + iter->pos = pos; + } else { + unsigned int pos = zbtScoreLeafScoreBound(leaf, score, exclusive); + if (pos == leaf->n.count) { + leaf = leaf->next; + if (leaf == NULL) goto notfound; + pos = 0; + } + iter->leaf = leaf; + iter->pos = pos; + } + if (rank) + *rank = zbtScoreRank(iter->leaf, iter->pos); + return 1; + +notfound: + iter->leaf = NULL; + iter->pos = 0; + return 0; +} + +/* Seek a lexical bound. Redis defines useful lexicographical ranges when all + * members have the same score. Under that condition the score tree is already + * ordered by member, so a normal (score, member) search finds the bound. + * shared.minstring and shared.maxstring represent the two unbounded ends. + */ +int zbtreeIteratorSeekLex(const zbtreeSet *zs, sds ele, int exclusive, + int reverse, zbtreeIterator *iter, + unsigned long *rank) +{ + if (zs->score_first == NULL) goto notfound; + + if (reverse) { + if (ele == shared.maxstring) { + int found = zbtreeIteratorStart(zs, 1, iter); + if (found && rank) *rank = zs->length - 1; + return found; + } + if (ele == shared.minstring) goto notfound; + } else { + if (ele == shared.minstring) { + int found = zbtreeIteratorStart(zs, 0, iter); + if (found && rank) *rank = 0; + return found; + } + if (ele == shared.maxstring) goto notfound; + } + + /* Use the common score to turn the lexical bound into a normal score-tree + * search instead of walking from one end of the set. */ + double score = zbtScoreLeafScore(zs->score_first, 0); + zbtScoreLeaf *leaf = + zbtScoreFindLeaf(zs, score, (unsigned char *)ele, sdslen(ele)); + unsigned int pos = + zbtScoreLeafLowerBound(leaf, score, (unsigned char *)ele, sdslen(ele)); + int equal = 0; + + if (pos < leaf->n.count) { + size_t len; + const unsigned char *found = zbtScoreLeafElement(leaf, pos, &len); + equal = zbtCompareElements(found, len, (unsigned char *)ele, + sdslen(ele)) == 0; + } + + if (reverse) { + if (!equal || exclusive) { + if (pos != 0) { + pos--; + } else { + leaf = leaf->prev; + if (leaf == NULL) goto notfound; + pos = leaf->n.count - 1; + } + } + } else if (equal && exclusive) { + pos++; + } + + if (!reverse && pos == leaf->n.count) { + leaf = leaf->next; + if (leaf == NULL) goto notfound; + pos = 0; + } + + iter->leaf = leaf; + iter->pos = pos; + if (rank) *rank = zbtScoreRank(leaf, pos); + return 1; + +notfound: + iter->leaf = NULL; + iter->pos = 0; + return 0; +} + +/* Position an iterator at the first element, or the last when reverse is + * non-zero. Return zero for an empty set. */ +int zbtreeIteratorStart(const zbtreeSet *zs, int reverse, + zbtreeIterator *iter) +{ + zbtScoreLeaf *leaf = reverse ? zs->score_last : zs->score_first; + iter->leaf = leaf; + iter->pos = leaf ? (reverse ? leaf->n.count - 1 : 0) : 0; + return leaf != NULL; +} + +/* Return the current member and score, then advance to the next logical + * position. The output pointers are optional. Return zero at the end. */ +int zbtreeIteratorNext(zbtreeIterator *iter, int reverse, + const unsigned char **ele, size_t *len, double *score) +{ + zbtScoreLeaf *leaf = iter->leaf; + if (leaf == NULL) return 0; + + unsigned int pos = iter->pos; + if (ele || len) { + size_t elelen; + const unsigned char *element = + zbtScoreLeafElement(leaf, pos, &elelen); + if (ele) *ele = element; + if (len) *len = elelen; + } + if (score) *score = zbtScoreLeafScore(leaf, pos); + + if (reverse) { + if (pos != 0) { + iter->pos--; + } else { + leaf = leaf->prev; + iter->leaf = leaf; + if (leaf) iter->pos = leaf->n.count - 1; + } + } else { + if (pos + 1 < leaf->n.count) { + iter->pos++; + } else { + leaf = leaf->next; + iter->leaf = leaf; + iter->pos = 0; + } + } + return 1; +} + +/* Return the rank where a score bound would be inserted. Complete children + * are counted on the way down, so there is no second walk back to the root. */ +static unsigned long zbtScoreRankForBound(const zbtreeSet *zs, double score, + int after_equal) +{ + zbtScoreNode *node = zs->score_root; + unsigned long rank = 0; + if (node == NULL) return 0; + + while (!node->isleaf) { + zbtScoreInner *inner = (zbtScoreInner *)node; + unsigned int pos = 0; + while (pos + 1 < inner->n.count && + (inner->max_score[pos] < score || + (after_equal && inner->max_score[pos] == score))) + { + rank += inner->child_count[pos++]; + } + node = inner->child[pos]; + } + + zbtScoreLeaf *leaf = (zbtScoreLeaf *)node; + return rank + zbtScoreLeafScoreBound(leaf, score, after_equal); +} + +/* Count a score range as the difference between its two insertion ranks. + * Inclusive infinities are already the first and final ranks and need no + * tree walk. */ +unsigned long zbtreeCountByScore(const zbtreeSet *zs, double min, int minex, + double max, int maxex) +{ + unsigned long first = !minex && isinf(min) && min < 0 ? 0 : + zbtScoreRankForBound(zs, min, minex); + unsigned long last = !maxex && isinf(max) && max > 0 ? zs->length : + zbtScoreRankForBound(zs, max, !maxex); + return last > first ? last - first : 0; +} + +/* Load eight leaf tags so that numeric byte i is tags[i] on every host. This + * is the lane order zbtIndexTagMask() and zbtIndexFirstTag() assume; bucket + * tags have it because they are built with shifts. Without the conversion a + * big endian host would select mirrored positions. */ +static inline uint64_t zbtLeafTagWord(const uint8_t *tags) { + uint64_t word; + memcpy(&word, tags, sizeof(word)); + return intrev64ifbe(word); +} + +/* Store the logical positions of the leaf tags that an index tag can stand + * for, and return their number. 'tags' is the physical array, which is in + * reverse order when 'reversed' is set. zbtIndexTag() folds leaf tag zero + * into one. Eight tags are compared at once. As with bucket tags, a mask bit + * can be set above a real match, so each candidate is checked against its + * byte. Only complete words are loaded: member records follow the tags. */ +static inline unsigned int zbtLeafTagMatches(const uint8_t *tags, + unsigned int count, int reversed, + uint8_t tag, + unsigned int *positions) +{ + unsigned int found = 0; + unsigned int words = count / 8; + unsigned int last = reversed ? count - 1 : 0; + for (unsigned int w = 0; w < words; w++) { + uint64_t word = zbtLeafTagWord(tags + w * 8); + uint64_t mask = zbtIndexTagMask(word, tag); + if (tag == 1) mask |= zbtIndexTagMask(word, 0); + while (mask) { + unsigned int pos = w * 8 + zbtIndexFirstTag(mask); + if (tags[pos] == tag || (tag == 1 && tags[pos] == 0)) + positions[found++] = reversed ? last - pos : pos; + mask &= mask - 1; + } + } + for (unsigned int pos = words * 8; pos < count; pos++) { + if (tags[pos] == tag || (tag == 1 && tags[pos] == 0)) + positions[found++] = reversed ? last - pos : pos; + } + return found; +} + +/* Return the members represented by one physical table slot. If home_end is + * nonzero, only return members whose home bucket is in [home_start,home_end). + * When several members of a leaf share the slot's tag, their slots look alike, + * so each such member is returned only by the first of them on its probe + * path. Otherwise every slot would return all of them. */ +static unsigned long zbtIndexScanSlot(const zbtreeSet *zs, + zbtIndexTable *table, + zbtIndexBucket *bucket, + unsigned int slot_pos, + unsigned long home_start, + unsigned long home_end, + zbtreeScanFunction *fn, + void *privdata) +{ + uint8_t tag = zbtIndexTags(bucket) >> (slot_pos * 8); + uint32_t id = zbtIndexGetId(table, bucket, slot_pos); + if (tag == 0 || id == ZBT_INDEX_DELETED_ID) return 0; + if (home_end) { + uint64_t bits = zbtIndexTagBits(tag); + unsigned long home; + for (home = home_start; home < home_end; home++) { + if ((*zbtIndexHomeTagsAt(table, home) & bits) == bits) break; + } + if (home == home_end) return 0; + } + zbtScoreLeaf *leaf = id < zs->next_score_leaf_id ? + zs->score_leaf_by_id[id] : NULL; + if (leaf == NULL || ZBT_IS_FREE_LEAF_ID(leaf)) { + serverAssert(zs->member_rehash && table == &zs->member_index); + return 0; + } + serverAssert(leaf->id == id); + if (zs->member_rehash) { + int in_new = table == &zs->member_rehash->table; + if (in_new != zbtIndexLeafMigrated(zs, leaf)) return 0; + } + + /* Match by the leaf's own stored tag bytes, without hashing. Every leaf + * tag is the member's real hash >> 24 and is only copied forward from an + * already correct leaf, so a match already implies zbtIndexTag(hash) == + * tag. */ + unsigned int positions[ZBT_SCORE_LEAF_MAX]; + unsigned int count = zbtLeafTagMatches(zbtScoreLeafHashTags(leaf), + leaf->n.count, leaf->reversed, + tag, positions); + serverAssert(count != 0); + + unsigned long emitted = 0; + for (unsigned int i = 0; i < count; i++) { + /* The real (expensive, re-hashed) member hash is only needed to + * disambiguate a genuine same-tag collision inside this leaf. Such + * members are returned at the first alike slot on their probe path, + * which the scan of their home group always visits, since the table + * does not change during a scan call. A single match needs no such + * check. */ + if (count > 1 || home_end) { + uint32_t hash = zbtScoreLeafHash(leaf, positions[i]); + unsigned long home = hash & (table->size - 1); + if (home_end && (home < home_start || home >= home_end)) + continue; + if (count > 1) { + zbtIndexBucket *first; + unsigned int first_pos; + /* Every member in the table is reachable from its home, or + * lookups would miss it too: the table is corrupt. */ + int found = zbtIndexTableFindReference(table, hash, id, + &first, &first_pos); + serverAssert(found); + if (first != bucket || first_pos != slot_pos) continue; + } + } + size_t len; + const unsigned char *ele = + zbtScoreLeafElement(leaf, positions[i], &len); + fn(privdata, ele, len, zbtScoreLeafScore(leaf, positions[i])); + emitted++; + } + /* With alike slots, one may return nothing: several members can share + * the same first slot. */ + serverAssert(home_end || count > 1 || emitted != 0); + return emitted; +} + +/* Return every member represented in 'table' whose home bucket is one of the + * buckets of scan group 'group'. Slots in the group's own buckets are returned + * without rehashing their members. That may repeat a member displaced from an + * earlier home, which SCAN allows. Members whose home is here but which were + * displaced further are in the probe run after the group. Same-tag references + * can exchange slots when a member changes score, but the slot standing for a + * member is always reachable from its home. A zero tag never returns to a + * table, so a run that was full when such a member was added is still full + * and the walk below cannot stop before reaching it. */ +static unsigned long zbtIndexScanGroup(const zbtreeSet *zs, + zbtIndexTable *table, uint64_t group, + zbtreeScanFunction *fn, + void *privdata) +{ + unsigned long home_start = group * ZBT_SCAN_BUCKETS_PER_STEP; + unsigned long home_end = home_start + ZBT_SCAN_BUCKETS_PER_STEP; + unsigned long emitted = 0; + + /* Groups are visited in an order unrelated to memory, and each slot names + * a different leaf. Start loading those leaves before reading them. */ + for (unsigned long index = home_start; index < home_end; index++) { + zbtIndexBucket *bucket = zbtIndexBucketAt(table, index); + for (unsigned int pos = 0; pos < ZBT_INDEX_BUCKET_ITEMS; pos++) { + if ((uint8_t)(zbtIndexTags(bucket) >> (pos * 8)) == 0) continue; + uint32_t id = zbtIndexGetId(table, bucket, pos); + if (id >= zs->next_score_leaf_id) continue; + zbtScoreLeaf *leaf = zs->score_leaf_by_id[id]; + if (leaf == NULL || ZBT_IS_FREE_LEAF_ID(leaf)) continue; + redis_prefetch_read(leaf); + redis_prefetch_read((unsigned char *)leaf + 64); + } + } + + for (unsigned long index = home_start; index < home_end; index++) { + zbtIndexBucket *bucket = zbtIndexBucketAt(table, index); + for (unsigned int pos = 0; pos < ZBT_INDEX_BUCKET_ITEMS; pos++) + emitted += zbtIndexScanSlot(zs, table, bucket, pos, 0, 0, + fn, privdata); + } + + unsigned long mask = table->size - 1; + unsigned long index = home_end - 1; + unsigned long remaining = table->size - ZBT_SCAN_BUCKETS_PER_STEP; + while (remaining-- && + !zbtIndexTagMask(zbtIndexTags(zbtIndexBucketAt(table, index)), 0)) + { + index = (index + 1) & mask; + zbtIndexBucket *bucket = zbtIndexBucketAt(table, index); + for (unsigned int pos = 0; pos < ZBT_INDEX_BUCKET_ITEMS; pos++) + emitted += zbtIndexScanSlot(zs, table, bucket, pos, + home_start, home_end, fn, privdata); + } + return emitted; +} + +/* Increment the bits of a cursor selected by 'mask', a run of low bits, in + * reverse binary order, as dictScan() does by reversing, adding one and + * reversing back. The carry runs from the highest bit down: set the highest + * zero bit and clear the ones above it. Bits outside the mask are dropped, + * and zero is returned after the last value. */ +static uint64_t zbtScanNext(uint64_t cursor, uint64_t mask) { + uint64_t missing = ~cursor & mask; + if (missing == 0) return 0; + uint64_t bit = UINT64_C(1) << (63 - __builtin_clzll(missing)); + return (cursor & (bit - 1)) | bit; +} + +/* Return the mask selecting a scan group of 'table'. */ +static uint64_t zbtScanGroupMask(const zbtIndexTable *table) { + uint64_t groups = table->size / ZBT_SCAN_BUCKETS_PER_STEP; + if (groups > ZBT_SCAN_MAX_GROUPS) groups = ZBT_SCAN_MAX_GROUPS; + return groups - 1; +} + +/* The cursor is a dictScan() reverse binary counter over scan groups. It + * depends only on member hashes, never on a table, leaf or slot, so it stays + * valid while tables grow, shrink, are cleaned, widened or rebuilt, and while + * members move between leaves and tables. Its value stays below the number + * of groups. + * + * Group g of a table with S buckets holds the members whose home bucket is in + * [4g, 4g+4), that is, whose H = hash >> 2 has H & (S/4 - 1) == g. Tables are + * powers of two, so groups nest like dictionary buckets: group g of size S + * becomes groups g and g + S/4 of size 2S. During a resize, each live member + * is represented in exactly one table: the new one if its leaf was copied. A + * step visits group 'cursor & m0' of the smaller table and, as dictScan() + * does, the groups of the larger table that expand it. + * + * Let R() reverse 64 bits, and let k0 and k1 be the bits of m0 and m1. A + * group of the smaller table holds the members with R(H) in [P, P + 2^(64-k0)), + * where P is R(cursor) with its low 64-k0 bits cleared. The larger table is + * visited from the group holding R(cursor) to the end of that interval, and + * the next cursor has R equal to P + 2^(64-k0), or wraps to zero. So each step + * returns every member with R(H) in [R(cursor), R(next)), whichever table + * represents it, and consecutive steps tile [0, 2^64). A member present + * during the whole iteration keeps its H and is returned by the step whose + * interval holds R(H). Every R(next) is a multiple of 2^(64-K), where 2^K is + * the largest group count used, so R strictly increases and an iteration ends + * after at most 2^K + 1 steps. Table sizes depend only on the number of + * members, so a set kept at a bounded size always finishes its scans, however + * much it changes in between. + * + * A cursor cannot be translated if a caller explicitly changes the object to + * another encoding. */ +uint64_t zbtreeScan(zbtreeSet *zs, uint64_t cursor, + unsigned long count, zbtreeScanFunction *fn, + void *privdata) +{ + if (zs->length == 0) return 0; + zbtIndexTable *small = &zs->member_index, *large = NULL; + if (zs->member_rehash) { + large = &zs->member_rehash->table; + if (large->size < small->size) { + zbtIndexTable *tmp = small; + small = large; + large = tmp; + } + } + uint64_t m0 = zbtScanGroupMask(small); + uint64_t m1 = large ? zbtScanGroupMask(large) : m0; + + /* Finish a whole step before returning the cursor. COUNT is a hint, as + * it is for a normal dictionary scan. */ + unsigned long emitted = 0; + do { + emitted += zbtIndexScanGroup(zs, small, cursor & m0, fn, privdata); + if (large == NULL) { + cursor = zbtScanNext(cursor, m0); + continue; + } + do { + emitted += zbtIndexScanGroup(zs, large, cursor & m1, fn, privdata); + cursor = zbtScanNext(cursor, m1); + } while (cursor & (m0 ^ m1)); + } while (cursor != 0 && emitted < count); + return cursor; +} + +#ifdef REDIS_TEST +#include "testhelp.h" + +/* Scan callback for zsetBtreeTest(): members are named "member:" for a + * known 0..N-1 range, so recover i and mark it seen instead of needing a + * general string-set to check scan completeness against. 'seen' has 'n' + * elements; the sscanf'd index is bounds-checked before use since it comes + * from the scanned data, not from a value the test already trusts. */ +typedef struct { + int *seen; + int n; +} zbtScanTestPrivdata; + +static void zbtScanTestMarkSeen(void *privdata, const unsigned char *ele, + size_t len, double score) +{ + UNUSED(score); + zbtScanTestPrivdata *pd = privdata; + char buf[32]; + serverAssert(len < sizeof(buf)); + memcpy(buf, ele, len); + buf[len] = '\0'; + int idx = -1; + serverAssert(sscanf(buf, "member:%d", &idx) == 1); + serverAssert(idx >= 0 && idx < pd->n); + pd->seen[idx]++; +} + +/* Insert members "member:".."member:" with scores from..to-1. + * Shared by every zsetBtreeTest() block that needs a populated set. */ +static void zbtTestPopulateRange(zbtreeSet *zs, int from, int to) { + char buf[32]; + for (int i = from; i < to; i++) { + snprintf(buf, sizeof(buf), "member:%d", i); + sds ele = sdsnew(buf); + double score; + zbtreeInsertPosition position; + int existing = zbtreeFindForAdd(zs, ele, &score, &position); + serverAssert(!existing); + zbtreeInsertNew(zs, (double)i, ele, &position); + sdsfree(ele); + } +} + +static void zbtTestPopulate(zbtreeSet *zs, int n) { + zbtTestPopulateRange(zs, 0, n); +} + +/* Helpers for the same-tag ZSCAN regression test. */ +static uint32_t zbtTestHash(const char *s) { + return (uint32_t)dictGenHashFunction((const unsigned char *)s, strlen(s)); +} + +static zbtScoreLeaf *zbtTestLeafOf(zbtreeSet *zs, const char *s) { + zbtScoreLeaf *leaf = NULL; + unsigned int pos; + serverAssert(zbtIndexFind(zs, zbtTestHash(s), (const unsigned char *)s, + strlen(s), NULL, &leaf, &pos, NULL, NULL)); + return leaf; +} + +static int zbtTestTagCount(zbtScoreLeaf *leaf, uint8_t tag) { + int count = 0; + for (unsigned int i = 0; i < leaf->n.count; i++) + if (zbtIndexTag(zbtScoreLeafHash(leaf, i)) == tag) count++; + return count; +} + +static int zbtTestSlotOf(zbtreeSet *zs, uint32_t hash, uint32_t id, + unsigned long *bucket_index, unsigned int *pos) +{ + zbtIndexBucket *bucket; + if (!zbtIndexTableFindReference(&zs->member_index, hash, id, &bucket, pos)) + return 0; + *bucket_index = ((unsigned char *)bucket - + (unsigned char *)zs->member_index.buckets) / + zbtIndexBucketBytes(&zs->member_index); + return 1; +} + +/* Reverse the 64 bits of a scan cursor, the R() of the zbtreeScan() proof. */ +static uint64_t zbtScanReverse(uint64_t v) { + v = ((v >> 1) & UINT64_C(0x5555555555555555)) | + ((v & UINT64_C(0x5555555555555555)) << 1); + v = ((v >> 2) & UINT64_C(0x3333333333333333)) | + ((v & UINT64_C(0x3333333333333333)) << 2); + v = ((v >> 4) & UINT64_C(0x0f0f0f0f0f0f0f0f)) | + ((v & UINT64_C(0x0f0f0f0f0f0f0f0f)) << 4); + v = ((v >> 8) & UINT64_C(0x00ff00ff00ff00ff)) | + ((v & UINT64_C(0x00ff00ff00ff00ff)) << 8); + v = ((v >> 16) & UINT64_C(0x0000ffff0000ffff)) | + ((v & UINT64_C(0x0000ffff0000ffff)) << 16); + return (v >> 32) | (v << 32); +} + +/* While the table does not change, a scan cursor has visited exactly the + * groups whose reversed index is below its own. Zero ends the iteration. */ +static int zbtTestGroupVisited(uint64_t cursor, uint64_t group) { + return cursor == 0 || zbtScanReverse(group) < zbtScanReverse(cursor); +} + +static void zbtTestSetScore(zbtreeSet *zs, const char *s, double score) { + sds ele = sdsnew(s); + double current; + zbtreeInsertPosition position; + serverAssert(zbtreeFindForAdd(zs, ele, ¤t, &position)); + zbtreeUpdateScore(zs, ele, score, &position); + sdsfree(ele); +} + +/* One member entry or member, as (leaf ID, folded tag), for comparing the + * member index with the score tree as multisets. */ +typedef struct { + uint32_t id; + uint8_t tag; +} zbtTestRef; + +static int zbtTestRefCompare(const void *a, const void *b) { + const zbtTestRef *x = a, *y = b; + if (x->id != y->id) return x->id < y->id ? -1 : 1; + return (int)x->tag - (int)y->tag; +} + +#define ZBT_TEST_CHECK(cond, ...) do { \ + if (!(cond)) { \ + printf(" check failed at line %d: ", __LINE__); \ + printf(__VA_ARGS__); \ + printf("\n"); \ + return -1; \ + } \ + } while (0) + +/* Check one score subtree and return its element count, or -1. */ +static long long zbtTestCheckNode(zbtScoreNode *node, zbtScoreInner *parent, + unsigned int index) +{ + ZBT_TEST_CHECK(node->parent == parent, "parent pointer"); + ZBT_TEST_CHECK(parent == NULL || node->parent_index == index, + "parent index"); + if (node->isleaf) { + ZBT_TEST_CHECK(node->subtree == node->count, "leaf subtree"); + return node->count; + } + zbtScoreInner *inner = (zbtScoreInner *)node; + ZBT_TEST_CHECK(inner->n.count >= 1 && + inner->n.count <= ZBT_SCORE_INNER_MAX, "inner count"); + long long total = 0; + for (unsigned int i = 0; i < inner->n.count; i++) { + zbtScoreNode *child = inner->child[i]; + long long sub = zbtTestCheckNode(child, inner, i); + if (sub < 0) return -1; + ZBT_TEST_CHECK(inner->child_count[i] == (uint64_t)sub && + child->subtree == (uint64_t)sub, "child count"); + ZBT_TEST_CHECK(inner->max_score[i] == zbtScoreNodeMaxScore(child), + "child maximum"); + total += sub; + } + ZBT_TEST_CHECK(inner->n.subtree == (uint64_t)total, "inner subtree"); + return total; +} + +/* Return the usable bytes of a score subtree and the external members its + * records own. */ +static size_t zbtTestNodeAllocSize(zbtScoreNode *node) { + size_t bytes = zmalloc_usable_size(node); + if (node->isleaf) { + zbtScoreLeaf *leaf = (zbtScoreLeaf *)node; + for (unsigned int i = 0; i < leaf->n.count; i++) { + unsigned char *record = zbtScoreLeafRecord(leaf, i); + if (record[0] != ZBT_PACKED_EXTERNAL) continue; + serverAssert(leaf->n.has_external); + bytes += zmalloc_usable_size(zbtExternalPointer(record)); + } + } else { + zbtScoreInner *inner = (zbtScoreInner *)node; + for (unsigned int i = 0; i < inner->n.count; i++) + bytes += zbtTestNodeAllocSize(inner->child[i]); + } + return bytes; +} + +/* Recompute alloc_size from every reachable allocation. A leaked, freed or + * doubly owned external member makes the two differ. */ +static size_t zbtTestAllocSize(zbtreeSet *zs) { + size_t bytes = zmalloc_usable_size(zs); + if (zs->score_root) bytes += zbtTestNodeAllocSize(zs->score_root); + if (zs->member_index.buckets) + bytes += zmalloc_usable_size(zs->member_index.buckets); + if (zs->member_rehash) { + bytes += zmalloc_usable_size(zs->member_rehash); + if (zs->member_rehash->table.buckets) + bytes += zmalloc_usable_size(zs->member_rehash->table.buckets); + } + if (zs->score_leaf_by_id) + bytes += zmalloc_usable_size(zs->score_leaf_by_id); + return bytes; +} + +/* Collect the entries of 'table' that stand for live members, and check the + * entries that do not. Return the number collected, or -1. */ +static long zbtTestTableRefs(zbtreeSet *zs, zbtIndexTable *table, + zbtTestRef *refs) +{ + int is_new = zs->member_rehash && table == &zs->member_rehash->table; + unsigned long used = 0, filled = 0; + long n = 0; + for (unsigned long b = 0; b < table->size; b++) { + zbtIndexBucket *bucket = zbtIndexBucketAt(table, b); + for (unsigned int pos = 0; pos < ZBT_INDEX_BUCKET_ITEMS; pos++) { + uint8_t tag = zbtIndexTags(bucket) >> (pos * 8); + if (tag == 0) continue; + filled++; + uint32_t id = zbtIndexGetId(table, bucket, pos); + if (id == ZBT_INDEX_DELETED_ID) continue; + used++; + zbtScoreLeaf *leaf = zbtScoreLeafById(zs, id); + /* Only the old table of a resize may keep stale entries. */ + int current = leaf && (!zs->member_rehash || + zbtIndexLeafMigrated(zs, leaf) == is_new); + ZBT_TEST_CHECK(current || (zs->member_rehash && !is_new), + "stale entry for leaf %u", id); + if (!current) continue; + refs[n].id = id; + refs[n].tag = tag; + n++; + } + } + ZBT_TEST_CHECK(table->used == used && table->filled == filled, + "table counters used=%lu/%lu filled=%lu/%lu", + table->used, used, table->filled, filled); + return n; +} + +/* Check that the score tree, the leaf IDs and the member index agree: every + * member is found through the table that represents it, the live entries of + * each table are exactly its members, and a narrow table that receives new + * leaf IDs cannot be given one it is unable to store. Return 1 if so. */ +static int zbtTestCheck(zbtreeSet *zs) { + if (zs->score_root && + zbtTestCheckNode(zs->score_root, NULL, 0) != (long long)zs->length) + return 0; + if (zbtTestAllocSize(zs) != zs->alloc_size) { + printf(" check failed: alloc_size %zu, reachable %zu\n", + zs->alloc_size, zbtTestAllocSize(zs)); + return 0; + } + zbtIndexTable *receiving = zs->member_rehash ? + &zs->member_rehash->table : &zs->member_index; + if (receiving->size && !receiving->wide_ids && + zs->next_score_leaf_id > UINT16_MAX) + { + printf(" check failed: narrow table with leaf IDs up to %u\n", + zs->next_score_leaf_id); + return 0; + } + + unsigned long slots = zbtIndexSlots(&zs->member_index); + if (zs->member_rehash) slots += zbtIndexSlots(&zs->member_rehash->table); + zbtTestRef *members = zmalloc(sizeof(*members) * (zs->length + 1) * 2); + zbtTestRef *entries = zmalloc(sizeof(*entries) * (slots + 1)); + int ok = 0; + for (int pass = 0; pass < (zs->member_rehash ? 2 : 1); pass++) { + zbtIndexTable *table = pass ? &zs->member_rehash->table : + &zs->member_index; + unsigned long count = 0, total = 0; + zbtScoreLeaf *prev = NULL; + double pscore = 0; + const unsigned char *pele = NULL; + size_t plen = 0; + for (zbtScoreLeaf *leaf = zs->score_first; leaf; leaf = leaf->next) { + if (leaf->prev != prev || leaf->n.count == 0 || + zbtScoreLeafById(zs, leaf->id) != leaf) + { + printf(" check failed: leaf list or ID of leaf %u\n", + leaf->id); + goto done; + } + int in_new = zs->member_rehash && zbtIndexLeafMigrated(zs, leaf); + for (unsigned int i = 0; i < leaf->n.count; i++) { + size_t len; + const unsigned char *ele = zbtScoreLeafElement(leaf, i, &len); + double score = zbtScoreLeafScore(leaf, i); + if (pele && zbtScoreCompare(pscore, pele, plen, + score, ele, len) >= 0) + { + printf(" check failed: order in leaf %u\n", leaf->id); + goto done; + } + pele = ele; plen = len; pscore = score; + total++; + if (in_new != pass) continue; + uint32_t hash = zbtScoreLeafHash(leaf, i); + zbtScoreLeaf *found = NULL; + unsigned int found_pos = 0; + double found_score; + if (zbtScoreLeafTag(leaf, i) != (uint8_t)(hash >> 24) || + !zbtIndexHomeMayContain(table, hash) || + !zbtIndexTableFind(zs, table, hash, ele, len, + &found_score, &found, &found_pos, + NULL, NULL, NULL, NULL) || + found != leaf || found_pos != i) + { + printf(" check failed: member %u/%u not indexed\n", + leaf->id, i); + goto done; + } + members[count].id = leaf->id; + members[count].tag = zbtIndexTag(hash); + count++; + } + prev = leaf; + } + if (prev != zs->score_last || total != zs->length) { + printf(" check failed: leaf list end or length\n"); + goto done; + } + long n = zbtTestTableRefs(zs, table, entries); + if (n != (long)count) { + if (n >= 0) printf(" check failed: %ld entries for %lu " + "members\n", n, count); + goto done; + } + qsort(members, count, sizeof(*members), zbtTestRefCompare); + qsort(entries, count, sizeof(*entries), zbtTestRefCompare); + for (unsigned long i = 0; i < count; i++) { + if (zbtTestRefCompare(&members[i], &entries[i]) != 0) { + printf(" check failed: leaf %u tag %u has an entry " + "mismatch\n", members[i].id, members[i].tag); + goto done; + } + } + } + ok = 1; +done: + zfree(members); + zfree(entries); + return ok; +} + +/* Make 'next' the next new leaf ID, as if the IDs below it belonged to other + * leaves. Reaching the leaf ID limits otherwise needs hundreds of megabytes. + * The skipped IDs stay NULL, so no entry or cursor can refer to them. */ +static void zbtTestSkipLeafIds(zbtreeSet *zs, uint32_t next) { + serverAssert(next >= zs->next_score_leaf_id); + if (next >= zs->score_leaf_cap) { + uint32_t cap = zs->score_leaf_cap ? zs->score_leaf_cap : 16; + while (cap <= next) cap *= 2; + size_t usable, old_usable = 0; + zs->score_leaf_by_id = zrealloc_usable(zs->score_leaf_by_id, + cap * sizeof(*zs->score_leaf_by_id), &usable, &old_usable); + memset(zs->score_leaf_by_id + zs->score_leaf_cap, 0, + (cap - zs->score_leaf_cap) * sizeof(*zs->score_leaf_by_id)); + zs->score_leaf_cap = cap; + zs->alloc_size += usable - old_usable; + } + zs->next_score_leaf_id = next; +} + +static uint64_t zbtTestRand(uint64_t *r) { + *r = *r * 6364136223846793005ULL + 1442695040888963407ULL; + return *r >> 17; +} + +/* Test members: "m::" padded to a chosen length. */ +static size_t zbtTestMember(char *buf, int i, size_t len) { + int n = snprintf(buf, 32, "m:%d:", i); + if (len < (size_t)n) len = n; + memset(buf + n, 'a' + i % 26, len - n); + return len; +} + +/* Insert or change one member like ZADD does. */ +static void zbtTestSet(zbtreeSet *zs, const char *buf, size_t len, + double score) +{ + sds ele = sdsnewlen(buf, len); + double cur; + zbtreeInsertPosition position; + if (!zbtreeFindForAdd(zs, ele, &cur, &position)) + zbtreeInsertNew(zs, score, ele, &position); + else if (cur != score) + zbtreeUpdateScore(zs, ele, score, &position); + sdsfree(ele); +} + +/* Return whether every member "m:" with alive[i] has score ref[i]. */ +static int zbtTestVerify(zbtreeSet *zs, int n, const double *ref, + const char *alive, const size_t *lens) +{ + char buf[1024]; + unsigned long count = 0; + for (int i = 0; i < n; i++) { + if (alive && !alive[i]) continue; + size_t len = zbtTestMember(buf, i, lens ? lens[i] : 0); + double score; + if (!zbtreeScoreRaw(zs, (unsigned char *)buf, len, &score) || + score != ref[i]) return 0; + count++; + } + return count == zbtreeLength(zs); +} + +static unsigned long zbtTestLeaves(zbtreeSet *zs) { + unsigned long leaves = 0; + for (zbtScoreLeaf *leaf = zs->score_first; leaf; leaf = leaf->next) + leaves++; + return leaves; +} + +/* Apply 'moves' random score changes to members 0..n-1 of an append-built + * set. Most land in another leaf. */ +static void zbtTestRandomMoves(zbtreeSet *zs, int n, double *ref, + const size_t *lens, long moves, uint64_t *r) +{ + char buf[1024]; + for (long k = 0; k < moves; k++) { + int i = (int)(zbtTestRand(r) % (uint64_t)n); + size_t len = zbtTestMember(buf, i, lens[i]); + double score = (double)(zbtTestRand(r) % (uint64_t)(n * 4)) + 0.5; + zbtTestSet(zs, buf, len, score); + ref[i] = score; + } +} + +/* Scan state of the bounded-churn tests: anchors are "anchor:". */ +typedef struct { + int *seen; + int anchors; +} zbtScanAnchorPrivdata; + +static void zbtScanTestMarkAnchor(void *privdata, const unsigned char *ele, + size_t len, double score) +{ + UNUSED(score); + zbtScanAnchorPrivdata *pd = privdata; + char buf[32]; + if (len >= sizeof(buf) || len < 7 || memcmp(ele, "anchor:", 7)) return; + memcpy(buf, ele, len); + buf[len] = '\0'; + int idx = atoi(buf + 7); + serverAssert(idx >= 0 && idx < pd->anchors); + pd->seen[idx]++; +} + +/* Return the number of groups the smaller current table has; a returned + * cursor must stay below it. */ +static uint64_t zbtTestScanGroups(zbtreeSet *zs) { + zbtIndexTable *t = &zs->member_index; + if (zs->member_rehash && zs->member_rehash->table.size < t->size) + t = &zs->member_rehash->table; + return zbtScanGroupMask(t) + 1; +} + +/* Return the number of groups of the larger current table. */ +static uint64_t zbtTestScanMaxGroups(zbtreeSet *zs) { + zbtIndexTable *t = &zs->member_index; + if (zs->member_rehash && zs->member_rehash->table.size > t->size) + t = &zs->member_rehash->table; + return zbtScanGroupMask(t) + 1; +} + +int zsetBtreeTest(int argc, char **argv, int flags) { + UNUSED(argc); + UNUSED(argv); + UNUSED(flags); + + printf("Testing B+ tree zbtreeReserve()\n"); + + { + zbtreeSet *zs = zbtreeCreate(); + size_t initial_size = zs->alloc_size; + uint32_t initial_revision = zs->member_revision; + test_cond("Zero reservation does not allocate an index", + zbtreeTryReserve(zs, 0) && zs->member_index.buckets == NULL && + zs->alloc_size == initial_size); + + /* Exercise count arithmetic, bucket-byte overflow, and an allocation + * rejected by ztrycalloc_usable's SIZE_MAX/2 guard on 64-bit builds. + * On 32-bit builds the last count fails before narrowing instead. */ + const uint64_t invalid_counts[] = { + UINT64_MAX, ULONG_MAX / 8, UINT64_C(1) << 60 + }; + for (unsigned int i = 0; i < sizeof(invalid_counts) / sizeof(*invalid_counts); i++) { + test_cond("Failed reservation leaves the fresh tree unchanged", + !zbtreeTryReserve(zs, invalid_counts[i]) && + zs->member_index.buckets == NULL && zs->member_index.size == 0 && + zs->alloc_size == initial_size && + zs->member_revision == initial_revision); + } + + test_cond("A failed reservation can be retried with a valid count", + zbtreeTryReserve(zs, 10000) && zbtTestCheck(zs)); + zbtTestPopulate(zs, 10000); + test_cond("TryReserve holds the requested members without rehashing", + zbtreeLength(zs) == 10000 && zs->member_rehash == NULL && zbtTestCheck(zs)); + size_t populated_size = zs->alloc_size; + test_cond("TryReserve leaves a populated tree intact", + zbtreeTryReserve(zs, UINT64_MAX) && + zs->alloc_size == populated_size && zbtTestCheck(zs)); + zbtreeFree(zs); + } + + { + /* Only a count that fills ZBT_INDEX_WIDE_ID_AT full leaves is sure to + * widen the table, so only such a count starts it wide. */ + const uint64_t wide_at = + (uint64_t)ZBT_SCORE_LEAF_MAX * ZBT_INDEX_WIDE_ID_AT; + zbtreeSet *zs = zbtreeCreate(); + test_cond("A reservation that cannot need wide leaf IDs stays narrow", + zbtreeTryReserve(zs, wide_at - 1) && !zs->member_index.wide_ids); + zbtreeFree(zs); + zs = zbtreeCreate(); + test_cond("A reservation that will need wide leaf IDs starts wide", + zbtreeTryReserve(zs, wide_at) && zs->member_index.wide_ids); + zbtTestPopulate(zs, 10000); + test_cond("A wide reserved index holds members with small leaf IDs", + zbtreeLength(zs) == 10000 && zs->member_rehash == NULL && + zbtTestCheck(zs)); + zbtreeFree(zs); + } + + { + zbtreeSet *zs = zbtreeCreate(); + zbtreeReserve(zs, 1); + zbtreeAppendBatch *batch = zbtreeAppendBatchCreate(zs); + char member[32]; + for (int i = 0; i < 1000; i++) { + int len = snprintf(member, sizeof(member), "member:%05d", i); + zbtreeAppendBatchAdd(batch, i, (unsigned char *)member, len); + } + zbtreeAppendBatchFinish(batch); + int valid = zbtreeLength(zs) == 1000; + for (int i = 0; i < 1000; i++) { + snprintf(member, sizeof(member), "member:%05d", i); + sds ele = sdsnew(member); + double score; + valid = valid && zbtreeScore(zs, ele, &score) && score == i; + sdsfree(ele); + } + test_cond("Batch append tolerates an underestimated reserve", valid); + zbtreeFree(zs); + } + + /* zbtreeReserve() on a fresh set sizes the member index directly, + * without ever starting an incremental member_rehash. */ + { + zbtreeSet *zs = zbtreeCreate(); + zbtreeReserve(zs, 10000); + test_cond("zbtreeReserve presizes an empty index", + zs->member_index.size > ZBT_INDEX_INITIAL_BUCKETS && + zs->member_rehash == NULL); + + zbtTestPopulate(zs, 10000); + test_cond("Reserved set still holds every inserted member", + zbtreeLength(zs) == 10000 && zs->member_rehash == NULL); + zbtreeFree(zs); + } + + /* zbtreeReserve() must be a no-op on a set that already has members or + * an index: it exists to presize a still-empty destination, not to + * silently resize a live one out from under its scan cursors. */ + { + zbtreeSet *zs = zbtreeCreate(); + sds ele = sdsnew("only-member"); + double score; + zbtreeInsertPosition position; + zbtreeFindForAdd(zs, ele, &score, &position); + zbtreeInsertNew(zs, 1.0, ele, &position); + sdsfree(ele); + unsigned long size_before = zs->member_index.size; + zbtreeReserve(zs, 1000000); + test_cond("zbtreeReserve is a no-op once the set is non-empty", + zs->member_index.size == size_before); + zbtreeFree(zs); + } + + printf("Testing B+ tree zbtIndexScanSlot() tag-only fast path\n"); + + /* A full ZSCAN pass (real cursor-driven, repeated zbtreeScan() calls + * exactly like the SCAN command does, not one big call) must still + * visit every live member -- checked against independent ground truth + * from ordered iteration, not just a count, since a scan bug could + * plausibly emit one member twice while missing another and still + * balance to the right total. members:N is dense enough that some + * leaves are very likely to hold two or more members whose hash tags + * collide, exercising the count>1 disambiguation path as well as the + * count==1 fast path. */ + { + const int N = 8000; + zbtreeSet *zs = zbtreeCreate(); + zbtreeReserve(zs, N); + zbtTestPopulate(zs, N); + + int *seen = zcalloc(sizeof(int) * N); + zbtScanTestPrivdata pd = {seen, N}; + uint64_t cursor = 0; + unsigned long calls = 0; + do { + cursor = zbtreeScan(zs, cursor, 10, zbtScanTestMarkSeen, &pd); + calls++; + /* A real client issues thousands of these for a set this size; + * this is just a sanity backstop against an infinite cursor. */ + serverAssert(calls < 100000); + } while (cursor != 0); + + int missed = 0; + for (int i = 0; i < N; i++) + if (seen[i] == 0) missed++; + test_cond("Cursor-driven scan visits every live member at least once", + missed == 0); + zfree(seen); + zbtreeFree(zs); + } + + printf("Testing B+ tree ZSCAN mid-incremental-rehash\n"); + + /* The fast path's "provably redundant" argument (see the comment in + * zbtIndexScanSlot()) rests on every leaf's stored tag always matching + * that member's real hash. zbtreeReserve() sidesteps member_rehash + * entirely, so the block above never actually exercises a scan while + * zs->member_rehash != NULL -- the split old/new-table walk in + * zbtreeScan() and the leaf-migration bookkeeping it depends on. + * + * Growing one member at a time from empty triggers a resize almost + * immediately (around 125 elements / 2 leaves given this build's load + * factor), and that first resize's own triggering insert also migrates + * 1-2 leaves via zbtIndexInsert()'s eager copy-on-insert -- for a set + * that small, that single call can migrate EVERY leaf, leaving the "old" + * table a near-empty husk that only the finalize step (never reached, + * since the test stops touching the tree) would free. A scan over that + * state exercises the new-table read path but not the old-table one -- + * confirmed by mutation-testing: disabling zbtreeScan()'s old-table + * bucket range entirely still passed a first version of this test. + * + * Fix: reserve a large capacity first (bypasses rehash, as usual), fill + * it to exactly that size, then keep inserting one at a time past it. + * The resize this triggers has hundreds of leaves to drain and still + * only migrates 1-2 per call, so by the time the loop below breaks, the + * old table genuinely holds the vast majority of live members -- not + * just the migration flag with nothing left to migrate. */ + { + const int RESERVED = 50000; + const int N = 100000; + zbtreeSet *zs = zbtreeCreate(); + zbtreeReserve(zs, RESERVED); + zbtTestPopulateRange(zs, 0, RESERVED); + serverAssert(zs->member_rehash == NULL); + + int inserted = RESERVED; + while (inserted < N) { + zbtTestPopulateRange(zs, inserted, inserted + 1); + inserted++; + /* See the comment above: a resize can start and fully drain + * within the same insert call while the tree is still small, so + * this loop doesn't stop on the first sighting of member_rehash + * -- it stops once a resize has genuinely outlived its own call, + * checked via the backstop assertion below. */ + if (zs->member_rehash != NULL && + zs->member_rehash->table.used < (unsigned long)zs->length) + break; + } + serverAssert(zs->member_rehash != NULL); + serverAssert(inserted < N); + /* The real point of this test: the new table must hold only a small + * minority of live members (not just "not literally all of them"), + * or the old-table read path in zbtreeScan()'s split walk goes + * essentially unexercised no matter how the loop above got here. At + * RESERVED=50000, only 1-2 leaves migrate via the triggering insert, + * so this holds by a wide margin (table.used stays in the low + * hundreds against a length above 50000) -- the loose bound below is + * a backstop against that margin eroding, not the actual number. */ + serverAssert(zs->member_rehash->table.used < (unsigned long)zs->length / 2); + + int *seen = zcalloc(sizeof(int) * inserted); + zbtScanTestPrivdata pd = {seen, inserted}; + uint64_t cursor = 0; + unsigned long calls = 0; + do { + cursor = zbtreeScan(zs, cursor, 10, zbtScanTestMarkSeen, &pd); + calls++; + serverAssert(calls < 100000); + } while (cursor != 0); + /* Documents an invariant, not a discriminating check: zbtreeScan() + * never calls zbtIndexFind()/zbtIndexRehashStep(), and this test makes + * no further insert/find calls after the loop above breaks, so nothing + * here could have driven the rehash to completion either way. Pins + * the assumption for a future reader, not a claim this could fail. */ + test_cond("member_rehash is still in progress after the scan", + zs->member_rehash != NULL); + + int missed = 0; + for (int i = 0; i < inserted; i++) + if (seen[i] == 0) missed++; + test_cond("Scan mid-incremental-rehash visits every live member", + missed == 0); + zfree(seen); + zbtreeFree(zs); + } + + printf("Testing B+ tree zbtreeInsertNewAppend() vs the general insert path\n"); + + /* zbtreeInsertNewAppend() takes a shortcut only valid when the caller + * guarantees strictly increasing (score, member) order -- exactly what + * zsetConvertAndExpand()'s SKIPLIST->BTREE loop feeds it. Build the same + * sorted, duplicate-score-heavy sequence two ways (the fast append path, + * and the general zbtreeFindForAdd()+zbtreeInsertNew() path every other + * caller uses) and require byte-identical trees: same length, same + * forward order, same reverse order. Spans several thousand elements so + * it crosses many leaf/inner-node boundaries, and groups of 20 members + * share each score so multi-child same-max_score runs (the case + * zbtScoreChild's member tie-break exists for) are exercised on both + * sides. */ + { + const int n = 4000, group_size = 20; + zbtreeSet *fast = zbtreeCreate(); + zbtreeSet *slow = zbtreeCreate(); + char buf[32]; + for (int i = 0; i < n; i++) { + snprintf(buf, sizeof(buf), "member:%05d", i); + double score = (double)(i / group_size); + + zbtreeInsertNewAppend(fast, score, (unsigned char *)buf, + strlen(buf)); + + sds ele = sdsnew(buf); + double existing_score; + zbtreeInsertPosition position; + int existing = zbtreeFindForAdd(slow, ele, &existing_score, + &position); + serverAssert(!existing); + zbtreeInsertNew(slow, score, ele, &position); + sdsfree(ele); + } + test_cond("Append and general insert build sets of the same length", + zbtreeLength(fast) == (unsigned long)n && + zbtreeLength(slow) == (unsigned long)n); + + zbtreeIterator fast_iter, slow_iter; + int ffound = zbtreeIteratorStart(fast, 0, &fast_iter); + int sfound = zbtreeIteratorStart(slow, 0, &slow_iter); + int forward_match = 1; + long forward_count = 0; + while (ffound && sfound) { + const unsigned char *fele, *sele; + size_t flen, slen; + double fscore, sscore; + ffound = zbtreeIteratorNext(&fast_iter, 0, &fele, &flen, &fscore); + sfound = zbtreeIteratorNext(&slow_iter, 0, &sele, &slen, &sscore); + if (ffound != sfound) { forward_match = 0; break; } + if (!ffound) break; + if (fscore != sscore || flen != slen || + memcmp(fele, sele, flen) != 0) + { + forward_match = 0; + break; + } + forward_count++; + } + test_cond("Forward iteration is byte-identical between the two trees", + forward_match && forward_count == n); + + ffound = zbtreeIteratorStart(fast, 1, &fast_iter); + sfound = zbtreeIteratorStart(slow, 1, &slow_iter); + int reverse_match = 1; + long reverse_count = 0; + while (ffound && sfound) { + const unsigned char *fele, *sele; + size_t flen, slen; + double fscore, sscore; + ffound = zbtreeIteratorNext(&fast_iter, 1, &fele, &flen, &fscore); + sfound = zbtreeIteratorNext(&slow_iter, 1, &sele, &slen, &sscore); + if (ffound != sfound) { reverse_match = 0; break; } + if (!ffound) break; + if (fscore != sscore || flen != slen || + memcmp(fele, sele, flen) != 0) + { + reverse_match = 0; + break; + } + reverse_count++; + } + test_cond("Reverse iteration is byte-identical between the two trees", + reverse_match && reverse_count == n); + + zbtreeFree(fast); + zbtreeFree(slow); + } + + printf("Testing B+ tree ZSCAN across a same-tag slot exchange\n"); + + /* Deterministic regression for a ZSCAN omission. X and Y have the same + * folded 8 bit index tag. Y's slot b1 lies on X's probe path before X's + * slot b2, and the two slots are in different scan groups. The cursor + * passes b1 while X and Y live in different leaves. X then moves into + * Y's leaf and out again. The second zbtIndexMove() relabels the first + * (tag, leaf) slot on X's path, b1, so X ends up represented only behind + * the cursor while b2 now stands for Y. X never leaves the set. */ + { + uint8_t seed[16] = {0}; + dictSetHashFunctionSeed(seed); + const int N = 7000; /* 1024 buckets, about 85% load, no resize. */ + zbtreeSet *zs = zbtreeCreate(); + zbtreeReserve(zs, N); + int *perm = zmalloc(sizeof(int) * N); + for (int i = 0; i < N; i++) perm[i] = i; + uint64_t rnd = 12345; + for (int i = N - 1; i > 0; i--) { + rnd = rnd * 6364136223846793005ULL + 1442695040888963407ULL; + int j = (int)((rnd >> 33) % (uint64_t)(i + 1)); + int tmp = perm[i]; perm[i] = perm[j]; perm[j] = tmp; + } + for (int k = 0; k < N; k++) + zbtTestPopulateRange(zs, perm[k], perm[k] + 1); + zfree(perm); + serverAssert(zs->member_rehash == NULL); + + zbtIndexTable *t = &zs->member_index; + unsigned long mask = t->size - 1; + int X = -1, Y = -1; + unsigned long b1 = 0, b2 = 0; + unsigned int p1 = 0, p2 = 0; + uint32_t ly_id = 0; + char xbuf[32], ybuf[32]; + + for (int x = 0; x < N && X < 0; x++) { + snprintf(xbuf, sizeof(xbuf), "member:%d", x); + uint32_t hx = zbtTestHash(xbuf); + uint8_t tag = zbtIndexTag(hx); + zbtScoreLeaf *lx = zbtTestLeafOf(zs, xbuf); + if (zbtTestTagCount(lx, tag) != 1) continue; + unsigned long sb; + unsigned int sp; + serverAssert(zbtTestSlotOf(zs, hx, lx->id, &sb, &sp)); + unsigned long home = hx & mask; + if (sb <= home) continue; /* Home slot, or a wrapped probe run. */ + for (unsigned long b = home; b < sb && X < 0; b++) { + /* The scan must visit b1's group strictly before b2's. */ + if (zbtScanReverse(b / ZBT_SCAN_BUCKETS_PER_STEP) >= + zbtScanReverse(sb / ZBT_SCAN_BUCKETS_PER_STEP)) continue; + zbtIndexBucket *bucket = zbtIndexBucketAt(t, b); + for (unsigned int p = 0; p < ZBT_INDEX_BUCKET_ITEMS; p++) { + uint8_t slot_tag = zbtIndexTags(bucket) >> (p * 8); + uint32_t id = zbtIndexGetId(t, bucket, p); + if (slot_tag != tag || id == ZBT_INDEX_DELETED_ID || + id == lx->id) continue; + zbtScoreLeaf *ly = zs->score_leaf_by_id[id]; + if (zbtTestTagCount(ly, tag) != 1) continue; + if (ly->n.count >= ZBT_SCORE_LEAF_MAX - 1 || + ly == zs->score_last) continue; + int yi = -1; + for (unsigned int i = 0; i < ly->n.count; i++) { + if (zbtIndexTag(zbtScoreLeafHash(ly, i)) != tag) + continue; + size_t len; + const unsigned char *e = + zbtScoreLeafElement(ly, i, &len); + char tmp[32]; + memcpy(tmp, e, len); + tmp[len] = '\0'; + serverAssert(sscanf(tmp, "member:%d", &yi) == 1); + } + if (yi < 0 || abs(yi - x) < 3) continue; + snprintf(ybuf, sizeof(ybuf), "member:%d", yi); + unsigned long yb; + unsigned int yp; + serverAssert(zbtTestSlotOf(zs, zbtTestHash(ybuf), id, + &yb, &yp)); + if (yb != b || yp != p) continue; + X = x; Y = yi; b1 = b; p1 = p; b2 = sb; p2 = sp; + ly_id = id; + break; + } + } + } + test_cond("Same-tag scan: found X/Y with same tag and overlapping probe runs", + X >= 0); + if (X >= 0) { + snprintf(xbuf, sizeof(xbuf), "member:%d", X); + snprintf(ybuf, sizeof(ybuf), "member:%d", Y); + printf(" X=%s Y=%s tag=%u b1=%lu/%u b2=%lu/%u leafY=%u\n", + xbuf, ybuf, zbtIndexTag(zbtTestHash(xbuf)), b1, p1, + b2, p2, ly_id); + + int *seen = zcalloc(sizeof(int) * N); + zbtScanTestPrivdata pd = {seen, N}; + uint64_t cursor = 0; + do { + cursor = zbtreeScan(zs, cursor, 1, zbtScanTestMarkSeen, &pd); + serverAssert(cursor != 0); + } while (!zbtTestGroupVisited(cursor, + b1 / ZBT_SCAN_BUCKETS_PER_STEP)); + serverAssert(!zbtTestGroupVisited(cursor, + b2 / ZBT_SCAN_BUCKETS_PER_STEP)); + test_cond("Same-tag scan: cursor is between b1 and b2", + seen[Y] > 0); + + /* Move X into Y's leaf: just below Y, far from X's neighbors. */ + zbtTestSetScore(zs, xbuf, (double)Y - 0.5); + test_cond("Same-tag scan: X now lives in Y's leaf", + zbtTestLeafOf(zs, xbuf)->id == ly_id); + zbtIndexBucket *s1 = zbtIndexBucketAt(t, b1); + zbtIndexBucket *s2 = zbtIndexBucketAt(t, b2); + test_cond("Same-tag scan: both slots now name Y's leaf", + zbtIndexGetId(t, s1, p1) == ly_id && + zbtIndexGetId(t, s2, p2) == ly_id); + + /* Move X out to the right edge. */ + zbtTestSetScore(zs, xbuf, 1e9); + uint32_t lr_id = zbtTestLeafOf(zs, xbuf)->id; + test_cond("Same-tag scan: X moved to another leaf", lr_id != ly_id); + printf(" after moves: slot b1 -> leaf %u, slot b2 -> leaf %u," + " X leaf %u\n", zbtIndexGetId(t, s1, p1), + zbtIndexGetId(t, s2, p2), lr_id); + + while (cursor != 0) + cursor = zbtreeScan(zs, cursor, 1, zbtScanTestMarkSeen, &pd); + int missed = 0; + for (int i = 0; i < N; i++) if (seen[i] == 0) missed++; + printf(" X seen %d time(s); total members missed: %d\n", + seen[X], missed); + test_cond("ZSCAN returns every member despite same-tag slot exchanges", + missed == 0); + zfree(seen); + } + zbtreeFree(zs); + } + + printf("Testing B+ tree memory after deleting most members\n"); + { + /* 319 byte members are the longest kept inside a leaf, so a leaf + * holds about a dozen. Build in random order, change and delete some + * members, reload in RDB order, then delete 95% spread evenly over + * all leaves. Deletion used to leave most surviving members alone in + * a 4 KB leaf: about five times the memory of a rebuilt set. */ + const int N = 60000; + char buf[319]; + memset(buf, 'x', sizeof(buf)); + int *order = zmalloc(sizeof(int) * N); + char *gone = zcalloc(N); + for (int i = 0; i < N; i++) order[i] = i; + uint64_t r = 654; + for (int i = N - 1; i > 0; i--) { + r = r * 6364136223846793005ULL + 1442695040888963407ULL; + int j = (int)((r >> 33) % (uint64_t)(i + 1)); + int t = order[i]; order[i] = order[j]; order[j] = t; + } +#define ZBT_TEST_SET(zs, i, score) do { \ + snprintf(buf, 13, "%012d", (i)); buf[12] = 'x'; \ + sds ele = sdsnewlen(buf, sizeof(buf)); \ + double cur; zbtreeInsertPosition position; \ + if (!zbtreeFindForAdd((zs), ele, &cur, &position)) \ + zbtreeInsertNew((zs), (score), ele, &position); \ + else if (cur != (score)) \ + zbtreeUpdateScore((zs), ele, (score), &position); \ + sdsfree(ele); \ + } while (0) +#define ZBT_TEST_DEL(zs, i) do { \ + snprintf(buf, 13, "%012d", (i)); buf[12] = 'x'; \ + sds ele = sdsnewlen(buf, sizeof(buf)); \ + serverAssert(zbtreeDelete((zs), ele)); \ + sdsfree(ele); gone[(i)] = 1; \ + } while (0) + zbtreeSet *zs = zbtreeCreate(); + for (int k = 0; k < N; k++) ZBT_TEST_SET(zs, order[k], order[k] % 997); + for (int k = 0; k < N; k++) { + int i = order[k]; + if (i % 5 == 0) ZBT_TEST_DEL(zs, i); + else if (i % 3 == 0) ZBT_TEST_SET(zs, i, -(double)i / 8); + } + /* Reload the way RDB does: descending order into a new set. */ + zbtreeSet *loaded = zbtreeCreate(); + zbtreeIterator it; + const unsigned char *e; + size_t len; + double score; + zbtreeIteratorStart(zs, 1, &it); + while (zbtreeIteratorNext(&it, 1, &e, &len, &score)) { + sds ele = sdsnewlen(e, len); + double cur; + zbtreeInsertPosition position; + serverAssert(!zbtreeFindForAdd(loaded, ele, &cur, &position)); + zbtreeInsertNew(loaded, score, ele, &position); + sdsfree(ele); + } + zbtreeFree(zs); + zs = loaded; + for (int i = 0; i < N; i++) + if (!gone[i] && i % 20 != 1) ZBT_TEST_DEL(zs, i); +#undef ZBT_TEST_SET +#undef ZBT_TEST_DEL + zbtreeSet *packed = zbtreeCreate(); + int intact = 1; + const unsigned char *pe = NULL; + size_t plen = 0; + double pscore = 0; + zbtreeIteratorStart(zs, 0, &it); + while (zbtreeIteratorNext(&it, 0, &e, &len, &score)) { + int i = atoi((const char *)e); + intact = intact && i % 20 == 1 && !gone[i] && + score == (i % 3 == 0 ? -(double)i / 8 : i % 997); + if (pe) intact = intact && + zbtScoreCompare(pscore, pe, plen, score, e, len) < 0; + pe = e; plen = len; pscore = score; + zbtreeInsertNewAppend(packed, score, e, len); + } + test_cond("Survivors are intact after deleting 95% of the members", + intact && zbtreeLength(zs) == (unsigned long)N / 20); + test_cond("Sparse leaves release memory after mass deletion", + zbtreeAllocSize(zs) < 2 * zbtreeAllocSize(packed)); + zbtreeFree(packed); + zbtreeFree(zs); + zfree(order); + zfree(gone); + } + + printf("Testing B+ tree memory when end deletions stop at one member\n"); + for (int reverse = 0; reverse <= 1; reverse++) { + zbtreeSet *zs = zbtreeCreate(); + char buf[319]; + memset(buf, 'x', sizeof(buf)); + for (int i = 0; i < 96; i++) { + snprintf(buf, 13, "%012d", i); buf[12] = 'x'; + zbtreeInsertNewAppend(zs, i, (unsigned char *)buf, sizeof(buf)); + } + for (int k = 0; k < 95; k++) { + int i = reverse ? 95 - k : k; + snprintf(buf, 13, "%012d", i); buf[12] = 'x'; + sds ele = sdsnewlen(buf, sizeof(buf)); + serverAssert(zbtreeDelete(zs, ele)); + sdsfree(ele); + } + test_cond("Stopping end deletions at one member releases the leaf's unused space", + zbtreeLength(zs) == 1 && zs->score_first == zs->score_last && + zmalloc_usable_size(zs->score_first) < 1024 && + zbtScoreLeafScore(zs->score_first, 0) == (reverse ? 0 : 95)); + zbtreeFree(zs); + } + + printf("Testing B+ tree member index shrinking during range deletions\n"); + { + /* Each range stays below a tenth of the set, so it never rebuilds + * the member index and only incremental resizes can shrink it. Range + * deletion looks no member up by name; unless it advances a resize + * as fast as ZREM would, the shrink cannot keep up with it. */ + const int N = 1000000, B = 3072, K = 96; + zbtreeSet *zs = zbtreeCreate(); + zbtTestPopulate(zs, N); + int bounded = 1; + for (int s = 0; s < N; s += B) { + zbtreeDeleteRangeByScore(zs, s + K, 0, s + B - 1, 0); + unsigned long buckets = zs->member_index.size; + if (zs->member_rehash) buckets += zs->member_rehash->table.size; + if (buckets > 16 * zbtIndexBucketsForElements(zbtreeLength(zs))) + bounded = 0; + } + test_cond("Range deletions shrink the member index as they go", + bounded && zbtTestCheck(zs) && + zbtreeLength(zs) == (unsigned long)(N + B - 1) / B * K); + zbtreeFree(zs); + } + + printf("Testing B+ tree member index clean-up close to the maximum load\n"); + { + /* Replacing members fills the table with tombstones until a clean-up + * resize starts. Here the set is then close to the maximum load of + * its table. Adding members must still let the copy proceed a few + * leaves at a time instead of finishing it within one insertion. */ + const unsigned long buckets = 4096; + const unsigned long max_live = buckets * ZBT_INDEX_BUCKET_ITEMS * + ZBT_INDEX_MAX_LOAD_NUM / ZBT_INDEX_MAX_LOAD_DEN; + char buf[32]; + int oldest = 0, next = (int)max_live - 32; + zbtreeSet *zs = zbtreeCreate(); + zbtTestPopulate(zs, next); + int sized = zs->member_index.size == buckets && !zs->member_rehash; + while (sized && !zs->member_rehash) { + snprintf(buf, sizeof(buf), "member:%d", oldest++); + sds ele = sdsnew(buf); + serverAssert(zbtreeDelete(zs, ele)); + sdsfree(ele); + zbtTestPopulateRange(zs, next, next + 1); + next++; + } + /* One insertion copies at most three leaves: one step in its lookup, + * its own leaf, and one step after it. A finished copy may be + * followed at once by a larger resize, so compare resize numbers. */ + int bounded = 1; + while (sized && zbtreeLength(zs) <= 2 * max_live) { + unsigned long pending = 0; + uint16_t resize_id = 0; + if (zs->member_rehash) { + resize_id = zs->member_rehash->resize_id; + for (zbtScoreLeaf *l = zs->score_first; l; l = l->next) + if (!zbtIndexLeafMigrated(zs, l)) pending++; + } + zbtTestPopulateRange(zs, next, next + 1); + next++; + if (pending > 3 && (!zs->member_rehash || + zs->member_rehash->resize_id != resize_id)) + bounded = 0; + } + test_cond("A clean-up resize next to the maximum load stays incremental", + sized && bounded && zbtTestCheck(zs)); + zbtreeFree(zs); + } + + printf("Testing B+ tree random deletes at both ends and in the middle\n"); + { + /* Mix pops from both ends, random deletes, inserts and score + * changes, with inline and external members. Compare everything + * with a plain array after every batch. */ + const int U = 3000; + double *ref = zmalloc(sizeof(double) * U); + char *alive = zcalloc(U); + char buf[512]; + uint64_t r = 4242; + int ok = 1; +#define ZBT_TEST_RAND() (r = r * 6364136223846793005ULL + \ + 1442695040888963407ULL, (unsigned)(r >> 33)) +#define ZBT_TEST_MEMBER(i) (snprintf(buf, 13, "%012d", (i)), \ + memset(buf + 12, 'a' + (i) % 26, \ + (size_t)(8 + ((i) * 37) % 400)), \ + (size_t)(20 + ((i) * 37) % 400)) + zbtreeSet *zs = zbtreeCreate(); + for (int round = 0; round < 60 && ok; round++) { + for (int op = 0; op < 2000; op++) { + unsigned kind = ZBT_TEST_RAND() % 10; + int i = (int)(ZBT_TEST_RAND() % U); + if (kind <= 1 && zbtreeLength(zs)) { + zbtreeIterator it; + const unsigned char *e; + size_t len; + double score; + int reverse = kind == 1; + serverAssert(zbtreeIteratorStart(zs, reverse, &it)); + serverAssert(zbtreeIteratorNext(&it, reverse, &e, &len, + &score)); + int m = atoi((const char *)e); + sds ele = sdsnewlen(e, len); + serverAssert(zbtreeDelete(zs, ele)); + sdsfree(ele); + ok = ok && alive[m] && ref[m] == score; + alive[m] = 0; + continue; + } + size_t len = ZBT_TEST_MEMBER(i); + sds ele = sdsnewlen(buf, len); + if (kind <= 4) { + ok = ok && zbtreeDelete(zs, ele) == alive[i]; + alive[i] = 0; + } else { + double score = (double)(ZBT_TEST_RAND() % 50); + double cur; + zbtreeInsertPosition position; + int found = zbtreeFindForAdd(zs, ele, &cur, &position); + ok = ok && found == alive[i]; + if (!found) + zbtreeInsertNew(zs, score, ele, &position); + else if (cur != score) + zbtreeUpdateScore(zs, ele, score, &position); + alive[i] = 1; + ref[i] = score; + } + sdsfree(ele); + } + unsigned long count = 0; + for (int i = 0; i < U; i++) { + if (!alive[i]) continue; + count++; + size_t len = ZBT_TEST_MEMBER(i); + double score; + ok = ok && zbtreeScoreRaw(zs, (unsigned char *)buf, len, + &score) && score == ref[i]; + } + zbtreeIterator it; + const unsigned char *e, *pe = NULL; + size_t len, plen = 0; + double score, pscore = 0; + unsigned long seen = 0; + zbtreeIteratorStart(zs, 0, &it); + while (zbtreeIteratorNext(&it, 0, &e, &len, &score)) { + if (pe) ok = ok && zbtScoreCompare(pscore, pe, plen, score, + e, len) < 0; + pe = e; plen = len; pscore = score; + seen++; + } + ok = ok && count == zbtreeLength(zs) && seen == count; + } +#undef ZBT_TEST_RAND +#undef ZBT_TEST_MEMBER + test_cond("Random pops, deletes and updates match a reference", ok); + zbtreeFree(zs); + zfree(ref); + zfree(alive); + } + + printf("Testing B+ tree first-leaf pop with mixed signed zero scores\n"); + { + /* +0.0 and -0.0 compare equal, so members order them. The last member + * of a leaf can then have the smaller ordered score. */ + const char *names[] = {"a", "b", "c", "d"}; + double in[] = {0.0, 0.0, 0.0, -0.0}; + zbtreeSet *zs = zbtreeCreate(); + for (int i = 0; i < 4; i++) { + sds ele = sdsnew(names[i]); + double cur; + zbtreeInsertPosition position; + serverAssert(!zbtreeFindForAdd(zs, ele, &cur, &position)); + zbtreeInsertNew(zs, in[i], ele, &position); + sdsfree(ele); + } + sds first = sdsnew("a"); + serverAssert(zbtreeDelete(zs, first)); + sdsfree(first); + int ok = zbtreeLength(zs) == 3; + for (int i = 1; i < 4; i++) { + sds ele = sdsnew(names[i]); + double score = 12345; + int found = zbtreeScore(zs, ele, &score); + printf(" %s found=%d score=%.17g signbit=%d\n", names[i], found, + score, signbit(score) != 0); + ok = ok && found && score == 0 && + (signbit(score) != 0) == (signbit(in[i]) != 0); + sdsfree(ele); + } + test_cond("Popping the first member keeps signed zero scores exact", ok); + zbtreeFree(zs); + } + + printf("Testing B+ tree leaf IDs and the narrow member index\n"); + + /* Score changes split leaves too. They used to do it without ever + * starting the widening, until a 16 bit leaf ID overflowed. Members are + * appended into full leaves, so most moves split one. */ + { + const int N = 20000; + char buf[64]; + uint64_t r = 1; + double *ref = zmalloc(sizeof(double) * N * 2); + size_t *lens = zcalloc(sizeof(size_t) * N * 2); + zbtreeSet *zs = zbtreeCreate(); + for (int i = 0; i < N; i++) { + size_t len = zbtTestMember(buf, i, 0); + zbtreeInsertNewAppend(zs, i * 2, (unsigned char *)buf, len); + ref[i] = i * 2; + } + serverAssert(!zs->member_index.wide_ids && !zs->member_rehash); + zbtTestSkipLeafIds(zs, ZBT_INDEX_WIDE_ID_AT - 4); + zbtTestRandomMoves(zs, N, ref, lens, 2000, &r); + int widening = zs->member_index.wide_ids || + (zs->member_rehash && + zs->member_rehash->table.wide_ids); + test_cond("Score changes start widening a narrow member index", + widening && zbtTestCheck(zs) && + zbtTestVerify(zs, N, ref, NULL, lens)); + zbtTestRandomMoves(zs, N, ref, lens, 20000, &r); + test_cond("Score changes finish widening the member index", + zs->member_index.wide_ids && !zs->member_rehash && + zbtTestCheck(zs) && zbtTestVerify(zs, N, ref, NULL, lens)); + zbtreeFree(zs); + + /* A narrow table next to the limit, as the score change path used + * to leave it: the next splits must not store an ID that does not + * fit. */ + zs = zbtreeCreate(); + for (int i = 0; i < N; i++) { + size_t len = zbtTestMember(buf, i, 0); + zbtreeInsertNewAppend(zs, i * 2, (unsigned char *)buf, len); + ref[i] = i * 2; + } + zbtTestSkipLeafIds(zs, UINT16_MAX - 4); + zbtTestRandomMoves(zs, N, ref, lens, 2000, &r); + test_cond("Score changes at the narrow ID limit widen first", + zbtTestCheck(zs) && zbtTestVerify(zs, N, ref, NULL, lens) && + zs->next_score_leaf_id > UINT16_MAX); + zbtreeFree(zs); + + /* A narrow resize that runs while IDs approach the limit. Released + * IDs are not reused during a resize, so its destination receives + * new IDs. It must end, and the widening start, in time: faster than + * usual from ZBT_INDEX_WIDE_ID_AT, at once next to the limit. */ + for (int late = 0; late <= 1; late++) { + zs = zbtreeCreate(); + zbtreeReserve(zs, N); + for (int i = 0; i < N; i++) { + size_t len = zbtTestMember(buf, i, 0); + zbtreeInsertNewAppend(zs, i * 2, (unsigned char *)buf, len); + ref[i] = i * 2; + } + int extra = N; + while (zs->member_rehash == NULL) { + size_t len = zbtTestMember(buf, extra, 0); + zbtreeInsertNewAppend(zs, extra * 2, (unsigned char *)buf, len); + ref[extra] = extra * 2; + extra++; + } + serverAssert(extra < N * 2); + serverAssert(!zs->member_rehash->table.wide_ids); + unsigned long leaves = zbtTestLeaves(zs); + zbtTestSkipLeafIds(zs, late ? UINT16_MAX - 4 : 40000); + long moves = 0; + int checked = 1; + while (moves < 5000 && + !(zs->member_rehash ? zs->member_rehash->table.wide_ids : + zs->member_index.wide_ids)) + { + zbtTestRandomMoves(zs, N, ref, lens, 1, &r); + moves++; + if (moves % 16 == 0) checked = checked && zbtTestCheck(zs); + } + printf(" narrow resize with %lu leaves, IDs from %s: widening " + "after %ld moves\n", leaves, late ? "65531" : "40000", + moves); + test_cond(late ? "A narrow resize next to the ID limit finishes " + "before a split" : + "A narrow resize past the widening point " + "finishes early", + moves <= (late ? 2 : (long)leaves / 3 + 2) && checked && + zbtTestCheck(zs) && + zbtTestVerify(zs, extra, ref, NULL, lens)); + zbtreeFree(zs); + } + zfree(ref); + zfree(lens); + } + + /* A new member inserted into the right half of a full leaf that the + * widening has not copied yet. The split used to relabel the old narrow + * entries of the moved members with the new, too wide, leaf ID. */ + { + const int N = 30000; + char buf[64]; + zbtreeSet *zs = zbtreeCreate(); + for (int i = 0; i < N; i++) { + size_t len = zbtTestMember(buf, i, 0); + zbtreeInsertNewAppend(zs, i * 2, (unsigned char *)buf, len); + } + zbtTestSkipLeafIds(zs, 40000); + size_t len = zbtTestMember(buf, N, 0); + zbtTestSet(zs, buf, len, N * 2); + serverAssert(zs->member_rehash && zs->member_rehash->table.wide_ids && + !zs->member_index.wide_ids); + zbtTestSkipLeafIds(zs, 70000); + zbtScoreLeaf *leaf = zs->score_last->prev->prev; + serverAssert(leaf->n.count == ZBT_SCORE_LEAF_MAX && + !zbtIndexLeafMigrated(zs, leaf)); + double score = zbtScoreLeafScore(leaf, 80) + 1; + len = zbtTestMember(buf, N + 1, 0); + zbtTestSet(zs, buf, len, score); + zbtScoreLeaf *found = zbtTestLeafOf(zs, buf); + double check; + test_cond("A split during a late widening keeps old entries narrow", + found->id >= 70000 && zbtIndexLeafMigrated(zs, found) && + zbtreeScoreRaw(zs, (unsigned char *)buf, len, &check) && + check == score && zbtTestCheck(zs) && + zbtreeLength(zs) == (unsigned long)N + 2); + zbtreeFree(zs); + } + + /* A set emptied by single deletions keeps its leaf ID counter and reuses + * released IDs. A new table must be wide enough for them. */ + { + char buf[64]; + zbtreeSet *zs = zbtreeCreate(); + for (int i = 0; i < 2000; i++) { + size_t len = zbtTestMember(buf, i, 0); + zbtreeInsertNewAppend(zs, i, (unsigned char *)buf, len); + } + zbtTestSkipLeafIds(zs, 70000); + for (int i = 2000; i < 6000; i++) { + size_t len = zbtTestMember(buf, i, 0); + zbtreeInsertNewAppend(zs, i, (unsigned char *)buf, len); + } + for (int i = 0; i < 6000; i++) { + size_t len = zbtTestMember(buf, i, 0); + sds ele = sdsnewlen(buf, len); + serverAssert(zbtreeDelete(zs, ele)); + sdsfree(ele); + } + serverAssert(zs->length == 0 && zs->member_index.size == 0 && + zs->free_score_leaf_id > 70000); + size_t len = zbtTestMember(buf, 0, 0); + zbtTestSet(zs, buf, len, 1); + test_cond("An emptied set reuses a high leaf ID with a wide table", + zs->score_first->id >= 70000 && zs->member_index.wide_ids && + zbtTestCheck(zs) && zbtreeLength(zs) == 1); + zbtreeFree(zs); + } + + printf("Testing B+ tree leaf density under score changes\n"); + + /* Score changes used to leave their old leaves behind without ever + * joining or shrinking them. Start from densely appended leaves and + * compare the leaf count with the fewest leaves that could hold the + * members. Mixed lengths are random per member, so neighbors do not have + * similar sizes. */ + { + static const struct { + const char *name; + size_t min, max; + } profiles[] = { + {"319 byte", 319, 319}, + {"short", 1, 1}, + {"mixed 1-319 byte", 1, 319}, + {"mixed 1-600 byte", 1, 600}, + }; + const int N = 10000; + for (size_t p = 0; p < sizeof(profiles) / sizeof(profiles[0]); p++) { + char buf[1024]; + uint64_t r = 77 + p; + double *ref = zmalloc(sizeof(double) * N); + size_t *lens = zmalloc(sizeof(size_t) * N); + zbtreeSet *zs = zbtreeCreate(); + for (int i = 0; i < N; i++) { + lens[i] = profiles[p].min + zbtTestRand(&r) % + (profiles[p].max - profiles[p].min + 1); + size_t len = zbtTestMember(buf, i, lens[i]); + zbtreeInsertNewAppend(zs, i, (unsigned char *)buf, len); + ref[i] = i; + } + unsigned long fresh = zbtTestLeaves(zs); + zbtTestRandomMoves(zs, N, ref, lens, (long)N * 20, &r); + size_t live = 0, external = 0; + unsigned long leaves = 0; + for (zbtScoreLeaf *leaf = zs->score_first; leaf; + leaf = leaf->next) + { + live += zbtScoreLeafLiveBytes(leaf); + leaves++; + for (unsigned int i = 0; i < leaf->n.count; i++) { + unsigned char *record = zbtScoreLeafRecord(leaf, i); + if (record[0] == ZBT_PACKED_EXTERNAL) + external += zbtExternalLength(record); + } + } + unsigned long fewest = (N + ZBT_SCORE_LEAF_MAX - 1) / + ZBT_SCORE_LEAF_MAX; + unsigned long by_bytes = (live + ZBT_SCORE_LEAF_BYTES - 1) / + ZBT_SCORE_LEAF_BYTES; + if (by_bytes > fewest) fewest = by_bytes; + printf(" %s members: %lu leaves fresh, %lu after moves, " + "fewest possible %lu (%.2fx), memory %.2f of live\n", + profiles[p].name, fresh, leaves, fewest, + (double)leaves / fewest, + (double)zbtreeAllocSize(zs) / (live + external)); + char title[128]; + snprintf(title, sizeof(title), + "Score changes keep %s leaves dense", profiles[p].name); + test_cond(title, leaves <= fewest * 2 && zbtTestCheck(zs) && + zbtTestVerify(zs, N, ref, NULL, lens)); + zbtreeFree(zs); + zfree(ref); + zfree(lens); + } + } + + /* A leaf of long members next to a leaf of short ones. Estimating the + * neighbor from the long members says that they cannot be joined, but + * once moves leave the long leaf sparse it must check the real sizes. */ + { + char buf[1024]; + zbtreeSet *zs = zbtreeCreate(); + for (int i = 0; i < ZBT_SCORE_LEAF_MAX; i++) { + size_t len = zbtTestMember(buf, i, 0); + zbtreeInsertNewAppend(zs, i, (unsigned char *)buf, len); + } + for (int i = ZBT_SCORE_LEAF_MAX; i < ZBT_SCORE_LEAF_MAX + 40; i++) { + size_t len = zbtTestMember(buf, i, 319); + zbtreeInsertNewAppend(zs, i, (unsigned char *)buf, len); + } + for (int i = 10; i < 46; i++) { + size_t len = zbtTestMember(buf, i, 0); + sds ele = sdsnewlen(buf, len); + serverAssert(zbtreeDelete(zs, ele)); + sdsfree(ele); + } + zbtScoreLeaf *shortleaf = zs->score_first; + zbtScoreLeaf *longleaf = shortleaf->next; + unsigned int shorts = shortleaf->n.count, longs = longleaf->n.count; + /* Moving all but four long members makes two quarter-page checks. */ + serverAssert(shorts == ZBT_SCORE_LEAF_MAX - 36 && longs >= 11); + int first_long = ZBT_SCORE_LEAF_MAX; + for (unsigned int k = 0; k < longs - 4; k++) { + size_t len = zbtTestMember(buf, first_long + k, 319); + zbtTestSet(zs, buf, len, 1e9 + k); + } + test_cond("A sparse leaf joins a neighbor with much shorter members", + zs->score_first->n.count == shorts + 4 && zbtTestCheck(zs)); + zbtreeFree(zs); + } + + /* Join and rebuild the old leaves of moved members while a resize copies + * leaves between tables. Most members are external, so joins also move + * allocation ownership. Moves drain the last leaves, which the resize + * copies last, into random leaves on both sides of its progress. */ + { + const int N = 40000; + char buf[1024]; + uint64_t r = 99; + double *ref = zmalloc(sizeof(double) * N); + size_t *lens = zmalloc(sizeof(size_t) * N); + zbtreeSet *zs = zbtreeCreate(); + zbtreeReserve(zs, N / 2); + int n = 0; + while (n < N && zs->member_rehash == NULL) { + lens[n] = 200 + zbtTestRand(&r) % 400; + size_t len = zbtTestMember(buf, n, lens[n]); + zbtreeInsertNewAppend(zs, n, (unsigned char *)buf, len); + ref[n] = n; + n++; + } + serverAssert(zs->member_rehash != NULL); + unsigned long leaves_before = zbtTestLeaves(zs); + int ok = 1, during = 0; + for (int k = 0; k < 3000 && ok; k++) { + int i = n - 1 - (int)(zbtTestRand(&r) % (uint64_t)(n / 32)); + size_t len = zbtTestMember(buf, i, lens[i]); + double score = (double)(zbtTestRand(&r) % (uint64_t)n) + 0.5; + zbtTestSet(zs, buf, len, score); + ref[i] = score; + if (zs->member_rehash) during++; + if (k % 250 == 0) ok = zbtTestCheck(zs); + } + printf(" %d of 3000 moves during the resize, leaves %lu -> %lu\n", + during, leaves_before, zbtTestLeaves(zs)); + test_cond("Score change maintenance during a resize keeps ownership", + ok && during > 300 && zbtTestCheck(zs) && + zbtTestVerify(zs, n, ref, NULL, lens)); + zbtreeFree(zs); + zfree(ref); + zfree(lens); + } + + printf("Testing B+ tree leaf tag lanes\n"); + + /* ZSCAN compares eight leaf tags at once. Byte i of the array must be + * lane i of the word on every host, as for the bucket tags, or a big + * endian host selects mirrored positions. */ + { + uint64_t r = 5; + int ok = 1; + for (int k = 0; k < 1000 && ok; k++) { + uint8_t bytes[8]; + uint64_t expected = 0; + for (int i = 0; i < 8; i++) { + bytes[i] = zbtTestRand(&r); + expected |= (uint64_t)bytes[i] << (i * 8); + } + ok = zbtLeafTagWord(bytes) == expected; + } + test_cond("Leaf tag words put byte i in lane i", ok); + + /* Few distinct tags make repeats, folded zero tags, and the false + * candidates above a real match that every word comparison must + * reject. */ + ok = 1; + for (int k = 0; k < 2000 && ok; k++) { + uint8_t tags[ZBT_SCORE_LEAF_MAX]; + unsigned int count = 1 + zbtTestRand(&r) % ZBT_SCORE_LEAF_MAX; + static const uint8_t pool[] = {0, 1, 2, 3, 0x80, 0x81, 0xfe, 0xff}; + for (unsigned int i = 0; i < count; i++) + tags[i] = k % 2 ? pool[zbtTestRand(&r) % sizeof(pool)] : + (uint8_t)zbtTestRand(&r); + int reversed = k % 3 == 0; + for (unsigned int tag = 1; tag <= 255 && ok; tag++) { + unsigned int got[ZBT_SCORE_LEAF_MAX], want[ZBT_SCORE_LEAF_MAX]; + unsigned int found = zbtLeafTagMatches(tags, count, reversed, + tag, got); + unsigned int expected = 0; + for (unsigned int i = 0; i < count; i++) { + if (tags[i] == tag || (tag == 1 && tags[i] == 0)) + want[expected++] = reversed ? count - 1 - i : i; + } + ok = found == expected && + memcmp(got, want, found * sizeof(got[0])) == 0; + } + } + test_cond("Leaf tag matching agrees with a byte by byte search", ok); + } + + printf("Testing B+ tree ZSCAN while member tables keep changing\n"); + + /* SCAN must finish while a set stays bounded, however often its tables + * are replaced. The old cursor named one table and restarted whenever it + * changed. Between every call, temporary members make the index grow, + * shrink, clean tombstones, widen, or be rebuilt, while a few anchors stay + * in the set. The proof in zbtreeScan() bounds the calls by the largest + * group count plus one. */ + { + const int A = 300, T = 2000; + char buf[64]; + uint64_t r = 3; + zbtreeSet *zs = zbtreeCreate(); + for (int i = 0; i < A; i++) { + int len = snprintf(buf, sizeof(buf), "anchor:%d", i); + zbtTestSet(zs, buf, len, i); + } + int *seen = zcalloc(sizeof(int) * A); + zbtScanAnchorPrivdata pd = {seen, A}; + uint64_t cursor = 0, max_groups = 0; + unsigned long calls = 0, replaced = 0; + int small_cursor = 1, temps = 0, next_temp = 0, widened = 0; + void *buckets = zs->member_index.buckets; + do { + cursor = zbtreeScan(zs, cursor, 10, zbtScanTestMarkAnchor, &pd); + calls++; + small_cursor = small_cursor && cursor < zbtTestScanGroups(zs); + switch (calls % 5) { + case 0: /* Grow. */ + for (int i = 0; i < T; i++, next_temp++, temps++) { + int len = snprintf(buf, sizeof(buf), "temp:%d", next_temp); + zbtTestSet(zs, buf, len, 1e6 + next_temp); + } + break; + case 1: /* Replace temporaries, which leaves tombstones. */ + for (int i = 0; i < temps && i < T; i++, next_temp++) { + int len = snprintf(buf, sizeof(buf), "temp:%d", + next_temp - temps); + sds ele = sdsnewlen(buf, len); + serverAssert(zbtreeDelete(zs, ele)); + sdsfree(ele); + len = snprintf(buf, sizeof(buf), "temp:%d", next_temp); + zbtTestSet(zs, buf, len, 1e6 + next_temp); + } + break; + case 2: /* Shrink by single deletions. */ + for (; temps > 0; temps--) { + int len = snprintf(buf, sizeof(buf), "temp:%d", + next_temp - temps); + sds ele = sdsnewlen(buf, len); + serverAssert(zbtreeDelete(zs, ele)); + sdsfree(ele); + } + break; + case 3: /* Grow, then rebuild the index by a large range delete. */ + for (int i = 0; i < T; i++, next_temp++) { + int len = snprintf(buf, sizeof(buf), "temp:%d", next_temp); + zbtTestSet(zs, buf, len, 1e6 + next_temp); + } + zbtreeDeleteRangeByScore(zs, 1e6, 0, 1e300, 0); + break; + case 4: /* Move anchors and temporaries; widen once. */ + if (calls == 4) zbtTestSkipLeafIds(zs, ZBT_INDEX_WIDE_ID_AT); + for (int i = 0; i < 200; i++) { + int a = zbtTestRand(&r) % A; + int len = snprintf(buf, sizeof(buf), "anchor:%d", a); + zbtTestSet(zs, buf, len, + (double)(zbtTestRand(&r) % (A * 4))); + } + break; + } + if (zs->member_index.buckets != buckets) replaced++; + buckets = zs->member_index.buckets; + widened = widened || zs->member_index.wide_ids || + (zs->member_rehash && zs->member_rehash->table.wide_ids); + if (zbtTestScanMaxGroups(zs) > max_groups) + max_groups = zbtTestScanMaxGroups(zs); + } while (cursor != 0 && calls < 100000); + int missed = 0; + for (int i = 0; i < A; i++) if (seen[i] == 0) missed++; + printf(" %lu calls, %lu table replacements, up to %llu groups, " + "%d anchors missed\n", calls, replaced, + (unsigned long long)max_groups, missed); + test_cond("ZSCAN finishes while member tables keep being replaced", + cursor == 0 && calls <= max_groups + 1 && + replaced >= calls / 3 && widened); + test_cond("ZSCAN returns every anchor across table replacements", + missed == 0 && zbtTestCheck(zs)); + test_cond("ZSCAN cursors stay below the group count", small_cursor); + zfree(seen); + zbtreeFree(zs); + } + + /* The cursor step must be dictScan()'s reverse, add one, reverse. */ + { + uint64_t r = 17; + int ok = 1; + for (int k = 0; k < 200000 && ok; k++) { + uint64_t mask = (UINT64_C(1) << (k % 31)) - 1; + uint64_t cursor = zbtTestRand(&r) ^ (zbtTestRand(&r) << 40); + if (k % 3 == 0) cursor &= mask; + uint64_t expected = + zbtScanReverse(zbtScanReverse(cursor | ~mask) + 1); + ok = zbtScanNext(cursor, mask) == expected; + } + test_cond("ZSCAN cursor steps match reversed binary increments", ok); + } + + /* Any cursor value is accepted. Bits beyond the current table are dropped + * as dictScan() does, so the scan still ends within the group count. */ + { + zbtreeSet *zs = zbtreeCreate(); + zbtTestPopulate(zs, 5000); + uint64_t groups = zbtTestScanGroups(zs); + static const uint64_t starts[] = { + UINT64_MAX, UINT64_C(1) << 63, UINT64_C(1) << 40, + (UINT64_C(12345) << 32) | 7, UINT64_C(0xffffffff), 1, + }; + int ok = 1; + for (size_t s = 0; s < sizeof(starts) / sizeof(starts[0]); s++) { + int *seen = zcalloc(sizeof(int) * 5000); + zbtScanTestPrivdata pd = {seen, 5000}; + uint64_t cursor = starts[s]; + unsigned long calls = 0; + do { + cursor = zbtreeScan(zs, cursor, 1, zbtScanTestMarkSeen, &pd); + ok = ok && cursor < groups; + calls++; + } while (cursor != 0 && calls <= groups); + ok = ok && cursor == 0; + zfree(seen); + } + test_cond("ZSCAN accepts any cursor and still finishes", ok); + + /* The smallest table has a single group: one call returns all. */ + zbtreeSet *tiny = zbtreeCreate(); + zbtTestPopulate(tiny, 3); + int *seen = zcalloc(sizeof(int) * 3); + zbtScanTestPrivdata pd = {seen, 3}; + uint64_t cursor = zbtreeScan(tiny, 0, 1, zbtScanTestMarkSeen, &pd); + test_cond("ZSCAN of a one-group table finishes in one call", + tiny->member_index.size == ZBT_INDEX_INITIAL_BUCKETS && + cursor == 0 && seen[0] && seen[1] && seen[2]); + zfree(seen); + zbtreeFree(tiny); + + /* Two members of one leaf with the same tag and home bucket have + * alike slots, each of which could stand for either member. A full + * scan must still return each of them once. */ + int a = -1, b = -1; + char abuf[32], bbuf[32]; + for (int j = 1; j < 4096 && a < 0; j++) { + snprintf(bbuf, sizeof(bbuf), "member:%d", j); + uint32_t hb = zbtTestHash(bbuf); + for (int i = 0; i < j; i++) { + snprintf(abuf, sizeof(abuf), "member:%d", i); + uint32_t ha = zbtTestHash(abuf); + if (zbtIndexTag(ha) == zbtIndexTag(hb) && + ((ha ^ hb) & (ZBT_INDEX_INITIAL_BUCKETS - 1)) == 0) + { + a = i; + b = j; + break; + } + } + } + test_cond("ZSCAN test found two members with alike slots", a >= 0); + if (a >= 0) { + zbtreeSet *alike = zbtreeCreate(); + zbtTestPopulateRange(alike, a, a + 1); + zbtTestPopulateRange(alike, b, b + 1); + seen = zcalloc(sizeof(int) * (b + 1)); + pd = (zbtScanTestPrivdata){seen, b + 1}; + cursor = zbtreeScan(alike, 0, 100, zbtScanTestMarkSeen, &pd); + test_cond("ZSCAN returns members with alike slots once", + zbtTestLeafOf(alike, abuf) == zbtTestLeafOf(alike, bbuf) && + cursor == 0 && seen[a] == 1 && seen[b] == 1); + zfree(seen); + zbtreeFree(alike); + } + zbtreeFree(zs); + } + + return 0; +} +#endif diff --git a/src/zset_btree.h b/src/zset_btree.h new file mode 100644 index 0000000000..2b3bed38ff --- /dev/null +++ b/src/zset_btree.h @@ -0,0 +1,141 @@ +/* zset_btree.h -- memory efficient sorted set implementation. + * + * Copyright (c) 2026-Present, Redis Ltd. + * All rights reserved. + * + * Licensed under your choice of (a) the Redis Source Available License 2.0 + * (RSALv2); or (b) the Server Side Public License v1 (SSPLv1); or (c) the + * GNU Affero General Public License v3 (AGPLv3). + */ + +#ifndef REDIS_ZSET_BTREE_H +#define REDIS_ZSET_BTREE_H + +#include "sds.h" +#include +#include + +typedef struct zbtreeSet zbtreeSet; + +/* A ZADD lookup records both index positions. An insertion can reuse the hash + * position after a miss, while an update can reuse the score position after + * a hit. The contents are private to zset_btree.c. */ +typedef struct { + void *bucket; + unsigned int pos; + void *score_leaf; + unsigned int score_pos; + uint32_t hash; + uint32_t revision; +} zbtreeInsertPosition; + +/* Iterators point inside score leaves. They and the member bytes returned by + * zbtreeIteratorNext() remain valid only until the set is modified. */ +typedef struct { + void *leaf; + unsigned int pos; +} zbtreeIterator; + +/* This state exists only while active defragmentation is visiting a large + * sorted set. It is kept by the defragmenter, not in every zset. */ +typedef struct { + zbtreeSet *tree; + uint32_t score_leaf_id; + uint16_t score_pos; + uint8_t leaf_moved; +} zbtreeDefragState; + +typedef void *zbtreeDefragAllocFunction(void *ptr); +typedef void zbtreeDefragFreeFunction(void *ptr); +typedef void zbtreeScanFunction(void *privdata, const unsigned char *ele, + size_t len, double score); + +/* Creation, release, and memory accounting. */ +zbtreeSet *zbtreeCreate(void); +zbtreeSet *zbtreeDup(const zbtreeSet *zs); +void zbtreeFree(zbtreeSet *zs); +unsigned long zbtreeLength(const zbtreeSet *zs); +size_t zbtreeAllocSize(const zbtreeSet *zs); +void zbtreeDismissMemory(zbtreeSet *zs); +void zbtreeReserve(zbtreeSet *zs, unsigned long count); +int zbtreeTryReserve(zbtreeSet *zs, uint64_t count); + +/* Active defragmentation. Start may move the set itself and therefore returns + * its current address. Step moves at most one leaf, one external member, and + * the inner pages that finish at that leaf. It returns one when complete. */ +zbtreeSet *zbtreeDefragStart(zbtreeSet *zs, + zbtreeDefragAllocFunction *defrag_alloc); +void zbtreeDefragStateReset(zbtreeDefragState *state); +int zbtreeDefragStep(zbtreeSet *zs, zbtreeDefragState *state, + zbtreeDefragAllocFunction *defrag_alloc, + zbtreeDefragFreeFunction *defrag_free, + unsigned long *scanned); + +/* Lookup and modification. InsertNew* require an absent member. UpdateScore + * requires the position returned by a successful FindForAdd. Delete returns + * zero when the member is absent. */ +int zbtreeScore(zbtreeSet *zs, sds ele, double *score); +int zbtreeScoreRaw(zbtreeSet *zs, const unsigned char *ele, + size_t elelen, double *score); +/* Look up up to 16 scores. Bit i of the result says whether scores[i] exists. */ +#define ZBT_SCORE_BATCH_SIZE 16 +unsigned int zbtreeScores(zbtreeSet *zs, const sds *members, unsigned int count, + double *scores); +int zbtreeFindForAdd(zbtreeSet *zs, sds ele, double *score, + zbtreeInsertPosition *position); +long zbtreeRank(zbtreeSet *zs, sds ele, int reverse, double *score); +void zbtreeInsertNew(zbtreeSet *zs, double score, sds ele, + const zbtreeInsertPosition *position); +void zbtreeInsertNewRaw(zbtreeSet *zs, double score, const unsigned char *ele, + size_t elelen, const zbtreeInsertPosition *position); +void zbtreeInsertNewAppend(zbtreeSet *zs, double score, + const unsigned char *ele, size_t elelen); +void zbtreeInsertNewAppendWithHash(zbtreeSet *zs, double score, + const unsigned char *ele, size_t elelen, + uint32_t hash); +/* Buffers up to one leaf's worth of zbtreeInsertNewAppend()-style inserts and + * flushes them with a single bulk leaf build instead of one insert per + * element -- see the comment above its implementation in zset_btree.c for + * why. Same precondition as zbtreeInsertNewAppend(): every added element + * must sort after everything already in the tree and be known absent. */ +typedef struct zbtreeAppendBatch zbtreeAppendBatch; +zbtreeAppendBatch *zbtreeAppendBatchCreate(zbtreeSet *zs); +void zbtreeAppendBatchAdd(zbtreeAppendBatch *b, double score, + const unsigned char *ele, size_t elelen); +void zbtreeAppendBatchAddWithHash(zbtreeAppendBatch *b, double score, + const unsigned char *ele, size_t elelen, + uint32_t hash); +void zbtreeAppendBatchFinish(zbtreeAppendBatch *b); +int zbtreeDelete(zbtreeSet *zs, sds ele); +void zbtreeDeleteEdge(zbtreeSet *zs, int reverse); +void zbtreeUpdateScore(zbtreeSet *zs, sds ele, double score, + const zbtreeInsertPosition *position); +unsigned long zbtreeDeleteRangeByRank(zbtreeSet *zs, unsigned long start, + unsigned long end); +unsigned long zbtreeDeleteRangeByScore(zbtreeSet *zs, double min, int minex, + double max, int maxex); + +/* Ordered access. A successful seek leaves the iterator on its result; Next + * returns that result and advances in the requested direction. Its member, + * length, and score output pointers may be NULL. */ +int zbtreeIteratorSeekRank(const zbtreeSet *zs, unsigned long rank, + zbtreeIterator *iter); +int zbtreeIteratorSeekScore(const zbtreeSet *zs, double score, int exclusive, + int reverse, zbtreeIterator *iter, + unsigned long *rank); +int zbtreeIteratorSeekLex(const zbtreeSet *zs, sds ele, int exclusive, + int reverse, zbtreeIterator *iter, + unsigned long *rank); +int zbtreeIteratorStart(const zbtreeSet *zs, int reverse, + zbtreeIterator *iter); +int zbtreeIteratorNext(zbtreeIterator *iter, int reverse, + const unsigned char **ele, size_t *len, double *score); +unsigned long zbtreeCountByScore(const zbtreeSet *zs, double min, int minex, + double max, int maxex); + +/* Incremental unordered access. The cursor is an opaque hash position. */ +uint64_t zbtreeScan(zbtreeSet *zs, uint64_t cursor, + unsigned long count, zbtreeScanFunction *fn, + void *privdata); + +#endif diff --git a/tests/integration/convert-ziplist-zset-on-load.tcl b/tests/integration/convert-ziplist-zset-on-load.tcl index 0fbb20138b..08686962a3 100644 --- a/tests/integration/convert-ziplist-zset-on-load.tcl +++ b/tests/integration/convert-ziplist-zset-on-load.tcl @@ -16,10 +16,10 @@ start_server [list overrides [list "dir" $server_path "dbfilename" "zset-ziplist exec cp -f tests/assets/zset-ziplist.rdb $server_path start_server [list overrides [list "dir" $server_path "dbfilename" "zset-ziplist.rdb" "zset-max-ziplist-entries" 1]] { - test "RDB load ziplist zset: converts to skiplist when zset-max-ziplist-entries is exceeded" { + test "RDB load ziplist zset: converts to btree when zset-max-ziplist-entries is exceeded" { r select 0 - assert_encoding skiplist zset + assert_encoding btree zset assert_equal 2 [r zcard zset] assert_match {one 1 two 2} [r zrange zset 0 -1 withscores] } diff --git a/tests/integration/corrupt-dump.tcl b/tests/integration/corrupt-dump.tcl index 42fa31bebb..c265717c4c 100644 --- a/tests/integration/corrupt-dump.tcl +++ b/tests/integration/corrupt-dump.tcl @@ -18,6 +18,63 @@ set run_oom_tests [expr {($arch_name == "x86_64" || $arch_name == "aarch64") && set corrupt_payload_7445 "\x0E\x01\x1D\x1D\x00\x00\x00\x16\x00\x00\x00\x03\x00\x00\x04\x43\x43\x43\x43\x06\x04\x42\x42\x42\x42\x06\x3F\x41\x41\x41\x41\xFF\x09\x00\x88\xA5\xCA\xA8\xC5\x41\xF4\x35" +test {corrupt payload: zset reservation rejects oversized counts} { + start_server {overrides {loglevel verbose}} { + r debug set-skip-checksum-validation 1 + r set target preserved + # Both legacy ZSET and ZSET_2 use the same 64-bit count. These values + # exercise sizing overflow and allocator rejection without attempting + # a large allocation. The footer has RDB version 9 and a dummy CRC. + foreach type {03 05} { + foreach count {fffffffffffffffe 1fffffffffffffff 1000000000000000} { + set payload [binary format H* "${type}81${count}09000000000000000000"] + assert_error {*Bad data format*} {r restore target 0 $payload replace} + assert_equal preserved [r get target] + assert_equal PONG [r ping] + } + } + verify_log_message 0 "*Unable to reserve zset index*" 0 + } +} + +test {Zset listpack scores are validated before loading or conversion} { + start_server {} { + r zadd score-source 1.5 a 2.5 b + assert_encoding listpack score-source + set payload [r dump score-source] + # Mutate only the score text; bypass the now-stale DUMP checksum so + # the test exercises semantic validation of a well-formed listpack. + r debug set-skip-checksum-validation 1 + # Without sanitization a listpack that stays one is trusted, as for + # other types. A conversion still checks every score first. + foreach {sanitize limit} {yes 128 yes 1 no 1} { + r config set sanitize-dump-payload $sanitize + r config set zset-max-listpack-entries $limit + foreach score {nan NaN abc 1.x} { + set invalid [string map [list 1.5 $score] $payload] + assert_error {*Bad data format*} {r restore invalid-score 0 $invalid} + assert_equal 0 [r exists invalid-score] + assert_equal PONG [r ping] + } + # Both infinities are legal and must remain reloadable. + foreach score {inf -inf} { + # -inf needs four bytes, so use ZADD's genuine payload. + r del score-valid + r zadd score-valid $score a 2.5 b + set valid [r dump score-valid] + r restore score-roundtrip 0 $valid replace + assert_equal $score [r zscore score-roundtrip a] + } + } + r debug set-skip-checksum-validation 0 + r config set sanitize-dump-payload yes + r config set zset-max-listpack-entries 128 + r debug reload + assert_equal PONG [r ping] + assert_equal 0 [r exists invalid-score] + } +} + test {corrupt payload: #7445 - with sanitize} { start_server [list overrides [list loglevel verbose use-exit-on-panic yes crash-memcheck-enabled no] ] { r config set sanitize-dump-payload yes @@ -1214,4 +1271,3 @@ test {corrupt payload: stream XTRIM with non-minimally encoded flags} { } } ;# tags - diff --git a/tests/integration/dismiss-mem.tcl b/tests/integration/dismiss-mem.tcl index 15ea006b3e..fde9fa111a 100644 --- a/tests/integration/dismiss-mem.tcl +++ b/tests/integration/dismiss-mem.tcl @@ -32,10 +32,16 @@ start_server {tags {"dismiss external:skip needs:debug"}} { } # zset - r zadd bigzset1 1.0 $bigstr ; # skiplist encoding + r zadd bigzset1 1.0 $bigstr ; # B+ tree, external member + assert_encoding btree bigzset1 for {set i 0} {$i < 128} {incr i} { - r zadd bigzset2 1.0 $64bytes ; # ziplist encoding + r zadd bigzset2 1.0 [format "%064d" $i] } + assert_encoding listpack bigzset2 + for {set i 0} {$i < 512} {incr i} { + r zadd bigzset3 $i "${64bytes}:$i" + } + assert_encoding btree bigzset3 # hash r hset bighash1 field1 $bigstr ; # hash encoding diff --git a/tests/modules/scan.c b/tests/modules/scan.c index 1723d3097a..6886cc6ddc 100644 --- a/tests/modules/scan.c +++ b/tests/modules/scan.c @@ -103,6 +103,52 @@ int scan_key(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) return REDISMODULE_OK; } +typedef struct { + RedisModuleCtx *ctx; + RedisModuleString *first; + size_t count; +} scan_key_delete_pd; + +void scan_key_delete_callback(RedisModuleKey *key, RedisModuleString *field, + RedisModuleString *value, void *privdata) +{ + REDISMODULE_NOT_USED(value); + scan_key_delete_pd *pd = privdata; + + RedisModule_ReplyWithString(pd->ctx, field); + pd->count++; + if (pd->count == 1) { + pd->first = field; + RedisModule_RetainString(pd->ctx, field); + } else if (pd->count == 2) { + int deleted; + RedisModule_ZsetRem(key, pd->first, &deleted); + RedisModule_FreeString(pd->ctx, pd->first); + pd->first = NULL; + } +} + +/* Scan a sorted set while deleting the first member from the second callback. + * This exercises the documented guarantee that deleting a returned element + * does not hide elements that remain present throughout the scan. */ +int scan_key_delete(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) +{ + if (argc != 2) return RedisModule_WrongArity(ctx); + + RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_READ | REDISMODULE_WRITE); + if (!key) return RedisModule_ReplyWithError(ctx, "not found"); + + scan_key_delete_pd pd = {.ctx = ctx}; + RedisModule_ReplyWithArray(ctx, REDISMODULE_POSTPONED_ARRAY_LEN); + RedisModuleScanCursor *cursor = RedisModule_ScanCursorCreate(); + while (RedisModule_ScanKey(key, cursor, scan_key_delete_callback, &pd)); + RedisModule_ScanCursorDestroy(cursor); + if (pd.first) RedisModule_FreeString(ctx, pd.first); + RedisModule_ReplySetArrayLength(ctx, pd.count); + RedisModule_CloseKey(key); + return REDISMODULE_OK; +} + int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { REDISMODULE_NOT_USED(argv); REDISMODULE_NOT_USED(argc); @@ -115,7 +161,8 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) if (RedisModule_CreateCommand(ctx, "scan.scan_key", scan_key, "", 0, 0, 0) == REDISMODULE_ERR) return REDISMODULE_ERR; + if (RedisModule_CreateCommand(ctx, "scan.scan_key_delete", scan_key_delete, "", 0, 0, 0) == REDISMODULE_ERR) + return REDISMODULE_ERR; + return REDISMODULE_OK; } - - diff --git a/tests/modules/zset.c b/tests/modules/zset.c index adb318d176..7ba5ff8b7b 100644 --- a/tests/modules/zset.c +++ b/tests/modules/zset.c @@ -1,6 +1,7 @@ #include "redismodule.h" #include #include +#include /* ZSET.REM key element * @@ -69,6 +70,112 @@ int zset_incrby(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { return RedisModule_ReplyWithError(ctx, "ERR ZsetIncrby failed"); } +/* Range APIs take doubles directly, so expose NaN bounds even though + * RedisModule_StringToDouble() rejects them. */ +static int zset_range_bound(RedisModuleString *arg, double *value) { + size_t len; + const char *str = RedisModule_StringPtrLen(arg, &len); + if (len == 3 && !strncasecmp(str, "nan", 3)) { + *value = NAN; + return REDISMODULE_OK; + } + if (len == 4 && !strncasecmp(str, "-nan", 4)) { + *value = -NAN; + return REDISMODULE_OK; + } + return RedisModule_StringToDouble(arg, value); +} + +/* ZSET.RANGE key score first|last min max minex maxex + * ZSET.RANGE key lex first|last min max + * + * Exposes the module range iterator to the test suite. The reply contains + * member, score pairs in the direction selected by first or last. */ +int zset_range(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { + if (argc != 6 && argc != 8) return RedisModule_WrongArity(ctx); + RedisModule_AutoMemory(ctx); + + size_t mode_len, direction_len; + const char *mode = RedisModule_StringPtrLen(argv[2], &mode_len); + const char *direction = RedisModule_StringPtrLen(argv[3], &direction_len); + int reverse; + if (direction_len == 5 && !strncasecmp(direction, "first", 5)) + reverse = 0; + else if (direction_len == 4 && !strncasecmp(direction, "last", 4)) + reverse = 1; + else + return RedisModule_ReplyWithError(ctx, "ERR invalid direction"); + + RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_READ); + int status; + if (mode_len == 5 && !strncasecmp(mode, "score", 5) && argc == 8) { + double min, max; + long long minex, maxex; + if (zset_range_bound(argv[4], &min) == REDISMODULE_ERR || + zset_range_bound(argv[5], &max) == REDISMODULE_ERR || + RedisModule_StringToLongLong(argv[6], &minex) == REDISMODULE_ERR || + RedisModule_StringToLongLong(argv[7], &maxex) == REDISMODULE_ERR || + (minex != 0 && minex != 1) || (maxex != 0 && maxex != 1)) + return RedisModule_ReplyWithError(ctx, "ERR invalid score range"); + status = reverse ? + RedisModule_ZsetLastInScoreRange(key, min, max, minex, maxex) : + RedisModule_ZsetFirstInScoreRange(key, min, max, minex, maxex); + } else if (mode_len == 3 && !strncasecmp(mode, "lex", 3) && argc == 6) { + status = reverse ? + RedisModule_ZsetLastInLexRange(key, argv[4], argv[5]) : + RedisModule_ZsetFirstInLexRange(key, argv[4], argv[5]); + } else { + return RedisModule_ReplyWithError(ctx, "ERR invalid range mode"); + } + if (status == REDISMODULE_ERR) + return RedisModule_ReplyWithError(ctx, "ERR range setup failed"); + + RedisModule_ReplyWithArray(ctx, REDISMODULE_POSTPONED_ARRAY_LEN); + long count = 0; + while (!RedisModule_ZsetRangeEndReached(key)) { + double score; + RedisModuleString *member = RedisModule_ZsetRangeCurrentElement(key, &score); + if (member == NULL) break; + RedisModule_ReplyWithString(ctx, member); + RedisModule_ReplyWithDouble(ctx, score); + RedisModule_FreeString(ctx, member); + count += 2; + if (reverse) + RedisModule_ZsetRangePrev(key); + else + RedisModule_ZsetRangeNext(key); + } + RedisModule_ZsetRangeStop(key); + RedisModule_ReplySetArrayLength(ctx, count); + return REDISMODULE_OK; +} + +/* Invalidate the rank saved by a B+ tree range iterator. The iterator API + * should report the end of the range rather than aborting Redis. */ +int zset_range_mutate(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { + if (argc != 2) return RedisModule_WrongArity(ctx); + RedisModule_AutoMemory(ctx); + + RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_READ | REDISMODULE_WRITE); + if (RedisModule_ZsetLastInScoreRange(key, REDISMODULE_NEGATIVE_INFINITE, + REDISMODULE_POSITIVE_INFINITE, 0, 0) == REDISMODULE_ERR) + return RedisModule_ReplyWithError(ctx, "ERR range setup failed"); + + /* Keep one member so the open key handle itself remains valid. */ + RedisModuleCallReply *reply = RedisModule_Call(ctx, "ZREMRANGEBYRANK", "sll", + argv[1], (long long)0, (long long)-2); + if (reply == NULL || RedisModule_CallReplyType(reply) == REDISMODULE_REPLY_ERROR) + return RedisModule_ReplyWithError(ctx, "ERR range mutation failed"); + + RedisModuleString *member = RedisModule_ZsetRangeCurrentElement(key, NULL); + if (member != NULL || RedisModule_ZsetRangePrev(key) != 0 || + !RedisModule_ZsetRangeEndReached(key)) + return RedisModule_ReplyWithError(ctx, "ERR invalid iterator state"); + + RedisModule_ZsetRangeStop(key); + return RedisModule_ReplyWithSimpleString(ctx, "OK"); +} + /* Structure to hold data for the delall scan callback */ typedef struct { RedisModuleCtx *ctx; @@ -172,6 +279,14 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) 1, 1, 1) == REDISMODULE_ERR) return REDISMODULE_ERR; + if (RedisModule_CreateCommand(ctx, "zset.range", zset_range, "readonly", + 1, 1, 1) == REDISMODULE_ERR) + return REDISMODULE_ERR; + + if (RedisModule_CreateCommand(ctx, "zset.range-mutate", zset_range_mutate, "write", + 1, 1, 1) == REDISMODULE_ERR) + return REDISMODULE_ERR; + if (RedisModule_CreateCommand(ctx, "zset.delall", zset_delall, "write touches-arbitrary-keys", 0, 0, 0) == REDISMODULE_ERR) return REDISMODULE_ERR; diff --git a/tests/unit/aofrw.tcl b/tests/unit/aofrw.tcl index 0a13783b54..7b8fd9fed5 100644 --- a/tests/unit/aofrw.tcl +++ b/tests/unit/aofrw.tcl @@ -166,7 +166,7 @@ start_server {tags {"aofrw external:skip debug_defrag:skip"} overrides {aof-use- } foreach d {string int} { - foreach e {listpack skiplist} { + foreach e {listpack btree} { test "AOF rewrite of zset with $e encoding, $d data" { r flushall if {$e eq {listpack}} {set len 10} else {set len 1000} @@ -191,6 +191,28 @@ start_server {tags {"aofrw external:skip debug_defrag:skip"} overrides {aof-use- } } + test {AOF rewrite of zset with external members} { + r flushall + set members {} + for {set j 0} {$j < 600} {incr j} { + set len [lindex {319 320 321 4096} [expr {$j % 4}]] + set prefix [format "%04d:" $j] + set member "$prefix[string repeat x [expr {$len-[string length $prefix]}]]" + lappend members $member + r zadd key $j $member + } + assert_encoding btree key + set expected [r zrange key 0 -1 withscores] + + r bgrewriteaof + waitForBgrewriteaof r + r debug loadaof + + assert_encoding btree key + assert_equal $expected [r zrange key 0 -1 withscores] + assert_equal 599 [r zscore key [lindex $members 599]] + } + test "AOF rewrite functions" { r flushall r FUNCTION LOAD {#!lua name=test diff --git a/tests/unit/keyspace.tcl b/tests/unit/keyspace.tcl index b46b36b82e..b22c514e1e 100644 --- a/tests/unit/keyspace.tcl +++ b/tests/unit/keyspace.tcl @@ -307,14 +307,14 @@ foreach {type large} [array get largevalue] { assert_equal $digest [debug_digest_value newzset1{t}] } - test {COPY basic usage for skiplist sorted set} { + test {COPY basic usage for btree sorted set} { r del zset2{t} newzset2{t} set original_max [lindex [r config get zset-max-ziplist-entries] 1] r config set zset-max-ziplist-entries 0 for {set j 0} {$j < 130} {incr j} { r zadd zset2{t} [randomInt 50] ele-[randomInt 10] } - assert_encoding skiplist zset2{t} + assert_encoding btree zset2{t} r copy zset2{t} newzset2{t} set digest [debug_digest_value zset2{t}] assert_equal $digest [debug_digest_value newzset2{t}] @@ -325,6 +325,35 @@ foreach {type large} [array get largevalue] { r config set zset-max-ziplist-entries $original_max } + test {COPY duplicates all btree pages while its member table grows} { + r del zset-copy-source{t} zset-copy-dest{t} + set original_max [lindex [r config get zset-max-listpack-entries] 1] + r config set zset-max-listpack-entries 0 + + # The last insertion starts a member-table resize and leaves many + # score pages still to copy into the new table. + for {set j 0} {$j < 3969} {incr j} { + r zadd zset-copy-source{t} [expr {$j * 17 % 5003}] \ + [format "member:%05d:abcdefghijklmnop" $j] + } + assert_encoding btree zset-copy-source{t} + r copy zset-copy-source{t} zset-copy-dest{t} + set digest [debug_digest_value zset-copy-source{t}] + assert_equal $digest [debug_digest_value zset-copy-dest{t}] + + # Free the source before exercising every copied lookup and rank. + r del zset-copy-source{t} + assert_equal $digest [debug_digest_value zset-copy-dest{t}] + for {set j 0} {$j < 3969} {incr j 31} { + set member [format "member:%05d:abcdefghijklmnop" $j] + assert {[r zscore zset-copy-dest{t} $member] ne {}} + assert {[r zrank zset-copy-dest{t} $member] ne {}} + r zincrby zset-copy-dest{t} 0.5 $member + } + assert_equal 3969 [r zcard zset-copy-dest{t}] + r config set zset-max-listpack-entries $original_max + } + test {COPY basic usage for listpack hash} { r config set hash-max-listpack-entries 512 r del hash1{t} newhash1{t} diff --git a/tests/unit/moduleapi/scan.tcl b/tests/unit/moduleapi/scan.tcl index aab1418f77..ea62dbf7b5 100644 --- a/tests/unit/moduleapi/scan.tcl +++ b/tests/unit/moduleapi/scan.tcl @@ -95,13 +95,28 @@ start_server {tags {"modules external:skip"}} { lsort [r scan.scan_key zz] } {{f1 1} {f2 2}} - test {Module scan zset skiplist} { + test {Module scan zset btree} { r config set zset-max-ziplist-entries 2 r zadd zz 3 f3 - assert_encoding skiplist zz + assert_encoding btree zz lsort [r scan.scan_key zz] } {{f1 1} {f2 2} {f3 3}} + test {Module scan zset btree tolerates deleting a returned member} { + r del zz + for {set i 0} {$i < 100} {incr i} { + r zadd zz $i [format "member:%03d" $i] + } + assert_encoding btree zz + + set visited [lsort -unique [r scan.scan_key_delete zz]] + set remaining [r zrange zz 0 -1] + assert_equal 99 [llength $remaining] + foreach member $remaining { + assert {$member in $visited} + } + } + test {Module scan set intset} { r sadd ss 1 2 assert_encoding intset ss @@ -124,4 +139,4 @@ start_server {tags {"modules external:skip"}} { test "Unload the module - scan" { assert_equal {OK} [r module unload scan] } -} \ No newline at end of file +} diff --git a/tests/unit/moduleapi/zset.tcl b/tests/unit/moduleapi/zset.tcl index f2a324c146..fb13c7b60e 100644 --- a/tests/unit/moduleapi/zset.tcl +++ b/tests/unit/moduleapi/zset.tcl @@ -34,6 +34,73 @@ start_server {tags {"modules external:skip"}} { assert_equal {hello 100} [r zrange k 0 -1 withscores] } + test {Module zset range iteration with btree encoding} { + set old_max_entries [lindex [r config get zset-max-listpack-entries] 1] + r config set zset-max-listpack-entries 0 + + r del k + r zadd k 1 a 2 b 2 c 3 d 4 e + assert_encoding btree k + assert_equal {d 3 e 4} [r zset.range k score first 2 4 1 0] + assert_equal {e 4 d 3} [r zset.range k score last 2 4 1 0] + assert_equal {} [r zset.range k score first 10 20 0 0] + + r del k + r zadd k 0 alpha 0 beta 0 delta 0 gamma + assert_encoding btree k + assert_equal {beta 0 delta 0} \ + [r zset.range k lex first \[beta \(gamma] + assert_equal {delta 0 beta 0} \ + [r zset.range k lex last \[beta \(gamma] + assert_equal {} [r zset.range k lex first \[z +] + + r config set zset-max-listpack-entries $old_max_entries + } + + foreach encoding {listpack btree} { + foreach count {3 1000} { + test "Module zset NaN score bounds are empty - $encoding, $count members" { + set old_max_entries [lindex [r config get zset-max-listpack-entries] 1] + r config set zset-max-listpack-entries [expr {$encoding eq "listpack" ? 2000 : 0}] + r del k + set elements {} + for {set j 1} {$j <= $count} {incr j} { + lappend elements $j member:$j + } + r zadd k {*}$elements + assert_encoding $encoding k + + foreach direction {first last} { + foreach {min max} {NaN 10 0 NaN NaN NaN -NaN 10 0 -NaN -NaN -NaN} { + foreach minex {0 1} { + foreach maxex {0 1} { + assert_equal {} [r zset.range k score $direction $min $max $minex $maxex] + } + } + } + } + assert_equal {member:1 1 member:2 2 member:3 3} \ + [r zset.range k score first -inf 3 0 0] + assert_equal {member:3 3 member:2 2 member:1 1} \ + [r zset.range k score last -inf 3 0 0] + r config set zset-max-listpack-entries $old_max_entries + } + } + } + + test {Module zset range iteration ends after its saved rank is invalidated} { + set old_max_entries [lindex [r config get zset-max-listpack-entries] 1] + r config set zset-max-listpack-entries 0 + + r del k + r zadd k 1 a 2 b 3 c + assert_encoding btree k + assert_equal OK [r zset.range-mutate k] + assert_equal {c} [r zrange k 0 -1] + + r config set zset-max-listpack-entries $old_max_entries + } + test {Module zset - KEYSIZES is updated as expected (like test at hash.tcl)} { proc run_cmd_verify_hist {cmd expOutput {retries 1}} { proc K {} {return [string map { "db0_distrib_zsets_items" "db0_ZSET" "# Keysizes" "" " " "" "\n" "" "\r" "" } [r info keysizes] ]} diff --git a/tests/unit/scan.tcl b/tests/unit/scan.tcl index e12b482e71..4d5ea56f1d 100644 --- a/tests/unit/scan.tcl +++ b/tests/unit/scan.tcl @@ -308,7 +308,7 @@ proc test_scan {type} { } } - foreach enc {listpack skiplist} { + foreach enc {listpack btree} { test "{$type} ZSCAN with encoding $enc" { # Create the Sorted Set r del zset @@ -348,6 +348,293 @@ proc test_scan {type} { } } + if {$type eq {standalone}} { + test {B+ tree ZSCAN small and streaming replies preserve framing} { + foreach proto {2 3} { + r hello $proto + foreach size {1 31 32 33 129} { + r del zset + for {set j 0} {$j < 129} {incr j} { + r zadd zset 1.25 member:$j + } + for {set j $size} {$j < 129} {incr j} { + r zrem zset member:$j + } + assert_encoding btree zset + r readraw 1 + assert_equal {*2} [r zscan zset 0 count 10000] + assert_equal {$1} [r read] + assert_equal 0 [r read] + set header [r read] + assert_equal {*} [string index $header 0] + set fields [string range $header 1 end] + assert {$fields % 2 == 0} + set seen {} + for {set j 0} {$j < $fields / 2} {incr j} { + assert_match {$*} [r read] + dict set seen [r read] 1 + assert_equal {$4} [r read] + assert_equal 1.25 [r read] + } + assert_equal $size [dict size $seen] + r readraw 0 + assert_equal {0 {}} [r zscan zset 0 count 10000 match absent:*] + assert_equal PONG [r ping] + set rd [redis_deferring_client] + $rd hello $proto + $rd read + $rd client reply off + $rd zscan zset 0 count 10000 + $rd client reply on + assert_equal OK [$rd read] + $rd client reply skip + $rd zscan zset 0 count 10000 + $rd ping + assert_equal PONG [$rd read] + $rd close + } + } + r hello 2 + set _ {} + } {} {resp3} + + test {B+ tree ZSCAN headers preserve queued reply order} { + r del zset{t} + for {set j 0} {$j < 129} {incr j} { + r zadd zset{t} 1.25 member:$j + } + assert_encoding btree zset{t} + foreach proto {2 3} { + r hello $proto + foreach size {0 1024 65536} { + set preceding [string repeat x $size] + r set preceding{t} $preceding + r multi + r getrange preceding{t} 0 -1 + r zscan zset{t} 0 count 10000 match absent:* + r zscan zset{t} 0 count 10000 + r ping + set replies [r exec] + assert_equal $preceding [lindex $replies 0] + assert_equal {0 {}} [lindex $replies 1] + assert_equal 0 [lindex $replies 2 0] + set seen {} + foreach {member score} [lindex $replies 2 1] { + assert_equal 1.25 $score + dict set seen $member 1 + } + assert_equal 129 [dict size $seen] + assert_equal PONG [lindex $replies 3] + } + } + r hello 2 + set _ {} + } {} {resp3} + + test {B+ tree ZSCAN keeps scores as bulk strings in RESP3} { + r del zset + set elements {} + for {set j 0} {$j < 129} {incr j} { + set score [expr {$j == 0 ? 3.3 : 1.25}] + lappend elements $score member:$j + } + r zadd zset {*}$elements + assert_encoding btree zset + + r hello 3 + r readraw 1 + assert_equal {*2} [r zscan zset 0 count 10000] + assert_equal {$1} [r read] + assert_equal 0 [r read] + set array_header [r read] + assert_equal {*} [string index $array_header 0] + set fields [string range $array_header 1 end] + assert {$fields % 2 == 0} + set seen {} + for {set j 0} {$j < $fields / 2} {incr j} { + assert_match {$*} [r read] + set member [r read] + set score_header [r read] + set score [r read] + dict set seen $member 1 + if {$member eq {member:0}} { + assert_equal {$3} $score_header + assert_equal 3.3 $score + } else { + assert_equal {$4} $score_header + assert_equal 1.25 $score + } + } + assert_equal 129 [dict size $seen] + r readraw 0 + r hello 2 + set _ {} + } {} {resp3} + } + + test "{$type} ZSCAN does not skip members after deleting returned members" { + r del zset + set elements {} + for {set j 0} {$j < 1000} {incr j} { + lappend elements $j key:$j + } + r zadd zset {*}$elements + assert_encoding btree zset + + lassign [r zscan zset 0 count 10] cursor items + set seen {} + foreach {member score} $items { + lappend seen $member + r zrem zset $member + } + while {$cursor != 0} { + lassign [r zscan zset $cursor count 10] cursor items + foreach {member score} $items { + lappend seen $member + } + } + + assert_equal 1000 [llength [lsort -unique $seen]] + } + + test "{$type} ZSCAN does not skip members changed between calls" { + r del zset + set elements {} + set members {} + for {set j 0} {$j < 2048} {incr j} { + set member [format "stable:%05d" $j] + lappend elements $j $member + lappend members $member + } + r zadd zset {*}$elements + assert_encoding btree zset + + set cursor 0 + set seen {} + set calls 0 + while 1 { + lassign [r zscan zset $cursor count 1] cursor items + foreach {member score} $items { + lappend seen $member + } + if {$cursor == 0} break + incr calls + + # Moving a member can change its score leaf, but it remains part of + # the set for the complete scan and must eventually be returned. + set member [lindex $members [expr {($calls * 31) % 2048}]] + r zadd zset [expr {2048 + $calls}] $member + } + + assert_equal 2048 [llength [lsort -unique $seen]] + } + + test "{$type} ZSCAN does not skip members while its table grows" { + r del zset + set elements {} + for {set j 0} {$j < 1984} {incr j} { + lappend elements $j [format "stable:%05d" $j] + } + r zadd zset {*}$elements + assert_encoding btree zset + + set cursor 0 + set seen {} + set calls 0 + while 1 { + lassign [r zscan zset $cursor count 1] cursor items + foreach {member score} $items { + if {[string match "stable:*" $member]} { + lappend seen $member + } + } + if {$cursor == 0} break + + # A 256-bucket narrow table has 2048 slots and grows when the + # 1985th member is added. Further additions advance the copy. + incr calls + r zadd zset [expr {1984 + $calls}] [format "added:%05d" $calls] + } + + assert_equal 1984 [llength [lsort -unique $seen]] + } + + test "{$type} ZSCAN finishes while a bounded set keeps replacing its tables" { + # Between calls the set grows and is cut back by a large range + # deletion, which installs new member tables every time. A cursor + # tied to one table restarted forever. This one must finish and + # return every member that stays in the set. + r del zset + set elements {} + for {set j 0} {$j < 200} {incr j} { + lappend elements $j [format "anchor:%03d" $j] + } + r zadd zset {*}$elements + assert_encoding btree zset + set temps {} + for {set j 0} {$j < 3000} {incr j} { + lappend temps [expr {1000000 + $j}] temp:$j + } + + set cursor 0 + set seen {} + set calls 0 + while 1 { + lassign [r zscan zset $cursor count 10] cursor items + foreach {member score} $items { + if {[string match "anchor:*" $member]} { + lappend seen $member + } + } + # Clients that keep cursors as doubles need them below 2^53. + assert {$cursor < 9007199254740992} + if {$cursor == 0} break + incr calls + assert {$calls < 1000} + r zadd zset {*}$temps + r zremrangebyscore zset 1000000 +inf + } + assert_equal 200 [llength [lsort -unique $seen]] + } + + test "{$type} ZSCAN finishes under constant size churn" { + # Replacing members leaves tombstones, and their cleanup installs a + # same-size table every few hundred calls, before a complete scan. + r del zset + set elements {} + for {set j 0} {$j < 5000} {incr j} { + lappend elements $j [format "anchor:%05d" $j] + lappend elements $j churn:$j + } + r zadd zset {*}$elements + assert_encoding btree zset + + set cursor 0 + set seen {} + set calls 0 + set next 5000 + while 1 { + lassign [r zscan zset $cursor count 10] cursor items + foreach {member score} $items { + if {[string match "anchor:*" $member]} { + lappend seen $member + } + } + if {$cursor == 0} break + incr calls + assert {$calls < 2000} + set add {} + set rem {} + for {set j 0} {$j < 200} {incr j; incr next} { + lappend add $next churn:$next + lappend rem churn:[expr {$next - 5000}] + } + r zadd zset {*}$add + r zrem zset {*}$rem + } + assert_equal 5000 [llength [lsort -unique $seen]] + } + test "{$type} SCAN guarantees check under write load" { r flushdb populate 100 @@ -425,7 +712,7 @@ proc test_scan {type} { assert {$first_score != 0} } - foreach enc {listpack skiplist} { + foreach enc {listpack btree} { test "{$type} ZSCAN scores match ZSCORE with encoding $enc" { r del mykey # 9.8813129168249309e-323 is the denormal from the #2175 test above; @@ -435,7 +722,7 @@ proc test_scan {type} { r zadd mykey 3.3 a 1 b 1.25 c -0.1 d 1e100 e 9.8813129168249309e-323 f \ inf g -inf h -1.2345678901234567e23 i 1152921504606846976 j \ 123456789.12345679 k - if {$enc eq {skiplist}} { + if {$enc eq {btree}} { # Push the set past zset-max-listpack-entries (default 128). set elements {} for {set j 0} {$j < 200} {incr j} { @@ -482,9 +769,9 @@ proc test_scan {type} { lappend elements $j $member } r zadd mykey {*}$elements - assert_encoding skiplist mykey + assert_encoding btree mykey - # Exceed the stack vector's capacity and delete the nodes before + # Exercise a large streaming reply and delete the nodes before # EXEC flushes the reply. The output buffer must own the returned bytes. r multi r zscan mykey 0 COUNT 1000 diff --git a/tests/unit/type/zset.tcl b/tests/unit/type/zset.tcl index 7e12048e48..e92f8db217 100644 --- a/tests/unit/type/zset.tcl +++ b/tests/unit/type/zset.tcl @@ -6,6 +6,44 @@ start_server {tags {"zset"}} { } } + test {Small ZSET_2 loads preserve signed zero when a long member needs a B+tree} { + set old_limit [lindex [r config get zset-max-listpack-entries] 1] + set longmember [string repeat L 65] + foreach count {2 128 129} { + foreach longscore {-1 1} { + r del zero-source{t} zero-restored{t} + r config set zset-max-listpack-entries 0 + r zadd zero-source{t} $longscore $longmember + for {set i 1} {$i < $count} {incr i} { + r zadd zero-source{t} -0 [format m%03d $i] + } + set zero_expected [r zrange zero-source{t} 0 -1 withscores] + set payload [r dump zero-source{t}] + r config set zset-max-listpack-entries 128 + r restore zero-restored{t} 0 $payload + assert_encoding btree zero-restored{t} + assert_equal $zero_expected [r zrange zero-restored{t} 0 -1 withscores] + assert_equal $payload [r dump zero-restored{t}] + } + } + r config set zset-max-listpack-entries $old_limit + r del zero-source{t} zero-restored{t} + } + + test {Small ZSET_2 loads still select listpack after all members are known} { + set old_limit [lindex [r config get zset-max-listpack-entries] 1] + r del small-load-source{t} small-load-restored{t} + r config set zset-max-listpack-entries 0 + r zadd small-load-source{t} 1 a 2 b 3 c + set payload [r dump small-load-source{t}] + r config set zset-max-listpack-entries 128 + r restore small-load-restored{t} 0 $payload + assert_encoding listpack small-load-restored{t} + assert_equal {a 1 b 2 c 3} [r zrange small-load-restored{t} 0 -1 withscores] + r config set zset-max-listpack-entries $old_limit + r del small-load-source{t} small-load-restored{t} + } + # A helper function to verify either ZPOP* or ZMPOP* response. proc verify_pop_response {pop res zpop_expected_response zmpop_expected_response} { if {[string match "*ZM*" $pop]} { @@ -102,7 +140,7 @@ start_server {tags {"zset"}} { if {$encoding == "listpack"} { r config set zset-max-ziplist-entries 128 r config set zset-max-ziplist-value 64 - } elseif {$encoding == "skiplist"} { + } elseif {$encoding == "btree"} { r config set zset-max-ziplist-entries 0 r config set zset-max-ziplist-value 0 } else { @@ -550,6 +588,10 @@ start_server {tags {"zset"}} { assert_equal {f e d} [r zrevrangebyscore zset 6 3] assert_equal {g f e} [r zrevrangebyscore zset +inf 4] assert_equal 3 [r zcount zset 0 3] + assert_equal 7 [r zcount zset -inf +inf] + assert_equal 6 [r zcount zset (-inf +inf] + assert_equal 6 [r zcount zset -inf (+inf] + assert_equal 5 [r zcount zset (-inf (+inf] # exclusive range assert_equal {b} [r zrangebyscore zset (-inf (2] @@ -1231,6 +1273,18 @@ start_server {tags {"zset"}} { } } + test {ZPOPMIN keeps mixed signed zero scores exact in a B+tree leaf} { + set old_limit [lindex [r config get zset-max-listpack-entries] 1] + r config set zset-max-listpack-entries 0 + r del signed-zero-pop + r zadd signed-zero-pop 0 a 0 b 0 c -0 d + assert_encoding btree signed-zero-pop + assert_equal {a 0} [r zpopmin signed-zero-pop] + assert_equal {b 0 c 0 d -0} [r zrange signed-zero-pop 0 -1 withscores] + r del signed-zero-pop + r config set zset-max-listpack-entries $old_limit + } + foreach {popmin popmax} {ZPOPMIN ZPOPMAX ZMPOP_MIN ZMPOP_MAX} { test "Basic $popmin/$popmax with a single key - $encoding" { r del zset @@ -1338,7 +1392,192 @@ start_server {tags {"zset"}} { } basics listpack - basics skiplist + basics btree + + test {ZSET with a large member stays in the B+ tree encoding} { + set original_max [lindex [r config get zset-max-listpack-entries] 1] + r config set zset-max-listpack-entries 0 + r del large-member-zset{t} + r zadd large-member-zset{t} 1 short + assert_encoding btree large-member-zset{t} + + set large [string repeat x 4096] + r zadd large-member-zset{t} 2 $large + assert_encoding btree large-member-zset{t} + assert_equal {short 1} [r zrange large-member-zset{t} 0 0 withscores] + assert_equal 2 [r zscore large-member-zset{t} $large] + + # Exercise lookup, ordering, score changes, and scans while the long + # member is stored outside its score leaf. + r zadd large-member-zset{t} 1.5 middle 3 tail + assert_equal 2 [r zrank large-member-zset{t} $large] + assert_equal 1 [r zrevrank large-member-zset{t} $large] + assert_equal {middle 1.5} \ + [r zrange large-member-zset{t} 1.5 1.5 byscore withscores] + assert_equal 0 [r zadd large-member-zset{t} 0 $large] + assert_equal $large [r zrange large-member-zset{t} 0 0] + + set scanned [lindex [r zscan large-member-zset{t} 0 count 100] 1] + set scanned_members {} + foreach {member score} $scanned { + lappend scanned_members $member + } + assert_equal [lsort [list short middle tail $large]] \ + [lsort $scanned_members] + + r del large-member-other{t} + r zadd large-member-other{t} 5 $large 7 other + assert_encoding btree large-member-other{t} + assert_equal [list $large 5] \ + [r zinter 2 large-member-zset{t} large-member-other{t} withscores] + + r del large-member-copy{t} + assert_equal 1 [r copy large-member-zset{t} large-member-copy{t}] + assert_encoding btree large-member-copy{t} + assert_equal [r zrange large-member-zset{t} 0 -1 withscores] \ + [r zrange large-member-copy{t} 0 -1 withscores] + assert_equal 1 [r zrem large-member-zset{t} middle] + assert_equal 1 [r zremrangebyscore large-member-zset{t} 3 3] + assert_equal [list $large 0] [r zpopmin large-member-zset{t}] + assert_equal 0 [r zscore large-member-copy{t} $large] + + r config set zset-max-listpack-entries $original_max + } + + test {ZSET B+ tree handles members at the external storage boundary} { + set original_max [lindex [r config get zset-max-listpack-entries] 1] + r config set zset-max-listpack-entries 0 + r del member-storage-boundary + + set members {} + for {set j 0} {$j < 400} {incr j} { + set len [lindex {319 320 321} [expr {$j % 3}]] + set prefix [format "%04d:" $j] + set member "$prefix[string repeat x [expr {$len-[string length $prefix]}]]" + assert_equal $len [string length $member] + lappend members $member + r zadd member-storage-boundary $j $member + } + assert_encoding btree member-storage-boundary + assert_equal 400 [r zcard member-storage-boundary] + + # Move inline and external members across many leaves, then remove + # ranges which force leaves to be rebuilt and joined. + for {set j 0} {$j < 400} {incr j 17} { + assert_equal 0 [r zadd member-storage-boundary [expr {1000-$j}] \ + [lindex $members $j]] + } + assert_equal 1000 [r zscore member-storage-boundary [lindex $members 0]] + assert_equal 1 [r zrem member-storage-boundary [lindex $members 319]] + assert_equal 100 [r zremrangebyrank member-storage-boundary 50 149] + assert_equal 299 [r zcard member-storage-boundary] + + foreach j {1 2 202 398 399} { + set member [lindex $members $j] + if {$j != 319} { + assert_equal $j [r zscore member-storage-boundary $member] + } + } + + r config set zset-max-listpack-entries $original_max + } + + test {RDB preserves a B+ tree containing a large member} { + set original_max [lindex [r config get zset-max-listpack-entries] 1] + r config set zset-max-listpack-entries 0 + r del large-member-rdb{t} + set large [string repeat x 4096] + r zadd large-member-rdb{t} 1 short 2 $large + assert_encoding btree large-member-rdb{t} + + r debug reload + assert_encoding btree large-member-rdb{t} + assert_equal [list short 1 $large 2] \ + [r zrange large-member-rdb{t} 0 -1 withscores] + r config set zset-max-listpack-entries $original_max + } {OK} {needs:debug} + + test {ZSET btree byte-full edge pages preserve order} { + set original_max [lindex [r config get zset-max-listpack-entries] 1] + r config set zset-max-listpack-entries 0 + r del edge-right{t} edge-left{t} + + for {set j 0} {$j < 200} {incr j} { + set member [format "%04d%s" $j [string repeat x 296]] + r zadd edge-right{t} $j $member + r zadd edge-left{t} [expr {-$j}] $member + } + + set first [format "%04d%s" 0 [string repeat x 296]] + set last [format "%04d%s" 199 [string repeat x 296]] + assert_encoding btree edge-right{t} + assert_encoding btree edge-left{t} + assert_equal $first [r zrange edge-right{t} 0 0] + assert_equal $last [r zrange edge-right{t} -1 -1] + assert_equal $last [r zrange edge-left{t} 0 0] + assert_equal $first [r zrange edge-left{t} -1 -1] + + assert_equal 150 [r zremrangebyrank edge-right{t} 25 174] + assert_equal [r zrange edge-right{t} 0 -1] \ + [lreverse [r zrevrange edge-right{t} 0 -1]] + + r config set zset-max-listpack-entries $original_max + } + + test {ZSET btree large range deletion rebuilds a usable member index} { + set original_max [lindex [r config get zset-max-listpack-entries] 1] + r config set zset-max-listpack-entries 0 + r del delete-rank{t} delete-score{t} + + set args {} + for {set j 0} {$j < 2048} {incr j} { + set member [format "member:%04d:%s" $j [string repeat x 372]] + lappend args $j $member + } + assert_equal 2048 [r zadd delete-rank{t} {*}$args] + assert_equal 1 [r copy delete-rank{t} delete-score{t}] + + assert_equal 1025 [r zremrangebyrank delete-rank{t} 512 1536] + assert_equal 1025 [r zremrangebyscore delete-score{t} 512 1536] + assert_equal [r zrange delete-rank{t} 0 -1 withscores] \ + [r zrange delete-score{t} 0 -1 withscores] + assert_equal 1023 [r zcard delete-rank{t}] + set member511 [format "member:%04d:%s" 511 [string repeat x 372]] + set member512 [format "member:%04d:%s" 512 [string repeat x 372]] + set member1537 [format "member:%04d:%s" 1537 [string repeat x 372]] + assert_equal 511 [r zrank delete-rank{t} $member511] + assert_equal 512 [r zrank delete-rank{t} $member1537] + assert_equal {} [r zscore delete-rank{t} $member512] + + set member0 [format "member:%04d:%s" 0 [string repeat x 372]] + set member2047 [format "member:%04d:%s" 2047 [string repeat x 372]] + foreach key {delete-rank{t} delete-score{t}} { + assert_equal 1 [r zadd $key 1000 new-member] + assert_equal 0 [r zadd $key 3000 $member0] + assert_equal 1 [r zrem $key $member2047] + } + assert_equal [r zrange delete-rank{t} 0 -1 withscores] \ + [r zrange delete-score{t} 0 -1 withscores] + + r config set zset-max-listpack-entries $original_max + } + + test {ZSET btree range deletion can empty a medium-sized index} { + set original_max [lindex [r config get zset-max-listpack-entries] 1] + r config set zset-max-listpack-entries 0 + r del delete-all + + set args {} + for {set j 0} {$j < 500} {incr j} { + lappend args $j member:$j + } + assert_equal 500 [r zadd delete-all {*}$args] + assert_encoding btree delete-all + assert_equal 500 [r zremrangebyrank delete-all 0 -1] + assert_equal 0 [r exists delete-all] + + r config set zset-max-listpack-entries $original_max + } test "ZPOP/ZMPOP against wrong type" { r set foo{t} bar @@ -1634,6 +1873,56 @@ start_server {tags {"zset"}} { } } + test {ZUNIONSTORE btree-encoded result matches aggregate-max with score collisions and weights} { + # Result is forced past zset-max-listpack-entries so it lands on the + # BTREE bulk-build fast path, with heavy score collisions (only 50 + # distinct scores across thousands of members) to exercise the + # (score, member) tie-break in the same order the B+tree itself + # expects. + set orig_listpack_entries [lindex [r config get zset-max-listpack-entries] 1] + r config set zset-max-listpack-entries 8 + r del u1{t} u2{t} udest{t} + array set expected {} + set cmd1 [list r zadd u1{t}] + for {set j 0} {$j < 300} {incr j} { + set m [format "m%05d" $j] + set sc [expr {$j % 20}] + lappend cmd1 $sc $m + set expected($m) $sc + } + {*}$cmd1 + set cmd2 [list r zadd u2{t}] + for {set j 150} {$j < 450} {incr j} { + set m [format "m%05d" $j] + set sc [expr {($j % 20) * 2}] + lappend cmd2 $sc $m + if {[info exists expected($m)]} { + if {$sc > $expected($m)} { set expected($m) $sc } + } else { + set expected($m) $sc + } + } + {*}$cmd2 + r zunionstore udest{t} 2 u1{t} u2{t} aggregate max + assert_encoding btree udest{t} + assert_equal [array size expected] [r zcard udest{t}] + set prevscore {} + set prevele {} + foreach {ele score} [r zrange udest{t} 0 -1 withscores] { + assert_equal $expected($ele) $score + if {$prevscore ne {}} { + if {$score == $prevscore} { + assert {[string compare $ele $prevele] >= 0} + } else { + assert {$score >= $prevscore} + } + } + set prevscore $score + set prevele $ele + } + r config set zset-max-listpack-entries $orig_listpack_entries + } + test "ZUNIONSTORE/ZINTERSTORE/ZDIFFSTORE error if using WITHSCORES " { assert_error "*ERR*syntax*" {r zunionstore foo{t} 2 zsetd{t} zsetf{t} withscores} assert_error "*ERR*syntax*" {r zinterstore foo{t} 2 zsetd{t} zsetf{t} withscores} @@ -1678,6 +1967,32 @@ start_server {tags {"zset"}} { assert_match {*ERR*wrong*number*arg*} $e } + test {ZMSCORE B+tree batches preserve order, duplicates and missing members} { + r del mscore-batches + set batch_args {} + for {set i 0} {$i < 300} {incr i} { + lappend batch_args [expr {$i - 150}] member:$i + } + r zadd mscore-batches {*}$batch_args + assert_encoding btree mscore-batches + foreach count {1 2 15 16 17 31 32 33 63 64 65 100} { + set batch_members {} + set batch_scores {} + for {set i 0} {$i < $count} {incr i} { + if {$i % 3 == 0} { + lappend batch_members missing:$i + lappend batch_scores {} + } else { + set member [expr {($i * 37) % 23}] + lappend batch_members member:$member + lappend batch_scores [expr {$member - 150}] + } + } + assert_equal $batch_scores [r zmscore mscore-batches {*}$batch_members] + } + r del mscore-batches + } + test "ZSET commands don't accept the empty strings as valid score" { assert_error "*not*float*" {r zadd myzset "" abc} } @@ -1700,7 +2015,7 @@ start_server {tags {"zset"}} { r config set zset-max-ziplist-entries 256 r config set zset-max-ziplist-value 64 set elements 128 - } elseif {$encoding == "skiplist"} { + } elseif {$encoding == "btree"} { r config set zset-max-ziplist-entries 0 r config set zset-max-ziplist-value 0 if {$::accurate} {set elements 1000} else {set elements 100} @@ -1764,7 +2079,7 @@ start_server {tags {"zset"}} { test "ZSCORE 17-19 significant digit mantissas (widened fast path) - $encoding" { # Exercise the widened fast_float_strtod path that handles # mantissas > 2^53 (via __uint128_t arithmetic). ZADD/ZSCORE - # must round-trip bit-exactly through the listpack/skiplist + # must round-trip bit-exactly through both zset encodings # encoding (parse on ingest, parse again on retrieval). Each # input string below parses to a specific IEEE double whose # canonical string representation is itself, so `expr` in Tcl @@ -2028,7 +2343,7 @@ start_server {tags {"zset"}} { } } - test "ZSETs skiplist implementation backlink consistency test - $encoding" { + test "ZSET forward/reverse consistency test - $encoding" { set diff 0 for {set j 0} {$j < $elements} {incr j} { r zadd myzset [expr rand()] "Element-$j" @@ -2046,7 +2361,7 @@ start_server {tags {"zset"}} { assert_equal 0 $diff } - test "ZSETs ZRANK augmented skip list stress testing - $encoding" { + test "ZSET ZRANK stress testing - $encoding" { set err {} r del myzset for {set k 0} {$k < 2000} {incr k} { @@ -2214,7 +2529,7 @@ start_server {tags {"zset"}} { tags {"slow"} { stresses listpack - stresses skiplist + stresses btree } test "BZPOP/BZMPOP against wrong type" { @@ -2362,7 +2677,7 @@ start_server {tags {"zset"}} { $rd2 close } {0} {cluster:skip} - test {ZSET skiplist order consistency when elements are moved} { + test {ZSET order consistency when elements are moved} { set original_max [lindex [r config get zset-max-ziplist-entries] 1] r config set zset-max-ziplist-entries 0 for {set times 0} {$times < 10} {incr times} { @@ -2489,13 +2804,13 @@ start_server {tags {"zset"}} { r config set zset-max-listpack-entries 0 r del z1{t} z2{t} r zadd z1{t} 1 a - assert_encoding skiplist z1{t} + assert_encoding btree z1{t} assert_equal 1 [r zrangestore z2{t} z1{t} 0 -1] - assert_encoding skiplist z2{t} + assert_encoding btree z2{t} r config set zset-max-listpack-entries $original_max } - test {ZRANGESTORE with zset-max-listpack-entries 1 dst key should use skiplist encoding} { + test {ZRANGESTORE with zset-max-listpack-entries 1 dst key should use btree encoding} { set original_max [lindex [r config get zset-max-listpack-entries] 1] r config set zset-max-listpack-entries 1 r del z1{t} z2{t} z3{t} @@ -2503,7 +2818,151 @@ start_server {tags {"zset"}} { assert_equal 1 [r zrangestore z2{t} z1{t} 0 0] assert_encoding listpack z2{t} assert_equal 2 [r zrangestore z3{t} z1{t} 0 1] - assert_encoding skiplist z3{t} + assert_encoding btree z3{t} + r config set zset-max-listpack-entries $original_max + } + + test {ZADD converts an existing listpack after lowering its limit} { + set conversion_limit [lindex [r config get zset-max-listpack-entries] 1] + r config set zset-max-listpack-entries 128 + r del conversion{t} + set conversion_args {} + for {set i 0} {$i < 40} {incr i} { + lappend conversion_args $i m$i + } + r zadd conversion{t} {*}$conversion_args + assert_encoding listpack conversion{t} + r config set zset-max-listpack-entries 1 + assert_equal 2 [r zadd conversion{t} 1 x 2 y] + assert_encoding btree conversion{t} + assert_equal 42 [r zcard conversion{t}] + foreach {score member} $conversion_args { + assert_equal $score [r zscore conversion{t} $member] + } + r del conversion{t} + r config set zset-max-listpack-entries $conversion_limit + } + + test {ZRANGESTORE bulk append during member-index rehash} { + set old_entries [lindex [r config get zset-max-listpack-entries] 1] + set old_value [lindex [r config get zset-max-listpack-value] 1] + r config set zset-max-listpack-entries 128 + r config set zset-max-listpack-value 64 + foreach size {65 127 319 320 4096} { + foreach count {40 60 100 1000} { + r del rehash-src{t} rehash-dst{t} + set elements {} + for {set j 0} {$j < $count} {incr j} { + set member [format %08d $j][string repeat x [expr {$size - 8}]] + lappend elements $j $member + } + r zadd rehash-src{t} {*}$elements + set rehash_expected [r zrange rehash-src{t} 0 -1 withscores] + foreach bounds {{0 -1} {-inf +inf BYSCORE}} { + assert_equal $count [r zrangestore rehash-dst{t} rehash-src{t} {*}$bounds] + assert_encoding btree rehash-dst{t} + assert_equal $rehash_expected [r zrange rehash-dst{t} 0 -1 withscores] + # Point reads advance and finish any remaining rehash. + foreach {member score} $rehash_expected { + assert_equal $score [r zscore rehash-dst{t} $member] + } + } + } + } + r del rehash-src{t} rehash-dst{t} + r config set zset-max-listpack-entries $old_entries + r config set zset-max-listpack-value $old_value + } + + test {ZRANGESTORE btree-encoded result matches expected slice, forward and REV, with score collisions} { + set original_max [lindex [r config get zset-max-listpack-entries] 1] + r config set zset-max-listpack-entries 8 + r del zsrc{t} zdst_fwd{t} zdst_rev{t} zdst_byscore{t} zdst_byscore_rev{t} zdst_lim{t} + set cmd [list r zadd zsrc{t}] + array set expected {} + for {set j 0} {$j < 500} {incr j} { + set m [format "member:%04d" $j] + set sc [expr {$j % 20}] + lappend cmd $sc $m + set expected($m) $sc + } + {*}$cmd + assert_encoding btree zsrc{t} + + # Ground truth: the full (score, member) ascending order, matching + # both zrangestore's own selection order and the B+tree's own order. + set full_sorted {} + foreach {ele score} [r zrange zsrc{t} 0 -1 withscores] { + lappend full_sorted [list $score $ele] + } + assert_equal 500 [llength $full_sorted] + + # Forward rank slice -- exercises the new descent-free append fast + # path directly (ascending direction, BTREE destination). + r zrangestore zdst_fwd{t} zsrc{t} 100 399 + assert_encoding btree zdst_fwd{t} + set got_fwd {} + foreach {ele score} [r zrange zdst_fwd{t} 0 -1 withscores] { + lappend got_fwd [list $score $ele] + } + assert_equal [lrange $full_sorted 100 399] $got_fwd + + # REV rank slice -- must keep using the original (slower but always + # correct) zsetAdd() path; the fast path is guarded off for reverse + # emission. Rank indices for REV count from the highest score down. + r zrangestore zdst_rev{t} zsrc{t} 100 399 REV + assert_encoding btree zdst_rev{t} + set full_desc [lreverse $full_sorted] + set expected_rev [lsort -command {apply {{a b} { + set sa [lindex $a 0]; set sb [lindex $b 0] + if {$sa != $sb} { return [expr {$sa < $sb ? -1 : 1}] } + return [string compare [lindex $a 1] [lindex $b 1]] + }}} [lrange $full_desc 100 399]] + set got_rev {} + foreach {ele score} [r zrange zdst_rev{t} 0 -1 withscores] { + lappend got_rev [list $score $ele] + } + assert_equal $expected_rev $got_rev + + # BYSCORE forward -- same fast path, different selection method + # feeding the same STORE result handler. + r zrangestore zdst_byscore{t} zsrc{t} 3 7 BYSCORE + assert_encoding btree zdst_byscore{t} + set expected_byscore {} + foreach pair $full_sorted { + set sc [lindex $pair 0] + if {$sc >= 3 && $sc <= 7} { lappend expected_byscore $pair } + } + set got_byscore {} + foreach {ele score} [r zrange zdst_byscore{t} 0 -1 withscores] { + lappend got_byscore [list $score $ele] + } + assert_equal $expected_byscore $got_byscore + + # BYSCORE REV -- fast path must NOT engage here either. + r zrangestore zdst_byscore_rev{t} zsrc{t} 7 3 BYSCORE REV + assert_encoding btree zdst_byscore_rev{t} + set got_byscore_rev {} + foreach {ele score} [r zrange zdst_byscore_rev{t} 0 -1 withscores] { + lappend got_byscore_rev [list $score $ele] + } + assert_equal $expected_byscore $got_byscore_rev + + # BYSCORE with LIMIT offset/count -- forward, still fast-path eligible. + r zrangestore zdst_lim{t} zsrc{t} 0 19 BYSCORE LIMIT 10 50 + assert_encoding btree zdst_lim{t} + set expected_lim {} + foreach pair $full_sorted { + set sc [lindex $pair 0] + if {$sc >= 0 && $sc <= 19} { lappend expected_lim $pair } + } + set expected_lim [lrange $expected_lim 10 59] + set got_lim {} + foreach {ele score} [r zrange zdst_lim{t} 0 -1 withscores] { + lappend got_lim [list $score $ele] + } + assert_equal $expected_lim $got_lim + r config set zset-max-listpack-entries $original_max } @@ -2540,7 +2999,7 @@ start_server {tags {"zset"}} { } } - foreach {type contents} "listpack {1 a 2 b 3 c} skiplist {1 a 2 b 3 [randstring 70 90 alpha]}" { + foreach {type contents} "listpack {1 a 2 b 3 c} btree {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 @@ -2599,7 +3058,7 @@ start_server {tags {"zset"}} { r readraw 0 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]} + btree {1 a 2 b 3 c 4 d 5 e 6 f 7 g 7 h 9 i 10 [randstring 70 90 alpha]} listpack {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] @@ -2776,7 +3235,7 @@ start_server {tags {"zset"}} { assert_encoding hashtable set_big{t} } - foreach zset_type {listpack skiplist} { + foreach zset_type {listpack btree} { r del zset_small{t} zset_big{t} if {$zset_type == "listpack"} { @@ -2784,12 +3243,12 @@ start_server {tags {"zset"}} { 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"} { + } elseif {$zset_type == "btree"} { 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} + assert_encoding btree zset_small{t} + assert_encoding btree zset_big{t} } # Test one key is big and one key is small separately. @@ -2850,9 +3309,98 @@ start_server {tags {"zset"}} { assert_encoding listpack myzset assert_equal $max_entries [r zcard myzset] assert_equal 1 [r zadd myzset 1 b] - assert_encoding skiplist myzset + assert_encoding btree myzset + assert_equal [expr {$max_entries + 1}] [r zcard myzset] + assert {[r zscore myzset 1] ne {}} r config set zset-max-listpack-entries $original_max } } + + if {$::debug_defrag} { + test "Active defrag relocates all B+ tree zset pages" { + r config set activedefrag no + r config set hz 100 + r config set active-defrag-threshold-lower 0 + r config set active-defrag-threshold-upper 1 + r config set active-defrag-ignore-bytes 1 + r config set active-defrag-cycle-min 75 + r config set active-defrag-cycle-max 75 + r config set active-defrag-max-scan-fields 10 + r config set zset-max-listpack-entries 0 + r del defrag-zset + + # The last insertion starts an index resize and leaves both hash + # table allocations alive for the forced defrag pass. + set elements 7937 + set args {} + for {set i 0} {$i < $elements} {incr i} { + set member [format "member:%05d" $i] + if {$i >= 5000 && $i % 100 == 0} { + append member [string repeat x 512] + } + lappend args $i $member + } + assert_equal $elements [r zadd defrag-zset {*}$args] + assert_encoding btree defrag-zset + set digest [debug_digest] + + r config resetstat + r config set activedefrag yes + wait_for_condition 1000 10 { + [s active_defrag_hits] > 1000 + } else { + fail "Active defrag did not relocate the B+ tree pages" + } + r config set activedefrag no + wait_for_condition 1000 10 { + [s active_defrag_running] == 0 + } else { + fail "Active defrag did not stop" + } + + assert_equal $elements [r zcard defrag-zset] + assert_equal 1234 [r zscore defrag-zset member:01234] + assert_equal 1234 [r zrank defrag-zset member:01234] + assert_equal {member:00000 member:00001 member:00002} \ + [r zrange defrag-zset 0 2] + + # Exercise the tree and member index after their allocations move. + assert_equal 0 [r zadd defrag-zset 99999 member:01234] + assert_equal [expr {$elements - 1}] \ + [r zrank defrag-zset member:01234] + assert_equal 0 [r zadd defrag-zset 1234 member:01234] + assert_equal 1 [r zrem defrag-zset member:04321] + assert_equal 1 [r zadd defrag-zset 4321 member:04321] + assert_equal $digest [debug_digest] + assert_equal OK [r save] + + # Small trees are moved immediately instead of being queued. + r del defrag-zset + r config set active-defrag-max-scan-fields 100 + assert_equal 5 [r zadd defrag-zset 1 one 2 two 3 three 4 four 5 five] + assert_encoding btree defrag-zset + set digest [debug_digest] + + r config resetstat + r config set activedefrag yes + wait_for_condition 1000 10 { + [s active_defrag_hits] > 4 + } else { + fail "Active defrag did not relocate the small B+ tree" + } + r config set activedefrag no + wait_for_condition 1000 10 { + [s active_defrag_running] == 0 + } else { + fail "Active defrag did not stop" + } + + assert_equal {one two three four five} \ + [r zrange defrag-zset 0 -1] + assert_equal 3 [r zscore defrag-zset three] + assert_equal $digest [debug_digest] + assert_equal OK [r save] + } + } }