diff options
| -rw-r--r-- | src/object.c | 2 | ||||
| -rw-r--r-- | src/t_set.c | 25 |
2 files changed, 20 insertions, 7 deletions
diff --git a/src/object.c b/src/object.c index a11ac1116..54ed57958 100644 --- a/src/object.c +++ b/src/object.c @@ -1036,6 +1036,8 @@ size_t objectComputeSize(robj *key, robj *o, size_t sample_size, int dbid) { if (samples) asize += (double)elesize/samples*dictSize(d); } else if (o->encoding == OBJ_ENCODING_INTSET) { asize = sizeof(*o)+zmalloc_size(o->ptr); + } else if (o->encoding == OBJ_ENCODING_LISTPACK) { + asize = sizeof(*o)+zmalloc_size(o->ptr); } else { serverPanic("Unknown set encoding"); } diff --git a/src/t_set.c b/src/t_set.c index 557ba35ce..6a4840676 100644 --- a/src/t_set.c +++ b/src/t_set.c @@ -969,6 +969,11 @@ void spopCommand(client *c) { * implementation for more info. */ #define SRANDMEMBER_SUB_STRATEGY_MUL 3 +/* If client is trying to ask for a very large number of random elements, + * queuing may consume an unlimited amount of memory, so we want to limit + * the number of randoms per time. */ +#define SRANDFIELD_RANDOM_SAMPLE_LIMIT 1000 + void srandmemberWithCountCommand(client *c) { long l; unsigned long count, size; @@ -1010,13 +1015,19 @@ void srandmemberWithCountCommand(client *c) { if (set->encoding == OBJ_ENCODING_LISTPACK && count > 1) { /* Specialized case for listpack, traversing it only once. */ - listpackEntry *entries = zmalloc(count * sizeof(listpackEntry)); - lpRandomEntries(set->ptr, count, entries); - for (unsigned long i = 0; i < count; i++) { - if (entries[i].sval) - addReplyBulkCBuffer(c, entries[i].sval, entries[i].slen); - else - addReplyBulkLongLong(c, entries[i].lval); + unsigned long limit, sample_count; + limit = count > SRANDFIELD_RANDOM_SAMPLE_LIMIT ? SRANDFIELD_RANDOM_SAMPLE_LIMIT : count; + listpackEntry *entries = zmalloc(limit * sizeof(listpackEntry)); + while (count) { + sample_count = count > limit ? limit : count; + count -= sample_count; + lpRandomEntries(set->ptr, sample_count, entries); + for (unsigned long i = 0; i < sample_count; i++) { + if (entries[i].sval) + addReplyBulkCBuffer(c, entries[i].sval, entries[i].slen); + else + addReplyBulkLongLong(c, entries[i].lval); + } } zfree(entries); return; |