1 /* 2 * Copyright (c) 2009-2012, Salvatore Sanfilippo <antirez at gmail dot com> 3 * All rights reserved. 4 * 5 * Redistribution and use in source and binary forms, with or without 6 * modification, are permitted provided that the following conditions are met: 7 * 8 * * Redistributions of source code must retain the above copyright notice, 9 * this list of conditions and the following disclaimer. 10 * * Redistributions in binary form must reproduce the above copyright 11 * notice, this list of conditions and the following disclaimer in the 12 * documentation and/or other materials provided with the distribution. 13 * * Neither the name of Redis nor the names of its contributors may be used 14 * to endorse or promote products derived from this software without 15 * specific prior written permission. 16 * 17 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 18 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 19 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 20 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE 21 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 22 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 23 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 24 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 25 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 26 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 27 * POSSIBILITY OF SUCH DAMAGE. 28 */ 29 30 #include "redis.h" 31 32 /*----------------------------------------------------------------------------- 33 * Set Commands 34 *----------------------------------------------------------------------------*/ 35 36 void sunionDiffGenericCommand(redisClient *c, robj **setkeys, int setnum, 37 robj *dstkey, int op); 38 39 /* Factory method to return a set that *can* hold "value". When the object has 40 * an integer-encodable value, an intset will be returned. Otherwise a regular 41 * hash table. */ 42 robj *setTypeCreate(robj *value) { 43 if (isObjectRepresentableAsLongLong(value,NULL) == REDIS_OK) 44 return createIntsetObject(); 45 return createSetObject(); 46 } 47 48 /* Add the specified value into a set. The function takes care of incrementing 49 * the reference count of the object if needed in order to retain a copy. 50 * 51 * If the value was already member of the set, nothing is done and 0 is 52 * returned, otherwise the new element is added and 1 is returned. */ 53 int setTypeAdd(robj *subject, robj *value) { 54 long long llval; 55 if (subject->encoding == REDIS_ENCODING_HT) { 56 if (dictAdd(subject->ptr,value,NULL) == DICT_OK) { 57 incrRefCount(value); 58 return 1; 59 } 60 } else if (subject->encoding == REDIS_ENCODING_INTSET) { 61 if (isObjectRepresentableAsLongLong(value,&llval) == REDIS_OK) { 62 uint8_t success = 0; 63 subject->ptr = intsetAdd(subject->ptr,llval,&success); 64 if (success) { 65 /* Convert to regular set when the intset contains 66 * too many entries. */ 67 if (intsetLen(subject->ptr) > server.set_max_intset_entries) 68 setTypeConvert(subject,REDIS_ENCODING_HT); 69 return 1; 70 } 71 } else { 72 /* Failed to get integer from object, convert to regular set. */ 73 setTypeConvert(subject,REDIS_ENCODING_HT); 74 75 /* The set *was* an intset and this value is not integer 76 * encodable, so dictAdd should always work. */ 77 redisAssertWithInfo(NULL,value, 78 dictAdd(subject->ptr,value,NULL) == DICT_OK); 79 incrRefCount(value); 80 return 1; 81 } 82 } else { 83 redisPanic("Unknown set encoding"); 84 } 85 return 0; 86 } 87 88 int setTypeRemove(robj *setobj, robj *value) { 89 long long llval; 90 if (setobj->encoding == REDIS_ENCODING_HT) { 91 if (dictDelete(setobj->ptr,value) == DICT_OK) { 92 if (htNeedsResize(setobj->ptr)) dictResize(setobj->ptr); 93 return 1; 94 } 95 } else if (setobj->encoding == REDIS_ENCODING_INTSET) { 96 if (isObjectRepresentableAsLongLong(value,&llval) == REDIS_OK) { 97 int success; 98 setobj->ptr = intsetRemove(setobj->ptr,llval,&success); 99 if (success) return 1; 100 } 101 } else { 102 redisPanic("Unknown set encoding"); 103 } 104 return 0; 105 } 106 107 int setTypeIsMember(robj *subject, robj *value) { 108 long long llval; 109 if (subject->encoding == REDIS_ENCODING_HT) { 110 return dictFind((dict*)subject->ptr,value) != NULL; 111 } else if (subject->encoding == REDIS_ENCODING_INTSET) { 112 if (isObjectRepresentableAsLongLong(value,&llval) == REDIS_OK) { 113 return intsetFind((intset*)subject->ptr,llval); 114 } 115 } else { 116 redisPanic("Unknown set encoding"); 117 } 118 return 0; 119 } 120 121 setTypeIterator *setTypeInitIterator(robj *subject) { 122 setTypeIterator *si = zmalloc(sizeof(setTypeIterator)); 123 si->subject = subject; 124 si->encoding = subject->encoding; 125 if (si->encoding == REDIS_ENCODING_HT) { 126 si->di = dictGetIterator(subject->ptr); 127 } else if (si->encoding == REDIS_ENCODING_INTSET) { 128 si->ii = 0; 129 } else { 130 redisPanic("Unknown set encoding"); 131 } 132 return si; 133 } 134 135 void setTypeReleaseIterator(setTypeIterator *si) { 136 if (si->encoding == REDIS_ENCODING_HT) 137 dictReleaseIterator(si->di); 138 zfree(si); 139 } 140 141 /* Move to the next entry in the set. Returns the object at the current 142 * position. 143 * 144 * Since set elements can be internally be stored as redis objects or 145 * simple arrays of integers, setTypeNext returns the encoding of the 146 * set object you are iterating, and will populate the appropriate pointer 147 * (eobj) or (llobj) accordingly. 148 * 149 * When there are no longer elements -1 is returned. 150 * Returned objects ref count is not incremented, so this function is 151 * copy on write friendly. */ 152 int setTypeNext(setTypeIterator *si, robj **objele, int64_t *llele) { 153 if (si->encoding == REDIS_ENCODING_HT) { 154 dictEntry *de = dictNext(si->di); 155 if (de == NULL) return -1; 156 *objele = dictGetKey(de); 157 *llele = -123456789; /* Not needed. Defensive. */ 158 } else if (si->encoding == REDIS_ENCODING_INTSET) { 159 if (!intsetGet(si->subject->ptr,si->ii++,llele)) 160 return -1; 161 *objele = NULL; /* Not needed. Defensive. */ 162 } else { 163 redisPanic("Wrong set encoding in setTypeNext"); 164 } 165 return si->encoding; 166 } 167 168 /* The not copy on write friendly version but easy to use version 169 * of setTypeNext() is setTypeNextObject(), returning new objects 170 * or incrementing the ref count of returned objects. So if you don't 171 * retain a pointer to this object you should call decrRefCount() against it. 172 * 173 * This function is the way to go for write operations where COW is not 174 * an issue as the result will be anyway of incrementing the ref count. */ 175 robj *setTypeNextObject(setTypeIterator *si) { 176 int64_t intele; 177 robj *objele; 178 int encoding; 179 180 encoding = setTypeNext(si,&objele,&intele); 181 switch(encoding) { 182 case -1: return NULL; 183 case REDIS_ENCODING_INTSET: 184 return createStringObjectFromLongLong(intele); 185 case REDIS_ENCODING_HT: 186 incrRefCount(objele); 187 return objele; 188 default: 189 redisPanic("Unsupported encoding"); 190 } 191 return NULL; /* just to suppress warnings */ 192 } 193 194 /* Return random element from a non empty set. 195 * The returned element can be a int64_t value if the set is encoded 196 * as an "intset" blob of integers, or a redis object if the set 197 * is a regular set. 198 * 199 * The caller provides both pointers to be populated with the right 200 * object. The return value of the function is the object->encoding 201 * field of the object and is used by the caller to check if the 202 * int64_t pointer or the redis object pointer was populated. 203 * 204 * When an object is returned (the set was a real set) the ref count 205 * of the object is not incremented so this function can be considered 206 * copy on write friendly. */ 207 int setTypeRandomElement(robj *setobj, robj **objele, int64_t *llele) { 208 if (setobj->encoding == REDIS_ENCODING_HT) { 209 dictEntry *de = dictGetRandomKey(setobj->ptr); 210 *objele = dictGetKey(de); 211 *llele = -123456789; /* Not needed. Defensive. */ 212 } else if (setobj->encoding == REDIS_ENCODING_INTSET) { 213 *llele = intsetRandom(setobj->ptr); 214 *objele = NULL; /* Not needed. Defensive. */ 215 } else { 216 redisPanic("Unknown set encoding"); 217 } 218 return setobj->encoding; 219 } 220 221 unsigned long setTypeSize(robj *subject) { 222 if (subject->encoding == REDIS_ENCODING_HT) { 223 return dictSize((dict*)subject->ptr); 224 } else if (subject->encoding == REDIS_ENCODING_INTSET) { 225 return intsetLen((intset*)subject->ptr); 226 } else { 227 redisPanic("Unknown set encoding"); 228 } 229 } 230 231 /* Convert the set to specified encoding. The resulting dict (when converting 232 * to a hash table) is presized to hold the number of elements in the original 233 * set. */ 234 void setTypeConvert(robj *setobj, int enc) { 235 setTypeIterator *si; 236 redisAssertWithInfo(NULL,setobj,setobj->type == REDIS_SET && 237 setobj->encoding == REDIS_ENCODING_INTSET); 238 239 if (enc == REDIS_ENCODING_HT) { 240 int64_t intele; 241 dict *d = dictCreate(&setDictType,NULL); 242 robj *element; 243 244 /* Presize the dict to avoid rehashing */ 245 dictExpand(d,intsetLen(setobj->ptr)); 246 247 /* To add the elements we extract integers and create redis objects */ 248 si = setTypeInitIterator(setobj); 249 while (setTypeNext(si,NULL,&intele) != -1) { 250 element = createStringObjectFromLongLong(intele); 251 redisAssertWithInfo(NULL,element, 252 dictAdd(d,element,NULL) == DICT_OK); 253 } 254 setTypeReleaseIterator(si); 255 256 setobj->encoding = REDIS_ENCODING_HT; 257 zfree(setobj->ptr); 258 setobj->ptr = d; 259 } else { 260 redisPanic("Unsupported set conversion"); 261 } 262 } 263 264 void saddCommand(redisClient *c) { 265 robj *set; 266 int j, added = 0; 267 268 set = lookupKeyWrite(c->db,c->argv[1]); 269 if (set == NULL) { 270 set = setTypeCreate(c->argv[2]); 271 dbAdd(c->db,c->argv[1],set); 272 } else { 273 if (set->type != REDIS_SET) { 274 addReply(c,shared.wrongtypeerr); 275 return; 276 } 277 } 278 279 for (j = 2; j < c->argc; j++) { 280 c->argv[j] = tryObjectEncoding(c->argv[j]); 281 if (setTypeAdd(set,c->argv[j])) added++; 282 } 283 if (added) { 284 signalModifiedKey(c->db,c->argv[1]); 285 notifyKeyspaceEvent(REDIS_NOTIFY_SET,"sadd",c->argv[1],c->db->id); 286 } 287 server.dirty += added; 288 addReplyLongLong(c,added); 289 } 290 291 void sremCommand(redisClient *c) { 292 robj *set; 293 int j, deleted = 0, keyremoved = 0; 294 295 if ((set = lookupKeyWriteOrReply(c,c->argv[1],shared.czero)) == NULL || 296 checkType(c,set,REDIS_SET)) return; 297 298 for (j = 2; j < c->argc; j++) { 299 if (setTypeRemove(set,c->argv[j])) { 300 deleted++; 301 if (setTypeSize(set) == 0) { 302 dbDelete(c->db,c->argv[1]); 303 keyremoved = 1; 304 break; 305 } 306 } 307 } 308 if (deleted) { 309 signalModifiedKey(c->db,c->argv[1]); 310 notifyKeyspaceEvent(REDIS_NOTIFY_SET,"srem",c->argv[1],c->db->id); 311 if (keyremoved) 312 notifyKeyspaceEvent(REDIS_NOTIFY_GENERIC,"del",c->argv[1], 313 c->db->id); 314 server.dirty += deleted; 315 } 316 addReplyLongLong(c,deleted); 317 } 318 319 void smoveCommand(redisClient *c) { 320 robj *srcset, *dstset, *ele; 321 srcset = lookupKeyWrite(c->db,c->argv[1]); 322 dstset = lookupKeyWrite(c->db,c->argv[2]); 323 ele = c->argv[3] = tryObjectEncoding(c->argv[3]); 324 325 /* If the source key does not exist return 0 */ 326 if (srcset == NULL) { 327 addReply(c,shared.czero); 328 return; 329 } 330 331 /* If the source key has the wrong type, or the destination key 332 * is set and has the wrong type, return with an error. */ 333 if (checkType(c,srcset,REDIS_SET) || 334 (dstset && checkType(c,dstset,REDIS_SET))) return; 335 336 /* If srcset and dstset are equal, SMOVE is a no-op */ 337 if (srcset == dstset) { 338 addReply(c,shared.cone); 339 return; 340 } 341 342 /* If the element cannot be removed from the src set, return 0. */ 343 if (!setTypeRemove(srcset,ele)) { 344 addReply(c,shared.czero); 345 return; 346 } 347 notifyKeyspaceEvent(REDIS_NOTIFY_SET,"srem",c->argv[1],c->db->id); 348 349 /* Remove the src set from the database when empty */ 350 if (setTypeSize(srcset) == 0) { 351 dbDelete(c->db,c->argv[1]); 352 notifyKeyspaceEvent(REDIS_NOTIFY_GENERIC,"del",c->argv[1],c->db->id); 353 } 354 signalModifiedKey(c->db,c->argv[1]); 355 signalModifiedKey(c->db,c->argv[2]); 356 server.dirty++; 357 358 /* Create the destination set when it doesn't exist */ 359 if (!dstset) { 360 dstset = setTypeCreate(ele); 361 dbAdd(c->db,c->argv[2],dstset); 362 } 363 364 /* An extra key has changed when ele was successfully added to dstset */ 365 if (setTypeAdd(dstset,ele)) { 366 server.dirty++; 367 notifyKeyspaceEvent(REDIS_NOTIFY_SET,"sadd",c->argv[2],c->db->id); 368 } 369 addReply(c,shared.cone); 370 } 371 372 void sismemberCommand(redisClient *c) { 373 robj *set; 374 375 if ((set = lookupKeyReadOrReply(c,c->argv[1],shared.czero)) == NULL || 376 checkType(c,set,REDIS_SET)) return; 377 378 c->argv[2] = tryObjectEncoding(c->argv[2]); 379 if (setTypeIsMember(set,c->argv[2])) 380 addReply(c,shared.cone); 381 else 382 addReply(c,shared.czero); 383 } 384 385 void scardCommand(redisClient *c) { 386 robj *o; 387 388 if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.czero)) == NULL || 389 checkType(c,o,REDIS_SET)) return; 390 391 addReplyLongLong(c,setTypeSize(o)); 392 } 393 394 /* Handle the "SPOP key <count>" variant. The normal version of the 395 * command is handled by the spopCommand() function itself. */ 396 397 /* How many times bigger should be the set compared to the remaining size 398 * for us to use the "create new set" strategy? Read later in the 399 * implementation for more info. */ 400 #define SPOP_MOVE_STRATEGY_MUL 5 401 402 void spopWithCountCommand(redisClient *c) { 403 long l; 404 unsigned long count, size; 405 robj *set; 406 407 /* Get the count argument */ 408 if (getLongFromObjectOrReply(c,c->argv[2],&l,NULL) != REDIS_OK) return; 409 if (l >= 0) { 410 count = (unsigned) l; 411 } else { 412 addReply(c,shared.outofrangeerr); 413 return; 414 } 415 416 /* Make sure a key with the name inputted exists, and that it's type is 417 * indeed a set. Otherwise, return nil */ 418 if ((set = lookupKeyReadOrReply(c,c->argv[1],shared.emptymultibulk)) 419 == NULL || checkType(c,set,REDIS_SET)) return; 420 421 /* If count is zero, serve an empty multibulk ASAP to avoid special 422 * cases later. */ 423 if (count == 0) { 424 addReply(c,shared.emptymultibulk); 425 return; 426 } 427 428 size = setTypeSize(set); 429 430 /* Generate an SPOP keyspace notification */ 431 notifyKeyspaceEvent(REDIS_NOTIFY_SET,"spop",c->argv[1],c->db->id); 432 server.dirty += count; 433 434 /* CASE 1: 435 * The number of requested elements is greater than or equal to 436 * the number of elements inside the set: simply return the whole set. */ 437 if (count >= size) { 438 /* We just return the entire set */ 439 sunionDiffGenericCommand(c,c->argv+1,1,NULL,REDIS_OP_UNION); 440 441 /* Delete the set as it is now empty */ 442 dbDelete(c->db,c->argv[1]); 443 notifyKeyspaceEvent(REDIS_NOTIFY_GENERIC,"del",c->argv[1],c->db->id); 444 445 /* Propagate this command as an DEL operation */ 446 rewriteClientCommandVector(c,2,shared.del,c->argv[1]); 447 signalModifiedKey(c->db,c->argv[1]); 448 server.dirty++; 449 return; 450 } 451 452 /* Case 2 and 3 require to replicate SPOP as a set of SERM commands. 453 * Prepare our replication argument vector. Also send the array length 454 * which is common to both the code paths. */ 455 robj *propargv[3]; 456 propargv[0] = createStringObject("SREM",4); 457 propargv[1] = c->argv[1]; 458 addReplyMultiBulkLen(c,count); 459 460 /* Common iteration vars. */ 461 robj *objele; 462 int encoding; 463 int64_t llele; 464 unsigned long remaining = size-count; /* Elements left after SPOP. */ 465 466 /* If we are here, the number of requested elements is less than the 467 * number of elements inside the set. Also we are sure that count < size. 468 * Use two different strategies. 469 * 470 * CASE 2: The number of elements to return is small compared to the 471 * set size. We can just extract random elements and return them to 472 * the set. */ 473 if (remaining*SPOP_MOVE_STRATEGY_MUL > count) { 474 while(count--) { 475 encoding = setTypeRandomElement(set,&objele,&llele); 476 if (encoding == REDIS_ENCODING_INTSET) { 477 objele = createStringObjectFromLongLong(llele); 478 } else { 479 incrRefCount(objele); 480 } 481 482 /* Return the element to the client and remove from the set. */ 483 addReplyBulk(c,objele); 484 setTypeRemove(set,objele); 485 486 /* Replicate/AOF this command as an SREM operation */ 487 propargv[2] = objele; 488 alsoPropagate(server.sremCommand,c->db->id,propargv,3, 489 REDIS_PROPAGATE_AOF|REDIS_PROPAGATE_REPL); 490 decrRefCount(objele); 491 } 492 } else { 493 /* CASE 3: The number of elements to return is very big, approaching 494 * the size of the set itself. After some time extracting random elements 495 * from such a set becomes computationally expensive, so we use 496 * a different strategy, we extract random elements that we don't 497 * want to return (the elements that will remain part of the set), 498 * creating a new set as we do this (that will be stored as the original 499 * set). Then we return the elements left in the original set and 500 * release it. */ 501 robj *newset = NULL; 502 503 /* Create a new set with just the remaining elements. */ 504 while(remaining--) { 505 encoding = setTypeRandomElement(set,&objele,&llele); 506 if (encoding == REDIS_ENCODING_INTSET) { 507 objele = createStringObjectFromLongLong(llele); 508 } else { 509 incrRefCount(objele); 510 } 511 if (!newset) newset = setTypeCreate(objele); 512 setTypeAdd(newset,objele); 513 setTypeRemove(set,objele); 514 decrRefCount(objele); 515 } 516 517 /* Assign the new set as the key value. */ 518 incrRefCount(set); /* Protect the old set value. */ 519 dbOverwrite(c->db,c->argv[1],newset); 520 521 /* Tranfer the old set to the client and release it. */ 522 setTypeIterator *si; 523 si = setTypeInitIterator(set); 524 while((encoding = setTypeNext(si,&objele,&llele)) != -1) { 525 if (encoding == REDIS_ENCODING_INTSET) { 526 objele = createStringObjectFromLongLong(llele); 527 } else { 528 incrRefCount(objele); 529 } 530 addReplyBulk(c,objele); 531 532 /* Replicate/AOF this command as an SREM operation */ 533 propargv[2] = objele; 534 alsoPropagate(server.sremCommand,c->db->id,propargv,3, 535 REDIS_PROPAGATE_AOF|REDIS_PROPAGATE_REPL); 536 537 decrRefCount(objele); 538 } 539 setTypeReleaseIterator(si); 540 decrRefCount(set); 541 } 542 543 /* Don't propagate the command itself even if we incremented the 544 * dirty counter. We don't want to propagate an SPOP command since 545 * we propagated the command as a set of SREMs operations using 546 * the alsoPropagate() API. */ 547 decrRefCount(propargv[0]); 548 preventCommandPropagation(c); 549 } 550 551 void spopCommand(redisClient *c) { 552 robj *set, *ele, *aux; 553 int64_t llele; 554 int encoding; 555 556 if (c->argc == 3) { 557 spopWithCountCommand(c); 558 return; 559 } else if (c->argc > 3) { 560 addReply(c,shared.syntaxerr); 561 return; 562 } 563 564 /* Make sure a key with the name inputted exists, and that it's type is 565 * indeed a set */ 566 if ((set = lookupKeyWriteOrReply(c,c->argv[1],shared.nullbulk)) == NULL || 567 checkType(c,set,REDIS_SET)) return; 568 569 /* Get a random element from the set */ 570 encoding = setTypeRandomElement(set,&ele,&llele); 571 572 /* Remove the element from the set */ 573 if (encoding == REDIS_ENCODING_INTSET) { 574 ele = createStringObjectFromLongLong(llele); 575 set->ptr = intsetRemove(set->ptr,llele,NULL); 576 } else { 577 incrRefCount(ele); 578 setTypeRemove(set,ele); 579 } 580 581 notifyKeyspaceEvent(REDIS_NOTIFY_SET,"spop",c->argv[1],c->db->id); 582 583 /* Replicate/AOF this command as an SREM operation */ 584 aux = createStringObject("SREM",4); 585 rewriteClientCommandVector(c,3,aux,c->argv[1],ele); 586 decrRefCount(ele); 587 decrRefCount(aux); 588 589 /* Add the element to the reply */ 590 addReplyBulk(c,ele); 591 592 /* Delete the set if it's empty */ 593 if (setTypeSize(set) == 0) { 594 dbDelete(c->db,c->argv[1]); 595 notifyKeyspaceEvent(REDIS_NOTIFY_GENERIC,"del",c->argv[1],c->db->id); 596 } 597 598 /* Set has been modified */ 599 signalModifiedKey(c->db,c->argv[1]); 600 server.dirty++; 601 } 602 603 /* handle the "SRANDMEMBER key <count>" variant. The normal version of the 604 * command is handled by the srandmemberCommand() function itself. */ 605 606 /* How many times bigger should be the set compared to the requested size 607 * for us to don't use the "remove elements" strategy? Read later in the 608 * implementation for more info. */ 609 #define SRANDMEMBER_SUB_STRATEGY_MUL 3 610 611 void srandmemberWithCountCommand(redisClient *c) { 612 long l; 613 unsigned long count, size; 614 int uniq = 1; 615 robj *set, *ele; 616 int64_t llele; 617 int encoding; 618 619 dict *d; 620 621 if (getLongFromObjectOrReply(c,c->argv[2],&l,NULL) != REDIS_OK) return; 622 if (l >= 0) { 623 count = (unsigned) l; 624 } else { 625 /* A negative count means: return the same elements multiple times 626 * (i.e. don't remove the extracted element after every extraction). */ 627 count = -l; 628 uniq = 0; 629 } 630 631 if ((set = lookupKeyReadOrReply(c,c->argv[1],shared.emptymultibulk)) 632 == NULL || checkType(c,set,REDIS_SET)) return; 633 size = setTypeSize(set); 634 635 /* If count is zero, serve it ASAP to avoid special cases later. */ 636 if (count == 0) { 637 addReply(c,shared.emptymultibulk); 638 return; 639 } 640 641 /* CASE 1: The count was negative, so the extraction method is just: 642 * "return N random elements" sampling the whole set every time. 643 * This case is trivial and can be served without auxiliary data 644 * structures. */ 645 if (!uniq) { 646 addReplyMultiBulkLen(c,count); 647 while(count--) { 648 encoding = setTypeRandomElement(set,&ele,&llele); 649 if (encoding == REDIS_ENCODING_INTSET) { 650 addReplyBulkLongLong(c,llele); 651 } else { 652 addReplyBulk(c,ele); 653 } 654 } 655 return; 656 } 657 658 /* CASE 2: 659 * The number of requested elements is greater than the number of 660 * elements inside the set: simply return the whole set. */ 661 if (count >= size) { 662 sunionDiffGenericCommand(c,c->argv+1,1,NULL,REDIS_OP_UNION); 663 return; 664 } 665 666 /* For CASE 3 and CASE 4 we need an auxiliary dictionary. */ 667 d = dictCreate(&setDictType,NULL); 668 669 /* CASE 3: 670 * The number of elements inside the set is not greater than 671 * SRANDMEMBER_SUB_STRATEGY_MUL times the number of requested elements. 672 * In this case we create a set from scratch with all the elements, and 673 * subtract random elements to reach the requested number of elements. 674 * 675 * This is done because if the number of requsted elements is just 676 * a bit less than the number of elements in the set, the natural approach 677 * used into CASE 3 is highly inefficient. */ 678 if (count*SRANDMEMBER_SUB_STRATEGY_MUL > size) { 679 setTypeIterator *si; 680 681 /* Add all the elements into the temporary dictionary. */ 682 si = setTypeInitIterator(set); 683 while((encoding = setTypeNext(si,&ele,&llele)) != -1) { 684 int retval = DICT_ERR; 685 686 if (encoding == REDIS_ENCODING_INTSET) { 687 retval = dictAdd(d,createStringObjectFromLongLong(llele),NULL); 688 } else { 689 retval = dictAdd(d,dupStringObject(ele),NULL); 690 } 691 redisAssert(retval == DICT_OK); 692 } 693 setTypeReleaseIterator(si); 694 redisAssert(dictSize(d) == size); 695 696 /* Remove random elements to reach the right count. */ 697 while(size > count) { 698 dictEntry *de; 699 700 de = dictGetRandomKey(d); 701 dictDelete(d,dictGetKey(de)); 702 size--; 703 } 704 } 705 706 /* CASE 4: We have a big set compared to the requested number of elements. 707 * In this case we can simply get random elements from the set and add 708 * to the temporary set, trying to eventually get enough unique elements 709 * to reach the specified count. */ 710 else { 711 unsigned long added = 0; 712 713 while(added < count) { 714 encoding = setTypeRandomElement(set,&ele,&llele); 715 if (encoding == REDIS_ENCODING_INTSET) { 716 ele = createStringObjectFromLongLong(llele); 717 } else { 718 ele = dupStringObject(ele); 719 } 720 /* Try to add the object to the dictionary. If it already exists 721 * free it, otherwise increment the number of objects we have 722 * in the result dictionary. */ 723 if (dictAdd(d,ele,NULL) == DICT_OK) 724 added++; 725 else 726 decrRefCount(ele); 727 } 728 } 729 730 /* CASE 3 & 4: send the result to the user. */ 731 { 732 dictIterator *di; 733 dictEntry *de; 734 735 addReplyMultiBulkLen(c,count); 736 di = dictGetIterator(d); 737 while((de = dictNext(di)) != NULL) 738 addReplyBulk(c,dictGetKey(de)); 739 dictReleaseIterator(di); 740 dictRelease(d); 741 } 742 } 743 744 void srandmemberCommand(redisClient *c) { 745 robj *set, *ele; 746 int64_t llele; 747 int encoding; 748 749 if (c->argc == 3) { 750 srandmemberWithCountCommand(c); 751 return; 752 } else if (c->argc > 3) { 753 addReply(c,shared.syntaxerr); 754 return; 755 } 756 757 if ((set = lookupKeyReadOrReply(c,c->argv[1],shared.nullbulk)) == NULL || 758 checkType(c,set,REDIS_SET)) return; 759 760 encoding = setTypeRandomElement(set,&ele,&llele); 761 if (encoding == REDIS_ENCODING_INTSET) { 762 addReplyBulkLongLong(c,llele); 763 } else { 764 addReplyBulk(c,ele); 765 } 766 } 767 768 int qsortCompareSetsByCardinality(const void *s1, const void *s2) { 769 return setTypeSize(*(robj**)s1)-setTypeSize(*(robj**)s2); 770 } 771 772 /* This is used by SDIFF and in this case we can receive NULL that should 773 * be handled as empty sets. */ 774 int qsortCompareSetsByRevCardinality(const void *s1, const void *s2) { 775 robj *o1 = *(robj**)s1, *o2 = *(robj**)s2; 776 777 return (o2 ? setTypeSize(o2) : 0) - (o1 ? setTypeSize(o1) : 0); 778 } 779 780 void sinterGenericCommand(redisClient *c, robj **setkeys, 781 unsigned long setnum, robj *dstkey) { 782 robj **sets = zmalloc(sizeof(robj*)*setnum); 783 setTypeIterator *si; 784 robj *eleobj, *dstset = NULL; 785 int64_t intobj; 786 void *replylen = NULL; 787 unsigned long j, cardinality = 0; 788 int encoding; 789 790 for (j = 0; j < setnum; j++) { 791 robj *setobj = dstkey ? 792 lookupKeyWrite(c->db,setkeys[j]) : 793 lookupKeyRead(c->db,setkeys[j]); 794 if (!setobj) { 795 zfree(sets); 796 if (dstkey) { 797 if (dbDelete(c->db,dstkey)) { 798 signalModifiedKey(c->db,dstkey); 799 server.dirty++; 800 } 801 addReply(c,shared.czero); 802 } else { 803 addReply(c,shared.emptymultibulk); 804 } 805 return; 806 } 807 if (checkType(c,setobj,REDIS_SET)) { 808 zfree(sets); 809 return; 810 } 811 sets[j] = setobj; 812 } 813 /* Sort sets from the smallest to largest, this will improve our 814 * algorithm's performance */ 815 qsort(sets,setnum,sizeof(robj*),qsortCompareSetsByCardinality); 816 817 /* The first thing we should output is the total number of elements... 818 * since this is a multi-bulk write, but at this stage we don't know 819 * the intersection set size, so we use a trick, append an empty object 820 * to the output list and save the pointer to later modify it with the 821 * right length */ 822 if (!dstkey) { 823 replylen = addDeferredMultiBulkLength(c); 824 } else { 825 /* If we have a target key where to store the resulting set 826 * create this key with an empty set inside */ 827 dstset = createIntsetObject(); 828 } 829 830 /* Iterate all the elements of the first (smallest) set, and test 831 * the element against all the other sets, if at least one set does 832 * not include the element it is discarded */ 833 si = setTypeInitIterator(sets[0]); 834 while((encoding = setTypeNext(si,&eleobj,&intobj)) != -1) { 835 for (j = 1; j < setnum; j++) { 836 if (sets[j] == sets[0]) continue; 837 if (encoding == REDIS_ENCODING_INTSET) { 838 /* intset with intset is simple... and fast */ 839 if (sets[j]->encoding == REDIS_ENCODING_INTSET && 840 !intsetFind((intset*)sets[j]->ptr,intobj)) 841 { 842 break; 843 /* in order to compare an integer with an object we 844 * have to use the generic function, creating an object 845 * for this */ 846 } else if (sets[j]->encoding == REDIS_ENCODING_HT) { 847 eleobj = createStringObjectFromLongLong(intobj); 848 if (!setTypeIsMember(sets[j],eleobj)) { 849 decrRefCount(eleobj); 850 break; 851 } 852 decrRefCount(eleobj); 853 } 854 } else if (encoding == REDIS_ENCODING_HT) { 855 /* Optimization... if the source object is integer 856 * encoded AND the target set is an intset, we can get 857 * a much faster path. */ 858 if (eleobj->encoding == REDIS_ENCODING_INT && 859 sets[j]->encoding == REDIS_ENCODING_INTSET && 860 !intsetFind((intset*)sets[j]->ptr,(long)eleobj->ptr)) 861 { 862 break; 863 /* else... object to object check is easy as we use the 864 * type agnostic API here. */ 865 } else if (!setTypeIsMember(sets[j],eleobj)) { 866 break; 867 } 868 } 869 } 870 871 /* Only take action when all sets contain the member */ 872 if (j == setnum) { 873 if (!dstkey) { 874 if (encoding == REDIS_ENCODING_HT) 875 addReplyBulk(c,eleobj); 876 else 877 addReplyBulkLongLong(c,intobj); 878 cardinality++; 879 } else { 880 if (encoding == REDIS_ENCODING_INTSET) { 881 eleobj = createStringObjectFromLongLong(intobj); 882 setTypeAdd(dstset,eleobj); 883 decrRefCount(eleobj); 884 } else { 885 setTypeAdd(dstset,eleobj); 886 } 887 } 888 } 889 } 890 setTypeReleaseIterator(si); 891 892 if (dstkey) { 893 /* Store the resulting set into the target, if the intersection 894 * is not an empty set. */ 895 int deleted = dbDelete(c->db,dstkey); 896 if (setTypeSize(dstset) > 0) { 897 dbAdd(c->db,dstkey,dstset); 898 addReplyLongLong(c,setTypeSize(dstset)); 899 notifyKeyspaceEvent(REDIS_NOTIFY_SET,"sinterstore", 900 dstkey,c->db->id); 901 } else { 902 decrRefCount(dstset); 903 addReply(c,shared.czero); 904 if (deleted) 905 notifyKeyspaceEvent(REDIS_NOTIFY_GENERIC,"del", 906 dstkey,c->db->id); 907 } 908 signalModifiedKey(c->db,dstkey); 909 server.dirty++; 910 } else { 911 setDeferredMultiBulkLength(c,replylen,cardinality); 912 } 913 zfree(sets); 914 } 915 916 void sinterCommand(redisClient *c) { 917 sinterGenericCommand(c,c->argv+1,c->argc-1,NULL); 918 } 919 920 void sinterstoreCommand(redisClient *c) { 921 sinterGenericCommand(c,c->argv+2,c->argc-2,c->argv[1]); 922 } 923 924 #define REDIS_OP_UNION 0 925 #define REDIS_OP_DIFF 1 926 #define REDIS_OP_INTER 2 927 928 void sunionDiffGenericCommand(redisClient *c, robj **setkeys, int setnum, 929 robj *dstkey, int op) { 930 robj **sets = zmalloc(sizeof(robj*)*setnum); 931 setTypeIterator *si; 932 robj *ele, *dstset = NULL; 933 int j, cardinality = 0; 934 int diff_algo = 1; 935 936 for (j = 0; j < setnum; j++) { 937 robj *setobj = dstkey ? 938 lookupKeyWrite(c->db,setkeys[j]) : 939 lookupKeyRead(c->db,setkeys[j]); 940 if (!setobj) { 941 sets[j] = NULL; 942 continue; 943 } 944 if (checkType(c,setobj,REDIS_SET)) { 945 zfree(sets); 946 return; 947 } 948 sets[j] = setobj; 949 } 950 951 /* Select what DIFF algorithm to use. 952 * 953 * Algorithm 1 is O(N*M) where N is the size of the element first set 954 * and M the total number of sets. 955 * 956 * Algorithm 2 is O(N) where N is the total number of elements in all 957 * the sets. 958 * 959 * We compute what is the best bet with the current input here. */ 960 if (op == REDIS_OP_DIFF && sets[0]) { 961 long long algo_one_work = 0, algo_two_work = 0; 962 963 for (j = 0; j < setnum; j++) { 964 if (sets[j] == NULL) continue; 965 966 algo_one_work += setTypeSize(sets[0]); 967 algo_two_work += setTypeSize(sets[j]); 968 } 969 970 /* Algorithm 1 has better constant times and performs less operations 971 * if there are elements in common. Give it some advantage. */ 972 algo_one_work /= 2; 973 diff_algo = (algo_one_work <= algo_two_work) ? 1 : 2; 974 975 if (diff_algo == 1 && setnum > 1) { 976 /* With algorithm 1 it is better to order the sets to subtract 977 * by decreasing size, so that we are more likely to find 978 * duplicated elements ASAP. */ 979 qsort(sets+1,setnum-1,sizeof(robj*), 980 qsortCompareSetsByRevCardinality); 981 } 982 } 983 984 /* We need a temp set object to store our union. If the dstkey 985 * is not NULL (that is, we are inside an SUNIONSTORE operation) then 986 * this set object will be the resulting object to set into the target key*/ 987 dstset = createIntsetObject(); 988 989 if (op == REDIS_OP_UNION) { 990 /* Union is trivial, just add every element of every set to the 991 * temporary set. */ 992 for (j = 0; j < setnum; j++) { 993 if (!sets[j]) continue; /* non existing keys are like empty sets */ 994 995 si = setTypeInitIterator(sets[j]); 996 while((ele = setTypeNextObject(si)) != NULL) { 997 if (setTypeAdd(dstset,ele)) cardinality++; 998 decrRefCount(ele); 999 } 1000 setTypeReleaseIterator(si); 1001 } 1002 } else if (op == REDIS_OP_DIFF && sets[0] && diff_algo == 1) { 1003 /* DIFF Algorithm 1: 1004 * 1005 * We perform the diff by iterating all the elements of the first set, 1006 * and only adding it to the target set if the element does not exist 1007 * into all the other sets. 1008 * 1009 * This way we perform at max N*M operations, where N is the size of 1010 * the first set, and M the number of sets. */ 1011 si = setTypeInitIterator(sets[0]); 1012 while((ele = setTypeNextObject(si)) != NULL) { 1013 for (j = 1; j < setnum; j++) { 1014 if (!sets[j]) continue; /* no key is an empty set. */ 1015 if (sets[j] == sets[0]) break; /* same set! */ 1016 if (setTypeIsMember(sets[j],ele)) break; 1017 } 1018 if (j == setnum) { 1019 /* There is no other set with this element. Add it. */ 1020 setTypeAdd(dstset,ele); 1021 cardinality++; 1022 } 1023 decrRefCount(ele); 1024 } 1025 setTypeReleaseIterator(si); 1026 } else if (op == REDIS_OP_DIFF && sets[0] && diff_algo == 2) { 1027 /* DIFF Algorithm 2: 1028 * 1029 * Add all the elements of the first set to the auxiliary set. 1030 * Then remove all the elements of all the next sets from it. 1031 * 1032 * This is O(N) where N is the sum of all the elements in every 1033 * set. */ 1034 for (j = 0; j < setnum; j++) { 1035 if (!sets[j]) continue; /* non existing keys are like empty sets */ 1036 1037 si = setTypeInitIterator(sets[j]); 1038 while((ele = setTypeNextObject(si)) != NULL) { 1039 if (j == 0) { 1040 if (setTypeAdd(dstset,ele)) cardinality++; 1041 } else { 1042 if (setTypeRemove(dstset,ele)) cardinality--; 1043 } 1044 decrRefCount(ele); 1045 } 1046 setTypeReleaseIterator(si); 1047 1048 /* Exit if result set is empty as any additional removal 1049 * of elements will have no effect. */ 1050 if (cardinality == 0) break; 1051 } 1052 } 1053 1054 /* Output the content of the resulting set, if not in STORE mode */ 1055 if (!dstkey) { 1056 addReplyMultiBulkLen(c,cardinality); 1057 si = setTypeInitIterator(dstset); 1058 while((ele = setTypeNextObject(si)) != NULL) { 1059 addReplyBulk(c,ele); 1060 decrRefCount(ele); 1061 } 1062 setTypeReleaseIterator(si); 1063 decrRefCount(dstset); 1064 } else { 1065 /* If we have a target key where to store the resulting set 1066 * create this key with the result set inside */ 1067 int deleted = dbDelete(c->db,dstkey); 1068 if (setTypeSize(dstset) > 0) { 1069 dbAdd(c->db,dstkey,dstset); 1070 addReplyLongLong(c,setTypeSize(dstset)); 1071 notifyKeyspaceEvent(REDIS_NOTIFY_SET, 1072 op == REDIS_OP_UNION ? "sunionstore" : "sdiffstore", 1073 dstkey,c->db->id); 1074 } else { 1075 decrRefCount(dstset); 1076 addReply(c,shared.czero); 1077 if (deleted) 1078 notifyKeyspaceEvent(REDIS_NOTIFY_GENERIC,"del", 1079 dstkey,c->db->id); 1080 } 1081 signalModifiedKey(c->db,dstkey); 1082 server.dirty++; 1083 } 1084 zfree(sets); 1085 } 1086 1087 void sunionCommand(redisClient *c) { 1088 sunionDiffGenericCommand(c,c->argv+1,c->argc-1,NULL,REDIS_OP_UNION); 1089 } 1090 1091 void sunionstoreCommand(redisClient *c) { 1092 sunionDiffGenericCommand(c,c->argv+2,c->argc-2,c->argv[1],REDIS_OP_UNION); 1093 } 1094 1095 void sdiffCommand(redisClient *c) { 1096 sunionDiffGenericCommand(c,c->argv+1,c->argc-1,NULL,REDIS_OP_DIFF); 1097 } 1098 1099 void sdiffstoreCommand(redisClient *c) { 1100 sunionDiffGenericCommand(c,c->argv+2,c->argc-2,c->argv[1],REDIS_OP_DIFF); 1101 } 1102 1103 void sscanCommand(redisClient *c) { 1104 robj *set; 1105 unsigned long cursor; 1106 1107 if (parseScanCursorOrReply(c,c->argv[2],&cursor) == REDIS_ERR) return; 1108 if ((set = lookupKeyReadOrReply(c,c->argv[1],shared.emptyscan)) == NULL || 1109 checkType(c,set,REDIS_SET)) return; 1110 scanGenericCommand(c,set,cursor); 1111 } 1112