From dbf95b4b3835bc0f6dd8a3909f65cea893271306 Mon Sep 17 00:00:00 2001
From: Hongze Zhang <hongze.zhang@intel.com>
Date: Wed, 24 Jul 2024 10:17:58 +0800
Subject: [PATCH] fixup

---
 cpp/core/jni/JniCommon.h             | 18 ++++++++++++------
 cpp/core/memory/AllocationListener.h |  2 +-
 cpp/core/memory/MemoryAllocator.cc   | 12 ++++++++++--
 cpp/core/memory/MemoryAllocator.h    |  6 ++----
 4 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/cpp/core/jni/JniCommon.h b/cpp/core/jni/JniCommon.h
index 8477f587ab17..1d784f3a5eda 100644
--- a/cpp/core/jni/JniCommon.h
+++ b/cpp/core/jni/JniCommon.h
@@ -403,9 +403,17 @@ class SparkAllocationListener final : public gluten::AllocationListener {
       env->CallLongMethod(jListenerGlobalRef_, jReserveMethod_, size);
       checkException(env);
     }
-    std::lock_guard<std::mutex> lock(mutex_);
     usedBytes_ += size;
-    peakBytes_ = std::max(usedBytes_, peakBytes_);
+    while (true) {
+      int64_t savedPeakBytes = peakBytes_;
+      if (usedBytes_ <= savedPeakBytes) {
+        break;
+      }
+      // usedBytes_ > savedPeakBytes, update peak
+      if (peakBytes_.compare_exchange_weak(savedPeakBytes, usedBytes_)) {
+        break;
+      }
+    }
   }
 
   int64_t currentBytes() override {
@@ -421,10 +429,8 @@ class SparkAllocationListener final : public gluten::AllocationListener {
   jobject jListenerGlobalRef_;
   const jmethodID jReserveMethod_;
   const jmethodID jUnreserveMethod_;
-  int64_t usedBytes_{0L};
-  int64_t peakBytes_{0L};
-
-  mutable std::mutex mutex_;
+  std::atomic_int64_t usedBytes_{0L};
+  std::atomic_int64_t peakBytes_{0L};
 };
 
 class BacktraceAllocationListener final : public gluten::AllocationListener {
diff --git a/cpp/core/memory/AllocationListener.h b/cpp/core/memory/AllocationListener.h
index a4e549cacc9d..41797641fe14 100644
--- a/cpp/core/memory/AllocationListener.h
+++ b/cpp/core/memory/AllocationListener.h
@@ -73,7 +73,7 @@ class BlockAllocationListener final : public AllocationListener {
   }
 
  private:
-  int64_t reserve(int64_t diff) {
+  inline int64_t reserve(int64_t diff) {
     std::lock_guard<std::mutex> lock(mutex_);
     usedBytes_ += diff;
     int64_t newBlockCount;
diff --git a/cpp/core/memory/MemoryAllocator.cc b/cpp/core/memory/MemoryAllocator.cc
index 4dd9de02827d..c637c6a9c13d 100644
--- a/cpp/core/memory/MemoryAllocator.cc
+++ b/cpp/core/memory/MemoryAllocator.cc
@@ -92,9 +92,17 @@ int64_t ListenableMemoryAllocator::peakBytes() const {
 
 void ListenableMemoryAllocator::updateUsage(int64_t size) {
   listener_->allocationChanged(size);
-  std::lock_guard<std::mutex> lock(mutex_);
   usedBytes_ += size;
-  peakBytes_ = std::max(peakBytes_, usedBytes_);
+  while (true) {
+    int64_t savedPeakBytes = peakBytes_;
+    if (usedBytes_ <= savedPeakBytes) {
+      break;
+    }
+    // usedBytes_ > savedPeakBytes, update peak
+    if (peakBytes_.compare_exchange_weak(savedPeakBytes, usedBytes_)) {
+      break;
+    }
+  }
 }
 
 bool StdMemoryAllocator::allocate(int64_t size, void** out) {
diff --git a/cpp/core/memory/MemoryAllocator.h b/cpp/core/memory/MemoryAllocator.h
index ffb932aa3efd..01271cc94673 100644
--- a/cpp/core/memory/MemoryAllocator.h
+++ b/cpp/core/memory/MemoryAllocator.h
@@ -72,10 +72,8 @@ class ListenableMemoryAllocator final : public MemoryAllocator {
   void updateUsage(int64_t size);
   MemoryAllocator* const delegated_;
   AllocationListener* const listener_;
-  int64_t usedBytes_{0L};
-  int64_t peakBytes_{0L};
-
-  mutable std::mutex mutex_;
+  std::atomic_int64_t usedBytes_{0L};
+  std::atomic_int64_t peakBytes_{0L};
 };
 
 class StdMemoryAllocator final : public MemoryAllocator {