void addReply(redisClient *c, robj *obj) {
if (listLength(c->reply) == 0 &&
(c->replstate == REDIS_REPL_NONE ||
c->replstate == REDIS_REPL_ONLINE) &&
aeCreateFileEvent(server.el, c->fd, AE_WRITABLE,
sendReplyToClient, c) == AE_ERR) return;
if (server.vm_enabled && obj->storage != REDIS_VM_MEMORY) {
obj = dupStringObject(obj);
obj->refcount = 0; /* getDecodedObject() will increment the refcount */
}
listAddNodeTail(c->reply,getDecodedObject(obj));
}
void srandmemberWithCountCommand(redisClient *c) {
long l;
unsigned long count, size;
int uniq = 1;
robj *set, *ele;
int64_t llele;
int encoding;
dict *d;
if (getLongFromObjectOrReply(c,c->argv[2],&l,NULL) != REDIS_OK) return;
if (l >= 0) {
count = (unsigned) l;
} else {
/* A negative count means: return the same elements multiple times
* (i.e. don't remove the extracted element after every extraction). */
count = -l;
uniq = 0;
}
if ((set = lookupKeyReadOrReply(c,c->argv[1],shared.emptymultibulk))
== NULL || checkType(c,set,REDIS_SET)) return;
size = setTypeSize(set);
/* If count is zero, serve it ASAP to avoid special cases later. */
if (count == 0) {
addReply(c,shared.emptymultibulk);
return;
}
/* CASE 1: The count was negative, so the extraction method is just:
* "return N random elements" sampling the whole set every time.
* This case is trivial and can be served without auxiliary data
* structures. */
if (!uniq) {
addReplyMultiBulkLen(c,count);
while(count--) {
encoding = setTypeRandomElement(set,&ele,&llele);
if (encoding == REDIS_ENCODING_INTSET) {
addReplyBulkLongLong(c,llele);
} else {
addReplyBulk(c,ele);
}
}
return;
}
/* CASE 2:
* The number of requested elements is greater than the number of
* elements inside the set: simply return the whole set. */
if (count >= size) {
sunionDiffGenericCommand(c,c->argv+1,1,NULL,REDIS_OP_UNION);
return;
}
/* For CASE 3 and CASE 4 we need an auxiliary dictionary. */
d = dictCreate(&setDictType,NULL);
/* CASE 3:
* The number of elements inside the set is not greater than
* SRANDMEMBER_SUB_STRATEGY_MUL times the number of requested elements.
* In this case we create a set from scratch with all the elements, and
* subtract random elements to reach the requested number of elements.
*
* This is done because if the number of requsted elements is just
* a bit less than the number of elements in the set, the natural approach
* used into CASE 3 is highly inefficient. */
if (count*SRANDMEMBER_SUB_STRATEGY_MUL > size) {
setTypeIterator *si;
/* Add all the elements into the temporary dictionary. */
si = setTypeInitIterator(set);
while((encoding = setTypeNext(si,&ele,&llele)) != -1) {
int retval = DICT_ERR;
if (encoding == REDIS_ENCODING_INTSET) {
retval = dictAdd(d,createStringObjectFromLongLong(llele),NULL);
} else if (ele->encoding == REDIS_ENCODING_RAW) {
retval = dictAdd(d,dupStringObject(ele),NULL);
} else if (ele->encoding == REDIS_ENCODING_INT) {
retval = dictAdd(d,
createStringObjectFromLongLong((long)ele->ptr),NULL);
}
redisAssert(retval == DICT_OK);
}
setTypeReleaseIterator(si);
redisAssert(dictSize(d) == size);
/* Remove random elements to reach the right count. */
while(size > count) {
dictEntry *de;
de = dictGetRandomKey(d);
dictDelete(d,dictGetKey(de));
size--;
}
}
/* CASE 4: We have a big set compared to the requested number of elements.
* In this case we can simply get random elements from the set and add
* to the temporary set, trying to eventually get enough unique elements
* to reach the specified count. */
else {
unsigned long added = 0;
while(added < count) {
encoding = setTypeRandomElement(set,&ele,&llele);
if (encoding == REDIS_ENCODING_INTSET) {
ele = createStringObjectFromLongLong(llele);
} else if (ele->encoding == REDIS_ENCODING_RAW) {
ele = dupStringObject(ele);
} else if (ele->encoding == REDIS_ENCODING_INT) {
ele = createStringObjectFromLongLong((long)ele->ptr);
}
/* Try to add the object to the dictionary. If it already exists
* free it, otherwise increment the number of objects we have
* in the result dictionary. */
if (dictAdd(d,ele,NULL) == DICT_OK)
added++;
else
decrRefCount(ele);
}
}
/* CASE 3 & 4: send the result to the user. */
{
dictIterator *di;
dictEntry *de;
addReplyMultiBulkLen(c,count);
di = dictGetIterator(d);
while((de = dictNext(di)) != NULL)
addReplyBulk(c,dictGetKey(de));
dictReleaseIterator(di);
dictRelease(d);
}
}
robj *dupStringObjectUnconstant(robj *o) {
if (o->constant) return o;
return dupStringObject(o);
}
