Merge pull request #4337 from dvirsky/module_ctx_flags

Module Context flags

3 files changed, 185 insertions, 5 deletions

diff --git a/src/module.c b/src/module.c
index fda68b273..ae76ae612 100644
--- a/src/module.c
+++ b/src/module.c
@@ -1262,6 +1262,74 @@ int RM_GetSelectedDb(RedisModuleCtx *ctx) {
     return ctx->client->db->id;
 }
 
+
+/* Return the current context's flags. The flags provide information on the 
+ * current request context (whether the client is a Lua script or in a MULTI),
+ * and about the Redis instance in general, i.e replication and persistence. 
+ * 
+ * The available flags are:
+ * 
+ *  * REDISMODULE_CTX_FLAGS_LUA: The command is running in a Lua script
+ * 
+ *  * REDISMODULE_CTX_FLAGS_MULTI: The command is running inside a transaction
+ * 
+ *  * REDISMODULE_CTX_FLAGS_MASTER: The Redis instance is a master
+ * 
+ *  * REDISMODULE_CTX_FLAGS_SLAVE: The Redis instance is a slave
+ * 
+ *  * REDISMODULE_CTX_FLAGS_READONLY: The Redis instance is read-only
+ * 
+ *  * REDISMODULE_CTX_FLAGS_CLUSTER: The Redis instance is in cluster mode
+ * 
+ *  * REDISMODULE_CTX_FLAGS_AOF: The Redis instance has AOF enabled
+ * 
+ *  * REDISMODULE_CTX_FLAGS_RDB: The instance has RDB enabled
+ * 
+ *  * REDISMODULE_CTX_FLAGS_MAXMEMORY:  The instance has Maxmemory set
+ * 
+ *  * REDISMODULE_CTX_FLAGS_EVICT:  Maxmemory is set and has an eviction
+ *    policy that may delete keys
+ */
+int RM_GetContextFlags(RedisModuleCtx *ctx) {
+    
+    int flags = 0;
+    /* Client specific flags */
+    if (ctx->client) {
+        if (ctx->client->flags & CLIENT_LUA) 
+         flags |= REDISMODULE_CTX_FLAGS_LUA;
+        if (ctx->client->flags & CLIENT_MULTI) 
+         flags |= REDISMODULE_CTX_FLAGS_MULTI;
+    }
+
+    if (server.cluster_enabled)
+        flags |= REDISMODULE_CTX_FLAGS_CLUSTER;
+    
+    /* Maxmemory and eviction policy */
+    if (server.maxmemory > 0) {
+        flags |= REDISMODULE_CTX_FLAGS_MAXMEMORY;
+        
+        if (server.maxmemory_policy != MAXMEMORY_NO_EVICTION)
+            flags |= REDISMODULE_CTX_FLAGS_EVICT;
+    }
+
+    /* Persistence flags */
+    if (server.aof_state != AOF_OFF)
+        flags |= REDISMODULE_CTX_FLAGS_AOF;
+    if (server.saveparamslen > 0)
+        flags |= REDISMODULE_CTX_FLAGS_RDB;
+
+    /* Replication flags */
+    if (server.masterhost == NULL) {
+        flags |= REDISMODULE_CTX_FLAGS_MASTER;
+    } else {
+        flags |= REDISMODULE_CTX_FLAGS_SLAVE;
+        if (server.repl_slave_ro)
+            flags |= REDISMODULE_CTX_FLAGS_READONLY;
+    }
+    
+    return flags;
+}
+
 /* Change the currently selected DB. Returns an error if the id
  * is out of range.
  *
@@ -3333,14 +3401,16 @@ void unblockClientFromModule(client *c) {
 RedisModuleBlockedClient *RM_BlockClient(RedisModuleCtx *ctx, RedisModuleCmdFunc reply_callback, RedisModuleCmdFunc timeout_callback, void (*free_privdata)(void*), long long timeout_ms) {
     client *c = ctx->client;
     int islua = c->flags & CLIENT_LUA;
+    int ismulti = c->flags & CLIENT_MULTI;
 
     c->bpop.module_blocked_handle = zmalloc(sizeof(RedisModuleBlockedClient));
     RedisModuleBlockedClient *bc = c->bpop.module_blocked_handle;
 
     /* We need to handle the invalid operation of calling modules blocking
-     * commands from Lua. We actually create an already aborted (client set to
-     * NULL) blocked client handle, and actually reply to Lua with an error. */
-    bc->client = islua ? NULL : c;
+     * commands from Lua or MULTI. We actually create an already aborted 
+     * (client set to NULL) blocked client handle, and actually reply with 
+     * an error. */
+    bc->client = (islua || ismulti) ? NULL : c;
     bc->module = ctx->module;
     bc->reply_callback = reply_callback;
     bc->timeout_callback = timeout_callback;
@@ -3351,9 +3421,10 @@ RedisModuleBlockedClient *RM_BlockClient(RedisModuleCtx *ctx, RedisModuleCmdFunc
     bc->dbid = c->db->id;
     c->bpop.timeout = timeout_ms ? (mstime()+timeout_ms) : 0;
 
-    if (islua) {
+    if (islua || ismulti) {
         c->bpop.module_blocked_handle = NULL;
-        addReplyError(c,"Blocking module command called from Lua script");
+        addReplyError(c, islua ? "Blocking module command called from Lua script" :
+                                 "Blocking module command called from transaction");
     } else {
         blockClient(c,BLOCKED_MODULE);
     }
@@ -3891,6 +3962,7 @@ void moduleRegisterCoreAPI(void) {
     REGISTER_API(IsKeysPositionRequest);
     REGISTER_API(KeyAtPos);
     REGISTER_API(GetClientId);
+    REGISTER_API(GetContextFlags);
     REGISTER_API(PoolAlloc);
     REGISTER_API(CreateDataType);
     REGISTER_API(ModuleTypeSetValue);
diff --git a/src/modules/testmodule.c b/src/modules/testmodule.c
index 8da45c0ea..a0d706fea 100644
--- a/src/modules/testmodule.c
+++ b/src/modules/testmodule.c
@@ -121,6 +121,81 @@ int TestStringPrintf(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
 }
 
 
+/* TEST.CTXFLAGS -- Test GetContextFlags. */
+int TestCtxFlags(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
+    REDISMODULE_NOT_USED(argc);
+    REDISMODULE_NOT_USED(argv);
+  
+    RedisModule_AutoMemory(ctx);
+  
+    int ok = 1;
+    const char *errString = NULL;
+  
+  #define FAIL(msg)    \
+    {                  \
+      ok = 0;          \
+      errString = msg; \
+      goto end;        \
+    }
+  
+    int flags = RedisModule_GetContextFlags(ctx);
+    if (flags == 0) {
+      FAIL("Got no flags");
+    }
+  
+    if (flags & REDISMODULE_CTX_FLAGS_LUA) FAIL("Lua flag was set");
+    if (flags & REDISMODULE_CTX_FLAGS_MULTI) FAIL("Multi flag was set");
+  
+    if (flags & REDISMODULE_CTX_FLAGS_AOF) FAIL("AOF Flag was set")
+    /* Enable AOF to test AOF flags */
+    RedisModule_Call(ctx, "config", "ccc", "set", "appendonly", "yes");
+    flags = RedisModule_GetContextFlags(ctx);
+    if (!(flags & REDISMODULE_CTX_FLAGS_AOF))
+      FAIL("AOF Flag not set after config set");
+  
+    if (flags & REDISMODULE_CTX_FLAGS_RDB) FAIL("RDB Flag was set");
+    /* Enable RDB to test RDB flags */
+    RedisModule_Call(ctx, "config", "ccc", "set", "save", "900 1");
+    flags = RedisModule_GetContextFlags(ctx);
+    if (!(flags & REDISMODULE_CTX_FLAGS_RDB))
+      FAIL("RDB Flag was not set after config set");
+  
+    if (!(flags & REDISMODULE_CTX_FLAGS_MASTER)) FAIL("Master flag was not set");
+    if (flags & REDISMODULE_CTX_FLAGS_SLAVE) FAIL("Slave flag was set");
+    if (flags & REDISMODULE_CTX_FLAGS_READONLY) FAIL("Read-only flag was set");
+    if (flags & REDISMODULE_CTX_FLAGS_CLUSTER) FAIL("Cluster flag was set");
+  
+    if (flags & REDISMODULE_CTX_FLAGS_MAXMEMORY) FAIL("Maxmemory flag was set");
+    ;
+    RedisModule_Call(ctx, "config", "ccc", "set", "maxmemory", "100000000");
+    flags = RedisModule_GetContextFlags(ctx);
+    if (!(flags & REDISMODULE_CTX_FLAGS_MAXMEMORY))
+      FAIL("Maxmemory flag was not set after config set");
+  
+    if (flags & REDISMODULE_CTX_FLAGS_EVICT) FAIL("Eviction flag was set");
+    RedisModule_Call(ctx, "config", "ccc", "set", "maxmemory-policy",
+                     "allkeys-lru");
+    flags = RedisModule_GetContextFlags(ctx);
+    if (!(flags & REDISMODULE_CTX_FLAGS_EVICT))
+      FAIL("Eviction flag was not set after config set");
+  
+  end:
+    /* Revert config changes */
+    RedisModule_Call(ctx, "config", "ccc", "set", "appendonly", "no");
+    RedisModule_Call(ctx, "config", "ccc", "set", "save", "");
+    RedisModule_Call(ctx, "config", "ccc", "set", "maxmemory", "0");
+    RedisModule_Call(ctx, "config", "ccc", "set", "maxmemory-policy", "noeviction");
+  
+    if (!ok) {
+      RedisModule_Log(ctx, "warning", "Failed CTXFLAGS Test. Reason: %s",
+                      errString);
+      return RedisModule_ReplyWithSimpleString(ctx, "ERR");
+    }
+  
+    return RedisModule_ReplyWithSimpleString(ctx, "OK");
+  }
+
+
 /* ----------------------------- Test framework ----------------------------- */
 
 /* Return 1 if the reply matches the specified string, otherwise log errors
@@ -188,6 +263,9 @@ int TestIt(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
     T("test.call","");
     if (!TestAssertStringReply(ctx,reply,"OK",2)) goto fail;
 
+    T("test.ctxflags","");
+    if (!TestAssertStringReply(ctx,reply,"OK",2)) goto fail;
+
     T("test.string.append","");
     if (!TestAssertStringReply(ctx,reply,"foobar",6)) goto fail;
 
@@ -229,6 +307,10 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc)
         TestStringPrintf,"write deny-oom",1,1,1) == REDISMODULE_ERR)
         return REDISMODULE_ERR;
 
+    if (RedisModule_CreateCommand(ctx,"test.ctxflags",
+        TestCtxFlags,"readonly",1,1,1) == REDISMODULE_ERR)
+        return REDISMODULE_ERR;
+
     if (RedisModule_CreateCommand(ctx,"test.it",
         TestIt,"readonly",1,1,1) == REDISMODULE_ERR)
         return REDISMODULE_ERR;
diff --git a/src/redismodule.h b/src/redismodule.h
index 7fc0fec40..8df203aba 100644
--- a/src/redismodule.h
+++ b/src/redismodule.h
@@ -58,6 +58,30 @@
 #define REDISMODULE_HASH_CFIELDS    (1<<2)
 #define REDISMODULE_HASH_EXISTS     (1<<3)
 
+/* Context Flags: Info about the current context returned by RM_GetContextFlags */
+
+/* The command is running in the context of a Lua script */
+#define REDISMODULE_CTX_FLAGS_LUA 0x0001
+/* The command is running inside a Redis transaction */
+#define REDISMODULE_CTX_FLAGS_MULTI 0x0002
+/* The instance is a master */
+#define REDISMODULE_CTX_FLAGS_MASTER 0x0004
+/* The instance is a slave */
+#define REDISMODULE_CTX_FLAGS_SLAVE 0x0008
+/* The instance is read-only (usually meaning it's a slave as well) */
+#define REDISMODULE_CTX_FLAGS_READONLY 0x0010
+/* The instance is running in cluster mode */
+#define REDISMODULE_CTX_FLAGS_CLUSTER 0x0020
+/* The instance has AOF enabled */
+#define REDISMODULE_CTX_FLAGS_AOF 0x0040 //
+/* The instance has RDB enabled */
+#define REDISMODULE_CTX_FLAGS_RDB 0x0080 //
+/* The instance has Maxmemory set */
+#define REDISMODULE_CTX_FLAGS_MAXMEMORY 0x0100
+/* Maxmemory is set and has an eviction policy that may delete keys */
+#define REDISMODULE_CTX_FLAGS_EVICT 0x0200 
+
+
 /* A special pointer that we can use between the core and the module to signal
  * field deletion, and that is impossible to be a valid pointer. */
 #define REDISMODULE_HASH_DELETE ((RedisModuleString*)(long)1)
@@ -183,6 +207,7 @@ int REDISMODULE_API_FUNC(RedisModule_HashGet)(RedisModuleKey *key, int flags, ..
 int REDISMODULE_API_FUNC(RedisModule_IsKeysPositionRequest)(RedisModuleCtx *ctx);
 void REDISMODULE_API_FUNC(RedisModule_KeyAtPos)(RedisModuleCtx *ctx, int pos);
 unsigned long long REDISMODULE_API_FUNC(RedisModule_GetClientId)(RedisModuleCtx *ctx);
+int REDISMODULE_API_FUNC(RedisModule_GetContextFlags)(RedisModuleCtx *ctx);
 void *REDISMODULE_API_FUNC(RedisModule_PoolAlloc)(RedisModuleCtx *ctx, size_t bytes);
 RedisModuleType *REDISMODULE_API_FUNC(RedisModule_CreateDataType)(RedisModuleCtx *ctx, const char *name, int encver, RedisModuleTypeMethods *typemethods);
 int REDISMODULE_API_FUNC(RedisModule_ModuleTypeSetValue)(RedisModuleKey *key, RedisModuleType *mt, void *value);
@@ -302,6 +327,7 @@ static int RedisModule_Init(RedisModuleCtx *ctx, const char *name, int ver, int
     REDISMODULE_GET_API(IsKeysPositionRequest);
     REDISMODULE_GET_API(KeyAtPos);
     REDISMODULE_GET_API(GetClientId);
+    REDISMODULE_GET_API(GetContextFlags);
     REDISMODULE_GET_API(PoolAlloc);
     REDISMODULE_GET_API(CreateDataType);
     REDISMODULE_GET_API(ModuleTypeSetValue);