From afbdca856f2a412aa6538ac3672a5b2bb995ac03 Mon Sep 17 00:00:00 2001
From: Xinyi Zou <zouxinyi02@gmail.com>
Date: Mon, 24 Jun 2024 19:56:20 +0800
Subject: [PATCH] 1

---
 be/src/olap/page_cache.cpp | 25 +++++++++----------------
 be/src/olap/page_cache.h   |  3 +--
 2 files changed, 10 insertions(+), 18 deletions(-)

diff --git a/be/src/olap/page_cache.cpp b/be/src/olap/page_cache.cpp
index fe0a99af34f737..1f0556f4642110 100644
--- a/be/src/olap/page_cache.cpp
+++ b/be/src/olap/page_cache.cpp
@@ -26,16 +26,14 @@
 namespace doris {
 template <typename TAllocator>
 PageBase<TAllocator>::PageBase(size_t b, bool use_cache, segment_v2::PageTypePB page_type)
-        : LRUCacheValueBase(),
-          _size(b),
-          _capacity(b),
-          _use_cache(use_cache),
-          _page_type(page_type) {
-    if (_use_cache) {
-        SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(
-                StoragePageCache::instance()->mem_tracker(_page_type));
-        _data = reinterpret_cast<char*>(TAllocator::alloc(_capacity, ALLOCATOR_ALIGNMENT_16));
+        : LRUCacheValueBase(), _size(b), _capacity(b) {
+    if (use_cache) {
+        _mem_tracker_by_allocator = StoragePageCache::instance()->mem_tracker(page_type);
     } else {
+        _mem_tracker_by_allocator = thread_context()->thread_mem_tracker_mgr->limiter_mem_tracker();
+    }
+    {
+        SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(_mem_tracker_by_allocator);
         _data = reinterpret_cast<char*>(TAllocator::alloc(_capacity, ALLOCATOR_ALIGNMENT_16));
     }
 }
@@ -44,13 +42,8 @@ template <typename TAllocator>
 PageBase<TAllocator>::~PageBase() {
     if (_data != nullptr) {
         DCHECK(_capacity != 0 && _size != 0);
-        if (_use_cache) {
-            SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(
-                    StoragePageCache::instance()->mem_tracker(_page_type));
-            TAllocator::free(_data, _capacity);
-        } else {
-            TAllocator::free(_data, _capacity);
-        }
+        SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(_mem_tracker_by_allocator);
+        TAllocator::free(_data, _capacity);
     }
 }
 
diff --git a/be/src/olap/page_cache.h b/be/src/olap/page_cache.h
index 23b3574a10a54a..09fc689959ce4c 100644
--- a/be/src/olap/page_cache.h
+++ b/be/src/olap/page_cache.h
@@ -60,8 +60,7 @@ class PageBase : private TAllocator, public LRUCacheValueBase {
     // Effective size, smaller than capacity, such as data page remove checksum suffix.
     size_t _size = 0;
     size_t _capacity = 0;
-    bool _use_cache;
-    segment_v2::PageTypePB _page_type;
+    std::shared_ptr<MemTrackerLimiter> _mem_tracker_by_allocator;
 };
 
 using DataPage = PageBase<Allocator<false>>;