Fix memory management in local context dict

Author: lantiga

src/dag.c

@@ -84,7 +84,7 @@ void *RedisAI_DagRunSession(RedisAI_RunInfo *rinfo) {
               const char *key_string = RedisModule_StringPtrLen(
                   currentOp->outkeys[outputNumber], NULL);
               const char *dictKey = RedisModule_Strdup(key_string);
-              AI_dictReplace(rinfo->dagTensorsContext, (void*)dictKey, tensor);
+              AI_dictReplace(rinfo->dagTensorsContext, (void*)dictKey, RAI_TensorGetShallowCopy(tensor));
             } else {
               RAI_SetError(currentOp->err, RAI_EMODELRUN,
                            "ERR output tensor on DAG's MODELRUN was null");
@@ -126,7 +126,7 @@ void *RedisAI_DagRunSession(RedisAI_RunInfo *rinfo) {
                         const char *key_string = RedisModule_StringPtrLen(
                                 currentOp->outkeys[outputNumber], NULL);
                         const char *dictKey = RedisModule_Strdup(key_string);
-                        AI_dictReplace(rinfo->dagTensorsContext, (void*)dictKey, tensor);
+                        AI_dictReplace(rinfo->dagTensorsContext, (void*)dictKey, RAI_TensorGetShallowCopy(tensor));
                     } else {
                         RAI_SetError(currentOp->err, RAI_EMODELRUN,
                                      "ERR output tensor on DAG's SCRIPTRUN was null");
@@ -329,7 +329,7 @@ int RAI_parseDAGLoadArgs(RedisModuleCtx *ctx, RedisModuleString **argv,
       }
       RedisModule_CloseKey(key);
       const char *dictKey = RedisModule_Strdup(arg_string);
-      AI_dictAdd(*localContextDict, (void*)dictKey, t);
+      AI_dictAdd(*localContextDict, (void*)dictKey, RAI_TensorGetShallowCopy(t));
       const char *keyspacePersistKey = RedisModule_Strdup(dictKey);
       AI_dictAdd(*loadedContextDict, (void*)keyspacePersistKey, (void *)1);
       number_loaded_keys++;

src/run_info.c

@@ -177,37 +177,6 @@ void RAI_FreeRunInfo(RedisModuleCtx *ctx, struct RedisAI_RunInfo *rinfo) {
   RAI_FreeError(rinfo->err);
 
   if (rinfo->dagTensorsContext) {
-    AI_dictIterator *iter = AI_dictGetSafeIterator(rinfo->dagTensorsContext);
-    AI_dictEntry *entry = AI_dictNext(iter);
-    RAI_Tensor *tensor = NULL;
-
-    while (entry) {
-      tensor = AI_dictGetVal(entry);
-      char *key = (char *)AI_dictGetKey(entry);
-
-      if (tensor && key != NULL) {
-        // if the key is persistent then we should not delete it
-        AI_dictEntry *persistent_entry =
-            AI_dictFind(rinfo->dagTensorsPersistentContext, key);
-        // if the key was loaded from the keyspace then we should not delete it
-        AI_dictEntry *loaded_entry =
-            AI_dictFind(rinfo->dagTensorsLoadedContext, key);
-
-        if (persistent_entry == NULL && loaded_entry == NULL) {
-          AI_dictDelete(rinfo->dagTensorsContext, key);
-        }
-
-        if (persistent_entry) {
-          AI_dictDelete(rinfo->dagTensorsPersistentContext, key);
-        }
-        if (loaded_entry) {
-          AI_dictDelete(rinfo->dagTensorsLoadedContext, key);
-        }
-      }
-      entry = AI_dictNext(iter);
-    }
-    AI_dictReleaseIterator(iter);
-
     AI_dictRelease(rinfo->dagTensorsContext);
     AI_dictRelease(rinfo->dagTensorsLoadedContext);
     AI_dictRelease(rinfo->dagTensorsPersistentContext);