Merge pull request #520 from RedisAI/Make_ref_count_atomic_read

Make atomic read for ref_count in DAG run_info free function.

Author: alonre24

src/background_workers.c

@@ -176,8 +176,7 @@ void *RedisAI_Run_ThreadMain(void *arg) {
                     if (timedOut == 1) {
                         queueEvict(run_queue_info->run_queue, item);
 
-                        int dagRefCount =
-                            __atomic_sub_fetch(rinfo->dagRefCount, 1, __ATOMIC_RELAXED);
+                        long long dagRefCount = RAI_DagRunInfoFreeShallowCopy(rinfo);
                         if (dagRefCount == 0 && rinfo->client) {
                             RedisModule_UnblockClient(rinfo->client, rinfo);
                         }
@@ -416,9 +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) {
-                        int dagRefCount =
-                            __atomic_sub_fetch(rinfo->dagRefCount, 1, __ATOMIC_RELAXED);
-
+                        long long dagRefCount = RAI_DagRunInfoFreeShallowCopy(rinfo);
                         if (dagRefCount == 0 && rinfo->client) {
                             RedisModule_UnblockClient(rinfo->client, rinfo);
                         }
@@ -436,12 +433,12 @@ void *RedisAI_Run_ThreadMain(void *arg) {
             int device_complete_after_run = RedisAI_DagDeviceComplete(batch_rinfo[0]);
             int dag_complete_after_run = RedisAI_DagComplete(batch_rinfo[0]);
 
-            int dagRefCount = -1;
+            long long dagRefCount = -1;
 
             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 = __atomic_sub_fetch(evicted_rinfo->dagRefCount, 1, __ATOMIC_RELAXED);
+                dagRefCount = RAI_DagRunInfoFreeShallowCopy(evicted_rinfo);
             }
 
             // If the DAG was complete, then it's time to unblock the client

src/dag.c

@@ -1231,10 +1231,8 @@ 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->dagRefCount = 1;
     rinfo_copies = array_append(rinfo_copies, rinfo);
 
     for (long long i = 1; i < ndevices; i++) {

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) = 0;
     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,20 +199,26 @@ void RAI_FreeDagOp(RedisModuleCtx *ctx, RAI_DagOp *dagOp) {
     }
 }
 
-void RAI_FreeRunInfo(RedisModuleCtx *ctx, struct RedisAI_RunInfo *rinfo) {
-    if (!rinfo) {
-        return;
+long long RAI_DagRunInfoFreeShallowCopy(RedisAI_RunInfo *rinfo) {
+    long long ref_count = __atomic_sub_fetch(rinfo->dagRefCount, 1, __ATOMIC_RELAXED);
+    RedisModule_Assert(ref_count >= 0 && "Tried to free the original RunInfo object");
+    if (rinfo->dagDeviceOps) {
+        array_free(rinfo->dagDeviceOps);
     }
-    if (*rinfo->dagRefCount > 0) {
-        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);
+    return ref_count;
+}
+
+void RAI_FreeRunInfo(RedisModuleCtx *ctx, struct RedisAI_RunInfo *rinfo) {
+    if (!rinfo) {
         return;
-    } else {
-        pthread_rwlock_destroy(rinfo->dagLock);
-        RedisModule_Free(rinfo->dagLock);
     }
+    long long ref_count = *rinfo->dagRefCount;
+    RedisModule_Assert(ref_count == 0);
+    pthread_rwlock_destroy(rinfo->dagLock);
+    RedisModule_Free(rinfo->dagLock);
 
     if (rinfo->dagTensorsContext) {
         AI_dictRelease(rinfo->dagTensorsContext);
@@ -225,10 +233,6 @@ void RAI_FreeRunInfo(RedisModuleCtx *ctx, struct RedisAI_RunInfo *rinfo) {
         array_free(rinfo->dagOps);
     }
 
-    if (rinfo->dagDeviceOps) {
-        array_free(rinfo->dagDeviceOps);
-    }
-
     if (rinfo->dagError) {
         RedisModule_Free(rinfo->dagError);
     }

src/run_info.h

@@ -103,6 +103,13 @@ int RAI_InitRunInfo(RedisAI_RunInfo **result);
 
 int RAI_ShallowCopyDagRunInfo(RedisAI_RunInfo **result, RedisAI_RunInfo *src);
 
+/**
+ * Frees the shallow copy of RedisAI_RunInfo pointed by rinfo.
+ * @param rinfo copy to be freed.
+ * @retval The ref_count of the rinfo object after freeing this copy.
+ */
+long long RAI_DagRunInfoFreeShallowCopy(RedisAI_RunInfo *rinfo);
+
 /**
  * Frees the memory allocated on RedisAI_RunInfo
  * @param ctx Context in which Redis modules operate