From 25335f5ab5c0f3dd448613bc0cf75ec6e58f6a5a Mon Sep 17 00:00:00 2001
From: Yang Yanchao <yangyanchao6@huawei.com>
Date: Wed, 6 Sep 2023 16:01:53 +0800
Subject: [PATCH] fix three issue

Check the return ptr of the malloc[#I7YV2X]
add hnid in gmemPrefetch [#I7XQMW]
free userData in stream to avoid async operations accessing illeagal address[#I7Z2RF]

Signed-off-by: Yang Yanchao <yangyanchao6@huawei.com>
---
 include/libgmem.h | 10 ++++++----
 src/libgmem.c     | 15 +++++++++++----
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/include/libgmem.h b/include/libgmem.h
index 4cd918b..ece13a3 100644
--- a/include/libgmem.h
+++ b/include/libgmem.h
@@ -62,17 +62,19 @@ int gmemFreeEager(unsigned long addr, size_t size, void *stream);
  * @args
  *     addr:    start address of the memory to be migrated
  *     length:  length of the memory to be migrated
+ *     hnid:	the numa id of the target device
  *     stream:  the stream where work queue used by operation
  *
  * The physical data mapped by [addr, addr + length) will migrate to
- * the device corresponding to streamid. The hnuma node of the device
- *  in the current context needs to be automatically obtained
- * and the prefetch instruction of hmadvise is invoked.
+ * the target device corresponding to streamid.
+ *
+ * If stream is NULL, perform synchronous operation; otherwise,
+ * perform asynchronous operation.
  *
  * The asynchronous operation is non-blocking, but is executed in
  * sequence with the tasks in the stream.
  */
-int gmemPrefetch(unsigned long addr, size_t length, void *stream);
+int gmemPrefetch(unsigned long addr, size_t length, int hnid, void *stream);
 
 /*
  * gmemGetNumaId - get the numaid of the current device
diff --git a/src/libgmem.c b/src/libgmem.c
index 9a34de0..59682af 100644
--- a/src/libgmem.c
+++ b/src/libgmem.c
@@ -37,6 +37,7 @@ int gmemAdvise(void *userData)
 	if(ret) {
 		printf("hmadvise failed: addr:%lx, size:%lx, behavior:%d, hnid:%d\n", msg.start, msg.len_in, msg.behavior, msg.hnid);
 	}
+	free(userData);
 	return ret;
 }
 
@@ -44,27 +45,33 @@ int gmemFreeEager(unsigned long addr, size_t length, void *stream)
 {
 	int ret = -1;
 	gm_msg_t userData = (gm_msg_t)malloc(sizeof(gm_msg));
+	if (userData == NULL) {
+		fprintf(stderr, "Failed to allocate memory for userData\n");
+		return -ENOMEM;
+	}
 	mix_userData(userData, addr, length, -1, MADV_DONTNEED);
 	if (!stream) {
 		ret = gmemAdvise(userData);
 	} else if(gmemSemantics.FreeEager != NULL) {
 		ret = gmemSemantics.FreeEager(userData, stream);
 	}
-	free(userData);
 	return ret;
 }
 
-int gmemPrefetch(unsigned long addr, size_t length, void *stream)
+int gmemPrefetch(unsigned long addr, size_t length, int hnid, void *stream)
 {
 	int ret = -1;
 	gm_msg_t userData = (gm_msg_t)malloc(sizeof(gm_msg));
-	mix_userData(userData, addr, length, gmemGetNumaId(), MADV_PREFETCH);
+	if (userData == NULL) {
+		fprintf(stderr, "Failed to allocate memory for userData\n");
+		return -ENOMEM;
+	}
+	mix_userData(userData, addr, length, hnid, MADV_PREFETCH);
 	if (!stream) {
 		ret = gmemAdvise(userData);
 	} else if(gmemSemantics.Prefetch != NULL) {
 		ret = gmemSemantics.Prefetch(userData, stream);
 	}
-	free(userData);
 	return ret;
 }
 
-- 
2.41.0.windows.3