[GLUTEN-5476] Add current_memory function ref (#488)

loneylee · kyligence-git · commit 1dfaf7ffeaa7 · 2024-05-07T23:02:45.000Z
(cherry picked from commit fe75cd3)
diff --git a/src/Common/CurrentMemoryTracker.cpp b/src/Common/CurrentMemoryTracker.cpp
@@ -39,6 +39,7 @@ using DB::current_thread;
 thread_local std::function<void(Int64, bool)> CurrentMemoryTracker::before_alloc = nullptr;
 
 thread_local std::function<void(Int64)> CurrentMemoryTracker::before_free  = nullptr;
+thread_local std::function<Int64()> CurrentMemoryTracker::current_memory = nullptr;
 
 AllocationTrace CurrentMemoryTracker::allocImpl(Int64 size, bool throw_if_memory_exceeded)
 {
diff --git a/src/Common/CurrentMemoryTracker.h b/src/Common/CurrentMemoryTracker.h
@@ -19,6 +19,7 @@ struct CurrentMemoryTracker
 
     static thread_local std::function<void(Int64, bool)> before_alloc;
     static thread_local std::function<void(Int64)> before_free;
+    static thread_local std::function<Int64()> current_memory;
 private:
     [[nodiscard]] static AllocationTrace allocImpl(Int64 size, bool throw_if_memory_exceeded);
 };

Original file line number	Diff line number	Diff line change
`@@ -39,6 +39,7 @@ using DB::current_thread;`
`39`	`39`	`thread_local std::function<void(Int64, bool)> CurrentMemoryTracker::before_alloc = nullptr;`
`40`	`40`
`41`	`41`	`thread_local std::function<void(Int64)> CurrentMemoryTracker::before_free = nullptr;`
	`42`	`+thread_local std::function<Int64()> CurrentMemoryTracker::current_memory = nullptr;`
`42`	`43`
`43`	`44`	`AllocationTrace CurrentMemoryTracker::allocImpl(Int64 size, bool throw_if_memory_exceeded)`
`44`	`45`	`{`