1 files changed, 98 insertions, 0 deletions

diff --git a/tests/modules/blockedclient.c b/tests/modules/blockedclient.c
index 3e97d6779..a2d7c6d00 100644
--- a/tests/modules/blockedclient.c
+++ b/tests/modules/blockedclient.c
@@ -1,3 +1,8 @@
+/* define macros for having usleep */
+#define _BSD_SOURCE
+#define _DEFAULT_SOURCE
+#include <unistd.h>
+
 #include "redismodule.h"
 #include <assert.h>
 #include <stdio.h>
@@ -5,6 +10,10 @@
 
 #define UNUSED(V) ((void) V)
 
+/* used to test processing events during slow bg operation */
+static volatile int g_slow_bg_operation = 0;
+static volatile int g_is_in_slow_bg_operation = 0;
+
 void *sub_worker(void *arg) {
     // Get Redis module context
     RedisModuleCtx *ctx = (RedisModuleCtx *)arg;
@@ -99,6 +108,16 @@ void *bg_call_worker(void *arg) {
     // Acquire GIL
     RedisModule_ThreadSafeContextLock(ctx);
 
+    // Test slow operation yielding
+    if (g_slow_bg_operation) {
+        g_is_in_slow_bg_operation = 1;
+        while (g_slow_bg_operation) {
+            RedisModule_Yield(ctx, REDISMODULE_YIELD_FLAG_CLIENTS, "Slow module operation");
+            usleep(1000);
+        }
+        g_is_in_slow_bg_operation = 0;
+    }
+
     // Call the command
     const char* cmd = RedisModule_StringPtrLen(bg->argv[1], NULL);
     RedisModuleCallReply* rep = RedisModule_Call(ctx, cmd, "v", bg->argv + 2, bg->argc - 2);
@@ -203,6 +222,73 @@ int do_fake_bg_true(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
     return REDISMODULE_OK;
 }
 
+
+/* this flag is used to work with busy commands, that might take a while
+ * and ability to stop the busy work with a different command*/
+static volatile int abort_flag = 0;
+
+int slow_fg_command(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
+    if (argc != 2) {
+        RedisModule_WrongArity(ctx);
+        return REDISMODULE_OK;
+    }
+    long long block_time = 0;
+    if (RedisModule_StringToLongLong(argv[1], &block_time) != REDISMODULE_OK) {
+        RedisModule_ReplyWithError(ctx, "Invalid integer value");
+        return REDISMODULE_OK;
+    }
+
+    uint64_t start_time = RedisModule_MonotonicMicroseconds();
+    /* when not blocking indefinitely, we don't process client commands in this test. */
+    int yield_flags = block_time? REDISMODULE_YIELD_FLAG_NONE: REDISMODULE_YIELD_FLAG_CLIENTS;
+    while (!abort_flag) {
+        RedisModule_Yield(ctx, yield_flags, "Slow module operation");
+        usleep(1000);
+        if (block_time && RedisModule_MonotonicMicroseconds() - start_time > (uint64_t)block_time)
+            break;
+    }
+
+    abort_flag = 0;
+    RedisModule_ReplyWithLongLong(ctx, 1);
+    return REDISMODULE_OK;
+}
+
+int stop_slow_fg_command(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
+    REDISMODULE_NOT_USED(argv);
+    REDISMODULE_NOT_USED(argc);
+    abort_flag = 1;
+    RedisModule_ReplyWithLongLong(ctx, 1);
+    return REDISMODULE_OK;
+}
+
+/* used to enable or disable slow operation in do_bg_rm_call */
+static int set_slow_bg_operation(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
+    if (argc != 2) {
+        RedisModule_WrongArity(ctx);
+        return REDISMODULE_OK;
+    }
+    long long ll;
+    if (RedisModule_StringToLongLong(argv[1], &ll) != REDISMODULE_OK) {
+        RedisModule_ReplyWithError(ctx, "Invalid integer value");
+        return REDISMODULE_OK;
+    }
+    g_slow_bg_operation = ll;
+    RedisModule_ReplyWithSimpleString(ctx, "OK");
+    return REDISMODULE_OK;
+}
+
+/* used to test if we reached the slow operation in do_bg_rm_call */
+static int is_in_slow_bg_operation(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
+    UNUSED(argv);
+    if (argc != 1) {
+        RedisModule_WrongArity(ctx);
+        return REDISMODULE_OK;
+    }
+
+    RedisModule_ReplyWithLongLong(ctx, g_is_in_slow_bg_operation);
+    return REDISMODULE_OK;
+}
+
 int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
     REDISMODULE_NOT_USED(argv);
     REDISMODULE_NOT_USED(argc);
@@ -223,5 +309,17 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc)
     if (RedisModule_CreateCommand(ctx, "do_fake_bg_true", do_fake_bg_true, "", 0, 0, 0) == REDISMODULE_ERR)
         return REDISMODULE_ERR;
 
+    if (RedisModule_CreateCommand(ctx, "slow_fg_command", slow_fg_command,"", 0, 0, 0) == REDISMODULE_ERR)
+        return REDISMODULE_ERR;
+
+    if (RedisModule_CreateCommand(ctx, "stop_slow_fg_command", stop_slow_fg_command,"allow-busy", 0, 0, 0) == REDISMODULE_ERR)
+        return REDISMODULE_ERR;
+
+    if (RedisModule_CreateCommand(ctx, "set_slow_bg_operation", set_slow_bg_operation, "allow-busy", 0, 0, 0) == REDISMODULE_ERR)
+        return REDISMODULE_ERR;
+
+    if (RedisModule_CreateCommand(ctx, "is_in_slow_bg_operation", is_in_slow_bg_operation, "allow-busy", 0, 0, 0) == REDISMODULE_ERR)
+        return REDISMODULE_ERR;
+
     return REDISMODULE_OK;
 }