Create shallow copies for every device and keep the original DAG runInfo in a designated field (so it won't be modified). Free the copy for every device that finishes, and use the original for calling OnFinish.

Author: alonre24

src/background_workers.c

@@ -175,10 +175,11 @@ void *RedisAI_Run_ThreadMain(void *arg) {
                     if (timedOut == 1) {
                         queueEvict(run_queue_info->run_queue, item);
 
+						RedisAI_RunInfo *orig = rinfo->orig_copy;
                         long long dagRefCount = RAI_DagRunInfoFreeShallowCopy(rinfo);
                         if (dagRefCount == 0) {
-                            RedisAI_OnFinishCtx finish_ctx = (RedisAI_RunInfo *)rinfo;
-                            rinfo->OnFinish(finish_ctx, rinfo->private_data);
+                            RedisAI_OnFinishCtx finish_ctx = (RedisAI_RunInfo *)orig;
+							orig->OnFinish(finish_ctx, orig->private_data);
                         }
 
                         queueItem *evicted_item = item;
@@ -413,10 +414,11 @@ void *RedisAI_Run_ThreadMain(void *arg) {
                     // If there was an error and the reference count for the dag
                     // has gone to zero and the client is still around, we unblock
                     if (dagError) {
+						RedisAI_RunInfo *orig = rinfo->orig_copy;
                         long long dagRefCount = RAI_DagRunInfoFreeShallowCopy(rinfo);
                         if (dagRefCount == 0) {
-                            RedisAI_OnFinishCtx finish_ctx = (RedisAI_RunInfo *)rinfo;
-                            rinfo->OnFinish(finish_ctx, rinfo->private_data);
+                            RedisAI_OnFinishCtx finish_ctx = (RedisAI_RunInfo *)orig;
+							orig->OnFinish(finish_ctx, orig->private_data);
                         }
                     } else {
                         rinfo->dagDeviceCompleteOpCount += 1;
@@ -433,22 +435,22 @@ void *RedisAI_Run_ThreadMain(void *arg) {
             int dag_complete_after_run = RedisAI_DagComplete(batch_rinfo[0]);
 
             long long dagRefCount = -1;
-
+			RedisAI_RunInfo *orig;
             if (device_complete == 1 || device_complete_after_run == 1) {
-                RedisAI_RunInfo *evicted_rinfo = (RedisAI_RunInfo *)(evicted_items[0]->value);
-                // We decrease and get the reference count for the DAG
+				RedisAI_RunInfo *evicted_rinfo = (RedisAI_RunInfo *)(evicted_items[0]->value);
+				orig = evicted_rinfo->orig_copy;
+                // We decrease and get the reference count for the DAG.
                 dagRefCount = RAI_DagRunInfoFreeShallowCopy(evicted_rinfo);
             }
 
             // If the DAG was complete, then it's time to unblock the client
             if (do_unblock == 1 || dag_complete_after_run == 1) {
-                RedisAI_RunInfo *evicted_rinfo = (RedisAI_RunInfo *)(evicted_items[0]->value);
 
                 // If the reference count for the DAG is zero and the client is still around,
                 // then we actually unblock the client
                 if (dagRefCount == 0) {
-                    RedisAI_OnFinishCtx finish_ctx = (RedisAI_RunInfo *)evicted_rinfo;
-                    evicted_rinfo->OnFinish(finish_ctx, evicted_rinfo->private_data);
+                    RedisAI_OnFinishCtx finish_ctx = (RedisAI_RunInfo *)orig;
+                    orig->OnFinish(finish_ctx, orig->private_data);
                 }
             }
 

src/dag.c

@@ -1302,7 +1302,7 @@ static bool DAG_InsertDAGToQueue(RedisAI_RunInfo *rinfo) {
 
 void DAG_ReplyAndUnblock(RedisAI_OnFinishCtx ctx, void *private_data) {
 
-    RedisAI_RunInfo *rinfo = (RedisAI_RunInfo *)private_data;
+    RedisAI_RunInfo *rinfo = (RedisAI_RunInfo *)ctx;
     if (rinfo->client)
         RedisModule_UnblockClient(rinfo->client, rinfo);
 }
@@ -1323,8 +1323,6 @@ int RedisAI_ProcessDagRunCommand(RedisModuleCtx *ctx, RedisModuleString **argv,
     // Block the client before adding rinfo to the run queues (sync call).
     rinfo->client = RedisModule_BlockClient(ctx, RedisAI_DagRun_Reply, NULL, RedisAI_FreeData, 0);
     RedisModule_SetDisconnectCallback(rinfo->client, RedisAI_Disconnected);
-    // Use the entire rinfo obj as the on finish's private data.
-    rinfo->private_data = rinfo;
     rinfo->OnFinish = DAG_ReplyAndUnblock;
     return DAG_InsertDAGToQueue(rinfo);
 }

src/run_info.c

@@ -130,6 +130,7 @@ int RAI_InitRunInfo(RedisAI_RunInfo **result) {
     rinfo->dagCompleteOpCount = RedisModule_Calloc(1, sizeof(long long));
     rinfo->dagDeviceOpCount = 0;
     rinfo->dagDeviceCompleteOpCount = 0;
+    rinfo->orig_copy = rinfo;
     pthread_rwlock_init(rinfo->dagLock, NULL);
     rinfo->timedOut = RedisModule_Calloc(1, sizeof(int));
     *result = rinfo;
@@ -206,9 +207,7 @@ long long RAI_DagRunInfoFreeShallowCopy(RedisAI_RunInfo *rinfo) {
     if (rinfo->dagDeviceOps) {
         array_free(rinfo->dagDeviceOps);
     }
-    // If this is the last run info copy we do not free it, the OnFinish callback may free.
-    if (ref_count > 0)
-        RedisModule_Free(rinfo);
+    RedisModule_Free(rinfo);
     return ref_count;
 }
 

src/run_info.h

@@ -79,7 +79,10 @@ typedef void (*RAI_OnFinishCB)(RedisAI_OnFinishCtx ctx, void *private_data);
  * Note that not all the context structure is always filled with actual values
  * but only the fields needed in a given operation.
  */
-typedef struct RedisAI_RunInfo {
+
+typedef struct RedisAI_RunInfo RedisAI_RunInfo;
+
+struct RedisAI_RunInfo {
     RedisModuleBlockedClient *client;
     int single_op_dag;
     int single_device_dag;
@@ -106,8 +109,9 @@ typedef struct RedisAI_RunInfo {
     int *timedOut;
     struct timeval queuingTime;
     RAI_OnFinishCB OnFinish;
+	RedisAI_RunInfo *orig_copy;
     void *private_data; // This is going to be sent to the OnFinish callback.
-} RedisAI_RunInfo;
+};
 
 /**
  * Allocate the memory and initialise the RedisAI_RunInfo.