fixup

apache · Jul 24, 2024 · dbf95b4 · dbf95b4
1 parent 8325fda
commit dbf95b4
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 13 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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 {