- Move the increase of the runInfo ref_count to inside the shallow copy function.

PR fixes.

Author: alonre24

src/background_workers.c

@@ -176,7 +176,7 @@ void *RedisAI_Run_ThreadMain(void *arg) {
                     if (timedOut == 1) {
                         queueEvict(run_queue_info->run_queue, item);
 
-                        long long dagRefCount = RAI_DagRunInfoDecreaseFetch(rinfo);
+                        long long dagRefCount = RAI_DagRunInfoFreeShallowCopy(rinfo);
                         if (dagRefCount == 0 && rinfo->client) {
                             RedisModule_UnblockClient(rinfo->client, rinfo);
                         }
@@ -415,7 +415,7 @@ 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) {
-                        long long dagRefCount = RAI_DagRunInfoDecreaseFetch(rinfo);
+                        long long dagRefCount = RAI_DagRunInfoFreeShallowCopy(rinfo);
                         if (dagRefCount == 0 && rinfo->client) {
                             RedisModule_UnblockClient(rinfo->client, rinfo);
                         }
@@ -438,7 +438,7 @@ void *RedisAI_Run_ThreadMain(void *arg) {
             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
-                dagRefCount = RAI_DagRunInfoDecreaseFetch(evicted_rinfo);
+                dagRefCount = RAI_DagRunInfoFreeShallowCopy(evicted_rinfo);
             }
 
             // If the DAG was complete, then it's time to unblock the client

src/dag.c

@@ -1230,9 +1230,6 @@ int RedisAI_DagRunSyntaxParser(RedisModuleCtx *ctx, RedisModuleString **argv, in
     }
 
     size_t ndevices = array_len(devices);
-
-    *rinfo->dagRefCount = ndevices;
-
     RedisAI_RunInfo **rinfo_copies = array_new(RedisAI_RunInfo *, ndevices);
     rinfo_copies = array_append(rinfo_copies, rinfo);
 

src/run_info.c

@@ -123,7 +123,8 @@ int RAI_InitRunInfo(RedisAI_RunInfo **result) {
     }
     rinfo->dagError = RedisModule_Calloc(1, sizeof(int));
     rinfo->dagLock = RedisModule_Alloc(sizeof(pthread_rwlock_t));
-    rinfo->dagRefCount = RedisModule_Calloc(1, sizeof(long long));
+    rinfo->dagRefCount = RedisModule_Alloc(sizeof(long long));
+    *rinfo->dagRefCount = 1;
     rinfo->dagOpCount = 0;
     rinfo->dagCompleteOpCount = RedisModule_Calloc(1, sizeof(long long));
     rinfo->dagDeviceOpCount = 0;
@@ -145,6 +146,7 @@ int RAI_ShallowCopyDagRunInfo(RedisAI_RunInfo **result, RedisAI_RunInfo *src) {
     if (!(rinfo->dagDeviceOps)) {
         return REDISMODULE_ERR;
     }
+    (*rinfo->dagRefCount)++;
     rinfo->dagDeviceOpCount = 0;
     rinfo->dagDeviceCompleteOpCount = 0;
     *result = rinfo;
@@ -197,7 +199,7 @@ void RAI_FreeDagOp(RedisModuleCtx *ctx, RAI_DagOp *dagOp) {
     }
 }
 
-long long RAI_DagRunInfoDecreaseFetch(RedisAI_RunInfo *rinfo) {
+long long RAI_DagRunInfoFreeShallowCopy(RedisAI_RunInfo *rinfo) {
     long long ref_count = __atomic_sub_fetch(rinfo->dagRefCount, 1, __ATOMIC_RELAXED);
     RedisModule_Assert(ref_count >= 0);
     if (rinfo->dagDeviceOps) {
@@ -213,7 +215,7 @@ void RAI_FreeRunInfo(RedisModuleCtx *ctx, struct RedisAI_RunInfo *rinfo) {
     if (!rinfo) {
         return;
     }
-    long long ref_count = __atomic_load_n(rinfo->dagRefCount, __ATOMIC_RELAXED);
+    long long ref_count = *rinfo->dagRefCount;
     RedisModule_Assert(ref_count == 0);
     pthread_rwlock_destroy(rinfo->dagLock);
     RedisModule_Free(rinfo->dagLock);

src/run_info.h

@@ -108,7 +108,7 @@ int RAI_ShallowCopyDagRunInfo(RedisAI_RunInfo **result, RedisAI_RunInfo *src);
  * @param rinfo copy to be freed.
  * @retval The ref_count of the rinfo object after freeing this copy.
  */
-long long RAI_DagRunInfoDecreaseFetch(RedisAI_RunInfo *rinfo);
+long long RAI_DagRunInfoFreeShallowCopy(RedisAI_RunInfo *rinfo);
 
 /**
  * Frees the memory allocated on RedisAI_RunInfo