From aea6e71ef82701e07177744e600e1ef20d60b7d0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Viktor=20S=C3=B6derqvist?= <viktor.soderqvist@est.tech>
Date: Fri, 5 Feb 2021 18:54:01 +0100
Subject: [PATCH] RM_ZsetRem: Delete key if empty (#8453)

Without this fix, RM_ZsetRem can leave empty sorted sets which are
not allowed to exist.

Removing from a sorted set while iterating seems to work (while
inserting causes failed assetions). RM_ZsetRangeEndReached is
modified to return 1 if the key doesn't exist, to terminate
iteration when the last element has been removed.
---
 runtest-moduleapi             |  1 +
 src/module.c                  |  2 ++
 tests/modules/Makefile        |  1 +
 tests/modules/zset.c          | 30 ++++++++++++++++++++++++++++++
 tests/unit/moduleapi/zset.tcl | 16 ++++++++++++++++
 5 files changed, 50 insertions(+)
 create mode 100644 tests/modules/zset.c
 create mode 100644 tests/unit/moduleapi/zset.tcl

diff --git a/runtest-moduleapi b/runtest-moduleapi
index e554226c1..53656fad7 100755
--- a/runtest-moduleapi
+++ b/runtest-moduleapi
@@ -32,5 +32,6 @@ $TCLSH tests/test_helper.tcl \
 --single unit/moduleapi/getkeys \
 --single unit/moduleapi/test_lazyfree \
 --single unit/moduleapi/defrag \
+--single unit/moduleapi/zset \
 --single unit/moduleapi/stream \
 "${@}"
diff --git a/src/module.c b/src/module.c
index b04595801..039465e1a 100644
--- a/src/module.c
+++ b/src/module.c
@@ -2612,6 +2612,7 @@ int RM_ZsetRem(RedisModuleKey *key, RedisModuleString *ele, int *deleted) {
     if (key->value && key->value->type != OBJ_ZSET) return REDISMODULE_ERR;
     if (key->value != NULL && zsetDel(key->value,ele->ptr)) {
         if (deleted) *deleted = 1;
+        moduleDelKeyIfEmpty(key);
     } else {
         if (deleted) *deleted = 0;
     }
@@ -2657,6 +2658,7 @@ void RM_ZsetRangeStop(RedisModuleKey *key) {
 
 /* Return the "End of range" flag value to signal the end of the iteration. */
 int RM_ZsetRangeEndReached(RedisModuleKey *key) {
+    if (!key->value || key->value->type != OBJ_ZSET) return 1;
     return key->u.zset.er;
 }
 
diff --git a/tests/modules/Makefile b/tests/modules/Makefile
index 73a9be3bc..8ea1d91a2 100644
--- a/tests/modules/Makefile
+++ b/tests/modules/Makefile
@@ -35,6 +35,7 @@ TEST_MODULES = \
     test_lazyfree.so \
     timer.so \
     defragtest.so \
+    zset.so \
     stream.so
 
 
diff --git a/tests/modules/zset.c b/tests/modules/zset.c
new file mode 100644
index 000000000..4806f6549
--- /dev/null
+++ b/tests/modules/zset.c
@@ -0,0 +1,30 @@
+#include "redismodule.h"
+
+/* ZSET.REM key element
+ *
+ * Removes an occurrence of an element from a sorted set. Replies with the
+ * number of removed elements (0 or 1).
+ */
+int zset_rem(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
+    if (argc != 3) return RedisModule_WrongArity(ctx);
+    RedisModule_AutoMemory(ctx);
+    int keymode = REDISMODULE_READ | REDISMODULE_WRITE;
+    RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], keymode);
+    int deleted;
+    if (RedisModule_ZsetRem(key, argv[2], &deleted) == REDISMODULE_OK)
+        return RedisModule_ReplyWithLongLong(ctx, deleted);
+    else
+        return RedisModule_ReplyWithError(ctx, "ERR ZsetRem failed");
+}
+
+int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
+    REDISMODULE_NOT_USED(argv);
+    REDISMODULE_NOT_USED(argc);
+    if (RedisModule_Init(ctx, "zset", 1, REDISMODULE_APIVER_1) ==
+        REDISMODULE_OK &&
+        RedisModule_CreateCommand(ctx, "zset.rem", zset_rem, "",
+                                  1, 1, 1) == REDISMODULE_OK)
+        return REDISMODULE_OK;
+    else
+        return REDISMODULE_ERR;
+}
diff --git a/tests/unit/moduleapi/zset.tcl b/tests/unit/moduleapi/zset.tcl
new file mode 100644
index 000000000..998d20765
--- /dev/null
+++ b/tests/unit/moduleapi/zset.tcl
@@ -0,0 +1,16 @@
+set testmodule [file normalize tests/modules/zset.so]
+
+start_server {tags {"modules"}} {
+    r module load $testmodule
+
+    test {Module zset rem} {
+        r del k
+        r zadd k 100 hello 200 world
+        assert_equal 1 [r zset.rem k hello]
+        assert_equal 0 [r zset.rem k hello]
+        assert_equal 1 [r exists k]
+        # Check that removing the last element deletes the key
+        assert_equal 1 [r zset.rem k world]
+        assert_equal 0 [r exists k]
+    }
+}