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