Fix Ut due to ClickHouse/ClickHouse#64423

apache · Jun 7, 2024 · 119bc89 · 119bc89
1 parent 75e840c
commit 119bc89
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 deletions.
diff --git a/...ickhouse/src/main/java/org/apache/gluten/memory/alloc/CHManagedCHReservationListener.java b/...ickhouse/src/main/java/org/apache/gluten/memory/alloc/CHManagedCHReservationListener.java
@@ -53,6 +53,10 @@ public void reserveOrThrow(long size) {
       if (!open) {
         return;
       }
+      if (size < 0) {
+        LOG.error(String.format("Negative size (%d) is passed to CHManagedCHReservationListener. It is a bug", size));
+        return;
+      }
       if (LOG.isDebugEnabled()) {
         LOG.debug(String.format("reserve memory size from native: %d", size));
       }
@@ -80,6 +84,10 @@ public long reserve(long size) {
       if (!open) {
         return 0L;
       }
+      if (size < 0) {
+        LOG.error(String.format("Negative size (%d) is passed to CHManagedCHReservationListener. It is a bug", size));
+        return 0L;
+      }
       if (LOG.isDebugEnabled()) {
         LOG.debug(String.format("reserve memory (without exception) size from native: %d", size));
       }

diff --git a/cpp-ch/local-engine/Common/CHUtil.h b/cpp-ch/local-engine/Common/CHUtil.h
@@ -214,11 +214,15 @@ class BackendFinalizerUtil
 class IgnoreMemoryTracker
 {
 public:
-    explicit IgnoreMemoryTracker(size_t limit_) : limit(limit_) { DB::CurrentThread::get().untracked_memory_limit += limit; }
-    ~IgnoreMemoryTracker() { DB::CurrentThread::get().untracked_memory_limit -= limit; }
+    explicit IgnoreMemoryTracker(size_t limit_)
+    {
+        orgianl_untracked_memory_limit = DB::CurrentThread::get().untracked_memory_limit;
+        DB::CurrentThread::get().untracked_memory_limit += limit_;
+    }
+    ~IgnoreMemoryTracker() { DB::CurrentThread::get().untracked_memory_limit = orgianl_untracked_memory_limit; }
 
 private:
-    size_t limit;
+    size_t orgianl_untracked_memory_limit;
 };
 
 class DateTimeUtil