From a6c0798b031c6d4e10eb63de0a36c364489abfcb Mon Sep 17 00:00:00 2001
From: Zhen Wang <643348094@qq.com>
Date: Wed, 9 Oct 2024 22:04:53 +0800
Subject: [PATCH] [CORE] Update AllocationListener usage after successfully
 releasing memory (#7396)

---
 cpp/core/memory/MemoryAllocator.cc | 41 ++++++++++++++++++++----------
 1 file changed, 28 insertions(+), 13 deletions(-)

diff --git a/cpp/core/memory/MemoryAllocator.cc b/cpp/core/memory/MemoryAllocator.cc
index e34c4c947962..dd5fb8e1974b 100644
--- 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;
 }