add memory_tracker_holder in TrackedMppDataPacket (#5840)

close #5830
pingcap · Sep 10, 2022 · 51b9ae4 · 51b9ae4
1 parent 6aa3b33
commit 51b9ae4
Showing 1 changed file with 10 additions and 4 deletions.
diff --git a/dbms/src/Flash/Mpp/TrackedMppDataPacket.h b/dbms/src/Flash/Mpp/TrackedMppDataPacket.h
@@ -46,17 +46,22 @@ inline size_t estimateAllocatedSize(const mpp::MPPDataPacket & data)
     return ret;
 }
 
+inline std::shared_ptr<MemoryTracker> getSharedPtrOfMemTracker(MemoryTracker * memory_tracker)
+{
+    return memory_tracker ? memory_tracker->shared_from_this() : nullptr;
+}
+
 struct MemTrackerWrapper
 {
     MemTrackerWrapper(size_t _size, MemoryTracker * memory_tracker)
-        : memory_tracker(memory_tracker)
+        : memory_tracker(getSharedPtrOfMemTracker(memory_tracker))
         , size(0)
     {
         alloc(_size);
     }
 
     explicit MemTrackerWrapper(MemoryTracker * memory_tracker)
-        : memory_tracker(memory_tracker)
+        : memory_tracker(getSharedPtrOfMemTracker(memory_tracker))
         , size(0)
     {}
 
@@ -88,7 +93,7 @@ struct MemTrackerWrapper
     {
         int bak_size = size;
         freeAll();
-        memory_tracker = new_memory_tracker;
+        memory_tracker = getSharedPtrOfMemTracker(new_memory_tracker);
         alloc(bak_size);
     }
     ~MemTrackerWrapper()
@@ -100,7 +105,8 @@ struct MemTrackerWrapper
     {
         free(size);
     }
-    MemoryTracker * memory_tracker;
+
+    std::shared_ptr<MemoryTracker> memory_tracker;
     size_t size = 0;
 };