diff --git a/src/module.c b/src/module.c index 31c882218..cb6ba1e8d 100644 --- a/src/module.c +++ b/src/module.c @@ -835,10 +835,95 @@ RedisModuleString *RM_ListPop(RedisModuleKey *key, int where) { * Key API for Sorted Set type * -------------------------------------------------------------------------- */ -int RM_ZsetAdd(RedisModuleKey *key, double score, RedisModuleString *ele) { +/* Conversion from/to public flags of the Modules API and our private flags, + * so that we have everything decoupled. */ +int RM_ZsetAddFlagsToCoreFlags(int flags) { + int retflags = 0; + if (flags & REDISMODULE_ZADD_XX) retflags |= ZADD_XX; + if (flags & REDISMODULE_ZADD_NX) retflags |= ZADD_NX; + return retflags; +} + +/* See previous function comment. */ +int RM_ZsetAddFlagsFromCoreFlags(int flags) { + int retflags = 0; + if (flags & ZADD_ADDED) retflags |= REDISMODULE_ZADD_ADDED; + if (flags & ZADD_UPDATED) retflags |= REDISMODULE_ZADD_UPDATED; + if (flags & ZADD_NOP) retflags |= REDISMODULE_ZADD_NOP; + return retflags; +} + +/* Add a new element into a sorted set, with the specified 'score'. + * If the element already exists, the score is updated. + * + * A new sorted set is created at value if the key is an empty open key + * setup for writing. + * + * Additional flags can be passed to the function via a pointer, the flags + * are both used to receive input and to communicate state when the function + * returns. 'flagsptr' can be NULL if no special flags are used. + * + * The input flags are: + * + * REDISMODULE_ZADD_XX: Element must already exist. Do nothing otherwise. + * REDISMODULE_ZADD_NX: Element must not exist. Do nothing otherwise. + * + * The output flags are: + * + * REDISMODULE_ZADD_ADDED: The new element was added to the sorted set. + * REDISMODULE_ZADD_UPDATED: The score of the element was updated. + * REDISMODULE_ZADD_NOP: No operation was performed because XX or NX flags. + * + * On success the function returns REDISMODULE_OK. On the following errors + * REDISMODULE_ERR is returned: + * + * - The key was not opened for writing. + * - The key is of the wrong type. + * - 'score' double value is not a number (NaN). + */ +int RM_ZsetAdd(RedisModuleKey *key, double score, RedisModuleString *ele, int *flagsptr) { + int flags = 0; if (!(key->mode & REDISMODULE_WRITE)) return REDISMODULE_ERR; - if (key->value == NULL) moduleCreateEmtpyKey(key,REDISMODULE_KEYTYPE_ZSET); if (key->value->type != OBJ_ZSET) return REDISMODULE_ERR; + if (key->value == NULL) moduleCreateEmtpyKey(key,REDISMODULE_KEYTYPE_ZSET); + if (flagsptr) flags = RM_ZsetAddFlagsToCoreFlags(*flagsptr); + if (zsetAdd(key->value,score,ele->ptr,&flags,NULL) == 0) { + if (flagsptr) *flagsptr = 0; + return REDISMODULE_ERR; + } + if (flagsptr) *flagsptr = RM_ZsetAddFlagsFromCoreFlags(flags); + return REDISMODULE_OK; +} + +/* This function works exactly like RM_ZsetAdd(), but instead of setting + * a new score, the score of the existing element is incremented, or if the + * element does not already exist, it is added assuming the old score was + * zero. + * + * The input and output flags, and the return value, have the same exact + * meaning, with the only difference that this function will return + * REDISMODULE_ERR even when 'score' is a valid double number, but adding it + * to the existing score resuts into a NaN (not a number) condition. + * + * This function has an additional field 'newscore', if not NULL is filled + * with the new score of the element after the increment, if no error + * is returned. */ +int RM_ZsetIncrby(RedisModuleKey *key, double score, RedisModuleString *ele, int *flagsptr, double *newscore) { + int flags = 0; + if (!(key->mode & REDISMODULE_WRITE)) return REDISMODULE_ERR; + if (key->value->type != OBJ_ZSET) return REDISMODULE_ERR; + if (key->value == NULL) moduleCreateEmtpyKey(key,REDISMODULE_KEYTYPE_ZSET); + if (flagsptr) flags = RM_ZsetAddFlagsToCoreFlags(*flagsptr); + if (zsetAdd(key->value,score,ele->ptr,&flags,newscore) == 0) { + if (flagsptr) *flagsptr = 0; + return REDISMODULE_ERR; + } + /* zsetAdd() may signal back that the resulting score is not a number. */ + if (flagsptr && (*flagsptr & ZADD_NAN)) { + *flagsptr = 0; + return REDISMODULE_ERR; + } + if (flagsptr) *flagsptr = RM_ZsetAddFlagsFromCoreFlags(flags); return REDISMODULE_OK; } diff --git a/src/redismodule.h b/src/redismodule.h index 0af014cf8..3ffa6c99d 100644 --- a/src/redismodule.h +++ b/src/redismodule.h @@ -39,6 +39,13 @@ /* Expire */ #define REDISMODULE_NO_EXPIRE -1 +/* Sorted set API flags. */ +#define REDISMODULE_ZADD_XX (1<<0) +#define REDISMODULE_ZADD_NX (1<<1) +#define REDISMODULE_ZADD_ADDED (1<<2) +#define REDISMODULE_ZADD_UPDATED (1<<3) +#define REDISMODULE_ZADD_NOP (1<<4) + /* Error messages. */ #define REDISMODULE_ERRORMSG_WRONGTYPE "WRONGTYPE Operation against a key holding the wrong kind of value"