diff options
| -rw-r--r-- | .gitignore | 1 | ||||
| -rw-r--r-- | src/module.c | 6 | ||||
| -rw-r--r-- | src/redismodule.h | 2 | ||||
| -rw-r--r-- | src/server.h | 2 | ||||
| -rw-r--r-- | src/t_zset.c | 31 | ||||
| -rw-r--r-- | tests/unit/type/zset.tcl | 58 |
6 files changed, 98 insertions, 2 deletions
diff --git a/.gitignore b/.gitignore index 721f1d74e..dd9d4b80e 100644 --- a/.gitignore +++ b/.gitignore @@ -37,3 +37,4 @@ Makefile.dep .ccls .ccls-cache/* compile_commands.json +redis.code-workspace diff --git a/src/module.c b/src/module.c index 02a3b1ddc..77468700e 100644 --- a/src/module.c +++ b/src/module.c @@ -2380,6 +2380,8 @@ int RM_ZsetAddFlagsToCoreFlags(int flags) { int retflags = 0; if (flags & REDISMODULE_ZADD_XX) retflags |= ZADD_XX; if (flags & REDISMODULE_ZADD_NX) retflags |= ZADD_NX; + if (flags & REDISMODULE_ZADD_GT) retflags |= ZADD_GT; + if (flags & REDISMODULE_ZADD_LT) retflags |= ZADD_LT; return retflags; } @@ -2406,6 +2408,10 @@ int RM_ZsetAddFlagsFromCoreFlags(int flags) { * * REDISMODULE_ZADD_XX: Element must already exist. Do nothing otherwise. * REDISMODULE_ZADD_NX: Element must not exist. Do nothing otherwise. + * REDISMODULE_ZADD_GT: If element exists, new score must be greater than the current score. + * Do nothing otherwise. Can optionally be combined with XX. + * REDISMODULE_ZADD_LT: If element exists, new score must be less than the current score. + * Do nothing otherwise. Can optionally be combined with XX. * * The output flags are: * diff --git a/src/redismodule.h b/src/redismodule.h index c0eedc221..00ca1f578 100644 --- a/src/redismodule.h +++ b/src/redismodule.h @@ -55,6 +55,8 @@ #define REDISMODULE_ZADD_ADDED (1<<2) #define REDISMODULE_ZADD_UPDATED (1<<3) #define REDISMODULE_ZADD_NOP (1<<4) +#define REDISMODULE_ZADD_GT (1<<5) +#define REDISMODULE_ZADD_LT (1<<6) /* Hash API flags. */ #define REDISMODULE_HASH_NONE 0 diff --git a/src/server.h b/src/server.h index c0ca89b8b..80df48881 100644 --- a/src/server.h +++ b/src/server.h @@ -1945,6 +1945,8 @@ void addACLLogEntry(client *c, int reason, int keypos, sds username); #define ZADD_INCR (1<<0) /* Increment the score instead of setting it. */ #define ZADD_NX (1<<1) /* Don't touch elements not already existing. */ #define ZADD_XX (1<<2) /* Only touch elements already existing. */ +#define ZADD_GT (1<<7) /* Only update existing when new scores are higher. */ +#define ZADD_LT (1<<8) /* Only update existing when new scores are lower. */ /* Output flags. */ #define ZADD_NOP (1<<3) /* Operation not performed because of conditionals.*/ diff --git a/src/t_zset.c b/src/t_zset.c index 4b28c24ef..2ef02ac8a 100644 --- a/src/t_zset.c +++ b/src/t_zset.c @@ -1283,6 +1283,10 @@ int zsetScore(robj *zobj, sds member, double *score) { * assume 0 as previous score. * ZADD_NX: Perform the operation only if the element does not exist. * ZADD_XX: Perform the operation only if the element already exist. + * ZADD_GT: Perform the operation on existing elements only if the new score is + * greater than the current score. + * ZADD_LT: Perform the operation on existing elements only if the new score is + * less than the current score. * * When ZADD_INCR is used, the new score of the element is stored in * '*newscore' if 'newscore' is not NULL. @@ -1317,6 +1321,8 @@ int zsetAdd(robj *zobj, double score, sds ele, int *flags, double *newscore) { int incr = (*flags & ZADD_INCR) != 0; int nx = (*flags & ZADD_NX) != 0; int xx = (*flags & ZADD_XX) != 0; + int gt = (*flags & ZADD_GT) != 0; + int lt = (*flags & ZADD_LT) != 0; *flags = 0; /* We'll return our response flags. */ double curscore; @@ -1348,7 +1354,12 @@ int zsetAdd(robj *zobj, double score, sds ele, int *flags, double *newscore) { } /* Remove and re-insert when score changed. */ - if (score != curscore) { + if (score != curscore && + /* LT? Only update if score is less than current. */ + (!lt || score < curscore) && + /* GT? Only update if score is greater than current. */ + (!gt || score > curscore)) + { zobj->ptr = zzlDelete(zobj->ptr,eptr); zobj->ptr = zzlInsert(zobj->ptr,ele,score); *flags |= ZADD_UPDATED; @@ -1393,7 +1404,12 @@ int zsetAdd(robj *zobj, double score, sds ele, int *flags, double *newscore) { } /* Remove and re-insert when score changes. */ - if (score != curscore) { + if (score != curscore && + /* LT? Only update if score is less than current. */ + (!lt || score < curscore) && + /* GT? Only update if score is greater than current. */ + (!gt || score > curscore)) + { znode = zslUpdateScore(zs->zsl,curscore,ele,score); /* Note that we did not removed the original element from * the hash table representing the sorted set, so we just @@ -1555,6 +1571,8 @@ void zaddGenericCommand(client *c, int flags) { else if (!strcasecmp(opt,"xx")) flags |= ZADD_XX; else if (!strcasecmp(opt,"ch")) flags |= ZADD_CH; else if (!strcasecmp(opt,"incr")) flags |= ZADD_INCR; + else if (!strcasecmp(opt,"gt")) flags |= ZADD_GT; + else if (!strcasecmp(opt,"lt")) flags |= ZADD_LT; else break; scoreidx++; } @@ -1564,6 +1582,8 @@ void zaddGenericCommand(client *c, int flags) { int nx = (flags & ZADD_NX) != 0; int xx = (flags & ZADD_XX) != 0; int ch = (flags & ZADD_CH) != 0; + int gt = (flags & ZADD_GT) != 0; + int lt = (flags & ZADD_LT) != 0; /* After the options, we expect to have an even number of args, since * we expect any number of score-element pairs. */ @@ -1580,6 +1600,13 @@ void zaddGenericCommand(client *c, int flags) { "XX and NX options at the same time are not compatible"); return; } + + if ((gt && nx) || (lt && nx) || (gt && lt)) { + addReplyError(c, + "GT, LT, and/or NX options at the same time are not compatible"); + return; + } + /* Note that XX is compatible with either GT or LT */ if (incr && elements > 1) { addReplyError(c, diff --git a/tests/unit/type/zset.tcl b/tests/unit/type/zset.tcl index d098840b0..e997bc66d 100644 --- a/tests/unit/type/zset.tcl +++ b/tests/unit/type/zset.tcl @@ -78,6 +78,46 @@ start_server {tags {"zset"}} { assert {[r zscore ztmp y] == 21} } + test "ZADD GT updates existing elements when new scores are greater" { + r del ztmp + r zadd ztmp 10 x 20 y 30 z + assert {[r zadd ztmp gt ch 5 foo 11 x 21 y 29 z] == 3} + assert {[r zcard ztmp] == 4} + assert {[r zscore ztmp x] == 11} + assert {[r zscore ztmp y] == 21} + assert {[r zscore ztmp z] == 30} + } + + test "ZADD LT updates existing elements when new scores are lower" { + r del ztmp + r zadd ztmp 10 x 20 y 30 z + assert {[r zadd ztmp lt ch 5 foo 11 x 21 y 29 z] == 2} + assert {[r zcard ztmp] == 4} + assert {[r zscore ztmp x] == 10} + assert {[r zscore ztmp y] == 20} + assert {[r zscore ztmp z] == 29} + } + + test "ZADD GT XX updates existing elements when new scores are greater and skips new elements" { + r del ztmp + r zadd ztmp 10 x 20 y 30 z + assert {[r zadd ztmp gt xx ch 5 foo 11 x 21 y 29 z] == 2} + assert {[r zcard ztmp] == 3} + assert {[r zscore ztmp x] == 11} + assert {[r zscore ztmp y] == 21} + assert {[r zscore ztmp z] == 30} + } + + test "ZADD LT XX updates existing elements when new scores are lower and skips new elements" { + r del ztmp + r zadd ztmp 10 x 20 y 30 z + assert {[r zadd ztmp lt xx ch 5 foo 11 x 21 y 29 z] == 1} + assert {[r zcard ztmp] == 3} + assert {[r zscore ztmp x] == 10} + assert {[r zscore ztmp y] == 20} + assert {[r zscore ztmp z] == 29} + } + test "ZADD XX and NX are not compatible" { r del ztmp catch {r zadd ztmp xx nx 10 x} err @@ -100,6 +140,24 @@ start_server {tags {"zset"}} { assert {[r zscore ztmp b] == 200} } + test "ZADD GT and NX are not compatible" { + r del ztmp + catch {r zadd ztmp gt nx 10 x} err + set err + } {ERR*} + + test "ZADD LT and NX are not compatible" { + r del ztmp + catch {r zadd ztmp lt nx 10 x} err + set err + } {ERR*} + + test "ZADD LT and GT are not compatible" { + r del ztmp + catch {r zadd ztmp lt gt 10 x} err + set err + } {ERR*} + test "ZADD INCR works like ZINCRBY" { r del ztmp r zadd ztmp 10 x 20 y 30 z |