RedisAI · lantiga · Aug 7, 2020 · Aug 5, 2020 · Aug 5, 2020 · Aug 5, 2020
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -95,7 +95,7 @@ commands:
 jobs:
   build-debian:
     docker:
-      - image: redisfab/rmbuilder:6.0.1-x64-buster
+      - image: redisfab/rmbuilder:6.0.5-x64-buster
     steps:
       - build-steps:
           platform: debian
@@ -112,7 +112,7 @@ jobs:
 
   coverage:
     docker:
-      - image: redisfab/rmbuilder:6.0.1-x64-buster
+      - image: redisfab/rmbuilder:6.0.5-x64-buster
     steps:
       - checkout
       - run:
@@ -136,7 +136,7 @@ jobs:
       - run:
           name: Test with coverage
           command: |
-            make -C opt test SHOW=1 COV=1
+            make -C opt test SHOW=1 COV=1 CLUSTER=1
             make -C opt cov-upload
           no_output_timeout: 20m
 
@@ -209,7 +209,7 @@ jobs:
       package:
         type: string
     docker:
-      - image: redisfab/rmbuilder:6.0.1-x64-buster
+      - image: redisfab/rmbuilder:6.0.5-x64-buster
     steps:
       - attach_workspace:
           at: workspace

diff --git a/opt/Makefile b/opt/Makefile
@@ -189,12 +189,14 @@ endif
 export GEN ?= 1
 export SLAVES ?= 1
 export AOF ?= 1
+export CLUSTER ?= 1
 
 test:
 	$(COVERAGE_RESET)
 	$(SHOW)\
 		DEVICE=$(DEVICE) \
 		MODULE=$(INSTALLED_TARGET) \
+		CLUSTER=$(CLUSTER) \
 		GEN=$(GEN) AOF=$(AOF) SLAVES=$(SLAVES) \
 		VALGRIND=$(VALGRIND) \
 		$(ROOT)/test/tests.sh

diff --git a/src/dag.c b/src/dag.c
@@ -510,7 +510,6 @@ int RAI_parseDAGLoadArgs(RedisModuleCtx *ctx, RedisModuleString **argv,
         ctx, "ERR invalid or negative value found in number of keys to LOAD");
     return -1;
   }
-
   int number_loaded_keys = 0;
   int separator_flag = 0;
   size_t argpos = 2;
@@ -522,8 +521,7 @@ int RAI_parseDAGLoadArgs(RedisModuleCtx *ctx, RedisModuleString **argv,
     } else {
       RAI_Tensor *t;
       RedisModuleKey *key;
-      const int status = RAI_GetTensorFromKeyspace(ctx, argv[argpos], &key, &t,
-                                                   REDISMODULE_READ);
+      const int status = RAI_GetTensorFromKeyspace(ctx, argv[argpos], &key, &t, REDISMODULE_READ);
       if (status == REDISMODULE_ERR) {
         RedisModule_Log(
             ctx, "warning",
@@ -532,9 +530,9 @@ int RAI_parseDAGLoadArgs(RedisModuleCtx *ctx, RedisModuleString **argv,
         return -1;
       }
       RedisModule_CloseKey(key);
-      char *dictKey = RedisModule_Alloc(strlen(arg_string) + 4);
+      char *dictKey = (char*) RedisModule_Alloc((strlen(arg_string) + 5)*sizeof(char));
       sprintf(dictKey, "%s%04d", arg_string, 1);
-      AI_dictAdd(*localContextDict, (void*)dictKey, t);
+      AI_dictAdd(*localContextDict, (void*)dictKey, (void *)t);
       AI_dictAdd(*loadedContextDict, (void*)dictKey, (void *)1);
       RedisModule_Free(dictKey);
       number_loaded_keys++;
@@ -587,10 +585,35 @@ int RAI_parseDAGPersistArgs(RedisModuleCtx *ctx, RedisModuleString **argv,
   return argpos;
 }
 
+int RedisAI_DagRun_IsKeysPositionRequest_ReportKeys(RedisModuleCtx *ctx,
+                                                    RedisModuleString **argv, int argc){
+    for (size_t argpos = 1; argpos < argc; argpos++){
+        const char *arg_string = RedisModule_StringPtrLen(argv[argpos], NULL);
+        if ( (!strcasecmp(arg_string, "LOAD") || !strcasecmp(arg_string, "PERSIST") ) && (argpos+1 < argc) ) {
+            long long n_keys;
+            argpos++;
+            const int retval = RedisModule_StringToLongLong(argv[argpos], &n_keys);
+            if(retval != REDISMODULE_OK){
+                return REDISMODULE_ERR;
+            }
+            argpos++;
+            if (n_keys > 0){
+                size_t last_persist_argpos = n_keys+argpos;
+                for (; argpos < last_persist_argpos &&  argpos < argc; argpos++){
+                    RedisModule_KeyAtPos(ctx, argpos);
+                }
+            }
+        }
+    }
+    return REDISMODULE_OK;
+}
+
 int RedisAI_DagRunSyntaxParser(RedisModuleCtx *ctx, RedisModuleString **argv,
                                  int argc, int dagMode) {
+  if (RedisModule_IsKeysPositionRequest(ctx)) {
+     return RedisAI_DagRun_IsKeysPositionRequest_ReportKeys(ctx, argv, argc);
+  }
   if (argc < 4) return RedisModule_WrongArity(ctx);
-
   RedisAI_RunInfo *rinfo = NULL;
   if (RAI_InitRunInfo(&rinfo) == REDISMODULE_ERR) {
     return RedisModule_ReplyWithError(
@@ -703,6 +726,7 @@ int RedisAI_DagRunSyntaxParser(RedisModuleCtx *ctx, RedisModuleString **argv,
 
   for (long long i=0; i<array_len(rinfo->dagOps); i++) {
     RAI_DagOp *currentOp = rinfo->dagOps[i];
+    if(currentOp==NULL) continue;
     int parse_result;
     switch (currentOp->commandType) {
       case REDISAI_DAG_CMD_TENSORSET:
@@ -818,7 +842,8 @@ int RedisAI_DagRunSyntaxParser(RedisModuleCtx *ctx, RedisModuleString **argv,
       } 
       int *instance = AI_dictGetVal(mangled_entry);
       RedisModuleString *mangled_key = RedisModule_CreateStringPrintf(ctx, "%s%04d", key, *instance);
-      AI_dictAdd(mangled_persisted, (void *)RedisModule_StringPtrLen(mangled_key, NULL), (void *)1);
+      const char* mangled_key_str = RedisModule_Strdup(RedisModule_StringPtrLen(mangled_key, NULL));
+      AI_dictAdd(mangled_persisted, (void *)mangled_key_str, (void *)1);
       entry = AI_dictNext(iter);
     }
     AI_dictReleaseIterator(iter);

diff --git a/src/dag.h b/src/dag.h
@@ -81,6 +81,21 @@ int RAI_parseDAGPersistArgs(RedisModuleCtx *ctx, RedisModuleString **argv,
                             int argc, AI_dict **localContextDict,
                             const char *chaining_operator);
 
+/**
+ * When a module command is called in order to obtain the position of
+ * keys, since it was flagged as "getkeys-api" during the registration,
+ * the command implementation checks for this special call using the
+ * RedisModule_IsKeysPositionRequest() API and uses this function in
+ * order to report keys.
+ * No real execution is done on this special call.
+ * @param ctx Context in which Redis modules operate
+ * @param argv Redis command arguments, as an array of strings
+ * @param argc Redis command number of arguments
+ * @return
+ */
+int RedisAI_DagRun_IsKeysPositionRequest_ReportKeys(RedisModuleCtx *ctx,
+                            RedisModuleString **argv, int argc);
+
 /**
  * DAGRUN and DAGRUN_RO parser, which reads the the sequence of
  * arguments and decides whether the sequence conforms to the syntax

diff --git a/src/model.c b/src/model.c
@@ -548,6 +548,18 @@ int RAI_ModelSerialize(RAI_Model *model, char **buffer, size_t *len, RAI_Error *
   return ret;
 }
 
+int RedisAI_ModelRun_IsKeysPositionRequest_ReportKeys(RedisModuleCtx *ctx,
+                                                    RedisModuleString **argv, int argc){
+  RedisModule_KeyAtPos(ctx, 1);
+  for (size_t argpos = 3; argpos < argc; argpos++){
+    const char *str = RedisModule_StringPtrLen(argv[argpos], NULL);
+    if (!strcasecmp(str, "OUTPUTS")) {
+      continue;
+    }
+    RedisModule_KeyAtPos(ctx,argpos);
+  }
+  return REDISMODULE_OK;
+}
 
 int RedisAI_Parse_ModelRun_RedisCommand(RedisModuleCtx *ctx,
                                         RedisModuleString **argv, int argc,
@@ -590,7 +602,6 @@ int RedisAI_Parse_ModelRun_RedisCommand(RedisModuleCtx *ctx,
       }
     }
   }
-
   if ((*mto)->inputs && array_len((*mto)->inputs) != ninputs) {
     RAI_SetError(
         error, RAI_EMODELRUN,

diff --git a/src/model.h b/src/model.h
@@ -201,6 +201,21 @@ int RAI_ModelSerialize(RAI_Model* model, char** buffer, size_t* len,
 int RAI_GetModelFromKeyspace(RedisModuleCtx* ctx, RedisModuleString* keyName,
                              RedisModuleKey** key, RAI_Model** model, int mode);
 
+/**
+ * When a module command is called in order to obtain the position of
+ * keys, since it was flagged as "getkeys-api" during the registration,
+ * the command implementation checks for this special call using the
+ * RedisModule_IsKeysPositionRequest() API and uses this function in
+ * order to report keys.
+ * No real execution is done on this special call.
+ * @param ctx Context in which Redis modules operate
+ * @param argv Redis command arguments, as an array of strings
+ * @param argc Redis command number of arguments
+ * @return
+ */
+int RedisAI_ModelRun_IsKeysPositionRequest_ReportKeys(RedisModuleCtx *ctx,
+                            RedisModuleString **argv, int argc);
+
 /**
  * Helper method to parse AI.MODELRUN arguments
  *

diff --git a/src/redisai.c b/src/redisai.c
@@ -535,20 +535,20 @@ int RedisAI_ModelScan_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv
  */
 int RedisAI_ModelRun_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv,
                                   int argc) {
+  if (RedisModule_IsKeysPositionRequest(ctx)) {
+    return RedisAI_ModelRun_IsKeysPositionRequest_ReportKeys(ctx, argv, argc);
+  }
   if (argc < 3) return RedisModule_WrongArity(ctx);
-
   RedisAI_RunInfo *rinfo = NULL;
   if (RAI_InitRunInfo(&rinfo) == REDISMODULE_ERR) {
     return RedisModule_ReplyWithError(ctx, "ERR Unable to allocate the memory and initialise the RedisAI_RunInfo structure");
   }
-
   RAI_Model *mto;
   RedisModuleKey *modelKey;
   const int status = RAI_GetModelFromKeyspace(ctx, argv[1], &modelKey, &mto, REDISMODULE_READ);
   if (status == REDISMODULE_ERR) {
       return REDISMODULE_ERR;
   }
-
   RedisModule_RetainString(NULL, argv[1]);
   rinfo->runkey = argv[1];
   rinfo->mctx = RAI_ModelRunCtxCreate(mto);
@@ -618,16 +618,17 @@ int RedisAI_ModelRun_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv,
 * AI.SCRIPTRUN <key> <function> INPUTS <input> [input ...] OUTPUTS <output> [output ...]
 */
 int RedisAI_ScriptRun_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
+  if (RedisModule_IsKeysPositionRequest(ctx)) {
+    return RedisAI_ScriptRun_IsKeysPositionRequest_ReportKeys(ctx, argv, argc);
+  }
   if (argc < 6) return RedisModule_WrongArity(ctx);
-
   RedisAI_RunInfo *rinfo = NULL;
   if (RAI_InitRunInfo(&rinfo) == REDISMODULE_ERR) {
     return RedisModule_ReplyWithError(
         ctx,
         "ERR Unable to allocate the memory and initialise the RedisAI_RunInfo"
         "structure");
   }
-
   RedisModule_RetainString(NULL, argv[1]);
   rinfo->runkey = argv[1];
 
@@ -1204,11 +1205,11 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc)
       == REDISMODULE_ERR)
     return REDISMODULE_ERR;
 
-  if (RedisModule_CreateCommand(ctx, "ai.dagrun", RedisAI_DagRun_RedisCommand, "write deny-oom", 3, 3, 1)
+  if (RedisModule_CreateCommand(ctx, "ai.dagrun", RedisAI_DagRun_RedisCommand, "write deny-oom getkeys-api", 3, 3, 1)
       == REDISMODULE_ERR)
     return REDISMODULE_ERR;
 
-  if (RedisModule_CreateCommand(ctx, "ai.dagrun_ro", RedisAI_DagRunRO_RedisCommand, "readonly", 3, 3, 1)
+  if (RedisModule_CreateCommand(ctx, "ai.dagrun_ro", RedisAI_DagRunRO_RedisCommand, "readonly getkeys-api", 3, 3, 1)
       == REDISMODULE_ERR)
     return REDISMODULE_ERR;
 

diff --git a/src/run_info.c b/src/run_info.c
@@ -179,6 +179,9 @@ int RAI_ShallowCopyDagRunInfo(RedisAI_RunInfo **result, RedisAI_RunInfo *src) {
 void RAI_FreeDagOp(RedisModuleCtx *ctx, RAI_DagOp *dagOp) {
   if (dagOp) {
     RAI_FreeError(dagOp->err);
+    if(dagOp->runkey){
+      RedisModule_FreeString(ctx,dagOp->runkey);
+    }
     if (dagOp->argv) {
       for (size_t i = 0; i < array_len(dagOp->argv); i++) {
         RedisModule_FreeString(ctx, dagOp->argv[i]);
@@ -203,14 +206,14 @@ void RAI_FreeDagOp(RedisModuleCtx *ctx, RAI_DagOp *dagOp) {
 
     if (dagOp->inkeys) {
       for (size_t i=0; i<array_len(dagOp->inkeys); i++) {
-        RedisModule_Free(dagOp->inkeys[i]);
+        RedisModule_FreeString(ctx,dagOp->inkeys[i]);
       }
       array_free(dagOp->inkeys);
     }
 
     if (dagOp->outkeys) {
       for (size_t i=0; i<array_len(dagOp->outkeys); i++) {
-        RedisModule_Free(dagOp->outkeys[i]);
+        RedisModule_FreeString(ctx,dagOp->outkeys[i]);
       }
       array_free(dagOp->outkeys);
     }
@@ -273,10 +276,6 @@ void RAI_FreeRunInfo(RedisModuleCtx *ctx, struct RedisAI_RunInfo *rinfo) {
       entry = AI_dictNext(iter);
     }
     AI_dictReleaseIterator(iter);
-
-    RedisModule_Free(rinfo->dagTensorsContext);
-    RedisModule_Free(rinfo->dagTensorsLoadedContext);
-    RedisModule_Free(rinfo->dagTensorsPersistedContext);
   }
 
   if (rinfo->dagOps) {

diff --git a/src/script.c b/src/script.c
@@ -266,6 +266,19 @@ int RAI_GetScriptFromKeyspace(RedisModuleCtx* ctx, RedisModuleString* keyName,
   return REDISMODULE_OK;
 }
 
+int RedisAI_ScriptRun_IsKeysPositionRequest_ReportKeys(RedisModuleCtx *ctx,
+                                                       RedisModuleString **argv, int argc){
+    RedisModule_KeyAtPos(ctx, 1);
+    for (size_t argpos = 4; argpos < argc; argpos++){
+        const char *str = RedisModule_StringPtrLen(argv[argpos], NULL);
+        if (!strcasecmp(str, "OUTPUTS")) {
+            continue;
+        }
+        RedisModule_KeyAtPos(ctx,argpos);
+    }
+    return REDISMODULE_OK;
+}
+
 /**
  * AI.SCRIPTRUN <key> <function> INPUTS <input> [input ...] OUTPUTS <output> [output ...]
  */

diff --git a/src/script.h b/src/script.h
@@ -170,6 +170,21 @@ int RAI_GetScriptFromKeyspace(RedisModuleCtx* ctx, RedisModuleString* keyName,
                               int mode);
 
 
+/**
+ * When a module command is called in order to obtain the position of
+ * keys, since it was flagged as "getkeys-api" during the registration,
+ * the command implementation checks for this special call using the
+ * RedisModule_IsKeysPositionRequest() API and uses this function in
+ * order to report keys.
+ * No real execution is done on this special call.
+ * @param ctx Context in which Redis modules operate
+ * @param argv Redis command arguments, as an array of strings
+ * @param argc Redis command number of arguments
+ * @return
+ */
+int RedisAI_ScriptRun_IsKeysPositionRequest_ReportKeys(RedisModuleCtx *ctx,
+                               RedisModuleString **argv, int argc);
+
 /**
  * Helper method to parse AI.SCRIPTRUN arguments
  *