[CORE] Update AllocationListener usage after successfully releasing m…

…emory (#7396)
apache · Oct 9, 2024 · a6c0798 · a6c0798
1 parent f0c643d
commit a6c0798
Showing 1 changed file with 28 additions and 13 deletions.
diff --git a/cpp/core/memory/MemoryAllocator.cc b/cpp/core/memory/MemoryAllocator.cc
@@ -50,12 +50,20 @@ bool ListenableMemoryAllocator::allocateAligned(uint64_t alignment, int64_t size
 
 bool ListenableMemoryAllocator::reallocate(void* p, int64_t size, int64_t newSize, void** out) {
   int64_t diff = newSize - size;
-  updateUsage(diff);
-  bool succeed = delegated_->reallocate(p, size, newSize, out);
-  if (!succeed) {
-    updateUsage(-diff);
+  if (diff >= 0) {
+    updateUsage(diff);
+    bool succeed = delegated_->reallocate(p, size, newSize, out);
+    if (!succeed) {
+      updateUsage(-diff);
+    }
+    return succeed;
+  } else {
+    bool succeed = delegated_->reallocate(p, size, newSize, out);
+    if (succeed) {
+      updateUsage(diff);
+    }
+    return succeed;
   }
-  return succeed;
 }
 
 bool ListenableMemoryAllocator::reallocateAligned(
@@ -65,19 +73,26 @@ bool ListenableMemoryAllocator::reallocateAligned(
     int64_t newSize,
     void** out) {
   int64_t diff = newSize - size;
-  updateUsage(diff);
-  bool succeed = delegated_->reallocateAligned(p, alignment, size, newSize, out);
-  if (!succeed) {
-    updateUsage(-diff);
+  if (diff >= 0) {
+    updateUsage(diff);
+    bool succeed = delegated_->reallocateAligned(p, alignment, size, newSize, out);
+    if (!succeed) {
+      updateUsage(-diff);
+    }
+    return succeed;
+  } else {
+    bool succeed = delegated_->reallocateAligned(p, alignment, size, newSize, out);
+    if (succeed) {
+      updateUsage(diff);
+    }
+    return succeed;
   }
-  return succeed;
 }
 
 bool ListenableMemoryAllocator::free(void* p, int64_t size) {
-  updateUsage(-size);
   bool succeed = delegated_->free(p, size);
-  if (!succeed) {
-    updateUsage(size);
+  if (succeed) {
+    updateUsage(-size);
   }
   return succeed;
 }