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