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